` with Tailwind classes (or `Typography.Text`) for text, `Card` from `antd`, etc. Note that `antd` has no `"yellow"` Tag color — use `"gold"` for amber/yellow.
- - The only exception is the Tremor Table component and its required Tremor Table sub components.
-
-2. **Use Common Components as much as possible**:
- - These are usually defined in the `common_components` directory
- - Use these components as much as possible and avoid building new components unless needed
-
-3. **Testing**:
- - The codebase uses **Vitest** and **React Testing Library**
- - **Query Priority Order**: Use query methods in this order: `getByRole`, `getByLabelText`, `getByPlaceholderText`, `getByText`, `getByTestId`
- - **Always use `screen`** instead of destructuring from `render()` (e.g., use `screen.getByText()` not `getByText`)
- - **Wrap user interactions in `act()`**: Always wrap `fireEvent` calls with `act()` to ensure React state updates are properly handled
- - **Use `query` methods for absence checks**: Use `queryBy*` methods (not `getBy*`) when expecting an element to NOT be present
- - **Test names must start with "should"**: All test names should follow the pattern `it("should ...")`
- - **Mock external dependencies**: Check `setupTests.ts` for global mocks and mock child components/networking calls as needed
- - **Structure tests properly**:
- - First test should verify the component renders successfully
- - Subsequent tests should focus on functionality and user interactions
- - Use `waitFor` for async operations that aren't already awaited
- - **Avoid using `querySelector`**: Prefer React Testing Library queries over direct DOM manipulation
-
-### IMPORTANT PATTERNS
-
-1. **Function/Tool Calling**:
- - LiteLLM standardizes tool calling across providers
- - OpenAI format is the standard, with transformations for other providers
- - See `litellm/llms/anthropic/chat/transformation.py` for complex tool handling
-
-2. **Streaming**:
- - All providers should support streaming where possible
- - Use consistent chunk formatting across providers
- - Handle both sync and async streaming
-
-3. **Error Handling**:
- - Use provider-specific exception classes
- - Maintain consistent error formats across providers
- - Include proper retry logic and fallback mechanisms
-
-4. **Configuration**:
- - Support both environment variables and programmatic configuration
- - Use `BaseConfig` classes for provider configurations
- - Allow dynamic parameter passing
-
-## PROXY SERVER (LLM GATEWAY)
-
-The proxy server is a critical component that provides:
-- Authentication and authorization
-- Rate limiting and budget management
-- Load balancing across multiple models/deployments
-- Observability and logging
-- Admin dashboard UI
-- Enterprise features
-
-Key files:
-- `litellm/proxy/proxy_server.py` - Main server implementation
-- `litellm/proxy/auth/` - Authentication logic
-- `litellm/proxy/management_endpoints/` - Admin API endpoints
-
-**Database (proxy)**: Use Prisma model methods (`prisma_client.db..upsert`, `.find_many`, `.find_unique`, etc.), not raw SQL (`execute_raw`/`query_raw`). See COMMON PITFALLS for details.
-
-## MCP (MODEL CONTEXT PROTOCOL) SUPPORT
-
-LiteLLM supports MCP for agent workflows:
-- MCP server integration for tool calling
-- Transformation between OpenAI and MCP tool formats
-- Support for external MCP servers (Zapier, Jira, Linear, etc.)
-- See `litellm/experimental_mcp_client/` and `litellm/proxy/_experimental/mcp_server/`
-
-## RUNNING SCRIPTS
-
-Use `uv run python script.py` to run Python scripts in the project environment (for non-test files).
-
-## GITHUB TEMPLATES
-
-When opening issues or pull requests, follow these templates:
-
-### Bug Reports (`.github/ISSUE_TEMPLATE/bug_report.yml`)
-- Describe what happened vs. expected behavior
-- Include relevant log output
-- Specify LiteLLM version
-- Indicate if you're part of an ML Ops team (helps with prioritization)
-
-### Feature Requests (`.github/ISSUE_TEMPLATE/feature_request.yml`)
-- Clearly describe the feature
-- Explain motivation and use case with concrete examples
-
-### Pull Requests (`.github/pull_request_template.md`)
-- Add at least 1 test in `tests/litellm/`
-- Ensure `make test-unit` passes
-
-
-## TESTING CONSIDERATIONS
-
-1. **Provider Tests**: Test against real provider APIs when possible
-2. **Proxy Tests**: Include authentication, rate limiting, and routing tests
-3. **Performance Tests**: Load testing for high-throughput scenarios
-4. **Integration Tests**: End-to-end workflows including tool calling
-
-## DOCUMENTATION
-
-- Keep documentation in sync with code changes
-- Update provider documentation when adding new providers
-- Include code examples for new features
-- Update changelog and release notes
-
-## SECURITY CONSIDERATIONS
-
-- Handle API keys securely
-- Validate all inputs, especially for proxy endpoints
-- Consider rate limiting and abuse prevention
-- Follow security best practices for authentication
-
-## ENTERPRISE FEATURES
-
-- Some features are enterprise-only
-- Check `enterprise/` directory for enterprise-specific code
-- Maintain compatibility between open-source and enterprise versions
-
-## COMMON PITFALLS TO AVOID
-
-1. **Breaking Changes**: LiteLLM has many users - avoid breaking existing APIs
-2. **Provider Specifics**: Each provider has unique quirks - handle them properly
-3. **Rate Limits**: Respect provider rate limits in tests
-4. **Memory Usage**: Be mindful of memory usage in streaming scenarios
-5. **Dependencies**: Keep dependencies minimal and well-justified
-6. **UI/Backend Contract Mismatch**: When adding a new entity type to the UI, always check whether the backend endpoint accepts a single value or an array. Match the UI control accordingly (single-select vs. multi-select) to avoid silently dropping user selections
-7. **Missing Tests for New Entity Types**: When adding a new entity type (e.g., in `EntityUsage`, `UsageViewSelect`), always add corresponding tests in the existing test files and update any icon/component mocks
-8. **Raw SQL in proxy DB code**: Do not use `execute_raw` or `query_raw` for proxy database access. Use Prisma model methods (e.g. `prisma_client.db.litellm_tooltable.upsert()`, `.find_many()`, `.find_unique()`) so behavior stays consistent with the schema, the client stays mockable in tests, and you avoid the pitfalls of hand-written SQL (parameter ordering, type casting, schema drift)
-
-8. **Do not hardcode model-specific flags**: Put model-specific capability flags in `model_prices_and_context_window.json` and read them via `get_model_info` (or existing helpers like `supports_reasoning`). This prevents users from needing to upgrade LiteLLM each time a new model supports a feature.
-
- **Example of BAD** (hardcoded model checks):
-
- ```python
- @staticmethod
- def _is_effort_supported_model(model: str) -> bool:
- """Check if the model supports the output_config.effort parameter..."""
- model_lower = model.lower()
- if AnthropicConfig._is_claude_4_6_model(model):
- return True
- return any(
- v in model_lower for v in ("opus-4-5", "opus_4_5", "opus-4.5", "opus_4.5")
- )
- ```
-
- **Example of GOOD** (config-driven or helper that reads from config):
-
- ```python
- if (
- "claude-3-7-sonnet" in model
- or AnthropicConfig._is_claude_4_6_model(model)
- or supports_reasoning(
- model=model,
- custom_llm_provider=self.custom_llm_provider,
- )
- ):
- ...
- ```
-
- Using helpers like `supports_reasoning` (which read from `model_prices_and_context_window.json` / `get_model_info`) allows future model updates to "just work" without code changes.
-
-9. **Never close HTTP/SDK clients on cache eviction**: Do not add `close()`, `aclose()`, or `create_task(close_fn())` inside `LLMClientCache._remove_key()` or any cache eviction path. Evicted clients may still be held by in-flight requests; closing them causes `RuntimeError: Cannot send a request, as the client has been closed.` in production after the cache TTL (1 hour) expires. Connection cleanup is handled at shutdown by `close_litellm_async_clients()`. See PR #22247 for the full incident history.
-
-## HELPFUL RESOURCES
-
-- Main documentation: https://docs.litellm.ai/ (source: [BerriAI/litellm-docs](https://github.com/BerriAI/litellm-docs))
-- Provider-specific docs: https://docs.litellm.ai/docs/providers/
-- Admin UI for testing proxy features
-
-## WHEN IN DOUBT
-
-- Follow existing patterns in the codebase
-- Check similar provider implementations
-- Ensure comprehensive test coverage
-- Update documentation appropriately
-- Consider backward compatibility impact
-
-## Cursor Cloud specific instructions
-
-### Environment
-
-- uv is installed in `~/.local/bin`; the update script ensures it is on `PATH`.
-- Python 3.12, Node 22 are pre-installed.
-- The project virtual environment lives under `.venv/`.
-
-### Running the proxy server
-
-Create a minimal config file and start the proxy:
-
-```yaml
-# config.yaml
-model_list:
- - model_name: fake-openai-endpoint
- litellm_params:
- model: openai/fake-model
- api_key: fake-key
- api_base: https://fake-api.example.com
-
-general_settings:
- master_key: sk-1234
-
-litellm_settings:
- drop_params: True
- telemetry: False
-```
-
-```bash
-uv run litellm --config config.yaml --port 4000
-```
-
-The proxy takes ~15-20 seconds to fully start (it runs Prisma migrations on boot). Wait for `/health` to return before sending requests. Without a PostgreSQL `DATABASE_URL`, the proxy connects to a default Neon dev database embedded in the `litellm-proxy-extras` package.
-
-### Running tests
-
-See `CLAUDE.md` and the `Makefile` for standard commands. Key notes:
-
-- `uv sync --group proxy-dev --extra proxy` installs the Prisma and proxy-side test dependencies used by the standard local workflow.
-- The `--timeout` pytest flag is NOT available; don't pass it.
-- Unit tests: `uv run pytest tests/test_litellm/ -x -vv -n 4`
-- **Before committing, always run `uv run black .` to format your code.** Black formatting is enforced in CI.
-- If `uv sync` fails because the lockfile is outdated, run `uv lock` and retry.
-
-### Lint
-
-```bash
-cd litellm && uv run ruff check .
-```
-
-Ruff is the primary fast linter. For the full lint suite (including mypy, black, circular imports), run `make lint` per `CLAUDE.md`.
-
-### UI Dashboard development
-
-- The UI is at `ui/litellm-dashboard/`. Run `npm run dev` from that directory for the Next.js dev server on port 3000.
-- The proxy at port 4000 serves a **pre-built** static UI from `litellm/proxy/_experimental/out/`. After making UI code changes, you must run `npm run build` in the dashboard directory and copy the output: `cp -r ui/litellm-dashboard/out/* litellm/proxy/_experimental/out/` for the proxy to serve the updated UI.
-- SVGs used as provider logos (loaded via `![]()
` tags) must NOT use `fill="currentColor"` — replace with an explicit color like `#000000` or use the `-color` variant from lobehub icons, since CSS color inheritance does not work inside `![]()
` elements.
-- Provider logos live in `ui/litellm-dashboard/public/assets/logos/` (source) and `litellm/proxy/_experimental/out/assets/logos/` (pre-built). Both locations must have the file for it to work in dev and proxy-served modes.
-- UI Vitest tests: `cd ui/litellm-dashboard && npx vitest run`
+Read @CLAUDE.md for coding guidelines
diff --git a/CLAUDE.md b/CLAUDE.md
index b9a336b8f40..3477b71a621 100644
--- a/CLAUDE.md
+++ b/CLAUDE.md
@@ -1,194 +1,70 @@
-# CLAUDE.md
+Do not write comments unless they are absolutely necessary to explain some very complex business logic. Please clean up if there are comments that are not absolutely necessary. Do not remove comments that are unrelated to the addition of the code of this PR
-This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
+Explanation: code comments are, in a way, a violation of DRY code. You must update logic in two locations to change the code and "hard to change" is literally the definition of tech debt. We should instead aim to write code that is intuitive to the reader, while being both easy to maintain and high performance
-## Confidentiality: Customer and Company Names in Code
+Don't assume that the existing code is correct or the right way of doing things / good coding patterns. In fact, there are a lot of bad coding practices, overly complex code, code smells, etc. If something doesn't look right, speak up. Feel free to break existing patterns or question weird existing code to make new code high quality, as in:
+- correct
+- secure
+- performant
+- readable
+- easy to maintain/change
+- modern
+In that order of importance
-The codebase is public. Before writing **any** third-party organization name into this repository — in source code, file or directory names, docstrings, comments, tests, fixtures, mock payloads, error messages, log lines, commit messages, or PR descriptions — pause and check:
+When adding new features, add meaningful tests. Don't add tests that don't check anything substantial and is there just to make the code coverage pass. Yes, code coverage is important, but I'd rather have no signal whether the code is working than tests that don't fail when code is broken. The goal is to have tests that would fail before the feature was added/if the code was mutated in a way that breaks the feature and succeed only when the feature is fully working. I should run mutation testing and see > 90% kill rate
-**Already in the codebase** (OpenAI, Anthropic, Google, Azure, Bedrock, Fireworks, and other established LLM providers / integrations) — fine to use. Quick check: `git grep -i ""` — if it returns hits in real code (not just your current diff), the name is established.
+Same thing for bug fixes. The tests should make it so that this specific bug can never happen again without failing tests (i.e., regression)
-**Anything else** — customers, prospects, partners, new vendor integrations, observability tools, infra vendors, or any organization name that does not already appear in the repo. STOP and surface it to the user. Ask for explicit consent before writing the name into any file, commit message, or PR description. Do not write it speculatively and clean up later. Do not substitute a placeholder and proceed. Do not assume it is safe because it "looks like" a public company. The user must approve first.
+When creating PRs, don't set base to `main`. `litellm_internal_staging` serves that purpose
-**What to do instead of a customer-specific reference:**
-- If you find yourself reaching for a customer name — real or fake — step back. The code shouldn't be customer-specific in the first place. Generalize the feature, or capture the customer motivation in internal docs (Notion / Linear / the internal staging PR description), never in the repo.
-- Frame changes by the capability they add, not the customer who asked for it ("add per-team Bedrock guardrail routing", not "add routing for $CUSTOMER").
-- Standard "fake value" markers (`example.com`, `localhost`, `127.0.0.1`, `test@example.com`) and abstract identifiers (`team_a`, `user_1`, `tenant_x`) are fine — those are not customer stand-ins.
+Always use @.github/pull_request_template.md as a guide for your PR body
-## Documentation
+Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, just tell me which URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, what fields to fill out, etc. along with the other commands to run in an ordered list, and I'll do it myself and post the screenshots after you make the PR
-Documentation lives in a separate repository: [BerriAI/litellm-docs](https://github.com/BerriAI/litellm-docs). It is served at [docs.litellm.ai](https://docs.litellm.ai). Do not create or edit documentation files in this repository — open doc PRs against `BerriAI/litellm-docs` instead.
+If you ever make public-facing PR descriptions, comments, issues, commit messages, etc., always follow these guidelines to sound less AI-y:
+- don't use emojis
+- don't use "—". Instead, reach for ";", ".", etc.
+- don't use the pattern "It's not X, it's Y", "You're not X, you're Y", etc.
+- don't use bulleted or numbered lists unless it would be nonsensical not to. Instead, prefer prose
-## Development Commands
+Don't hesitate to use values in .env to get needed API keys and other secrets, as long as you never add them to conversation history, commit them, or include them in GitHub issues / PRs
-### Installation
-- `make install-dev` - Install core development dependencies
-- `make install-proxy-dev` - Install proxy development dependencies with full feature set
-- `make install-test-deps` - Install the full local test environment and generate the Prisma client
+Run tests, format your code, and lint your code before each commit
-### Testing
-- `make test` - Run all tests
-- `make test-unit` - Run unit tests (tests/test_litellm) with 4 parallel workers
-- `make test-integration` - Run integration tests (excludes unit tests)
-- `pytest tests/` - Direct pytest execution
+Ask to commit and push your work when you're done (or if you're confident that your code is good and works, just do it)
-### Code Quality
-- `make lint` - Run all linting (Ruff, MyPy, Black, circular imports, import safety)
-- `make format` - Apply Black code formatting
-- `make lint-ruff` - Run Ruff linting only
-- `make lint-mypy` - Run MyPy type checking only
-- **Before committing, always run `uv run black .` to format your code.** Black formatting is enforced in CI.
+When you must use real LLM models to, for example, write e2e tests, write a QA runbook, etc., make sure to use the latest models (doesn't have to be smartest, can also be a modern small, fast one. No strong preference for smart vs fast here, just use something modern) as of the year and month of the current date. Do a web search as necessary to figure that out
-### Single Test Files
-- `uv run pytest tests/path/to/test_file.py -v` - Run specific test file
-- `uv run pytest tests/path/to/test_file.py::test_function -v` - Run specific test
+If you're an internal contributor, when creating a new PR, the typical flow is to branch off litellm_internal_staging and create a branch prefixed with litellm_. Do not create a branch prefixed with claude/ and generally do not have / in your branch names
-### Running Scripts
-- `uv run python script.py` - Run Python scripts (use for non-test files)
+Do not add `Co-Authored-By: Claude` or any Claude attribution to commit messages. Never use a `claude/` prefix or put a `/` in a branch name. Do not add "Generated with Claude Code" (or any similar attribution) to PR descriptions. Do not create a new PR/branch off the existing PR to fix/add something that is related and could've just been committed directly to the existing PR's branch
-### GitHub Issue & PR Templates
-When contributing to the project, use the appropriate templates:
+When working on a PR, keep the PR description in sync with new commits being made
-**Bug Reports** (`.github/ISSUE_TEMPLATE/bug_report.yml`):
-- Describe what happened vs. what you expected
-- Include relevant log output
-- Specify your LiteLLM version
+Monkeypatching attributes of a class to do testing is an anti-pattern. Prefer dependency-injecting things into classes. That way, at unit test time, you can pass a mocked dependency in
-**Feature Requests** (`.github/ISSUE_TEMPLATE/feature_request.yml`):
-- Describe the feature clearly
-- Explain the motivation and use case
+Do not put names of customers or customer company names in code, PRs, and issues. The codebase is public
-**Pull Requests** (`.github/pull_request_template.md`):
-- Add at least 1 test in `tests/litellm/`
-- Ensure `make test-unit` passes
+CI supply-chain safety: Never pipe a remote script into a shell (`curl ... | bash`, `wget ... | sh`); download the artifact to a file, verify its SHA-256 checksum, then install. Pin every external tool to a specific version with a full URL (not `latest` or `stable`). Verify checksums for all downloaded binaries, using the provider's official `.sha256` / `.sha256sum` sidecar when available. These rules apply to every download in CI
-## Architecture Overview
+## Think Before Coding
-LiteLLM is a unified interface for 100+ LLM providers with two main components:
+**Don't assume. Don't hide confusion. Surface tradeoffs.**
-### Core Library (`litellm/`)
-- **Main entry point**: `litellm/main.py` - Contains core completion() function
-- **Provider implementations**: `litellm/llms/` - Each provider has its own subdirectory
-- **Router system**: `litellm/router.py` + `litellm/router_utils/` - Load balancing and fallback logic
-- **Type definitions**: `litellm/types/` - Pydantic models and type hints
-- **Integrations**: `litellm/integrations/` - Third-party observability, caching, logging
-- **Caching**: `litellm/caching/` - Multiple cache backends (Redis, in-memory, S3, etc.)
+Before implementing:
+- State your assumptions explicitly. If uncertain, ask.
+- If multiple interpretations exist, present them. Don't pick silently.
+- If a simpler approach exists, say so. Push back when warranted.
+- If something is unclear, stop. Name what's confusing. Ask.
-### Proxy Server (`litellm/proxy/`)
-- **Main server**: `proxy_server.py` - FastAPI application
-- **Authentication**: `auth/` - API key management, JWT, OAuth2
-- **Database**: `db/` - Prisma ORM with PostgreSQL/SQLite support
-- **Management endpoints**: `management_endpoints/` - Admin APIs for keys, teams, models
-- **Pass-through endpoints**: `pass_through_endpoints/` - Provider-specific API forwarding
-- **Guardrails**: `guardrails/` - Safety and content filtering hooks
-- **UI Dashboard**: Served from `_experimental/out/` (Next.js build)
+## Simplicity First
-## Key Patterns
+**Minimum code that solves the problem. Nothing speculative.**
-### Provider Implementation
-- Providers inherit from base classes in `litellm/llms/base.py`
-- Each provider has transformation functions for input/output formatting
-- Support both sync and async operations
-- Handle streaming responses and function calling
+- No features beyond what was asked.
+- No abstractions for single-use code.
+- No "flexibility" or "configurability" that wasn't requested.
+- No error handling for impossible scenarios.
+- If you write 200 lines and it could be 50, rewrite it.
-### Error Handling
-- Provider-specific exceptions mapped to OpenAI-compatible errors
-- Fallback logic handled by Router system
-- Comprehensive logging through `litellm/_logging.py`
-
-### Configuration
-- YAML config files for proxy server (see `proxy/example_config_yaml/`)
-- Environment variables for API keys and settings
-- Database schema managed via Prisma (`proxy/schema.prisma`)
-
-## Development Notes
-
-### Code Style
-- Uses Black formatter, Ruff linter, MyPy type checker
-- Pydantic v2 for data validation
-- Async/await patterns throughout
-- Type hints required for all public APIs
-- **Avoid imports within methods** — place all imports at the top of the file (module-level). Inline imports inside functions/methods make dependencies harder to trace and hurt readability. The only exception is avoiding circular imports where absolutely necessary.
-- **Use dict spread for immutable copies** — prefer `{**original, "key": new_value}` over `dict(obj)` + mutation. The spread produces the final dict in one step and makes intent clear.
-- **Guard at resolution time** — when resolving an optional value through a fallback chain (`a or b or ""`), raise immediately if the resolved result being empty is an error. Don't pass empty strings or sentinel values downstream for the callee to deal with.
-- **Extract complex comprehensions to named helpers** — a set/dict comprehension that calls into the DB or manager (e.g. "which of these server IDs are OAuth2?") belongs in a named helper function, not inline in the caller.
-- **FastAPI parameter declarations** — mark required query/form params with `= Query(...)` / `= Form(...)` explicitly when other params in the same handler are optional. Mixing `str` (required) with `Optional[str] = None` in the same signature causes silent 422s when the required param is missing.
-
-### Testing Strategy
-- Unit tests in `tests/test_litellm/`
-- Integration tests for each provider in `tests/llm_translation/`
-- Proxy tests in `tests/proxy_unit_tests/`
-- Load tests in `tests/load_tests/`
-- **Always add tests when adding new entity types or features** — if the existing test file covers other entity types, add corresponding tests for the new one
-- **Keep monkeypatch stubs in sync with real signatures** — when a function gains a new optional parameter, update every `fake_*` / `stub_*` in tests that patch it to also accept that kwarg (even as `**kwargs`). Stale stubs fail with `unexpected keyword argument` and mask real bugs.
-- **Test all branches of name→ID resolution** — when adding server/resource lookup that resolves names to UUIDs, test: (1) name resolves and UUID is allowed, (2) name resolves but UUID is not allowed, (3) name does not resolve at all. The silent-fallback path is where access-control bugs hide.
-
-### UI / Backend Consistency
-- When wiring a new UI entity type to an existing backend endpoint, verify the backend API contract (single value vs. array, required vs. optional params) and ensure the UI controls match — e.g., use a single-select dropdown when the backend accepts a single value, not a multi-select
-
-### UI Component Library
-- **Always use `antd` for new UI components** — we are migrating off of `@tremor/react`. Do not introduce new `Badge`, `Text`, `Card`, `Grid`, `Title`, or other imports from `@tremor/react` in any new or modified file. Use `antd` equivalents: `Tag` for labels, `Typography.Text` / `Typography.Title` / `Typography.Paragraph` for textual content (avoid plain text-only ``, `
diff --git a/ui/litellm-dashboard/src/components/agents/agent_card_discovery.test.tsx b/ui/litellm-dashboard/src/components/agents/agent_card_discovery.test.tsx
new file mode 100644
index 00000000000..18b874c95ee
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/agents/agent_card_discovery.test.tsx
@@ -0,0 +1,299 @@
+import React from "react";
+import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
+import { screen, waitFor } from "@testing-library/react";
+import userEvent from "@testing-library/user-event";
+import { renderWithProviders } from "../../../tests/test-utils";
+import AgentCardDiscovery from "./agent_card_discovery";
+
+vi.mock("../networking", async () => {
+ const actual = await vi.importActual`, `
+
+ Using the connection details you entered above. We'll fetch:
+
+
+ Paste the upstream agent's base URL. We'll try{" "}
+
+
+
+ setManualUrl(e.target.value)}
+ onPressEnter={handleDiscover}
+ allowClear
+ disabled={loading}
+ />
+ :
+
+ )}
+
+ {error && (
+ setError(null)}
+ />
+ )}
+
+ {loading && !card && (
+
+ );
+};
+
+export default AgentCardDiscovery;
diff --git a/ui/litellm-dashboard/src/components/agents/agent_discovery_utils.test.ts b/ui/litellm-dashboard/src/components/agents/agent_discovery_utils.test.ts
new file mode 100644
index 00000000000..632fba1532c
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/agents/agent_discovery_utils.test.ts
@@ -0,0 +1,64 @@
+import { describe, it, expect } from "vitest";
+import {
+ selectionsFromSavedAgentCard,
+ selectionsFromUpstreamCard,
+ skillId,
+} from "./agent_discovery_utils";
+
+const upstreamCard = {
+ name: "Upstream Agent",
+ description: "Upstream description",
+ capabilities: { streaming: true },
+ skills: [
+ { id: "search", name: "Search", description: "Search the web" },
+ { id: "summarize", name: "Summarize", description: "Summarize docs" },
+ { id: "chat", name: "Chat", description: "General chat" },
+ ],
+};
+
+describe("selectionsFromSavedAgentCard", () => {
+ it("pre-selects only skills that exist in the saved DB card", () => {
+ const savedCard = {
+ name: "My Agent",
+ description: "Saved description",
+ capabilities: { streaming: false },
+ skills: [{ id: "search", name: "Search" }],
+ };
+
+ const result = selectionsFromSavedAgentCard(upstreamCard, savedCard);
+
+ expect(result.editedName).toBe("My Agent");
+ expect(result.editedDescription).toBe("Saved description");
+ expect(result.selectedCapabilities.streaming).toBe(false);
+ expect(result.selectedSkillIds.has(skillId(upstreamCard.skills![0], 0))).toBe(
+ true,
+ );
+ expect(
+ result.selectedSkillIds.has(skillId(upstreamCard.skills![1], 1)),
+ ).toBe(false);
+ expect(
+ result.selectedSkillIds.has(skillId(upstreamCard.skills![2], 2)),
+ ).toBe(false);
+ });
+
+ it("matches saved skills by name when id is missing", () => {
+ const savedCard = {
+ skills: [{ name: "Summarize" }],
+ };
+
+ const result = selectionsFromSavedAgentCard(upstreamCard, savedCard);
+
+ expect(
+ result.selectedSkillIds.has(skillId(upstreamCard.skills![1], 1)),
+ ).toBe(true);
+ expect(result.selectedSkillIds.size).toBe(1);
+ });
+});
+
+describe("selectionsFromUpstreamCard", () => {
+ it("selects all upstream skills for create flow", () => {
+ const result = selectionsFromUpstreamCard(upstreamCard);
+ expect(result.selectedSkillIds.size).toBe(3);
+ expect(result.editedName).toBe("Upstream Agent");
+ });
+});
diff --git a/ui/litellm-dashboard/src/components/agents/agent_discovery_utils.ts b/ui/litellm-dashboard/src/components/agents/agent_discovery_utils.ts
new file mode 100644
index 00000000000..bb3d79e259b
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/agents/agent_discovery_utils.ts
@@ -0,0 +1,176 @@
+import {
+ AgentCreateInfo,
+ DiscoveredAgentCard,
+ DiscoveryMode,
+} from "../networking";
+
+export interface DiscoveryRequestPlan {
+ url: string;
+ discovery_mode: DiscoveryMode;
+ params?: Record
+ Discover from agent URL
+
+
+
+ {isParentDriven ? (
+ <>
+
+ {discoveryRequest!.display_url || effectiveUrl || (
+
+ Fill in the fields above first
+
+ )}
+
+
+ :
+
+ ) : (
+ <>
+ /.well-known/agent-card.json,{" "}
+ /.well-known/agent.json, and /agent.json{" "}
+ in order.
+
+
+ )}
+
+ {card && (
+
+
+
+ Skills
+
+ {selectedSkillCount} / {skillCount} selected
+
+
+ }
+ >
+ {skillCount === 0 ? (
+
+
+
+ Capabilities
+
+
+
+ }
+ >
+
+
+
+
+ )}
+
+
+ Upstream card loaded
+ {card.version && v{card.version} }
+ {card.provider?.organization && (
+ {card.provider.organization}
+ )}
+
+
+
+
+
+
+
+
+ setEditedName(e.target.value)}
+ placeholder="Agent name"
+ />
+
+
+
+ setEditedDescription(e.target.value)}
+ rows={2}
+ placeholder="What this agent does"
+ />
+
+
+ {(card.skills ?? []).map((skill, idx) => {
+ const id = skillId(skill, idx);
+ const checked = selectedSkillIds.has(id);
+ return (
+
+ );
+ })}
+
+ )}
+
+ {ALLOWED_CAPABILITY_KEYS.map((key) => {
+ const upstreamHas = Boolean(card.capabilities?.[key]);
+ return (
+
+
+
+ setSelectedCapabilities((prev) => ({
+ ...prev,
+ [key]: checked,
+ }))
+ }
+ />
+
+ );
+ })}
+
+
+ {key}
+
+ {!upstreamHas && (
+
+ not advertised upstream
+
+ )}
+
+
+
+ )}
+
@@ -321,6 +399,7 @@ const AgentInfoView: React.FC = ({
{
+ setAppliedDiscoveredSelection(null);
setIsEditing(false);
fetchAgentInfo();
}}>
diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx
index 57e7d51123e..aa5e08d195f 100644
--- a/ui/litellm-dashboard/src/components/networking.tsx
+++ b/ui/litellm-dashboard/src/components/networking.tsx
@@ -6299,6 +6299,79 @@ export const createAgentCall = async (accessToken: string, agentData: any) => {
}
};
+export interface DiscoveredAgentCard {
+ protocolVersion?: string;
+ name?: string;
+ description?: string;
+ version?: string;
+ url?: string;
+ iconUrl?: string;
+ documentationUrl?: string;
+ defaultInputModes?: string[];
+ defaultOutputModes?: string[];
+ capabilities?: Record;
+ skills?: Array<{
+ id?: string;
+ name?: string;
+ description?: string;
+ tags?: string[];
+ examples?: string[];
+ [key: string]: any;
+ }>;
+ provider?: { organization?: string; url?: string };
+ [key: string]: any;
+}
+
+export interface DiscoverAgentCardResponse {
+ url: string;
+ agent_card: DiscoveredAgentCard;
+}
+
+/**
+ * How the backend should locate the upstream agent card.
+ *
+ * - ``well_known_fallback`` (default): pure A2A — try the three standard
+ * well-known paths under the base URL.
+ * - ``langgraph_platform``: LangGraph Platform — hits the canonical
+ * well-known path with an ``assistant_id`` query parameter, because
+ * LangGraph mounts one shared card endpoint per deployment.
+ */
+export type DiscoveryMode = "well_known_fallback" | "langgraph_platform";
+
+export interface DiscoverAgentCardOptions {
+ discovery_mode?: DiscoveryMode;
+ /** Mode-specific params. ``langgraph_platform`` requires ``assistant_id``. */
+ params?: Record;
+}
+
+export const discoverAgentCardCall = async (
+ accessToken: string,
+ url: string,
+ options?: DiscoverAgentCardOptions,
+): Promise => {
+ const endpoint = proxyBaseUrl ? `${proxyBaseUrl}/v1/a2a/discover` : `/v1/a2a/discover`;
+ const body: Record = { url };
+ if (options?.discovery_mode) body.discovery_mode = options.discovery_mode;
+ if (options?.params) body.params = options.params;
+
+ const response = await fetch(endpoint, {
+ method: "POST",
+ headers: {
+ [globalLitellmHeaderName]: `Bearer ${accessToken}`,
+ "Content-Type": "application/json",
+ },
+ body: JSON.stringify(body),
+ });
+
+ if (!response.ok) {
+ const errorData = await response.text();
+ handleError(errorData);
+ throw new Error(errorData);
+ }
+
+ return (await response.json()) as DiscoverAgentCardResponse;
+};
+
export const createGuardrailCall = async (accessToken: string, guardrailData: any) => {
try {
const url = proxyBaseUrl ? `${proxyBaseUrl}/guardrails` : `/guardrails`;
From c754c560ddfcb3374e695a4a4f2ae3c73a953f71 Mon Sep 17 00:00:00 2001
From: ryan-crabbe-berri
Date: Fri, 29 May 2026 21:30:45 -0700
Subject: [PATCH 17/44] fix(proxy): link passthrough success spans to the
SERVER root OTEL span (#29315)
* fix(proxy): link passthrough success spans to the SERVER root OTEL span
Passthrough requests never wired user_api_key_dict.parent_otel_span into the
logging metadata, so on success the litellm_request span orphaned into its own
trace and the "Received Proxy Server Request" root span was never ended. Setting
it once in _init_kwargs_for_pass_through_endpoint fixes both the non-streaming
and streaming paths, since update_environment_variables copies that metadata onto
the logging object's model_call_details, which is what the OTEL handler reads.
Resolves LIT-3443
* fix(proxy): set passthrough parent span after client metadata merge
Greptile flagged that litellm_parent_otel_span was assigned before the
_metadata.update() calls that merge request-body metadata, so a client body
mirroring the internal key could overwrite the real span with a JSON scalar
and null the fix for that request. Move the assignment after the merge and
add a regression test that fails on the old ordering.
* fix(proxy): also set user_api_key after client metadata merge
Per Greptile, user_api_key had the same clobber window as the parent span: a
passthrough request body mirroring the key could overwrite the authenticated
value in the logged metadata. Move it into the same post-merge block and add a
deterministic contract test asserting both internal keys resist client-supplied
metadata.
---
.../pass_through_endpoints.py | 8 +-
.../test_passthrough_parent_span.py | 333 ++++++++++++++++++
2 files changed, 339 insertions(+), 2 deletions(-)
create mode 100644 tests/test_litellm/integrations/open_telemetry/test_passthrough_parent_span.py
diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
index a9fb02a1c95..67af176c1eb 100644
--- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
+++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
@@ -543,8 +543,6 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils):
)
)
- _metadata["user_api_key"] = user_api_key_dict.api_key
-
litellm_metadata = litellm_params_in_body.pop("litellm_metadata", None)
metadata = litellm_params_in_body.pop("metadata", None)
if litellm_metadata:
@@ -557,6 +555,12 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils):
metadata=_metadata,
)
+ # Set internal keys after merging client-supplied metadata so a request
+ # body that mirrors them cannot clobber the authenticated key or the
+ # real parent span.
+ _metadata["user_api_key"] = user_api_key_dict.api_key
+ _metadata["litellm_parent_otel_span"] = user_api_key_dict.parent_otel_span
+
kwargs = {
"litellm_params": {
**litellm_params_in_body, # type: ignore
diff --git a/tests/test_litellm/integrations/open_telemetry/test_passthrough_parent_span.py b/tests/test_litellm/integrations/open_telemetry/test_passthrough_parent_span.py
new file mode 100644
index 00000000000..bd37b3e3c76
--- /dev/null
+++ b/tests/test_litellm/integrations/open_telemetry/test_passthrough_parent_span.py
@@ -0,0 +1,333 @@
+"""LIT-3443 — passthrough success spans must hang off the SERVER root span.
+
+_init_kwargs_for_pass_through_endpoint is the single place both passthrough
+paths get their logging metadata, and update_environment_variables copies that
+metadata onto the logging object's model_call_details — which is exactly what
+the OTEL success handler reads. So wiring the parent span in there once fixes
+both the non-streaming and streaming paths; the streaming handler rebuilds its
+kwargs from raw SSE bytes and never sees that metadata, but it doesn't need to.
+
+These tests drive the real passthrough logging code into the real OpenTelemetry
+success handler, capturing every span in an InMemorySpanExporter:
+
+ * non-streaming: _init_kwargs_for_pass_through_endpoint -> async_success_handler
+ * streaming: _route_streaming_logging_to_handler over real Anthropic SSE
+
+Before the fix the parent span is never wired in, so the litellm_request span
+orphans into its own trace and the SERVER root span is never ended. Each test
+asserts the SERVER root is exported (ended) and that nothing escapes into a
+foreign trace; the USE_OTEL_LITELLM_REQUEST_SPAN variants additionally assert
+the litellm_request child is parented to the SERVER root.
+"""
+
+import asyncio
+from datetime import datetime
+from typing import Optional, Tuple
+
+import pytest
+from starlette.requests import Request
+
+import litellm
+from litellm.integrations.opentelemetry import LITELLM_PROXY_REQUEST_SPAN_NAME
+from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
+ HttpPassThroughEndpointHelpers,
+)
+from litellm.proxy.pass_through_endpoints.streaming_handler import (
+ PassThroughStreamingHandler,
+)
+from litellm.proxy.pass_through_endpoints.success_handler import (
+ PassThroughEndpointLogging,
+)
+from litellm.types.passthrough_endpoints.pass_through_endpoints import (
+ EndpointType,
+ PassthroughStandardLoggingPayload,
+)
+from litellm.types.utils import Choices, Message, ModelResponse, Usage
+
+URL_ROUTE = "https://api.anthropic.com/v1/messages"
+MODEL = "claude-sonnet-4-5-20250929"
+
+
+@pytest.fixture
+def otel_success_callback(otel_with_exporter, monkeypatch):
+ """Register our in-memory OTEL instance where async_success_handler looks
+ for success callbacks (litellm._async_success_callback), so the real
+ logging path drives it."""
+ otel, exporter = otel_with_exporter
+ monkeypatch.setattr(litellm, "callbacks", [otel])
+ monkeypatch.setattr(litellm, "_async_success_callback", [otel])
+ return otel, exporter
+
+
+def _make_request() -> Request:
+ return Request(
+ {
+ "type": "http",
+ "method": "POST",
+ "path": "/anthropic/v1/messages",
+ "raw_path": b"/anthropic/v1/messages",
+ "query_string": b"",
+ "headers": [(b"content-type", b"application/json")],
+ "scheme": "http",
+ "server": ("testserver", 80),
+ "client": ("testclient", 50000),
+ }
+ )
+
+
+def _build_logging_obj_wired_to_root(
+ root_span, *, stream: bool, extra_body: Optional[dict] = None
+) -> Tuple[LiteLLMLoggingObj, dict, datetime]:
+ """Mirror pass_through_endpoints.py: build the logging object and run the
+ real _init_kwargs + update_environment_variables so the parent span lands
+ on model_call_details exactly the way production wires it."""
+ request = _make_request()
+ body = {"model": MODEL, "messages": [{"role": "user", "content": "hi"}]}
+ if extra_body:
+ body.update(extra_body)
+ user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", parent_otel_span=root_span)
+ start_time = datetime.now()
+ logging_obj = LiteLLMLoggingObj(
+ model="unknown",
+ messages=[{"role": "user", "content": "hi"}],
+ stream=stream,
+ call_type="pass_through_endpoint",
+ start_time=start_time,
+ litellm_call_id="lit-3443-call",
+ function_id="1245",
+ )
+ payload = PassthroughStandardLoggingPayload(
+ url=URL_ROUTE, request_body=body, request_method="POST"
+ )
+ kwargs = HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint(
+ request=request,
+ user_api_key_dict=user_api_key_dict,
+ passthrough_logging_payload=payload,
+ logging_obj=logging_obj,
+ _parsed_body=body,
+ litellm_call_id="lit-3443-call",
+ )
+ logging_obj.update_environment_variables(
+ model="unknown",
+ user="unknown",
+ optional_params={},
+ litellm_params=kwargs["litellm_params"],
+ call_type="pass_through_endpoint",
+ )
+ logging_obj.model_call_details["litellm_call_id"] = "lit-3443-call"
+ return logging_obj, kwargs, start_time
+
+
+def _model_response() -> ModelResponse:
+ resp = ModelResponse()
+ resp.model = MODEL
+ resp.choices = [Choices(message=Message(role="assistant", content="hi there"))]
+ resp.usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15)
+ return resp
+
+
+# Real Anthropic SSE stream (single text block) reused for the streaming path.
+STREAM_CHUNKS = [
+ "event: message_start",
+ 'data: {"type":"message_start","message":{"id":"msg_1","type":"message","role":"assistant","model":"claude-sonnet-4-5-20250929","content":[],"stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":17,"output_tokens":5}}}',
+ "event: content_block_start",
+ 'data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}',
+ "event: content_block_delta",
+ 'data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello world"}}',
+ "event: content_block_stop",
+ 'data: {"type":"content_block_stop","index":0}',
+ "event: message_delta",
+ 'data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":2}}',
+ "event: message_stop",
+ 'data: {"type":"message_stop"}',
+]
+
+
+def _assert_root_closed_and_no_orphan(exporter, root_span, where):
+ finished = exporter.get_finished_spans()
+ root_ctx = root_span.get_span_context()
+
+ server_spans = [s for s in finished if s.name == LITELLM_PROXY_REQUEST_SPAN_NAME]
+ assert server_spans, (
+ f"{where}: SERVER root span was never ended/exported — exporter saw "
+ f"{[s.name for s in finished]}"
+ )
+
+ foreign = [s for s in finished if s.context.trace_id != root_ctx.trace_id]
+ assert not foreign, (
+ f"{where}: span(s) orphaned into a foreign trace: "
+ f"{[(s.name, hex(s.context.trace_id)) for s in foreign]} "
+ f"(root trace={hex(root_ctx.trace_id)})"
+ )
+
+
+def _assert_child_parented_to_root(exporter, root_span, where):
+ finished = exporter.get_finished_spans()
+ root_ctx = root_span.get_span_context()
+ children = [
+ s
+ for s in finished
+ if s.name != LITELLM_PROXY_REQUEST_SPAN_NAME
+ and s.parent is not None
+ and s.parent.span_id == root_ctx.span_id
+ ]
+ assert children, (
+ f"{where}: no litellm_request child parented to the SERVER root — "
+ f"finished={[(s.name, s.parent and hex(s.parent.span_id)) for s in finished]}"
+ )
+ for child in children:
+ assert child.context.trace_id == root_ctx.trace_id, (
+ f"{where}: child {child.name} in trace {hex(child.context.trace_id)}, "
+ f"expected root trace {hex(root_ctx.trace_id)}"
+ )
+
+
+@pytest.mark.parametrize("use_request_span", [False, True])
+def test_non_streaming_passthrough_links_to_server_root(
+ otel_success_callback,
+ server_span_factory,
+ monkeypatch,
+ use_request_span,
+):
+ if use_request_span:
+ monkeypatch.setenv("USE_OTEL_LITELLM_REQUEST_SPAN", "true")
+ _otel, exporter = otel_success_callback
+ root = server_span_factory("/anthropic/v1/messages")
+
+ logging_obj, kwargs, start_time = _build_logging_obj_wired_to_root(
+ root, stream=False
+ )
+ end_time = datetime.now()
+ asyncio.run(
+ logging_obj.async_success_handler(
+ result=_model_response(),
+ start_time=start_time,
+ end_time=end_time,
+ cache_hit=False,
+ **kwargs,
+ )
+ )
+
+ where = f"non-streaming (use_request_span={use_request_span})"
+ _assert_root_closed_and_no_orphan(exporter, root, where)
+ if use_request_span:
+ _assert_child_parented_to_root(exporter, root, where)
+
+
+@pytest.mark.parametrize("use_request_span", [False, True])
+def test_streaming_passthrough_links_to_server_root(
+ otel_success_callback,
+ server_span_factory,
+ monkeypatch,
+ use_request_span,
+):
+ if use_request_span:
+ monkeypatch.setenv("USE_OTEL_LITELLM_REQUEST_SPAN", "true")
+ _otel, exporter = otel_success_callback
+ root = server_span_factory("/anthropic/v1/messages")
+
+ logging_obj, _kwargs, start_time = _build_logging_obj_wired_to_root(
+ root, stream=True
+ )
+ raw_bytes = ["\n".join(STREAM_CHUNKS).encode("utf-8")]
+ end_time = datetime.now()
+ asyncio.run(
+ PassThroughStreamingHandler._route_streaming_logging_to_handler(
+ litellm_logging_obj=logging_obj,
+ passthrough_success_handler_obj=PassThroughEndpointLogging(),
+ url_route=URL_ROUTE,
+ request_body={"model": MODEL, "stream": True},
+ endpoint_type=EndpointType.ANTHROPIC,
+ start_time=start_time,
+ raw_bytes=raw_bytes,
+ end_time=end_time,
+ )
+ )
+
+ where = f"streaming (use_request_span={use_request_span})"
+ _assert_root_closed_and_no_orphan(exporter, root, where)
+ if use_request_span:
+ _assert_child_parented_to_root(exporter, root, where)
+
+
+def test_client_body_metadata_cannot_clobber_parent_span(
+ otel_success_callback,
+ server_span_factory,
+ monkeypatch,
+):
+ """A passthrough request body whose metadata mirrors the internal
+ litellm_parent_otel_span key must not override the real parent span. The
+ internal span is wired after the client-metadata merge, so the SERVER root
+ still links and closes. With the old ordering the JSON scalar would win and
+ the litellm_request span would orphan."""
+ monkeypatch.setenv("USE_OTEL_LITELLM_REQUEST_SPAN", "true")
+ _otel, exporter = otel_success_callback
+ root = server_span_factory("/anthropic/v1/messages")
+
+ logging_obj, kwargs, start_time = _build_logging_obj_wired_to_root(
+ root,
+ stream=False,
+ extra_body={"metadata": {"litellm_parent_otel_span": "not-a-real-span"}},
+ )
+ end_time = datetime.now()
+ asyncio.run(
+ logging_obj.async_success_handler(
+ result=_model_response(),
+ start_time=start_time,
+ end_time=end_time,
+ cache_hit=False,
+ **kwargs,
+ )
+ )
+
+ where = "client-metadata-clobber"
+ _assert_root_closed_and_no_orphan(exporter, root, where)
+ _assert_child_parented_to_root(exporter, root, where)
+
+
+def test_init_kwargs_internal_keys_resist_client_metadata(server_span_factory):
+ """Deterministic contract test on _init_kwargs_for_pass_through_endpoint:
+ a request body whose metadata mirrors the internal user_api_key and
+ litellm_parent_otel_span keys must not override the authenticated values.
+ Pure dict assertion, no async or OTEL execution. Fails on the old ordering
+ where the client values were merged in last."""
+ real_span = server_span_factory("/anthropic/v1/messages")
+ user_api_key_dict = UserAPIKeyAuth(
+ api_key="sk-real-key", parent_otel_span=real_span
+ )
+ body = {
+ "model": MODEL,
+ "messages": [{"role": "user", "content": "hi"}],
+ "metadata": {
+ "user_api_key": "sk-SPOOFED",
+ "litellm_parent_otel_span": "not-a-real-span",
+ },
+ }
+ logging_obj = LiteLLMLoggingObj(
+ model="unknown",
+ messages=[{"role": "user", "content": "hi"}],
+ stream=False,
+ call_type="pass_through_endpoint",
+ start_time=datetime.now(),
+ litellm_call_id="lit-3443-clobber",
+ function_id="1245",
+ )
+ payload = PassthroughStandardLoggingPayload(
+ url=URL_ROUTE, request_body=body, request_method="POST"
+ )
+ kwargs = HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint(
+ request=_make_request(),
+ user_api_key_dict=user_api_key_dict,
+ passthrough_logging_payload=payload,
+ logging_obj=logging_obj,
+ _parsed_body=body,
+ litellm_call_id="lit-3443-clobber",
+ )
+ md = kwargs["litellm_params"]["metadata"]
+ # api_key is stored hashed on the auth object; the authenticated value must
+ # win over the client-supplied spoof.
+ assert md["user_api_key"] == user_api_key_dict.api_key
+ assert md["user_api_key"] != "sk-SPOOFED"
+ assert md["litellm_parent_otel_span"] is real_span
From 581c30f1e8567dab604d292039dae7f2bb7f2217 Mon Sep 17 00:00:00 2001
From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com>
Date: Fri, 29 May 2026 22:23:24 -0700
Subject: [PATCH 18/44] [internal copy of #29089] fix: duplicate claude code
traces (#29311)
---
litellm/litellm_core_utils/litellm_logging.py | 106 +++++++--
.../litellm_core_utils/streaming_handler.py | 20 +-
litellm/proxy/common_request_processing.py | 49 +---
.../streaming_handler.py | 22 +-
.../pass_through_endpoints/success_handler.py | 24 +-
.../test_unit_test_streaming.py | 117 ++++++++++
.../test_proxy_reject_logging.py | 19 +-
.../test_litellm_logging.py | 212 +++++++++++++++++-
.../test_streaming_handler.py | 16 +-
.../test_deferred_guardrail_logging.py | 103 ++++++++-
10 files changed, 566 insertions(+), 122 deletions(-)
diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py
index 97266096ef9..606d28585bc 100644
--- a/litellm/litellm_core_utils/litellm_logging.py
+++ b/litellm/litellm_core_utils/litellm_logging.py
@@ -1612,6 +1612,90 @@ class Logging(LiteLLMLoggingBaseClass):
) -> Optional[float]:
return self._response_cost_calculator(result=result, cache_hit=cache_hit)
+ @staticmethod
+ def _is_sync_litellm_request(litellm_params: dict) -> bool:
+ """True for sync SDK entrypoints (``completion``), false for async (``acompletion``, etc.)."""
+ return (
+ litellm_params.get(CallTypes.acompletion.value, False) is not True
+ and litellm_params.get(CallTypes.aresponses.value, False) is not True
+ and litellm_params.get(CallTypes.aembedding.value, False) is not True
+ and litellm_params.get(CallTypes.aimage_generation.value, False) is not True
+ and litellm_params.get(CallTypes.atranscription.value, False) is not True
+ )
+
+ def _is_assembled_stream_success(self, result=None) -> bool:
+ """Final assembled stream export (not a per-chunk success call).
+
+ Per-chunk callers pass a ``ModelResponseStream`` (or ``None``); the
+ final assembled response is any other non-``None`` value (typically a
+ ``ModelResponse``). Treating a chunk as the assembled response would
+ prematurely set the ``has_dispatched_final_stream_success`` dedup
+ guard and silently suppress the real final stream log.
+ """
+ if self.stream is not True:
+ return False
+ if result is not None and not isinstance(result, ModelResponseStream):
+ return True
+ return (
+ "async_complete_streaming_response" in self.model_call_details
+ or self.model_call_details.get("complete_streaming_response") is not None
+ )
+
+ async def dispatch_success_handlers(
+ self,
+ result=None,
+ start_time=None,
+ end_time=None,
+ cache_hit=None,
+ prefer_async_handlers: bool = False,
+ **kwargs,
+ ) -> None:
+ """Route success logging to async and/or sync handlers for this request.
+
+ ``prefer_async_handlers`` only bypasses the sync-SDK-only shortcut (e.g.
+ ``async for`` on a stream from ``completion()``). Legacy string callbacks
+ still run via ``executor.submit(success_handler)`` when configured.
+ """
+ from litellm.litellm_core_utils.thread_pool_executor import executor
+
+ if self._is_assembled_stream_success(result):
+ if self.model_call_details.get("has_dispatched_final_stream_success"):
+ return
+ self.model_call_details["has_dispatched_final_stream_success"] = True
+
+ litellm_params = self.model_call_details.get("litellm_params", {}) or {}
+ sync_sdk = self._is_sync_litellm_request(litellm_params)
+ passthrough = self.call_type == CallTypes.pass_through.value
+ if sync_sdk and not prefer_async_handlers and not passthrough:
+ self.success_handler(
+ result,
+ start_time=start_time,
+ end_time=end_time,
+ cache_hit=cache_hit,
+ **kwargs,
+ )
+ return
+
+ await self.async_success_handler(
+ result,
+ start_time=start_time,
+ end_time=end_time,
+ cache_hit=cache_hit,
+ **kwargs,
+ )
+
+ if not self._should_run_sync_callbacks_for_async_calls():
+ return
+
+ executor.submit(
+ self.success_handler,
+ result,
+ start_time=start_time,
+ end_time=end_time,
+ cache_hit=cache_hit,
+ **kwargs,
+ )
+
def should_run_logging(
self,
event_type: Literal[
@@ -2034,13 +2118,7 @@ class Logging(LiteLLMLoggingBaseClass):
standard_logging_object=kwargs.get("standard_logging_object", None),
)
litellm_params = self.model_call_details.get("litellm_params", {})
- is_sync_request = (
- litellm_params.get(CallTypes.acompletion.value, False) is not True
- and litellm_params.get(CallTypes.aresponses.value, False) is not True
- and litellm_params.get(CallTypes.aembedding.value, False) is not True
- and litellm_params.get(CallTypes.aimage_generation.value, False) is not True
- and litellm_params.get(CallTypes.atranscription.value, False) is not True
- )
+ is_sync_request = self._is_sync_litellm_request(litellm_params)
try:
## BUILD COMPLETE STREAMED RESPONSE
complete_streaming_response: Optional[
@@ -2496,9 +2574,11 @@ class Logging(LiteLLMLoggingBaseClass):
print_verbose(
"Logging Details LiteLLM-Async Success Call, cache_hit={}".format(cache_hit)
)
- if not self.should_run_logging(
+ if not self._is_assembled_stream_success(
+ result
+ ) and not self.should_run_logging(
event_type="async_success"
- ): # prevent double logging
+ ): # prevent double logging (non-streaming)
return
## CALCULATE COST FOR BATCH JOBS
@@ -2948,13 +3028,7 @@ class Logging(LiteLLMLoggingBaseClass):
): # prevent double logging
return
litellm_params = self.model_call_details.get("litellm_params", {})
- is_sync_request = (
- litellm_params.get(CallTypes.acompletion.value, False) is not True
- and litellm_params.get(CallTypes.aresponses.value, False) is not True
- and litellm_params.get(CallTypes.aembedding.value, False) is not True
- and litellm_params.get(CallTypes.aimage_generation.value, False) is not True
- and litellm_params.get(CallTypes.atranscription.value, False) is not True
- )
+ is_sync_request = self._is_sync_litellm_request(litellm_params)
try:
start_time, end_time = self._failure_handler_helper_fn(
diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py
index 4642201ca67..55042a733ed 100644
--- a/litellm/litellm_core_utils/streaming_handler.py
+++ b/litellm/litellm_core_utils/streaming_handler.py
@@ -1835,8 +1835,10 @@ class CustomStreamWrapper:
processed_chunk, None, None, cache_hit
)
)
- ## SYNC LOGGING
- self.logging_obj.success_handler(processed_chunk, None, None, cache_hit)
+ ## SYNC LOGGING — only for sync SDK entrypoints; async proxy paths export via async_success_handler
+ litellm_params = self.logging_obj.model_call_details.get("litellm_params", {})
+ if self.logging_obj._is_sync_litellm_request(litellm_params):
+ self.logging_obj.success_handler(processed_chunk, None, None, cache_hit)
def finish_reason_handler(self):
model_response = self.model_response_creator()
@@ -2231,23 +2233,19 @@ class CustomStreamWrapper:
cache_hit,
)
else:
+ # prefer_async_handlers routes CustomLogger to async_success_handler
+ # when consumers use ``async for`` on sync-SDK streams. Legacy string
+ # callbacks still run via executor.submit inside dispatch_success_handlers.
asyncio.create_task(
- self.logging_obj.async_success_handler(
+ self.logging_obj.dispatch_success_handlers(
complete_streaming_response,
cache_hit=cache_hit,
start_time=None,
end_time=None,
+ prefer_async_handlers=True,
)
)
- executor.submit(
- self.logging_obj.success_handler,
- complete_streaming_response,
- cache_hit=cache_hit,
- start_time=None,
- end_time=None,
- )
-
raise StopAsyncIteration # Re-raise StopIteration
else:
self.sent_last_chunk = True
diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py
index 6782208458e..36acd9653e8 100644
--- a/litellm/proxy/common_request_processing.py
+++ b/litellm/proxy/common_request_processing.py
@@ -1291,7 +1291,7 @@ class ProxyBaseLLMRequestProcessing:
# (ProxyLogging._fire_deferred_stream_logging) fires the
# closure after the full streaming pipeline finishes.
# The closure runs non-apply_guardrail hooks on the
- # assembled response, then fires both logging handlers.
+ # assembled response, then fires success logging.
# Only for CustomStreamWrapper — raw async generators from
# passthrough routes bypass CSW and would orphan the closure.
from litellm.litellm_core_utils.streaming_handler import (
@@ -1427,33 +1427,18 @@ class ProxyBaseLLMRequestProcessing:
logging_obj._on_deferred_stream_complete = None # type: ignore[union-attr]
try:
asyncio.create_task(
- logging_obj.async_success_handler(
+ logging_obj.dispatch_success_handlers(
response,
cache_hit=None,
start_time=None,
end_time=None,
+ prefer_async_handlers=True,
)
)
except Exception as e:
verbose_proxy_logger.exception(
"Error in orphaned streaming async logging: %s", e
)
- try:
- from litellm.litellm_core_utils.thread_pool_executor import (
- executor as _exc,
- )
-
- _exc.submit(
- logging_obj.success_handler,
- response,
- cache_hit=None,
- start_time=None,
- end_time=None,
- )
- except Exception as e:
- verbose_proxy_logger.exception(
- "Error in orphaned streaming sync logging: %s", e
- )
# Always return the client-requested model name (not provider-prefixed internal identifiers)
# for OpenAI-compatible responses.
@@ -1740,7 +1725,7 @@ class ProxyBaseLLMRequestProcessing:
) -> None:
"""
Run non-streaming post-call guardrail hooks on an assembled streaming
- response, then fire both async and sync logging handlers.
+ response, then fire success logging via ``dispatch_success_handlers``.
Called by ProxyLogging._fire_deferred_stream_logging after the full
streaming pipeline (including unified_guardrail end-of-stream blocks)
@@ -1756,8 +1741,6 @@ class ProxyBaseLLMRequestProcessing:
Extracted as a static method so tests can call the production
implementation directly rather than reimplementing the closure.
"""
- from litellm.litellm_core_utils.thread_pool_executor import executor
-
_response = assembled_response
try:
from litellm.proxy.proxy_server import llm_router as _global_llm_router
@@ -1816,31 +1799,23 @@ class ProxyBaseLLMRequestProcessing:
)
finally:
try:
+ # Proxy streaming always runs in async context and proxy spend
+ # logging is async-only; force async dispatch so DB/spend
+ # callbacks fire regardless of the call-type heuristic in
+ # _is_sync_litellm_request (which only recognizes a subset of
+ # async markers stored in litellm_params).
asyncio.create_task(
- captured_logging_obj.async_success_handler(
+ captured_logging_obj.dispatch_success_handlers(
_response,
cache_hit=cache_hit,
start_time=None,
end_time=None,
+ prefer_async_handlers=True,
)
)
except Exception as e:
verbose_proxy_logger.exception(
- "Error in deferred streaming async logging: %s",
- e,
- )
-
- try:
- executor.submit(
- captured_logging_obj.success_handler,
- _response,
- cache_hit=cache_hit,
- start_time=None,
- end_time=None,
- )
- except Exception as e:
- verbose_proxy_logger.exception(
- "Error in deferred streaming sync logging: %s",
+ "Error in deferred streaming success logging: %s",
e,
)
diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py
index 235a38b75f9..33a6b719280 100644
--- a/litellm/proxy/pass_through_endpoints/streaming_handler.py
+++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py
@@ -7,7 +7,6 @@ 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.thread_pool_executor import executor
from litellm.proxy._types import PassThroughEndpointLoggingResultValues
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType
@@ -145,25 +144,16 @@ class PassThroughStreamingHandler:
end_time=end_time,
model=model,
)
- await litellm_logging_obj.async_success_handler(
+ # Always reached from an async context (anthropic_messages,
+ # google_genai, and proxy pass-through stream tasks). prefer_async_handlers
+ # keeps async-only loggers running even when call_type isn't pass_through
+ # and litellm_params lacks an async flag (e.g. aanthropic_messages).
+ await litellm_logging_obj.dispatch_success_handlers(
result=standard_logging_response_object,
start_time=start_time,
end_time=end_time,
cache_hit=False,
- **kwargs,
- )
- if (
- litellm_logging_obj._should_run_sync_callbacks_for_async_calls()
- is False
- ):
- return
-
- executor.submit(
- litellm_logging_obj.success_handler,
- result=standard_logging_response_object,
- end_time=end_time,
- cache_hit=False,
- start_time=start_time,
+ prefer_async_handlers=True,
**kwargs,
)
except Exception as e:
diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py
index 0bc0183aa7c..292871bae67 100644
--- a/litellm/proxy/pass_through_endpoints/success_handler.py
+++ b/litellm/proxy/pass_through_endpoints/success_handler.py
@@ -11,7 +11,6 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import (
PassthroughStandardLoggingPayload,
)
from litellm.types.utils import StandardPassThroughResponseObject
-from litellm.utils import executor as thread_pool_executor
from .llm_provider_handlers.anthropic_passthrough_logging_handler import (
AnthropicPassthroughLoggingHandler,
@@ -94,19 +93,15 @@ class PassThroughEndpointLogging:
cache_hit: bool,
**kwargs,
):
- """Helper function to handle both sync and async logging operations"""
- # Submit to thread pool for sync logging
- thread_pool_executor.submit(
- logging_obj.success_handler,
- standard_logging_response_object,
- start_time,
- end_time,
- cache_hit,
- **kwargs,
- )
-
- # Handle async logging
- await logging_obj.async_success_handler(
+ """Log pass-through success via the shared async dispatch path."""
+ # Always reached from pass_through_async_success_handler, which runs in
+ # an async context. call_type is "pass_through_endpoint" here, so the
+ # passthrough guard in dispatch_success_handlers already forces the
+ # async handler to run; pass prefer_async_handlers explicitly to match
+ # the streaming sibling (_route_streaming_logging_to_handler) and keep
+ # async-only loggers (e.g. the proxy spend logger) firing regardless of
+ # how the call-type classification evolves.
+ await logging_obj.dispatch_success_handlers(
result=(
json.dumps(result)
if isinstance(result, dict)
@@ -115,6 +110,7 @@ class PassThroughEndpointLogging:
start_time=start_time,
end_time=end_time,
cache_hit=False,
+ prefer_async_handlers=True,
**kwargs,
)
diff --git a/tests/pass_through_unit_tests/test_unit_test_streaming.py b/tests/pass_through_unit_tests/test_unit_test_streaming.py
index 38b650121bd..63965320f2b 100644
--- a/tests/pass_through_unit_tests/test_unit_test_streaming.py
+++ b/tests/pass_through_unit_tests/test_unit_test_streaming.py
@@ -97,6 +97,123 @@ async def test_chunk_processor_yields_raw_bytes(endpoint_type, url_route):
), "Collected chunks do not match raw chunks"
+@pytest.mark.asyncio
+async def test_route_streaming_logging_runs_async_handler_for_sdk_passthrough():
+ """
+ SDK pass-through streaming (anthropic_messages, google generate_content) must run
+ the async success handler so async-only loggers record the assembled stream.
+
+ Regression for duplicate-trace dedupe: dispatch_success_handlers treated these as
+ sync SDK requests because call_type is not ``pass_through_endpoint`` and
+ litellm_params carries no ``acompletion`` flag, so only the sync success_handler
+ ran and CustomLogger.async_log_success_event never fired.
+ """
+ import time
+
+ from litellm.types.utils import CallTypes
+
+ logging_obj = LiteLLMLoggingObj(
+ model="claude-sonnet-4-5",
+ messages=[{"role": "user", "content": "hi"}],
+ stream=True,
+ call_type=CallTypes.anthropic_messages.value,
+ start_time=time.time(),
+ litellm_call_id="test-id",
+ function_id="fn",
+ )
+ logging_obj.model_call_details["litellm_params"] = {"anthropic_messages": True}
+
+ with (
+ patch.object(
+ PassThroughStreamingHandler,
+ "_build_passthrough_logging_result",
+ return_value=({"id": "slp"}, {}),
+ ),
+ patch.object(
+ logging_obj, "async_success_handler", new_callable=AsyncMock
+ ) as mock_async,
+ patch.object(
+ logging_obj, "success_handler", new_callable=MagicMock
+ ) as mock_sync,
+ patch.object(
+ logging_obj,
+ "_should_run_sync_callbacks_for_async_calls",
+ return_value=False,
+ ),
+ ):
+ await PassThroughStreamingHandler._route_streaming_logging_to_handler(
+ litellm_logging_obj=logging_obj,
+ passthrough_success_handler_obj=MagicMock(),
+ url_route="/v1/messages",
+ request_body={},
+ endpoint_type=EndpointType.ANTHROPIC,
+ start_time=datetime.now(),
+ raw_bytes=[],
+ end_time=datetime.now(),
+ )
+
+ mock_async.assert_awaited_once()
+ mock_sync.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_handle_logging_runs_async_handler_for_passthrough():
+ """
+ Non-streaming pass-through logging (_handle_logging) must always run the
+ async success handler so async-only loggers (e.g. the proxy spend logger)
+ record the request.
+
+ _handle_logging is only ever reached from pass_through_async_success_handler
+ (an async context), so it forces async dispatch via prefer_async_handlers.
+ This pins that contract independent of the call-type classification: even a
+ call_type that _is_sync_litellm_request would classify as sync (here
+ "completion" with no async marker in litellm_params) must still reach
+ async_success_handler. Without prefer_async_handlers=True the sync-only
+ branch would return early and async_log_success_event would never fire.
+ """
+ import time
+
+ from litellm.types.utils import CallTypes
+
+ logging_obj = LiteLLMLoggingObj(
+ model="claude-sonnet-4-5",
+ messages=[{"role": "user", "content": "hi"}],
+ stream=False,
+ call_type=CallTypes.completion.value,
+ start_time=time.time(),
+ litellm_call_id="test-id",
+ function_id="fn",
+ )
+ logging_obj.model_call_details["litellm_params"] = {}
+
+ handler = PassThroughEndpointLogging()
+
+ with (
+ patch.object(
+ logging_obj, "async_success_handler", new_callable=AsyncMock
+ ) as mock_async,
+ patch.object(
+ logging_obj, "success_handler", new_callable=MagicMock
+ ) as mock_sync,
+ patch.object(
+ logging_obj,
+ "_should_run_sync_callbacks_for_async_calls",
+ return_value=False,
+ ),
+ ):
+ await handler._handle_logging(
+ logging_obj=logging_obj,
+ standard_logging_response_object={"id": "slp"},
+ result="",
+ start_time=datetime.now(),
+ end_time=datetime.now(),
+ cache_hit=False,
+ )
+
+ mock_async.assert_awaited_once()
+ mock_sync.assert_not_called()
+
+
def test_convert_raw_bytes_to_str_lines():
"""
Test that the _convert_raw_bytes_to_str_lines method correctly converts raw bytes to a list of strings
diff --git a/tests/proxy_unit_tests/test_proxy_reject_logging.py b/tests/proxy_unit_tests/test_proxy_reject_logging.py
index 51a92fa3b4b..e0b575f4a71 100644
--- a/tests/proxy_unit_tests/test_proxy_reject_logging.py
+++ b/tests/proxy_unit_tests/test_proxy_reject_logging.py
@@ -95,6 +95,21 @@ router = Router(
)
+def _register_proxy_test_logger(callback_logger: testLogger) -> None:
+ """
+ Register the test logger on global callback lists.
+
+ ``function_setup`` dedupes by object identity; each parametrized case
+ constructs a new ``testLogger`` and must replace the global lists, not
+ only ``litellm.callbacks``.
+ """
+ litellm.callbacks = [callback_logger]
+ litellm.success_callback = [callback_logger]
+ litellm.failure_callback = [callback_logger]
+ litellm._async_success_callback = [callback_logger]
+ litellm._async_failure_callback = [callback_logger]
+
+
@pytest.mark.parametrize(
"route, body",
[
@@ -115,7 +130,7 @@ router = Router(
"/v1/embeddings",
{
"input": "The food was delicious and the waiter...",
- "model": "text-embedding-ada-002",
+ "model": "fake-model",
"encoding_format": "float",
},
),
@@ -133,7 +148,7 @@ async def test_chat_completion_request_with_redaction(route, body):
setattr(proxy_server, "llm_router", router)
_test_logger = testLogger()
- litellm.callbacks = [_test_logger]
+ _register_proxy_test_logger(_test_logger)
litellm.set_verbose = True
# Prepare the query string
diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py
index 07ab29c5231..b64cb7c6905 100644
--- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py
+++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py
@@ -1,6 +1,7 @@
import os
import sys
-from unittest.mock import MagicMock, patch
+import asyncio
+from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@@ -786,6 +787,211 @@ def test_success_handler_runs_sync_callbacks_for_sync_requests(logging_obj, call
dummy_logger.log_stream_event.assert_not_called()
+def test_is_sync_litellm_request():
+ assert LitellmLogging._is_sync_litellm_request({}) is True
+ assert LitellmLogging._is_sync_litellm_request({"acompletion": True}) is False
+
+
+@pytest.mark.asyncio
+async def test_dispatch_success_handlers_invokes_callbacks_once_for_final_stream(
+ logging_obj,
+):
+ """Second final-stream dispatch must not re-export (CSW + deferred guardrail paths)."""
+ import litellm
+ from litellm.integrations.custom_logger import CustomLogger
+
+ class MockCallback(CustomLogger):
+ pass
+
+ mock_callback = MockCallback()
+ original_async_callbacks = list(litellm._async_success_callback or [])
+ litellm._async_success_callback = [mock_callback]
+
+ result = ModelResponse(
+ id="resp-dedupe",
+ model="gpt-4o-mini",
+ choices=[
+ {
+ "message": {"role": "assistant", "content": "hi"},
+ "finish_reason": "stop",
+ "index": 0,
+ }
+ ],
+ usage={"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
+ )
+
+ try:
+ logging_obj.stream = True
+ logging_obj.model_call_details["litellm_params"] = {"acompletion": True}
+
+ with (
+ patch.object(
+ mock_callback, "async_log_success_event", new_callable=AsyncMock
+ ) as mock_async_log,
+ patch.object(mock_callback, "log_success_event") as mock_sync_log,
+ patch.object(
+ logging_obj,
+ "_success_handler_helper_fn",
+ return_value=(time.time(), time.time(), result),
+ ),
+ patch.object(
+ logging_obj,
+ "_get_assembled_streaming_response",
+ return_value=result,
+ ),
+ patch.object(
+ logging_obj,
+ "_should_run_sync_callbacks_for_async_calls",
+ return_value=True,
+ ),
+ ):
+ await logging_obj.dispatch_success_handlers(result=result)
+ await logging_obj.dispatch_success_handlers(result=result)
+
+ mock_async_log.assert_awaited_once()
+ mock_sync_log.assert_not_called()
+ finally:
+ litellm._async_success_callback = original_async_callbacks
+
+
+@pytest.mark.asyncio
+async def test_dispatch_success_handlers_sync_path_invokes_callback_once_for_final_stream(
+ logging_obj,
+):
+ """Sync dispatch path must also dedupe when dispatch is called twice."""
+ import litellm
+ from litellm.integrations.custom_logger import CustomLogger
+
+ class MockCallback(CustomLogger):
+ pass
+
+ mock_callback = MockCallback()
+ original_success_callbacks = list(litellm.success_callback or [])
+ litellm.success_callback = [mock_callback]
+
+ result = ModelResponse(
+ id="resp-sync-dedupe",
+ model="gpt-4o-mini",
+ choices=[
+ {
+ "message": {"role": "assistant", "content": "hi"},
+ "finish_reason": "stop",
+ "index": 0,
+ }
+ ],
+ usage={"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
+ )
+
+ try:
+ logging_obj.stream = True
+ logging_obj.model_call_details["litellm_params"] = {}
+
+ with (
+ patch.object(mock_callback, "log_success_event") as mock_sync_log,
+ patch.object(
+ mock_callback, "async_log_success_event", new_callable=AsyncMock
+ ) as mock_async_log,
+ patch.object(
+ logging_obj,
+ "_success_handler_helper_fn",
+ return_value=(time.time(), time.time(), result),
+ ),
+ patch.object(
+ logging_obj,
+ "_get_assembled_streaming_response",
+ return_value=result,
+ ),
+ ):
+ await logging_obj.dispatch_success_handlers(result=result)
+ await logging_obj.dispatch_success_handlers(result=result)
+
+ mock_sync_log.assert_called_once()
+ mock_async_log.assert_not_awaited()
+ finally:
+ litellm.success_callback = original_success_callbacks
+
+
+@pytest.mark.asyncio
+async def test_dispatch_prefer_async_handlers_runs_legacy_callbacks(
+ logging_obj,
+):
+ """``prefer_async_handlers`` must not skip executor.submit for string callbacks."""
+ result = ModelResponse(
+ id="resp-prefer-async",
+ model="gpt-4o-mini",
+ choices=[
+ {
+ "message": {"role": "assistant", "content": "hi"},
+ "finish_reason": "stop",
+ "index": 0,
+ }
+ ],
+ )
+
+ logging_obj.stream = True
+ logging_obj.model_call_details["litellm_params"] = {}
+
+ with (
+ patch.object(
+ logging_obj, "async_success_handler", new_callable=AsyncMock
+ ) as mock_async,
+ patch.object(
+ logging_obj, "success_handler", new_callable=MagicMock
+ ) as mock_sync,
+ patch.object(
+ logging_obj,
+ "_should_run_sync_callbacks_for_async_calls",
+ return_value=True,
+ ),
+ patch(
+ "litellm.litellm_core_utils.litellm_logging.executor.submit"
+ ) as mock_submit,
+ ):
+ await logging_obj.dispatch_success_handlers(
+ result=result,
+ prefer_async_handlers=True,
+ )
+
+ mock_async.assert_awaited_once()
+ mock_sync.assert_not_called()
+ mock_submit.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_dispatch_success_handlers_invokes_async_callback_for_pass_through(
+ logging_obj,
+):
+ """Pass-through must use async_success_handler (CustomLogger skips sync success_handler)."""
+ import litellm
+ from litellm.integrations.custom_logger import CustomLogger
+ from litellm.types.utils import CallTypes
+
+ class MockCallback(CustomLogger):
+ pass
+
+ mock_callback = MockCallback()
+ original_async_callbacks = list(litellm._async_success_callback or [])
+ litellm._async_success_callback = [mock_callback]
+
+ logging_obj.call_type = CallTypes.pass_through.value
+ logging_obj.stream = False
+ logging_obj.model_call_details["litellm_params"] = {}
+
+ try:
+ with (
+ patch.object(
+ mock_callback, "async_log_success_event", new_callable=AsyncMock
+ ) as mock_async_log,
+ patch.object(mock_callback, "log_success_event") as mock_sync_log,
+ ):
+ await logging_obj.dispatch_success_handlers(result={"id": "pt-1"})
+
+ mock_async_log.assert_awaited_once()
+ mock_sync_log.assert_not_called()
+ finally:
+ litellm._async_success_callback = original_async_callbacks
+
+
def test_success_handler_skips_guardrail_logging_hook_when_disabled(logging_obj):
"""Ensure CustomGuardrail logging_hook is skipped when should_run_guardrail is False."""
import datetime
@@ -1351,7 +1557,7 @@ async def test_e2e_generate_cold_storage_object_key_with_custom_logger_s3_path()
Test that _generate_cold_storage_object_key uses s3_path from custom logger instance.
"""
from datetime import datetime, timezone
- from unittest.mock import MagicMock, patch
+ from unittest.mock import AsyncMock, MagicMock, patch
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
@@ -1404,7 +1610,7 @@ async def test_e2e_generate_cold_storage_object_key_with_logger_no_s3_path():
Test that _generate_cold_storage_object_key falls back to empty s3_path when logger has no s3_path.
"""
from datetime import datetime, timezone
- from unittest.mock import MagicMock, patch
+ from unittest.mock import AsyncMock, MagicMock, patch
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py
index 49d3c51e340..63e2cb7f35c 100644
--- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py
+++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py
@@ -569,8 +569,6 @@ async def test_streaming_with_usage_and_logging(sync_mode: bool):
== final_usage_block
)
- print(mock_log_success_event.call_args.kwargs.keys())
-
def test_streaming_handler_with_stop_chunk(
initialized_custom_stream_wrapper: CustomStreamWrapper,
@@ -2036,23 +2034,19 @@ async def test_azure_streaming_role_preserved_with_include_usage(sync_mode: bool
chunks.append(chunk)
# The prompt_filter chunk should be forwarded with choices=[]
- assert len(chunks[0].choices) == 0, (
- f"Expected prompt_filter chunk with choices=[], got {len(chunks[0].choices)} choices"
- )
+ assert (
+ len(chunks[0].choices) == 0
+ ), f"Expected prompt_filter chunk with choices=[], got {len(chunks[0].choices)} choices"
# At least one chunk must have role='assistant' in its delta
has_role = any(
- len(c.choices) > 0
- and getattr(c.choices[0].delta, "role", None) == "assistant"
+ len(c.choices) > 0 and getattr(c.choices[0].delta, "role", None) == "assistant"
for c in chunks
)
assert has_role, (
"No chunk contained role='assistant' in delta (issue #24221). "
"Chunk deltas: "
- + str([
- c.choices[0].delta if c.choices else "no choices"
- for c in chunks
- ])
+ + str([c.choices[0].delta if c.choices else "no choices" for c in chunks])
)
diff --git a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py
index e10258c0829..e9ff193e044 100644
--- a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py
+++ b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py
@@ -18,7 +18,7 @@ import asyncio
import os
import sys
from typing import Any
-from unittest.mock import MagicMock, patch
+from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@@ -38,6 +38,24 @@ from litellm.types.guardrails import GuardrailEventHooks
# ---------------------------------------------------------------------------
+def _attach_mock_success_dispatch(mock_logging_obj, async_success_fn):
+ """Match production entrypoint: ``_run_deferred_stream_guardrails`` uses dispatch."""
+
+ async def dispatch_success_handlers(
+ result=None, start_time=None, end_time=None, cache_hit=None, **kwargs
+ ):
+ await async_success_fn(
+ result,
+ start_time=start_time,
+ end_time=end_time,
+ cache_hit=cache_hit,
+ **kwargs,
+ )
+
+ mock_logging_obj.dispatch_success_handlers = dispatch_success_handlers
+ mock_logging_obj.async_success_handler = async_success_fn
+
+
class PostCallGuardrail(CustomGuardrail):
"""A post-call guardrail."""
@@ -454,7 +472,7 @@ class TestDeferredStreamingClosure:
async def track_async_success(*args, **kwargs):
pass
- mock_logging_obj.async_success_handler = track_async_success
+ _attach_mock_success_dispatch(mock_logging_obj, track_async_success)
tracking_guardrail = TrackingGuardrail()
tracking_logger = TrackingLogger()
@@ -511,7 +529,7 @@ class TestDeferredStreamingClosure:
nonlocal logged_response
logged_response = args[0] if args else None
- mock_logging_obj.async_success_handler = track_async_success
+ _attach_mock_success_dispatch(mock_logging_obj, track_async_success)
class ModifyingGuardrail(CustomGuardrail):
def __init__(self):
@@ -573,7 +591,7 @@ class TestDeferredStreamingClosure:
nonlocal logging_called
logging_called = True
- mock_logging_obj.async_success_handler = track_async_success
+ _attach_mock_success_dispatch(mock_logging_obj, track_async_success)
guardrail = BlockingGuardrail()
@@ -621,7 +639,7 @@ class TestDeferredStreamingClosure:
async def track_async_success(*args, **kwargs):
pass
- mock_logging_obj.async_success_handler = track_async_success
+ _attach_mock_success_dispatch(mock_logging_obj, track_async_success)
guardrail = TransientErrorGuardrail()
@@ -656,7 +674,7 @@ class TestDeferredStreamingClosure:
nonlocal logged_response
logged_response = args[0] if args else None
- mock_logging_obj.async_success_handler = track_async_success
+ _attach_mock_success_dispatch(mock_logging_obj, track_async_success)
class TestGuardrail(CustomGuardrail):
def __init__(self):
@@ -739,7 +757,7 @@ class TestDeferredStreamingClosure:
async def track_async_success(*args, **kwargs):
pass
- mock_logging_obj.async_success_handler = track_async_success
+ _attach_mock_success_dispatch(mock_logging_obj, track_async_success)
guardrail = ApplyGuardrailType()
@@ -792,7 +810,7 @@ class TestDeferredStreamingClosure:
async def track_async_success(*args, **kwargs):
pass
- mock_logging_obj.async_success_handler = track_async_success
+ _attach_mock_success_dispatch(mock_logging_obj, track_async_success)
guardrail = IteratorHookGuardrail()
@@ -847,7 +865,7 @@ class TestDeferredStreamingClosure:
async def track_async_success(*args, **kwargs):
pass
- mock_logging_obj.async_success_handler = track_async_success
+ _attach_mock_success_dispatch(mock_logging_obj, track_async_success)
guardrail = InspectingGuardrail()
@@ -914,7 +932,7 @@ class TestDeferredStreamingClosure:
async def track_async_success(*args, **kwargs):
pass
- mock_logging_obj.async_success_handler = track_async_success
+ _attach_mock_success_dispatch(mock_logging_obj, track_async_success)
guardrail_a = TaggedGuardrail("guardrail-a")
guardrail_b = TaggedGuardrail("guardrail-b")
@@ -962,7 +980,7 @@ class TestDeferredStreamingClosure:
nonlocal logging_called
logging_called = True
- mock_logging_obj.async_success_handler = track_async_success
+ _attach_mock_success_dispatch(mock_logging_obj, track_async_success)
def exploding_merge(data, llm_router):
raise RuntimeError("Simulated init failure")
@@ -986,6 +1004,67 @@ class TestDeferredStreamingClosure:
logging_called is True
), "Logging must fire even when guardrail initialization raises"
+ @pytest.mark.asyncio
+ async def test_deferred_logging_forces_async_for_sync_classified_call_type(self):
+ """
+ Regression: proxy deferred streaming logging must reach the async success
+ handler (which runs the async-only DB/spend logger) even when the call
+ type is classified as a sync SDK request by _is_sync_litellm_request.
+
+ Without prefer_async_handlers=True, an async proxy stream whose
+ litellm_params lacks a recognized async marker would enter the sync
+ branch of dispatch_success_handlers and silently skip spend tracking.
+
+ Uses the real dispatch_success_handlers via the production
+ _run_deferred_stream_guardrails entrypoint.
+ """
+ import time
+
+ from litellm.litellm_core_utils.litellm_logging import (
+ Logging as LiteLLMLoggingObj,
+ )
+
+ logging_obj = LiteLLMLoggingObj(
+ model="gpt-4o-mini",
+ messages=[{"role": "user", "content": "hi"}],
+ stream=True,
+ call_type="completion", # not pass_through_endpoint
+ start_time=time.time(),
+ litellm_call_id="test-id",
+ function_id="fn",
+ )
+ # litellm_params with no recognized async marker -> classified sync.
+ logging_obj.model_call_details["litellm_params"] = {}
+ assert LiteLLMLoggingObj._is_sync_litellm_request({}) is True
+
+ with (
+ patch.object(
+ logging_obj, "async_success_handler", new_callable=AsyncMock
+ ) as mock_async,
+ patch.object(
+ logging_obj, "success_handler", new_callable=MagicMock
+ ) as mock_sync,
+ patch.object(
+ logging_obj,
+ "_should_run_sync_callbacks_for_async_calls",
+ return_value=False,
+ ),
+ patch("litellm.callbacks", [PostCallGuardrail()]),
+ ):
+ await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails(
+ captured_data={"model": "gpt-4o-mini", "metadata": {}},
+ captured_user_api_key_dict=UserAPIKeyAuth(api_key="test"),
+ captured_logging_obj=logging_obj,
+ assembled_response=MagicMock(),
+ cache_hit=False,
+ )
+
+ await asyncio.sleep(0)
+ await asyncio.sleep(0)
+
+ mock_async.assert_awaited_once()
+ mock_sync.assert_not_called()
+
# ---------------------------------------------------------------------------
# 7. _fire_deferred_stream_logging
@@ -1054,7 +1133,7 @@ class TestFireDeferredStreamLogging:
nonlocal logged_response
logged_response = args[0] if args else None
- mock_logging_obj.async_success_handler = track_async_success
+ _attach_mock_success_dispatch(mock_logging_obj, track_async_success)
class InfoWritingGuardrail(CustomGuardrail):
def __init__(self):
From d82eb33a60b72639eb9ce39cb7bdd1d7188ba28c Mon Sep 17 00:00:00 2001
From: Yassin Kortam
Date: Fri, 29 May 2026 23:15:27 -0700
Subject: [PATCH 19/44] feat(otel): typed semconv-aligned OpenTelemetry
instrumentation (#28909)
---
litellm/_service_logger.py | 103 +-
litellm/constants.py | 7 +
litellm/integrations/otel/README.md | 258 ++++
litellm/integrations/otel/__init__.py | 102 ++
litellm/integrations/otel/emitter.py | 175 +++
litellm/integrations/otel/logger.py | 495 ++++++++
litellm/integrations/otel/mappers/__init__.py | 58 +
litellm/integrations/otel/mappers/base.py | 36 +
litellm/integrations/otel/mappers/genai.py | 137 ++
litellm/integrations/otel/mappers/langfuse.py | 84 ++
.../integrations/otel/mappers/langtrace.py | 64 +
litellm/integrations/otel/mappers/legacy.py | 97 ++
.../otel/mappers/openinference.py | 128 ++
litellm/integrations/otel/mappers/utils.py | 76 ++
litellm/integrations/otel/mappers/weave.py | 48 +
litellm/integrations/otel/model/__init__.py | 0
litellm/integrations/otel/model/baggage.py | 76 ++
litellm/integrations/otel/model/config.py | 236 ++++
litellm/integrations/otel/model/metadata.py | 315 +++++
litellm/integrations/otel/model/payloads.py | 468 +++++++
litellm/integrations/otel/model/semconv.py | 201 +++
litellm/integrations/otel/model/spans.py | 203 +++
litellm/integrations/otel/model/utils.py | 103 ++
litellm/integrations/otel/mount.py | 130 ++
.../integrations/otel/plumbing/__init__.py | 0
litellm/integrations/otel/plumbing/context.py | 127 ++
litellm/integrations/otel/plumbing/metrics.py | 28 +
.../integrations/otel/plumbing/providers.py | 220 ++++
litellm/integrations/otel/plumbing/routing.py | 101 ++
litellm/integrations/otel/presets/__init__.py | 78 ++
litellm/integrations/otel/presets/agentops.py | 139 ++
litellm/integrations/otel/presets/arize.py | 75 ++
litellm/integrations/otel/presets/base.py | 25 +
litellm/integrations/otel/presets/langfuse.py | 43 +
.../integrations/otel/presets/langtrace.py | 22 +
litellm/integrations/otel/presets/levo.py | 24 +
litellm/integrations/otel/presets/phoenix.py | 48 +
litellm/integrations/otel/presets/utils.py | 16 +
litellm/integrations/otel/presets/weave.py | 43 +
litellm/integrations/otel/runtime.py | 38 +
litellm/litellm_core_utils/litellm_logging.py | 79 ++
litellm/proxy/auth/user_api_key_auth.py | 148 ++-
litellm/proxy/db/log_db_metrics.py | 37 +-
litellm/proxy/management_helpers/utils.py | 7 +
.../pass_through_endpoints.py | 9 +-
litellm/proxy/proxy_server.py | 52 +
litellm/proxy/utils.py | 35 +-
pyproject.toml | 3 +
.../integrations/otel/test_otel_v2_baggage.py | 172 +++
.../otel/test_otel_v2_components.py | 463 +++++++
..._v2_config_baggage_parenting_guardrails.py | 237 ++++
.../integrations/otel/test_otel_v2_dynamic.py | 131 ++
.../integrations/otel/test_otel_v2_emitter.py | 229 ++++
.../integrations/otel/test_otel_v2_logger.py | 1122 +++++++++++++++++
.../integrations/otel/test_otel_v2_mount.py | 151 +++
.../otel/test_otel_v2_multibackend.py | 89 ++
.../integrations/otel/test_otel_v2_presets.py | 122 ++
.../otel/test_otel_v2_sources_of_truth.py | 454 +++++++
.../otel/test_otel_v2_vendor_mappers.py | 196 +++
tests/test_litellm/proxy/test_proxy_utils.py | 47 +
tests/test_litellm/test_service_logger.py | 157 ++-
uv.lock | 50 +
62 files changed, 8506 insertions(+), 111 deletions(-)
create mode 100644 litellm/integrations/otel/README.md
create mode 100644 litellm/integrations/otel/__init__.py
create mode 100644 litellm/integrations/otel/emitter.py
create mode 100644 litellm/integrations/otel/logger.py
create mode 100644 litellm/integrations/otel/mappers/__init__.py
create mode 100644 litellm/integrations/otel/mappers/base.py
create mode 100644 litellm/integrations/otel/mappers/genai.py
create mode 100644 litellm/integrations/otel/mappers/langfuse.py
create mode 100644 litellm/integrations/otel/mappers/langtrace.py
create mode 100644 litellm/integrations/otel/mappers/legacy.py
create mode 100644 litellm/integrations/otel/mappers/openinference.py
create mode 100644 litellm/integrations/otel/mappers/utils.py
create mode 100644 litellm/integrations/otel/mappers/weave.py
create mode 100644 litellm/integrations/otel/model/__init__.py
create mode 100644 litellm/integrations/otel/model/baggage.py
create mode 100644 litellm/integrations/otel/model/config.py
create mode 100644 litellm/integrations/otel/model/metadata.py
create mode 100644 litellm/integrations/otel/model/payloads.py
create mode 100644 litellm/integrations/otel/model/semconv.py
create mode 100644 litellm/integrations/otel/model/spans.py
create mode 100644 litellm/integrations/otel/model/utils.py
create mode 100644 litellm/integrations/otel/mount.py
create mode 100644 litellm/integrations/otel/plumbing/__init__.py
create mode 100644 litellm/integrations/otel/plumbing/context.py
create mode 100644 litellm/integrations/otel/plumbing/metrics.py
create mode 100644 litellm/integrations/otel/plumbing/providers.py
create mode 100644 litellm/integrations/otel/plumbing/routing.py
create mode 100644 litellm/integrations/otel/presets/__init__.py
create mode 100644 litellm/integrations/otel/presets/agentops.py
create mode 100644 litellm/integrations/otel/presets/arize.py
create mode 100644 litellm/integrations/otel/presets/base.py
create mode 100644 litellm/integrations/otel/presets/langfuse.py
create mode 100644 litellm/integrations/otel/presets/langtrace.py
create mode 100644 litellm/integrations/otel/presets/levo.py
create mode 100644 litellm/integrations/otel/presets/phoenix.py
create mode 100644 litellm/integrations/otel/presets/utils.py
create mode 100644 litellm/integrations/otel/presets/weave.py
create mode 100644 litellm/integrations/otel/runtime.py
create mode 100644 tests/test_litellm/integrations/otel/test_otel_v2_baggage.py
create mode 100644 tests/test_litellm/integrations/otel/test_otel_v2_components.py
create mode 100644 tests/test_litellm/integrations/otel/test_otel_v2_config_baggage_parenting_guardrails.py
create mode 100644 tests/test_litellm/integrations/otel/test_otel_v2_dynamic.py
create mode 100644 tests/test_litellm/integrations/otel/test_otel_v2_emitter.py
create mode 100644 tests/test_litellm/integrations/otel/test_otel_v2_logger.py
create mode 100644 tests/test_litellm/integrations/otel/test_otel_v2_mount.py
create mode 100644 tests/test_litellm/integrations/otel/test_otel_v2_multibackend.py
create mode 100644 tests/test_litellm/integrations/otel/test_otel_v2_presets.py
create mode 100644 tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py
create mode 100644 tests/test_litellm/integrations/otel/test_otel_v2_vendor_mappers.py
diff --git a/litellm/_service_logger.py b/litellm/_service_logger.py
index 1a3be203fec..5531c418799 100644
--- a/litellm/_service_logger.py
+++ b/litellm/_service_logger.py
@@ -24,6 +24,22 @@ else:
UserAPIKeyAuth = Any
+def _get_otel_v2_class() -> Optional[type]:
+ """Return the ``OpenTelemetryV2`` class, or ``None`` if the OTel SDK is absent.
+
+ Imported lazily: ``litellm.integrations.otel.logger`` imports the OpenTelemetry
+ SDK at module scope, so importing it eagerly would break installs without the
+ SDK. The V2 logger only exists when ``LITELLM_OTEL_V2`` is enabled (which
+ requires the SDK), so a failed import simply means "no V2 logger in play".
+ """
+ try:
+ from litellm.integrations.otel.logger import OpenTelemetryV2
+
+ return OpenTelemetryV2
+ except Exception:
+ return None
+
+
class ServiceLogging(CustomLogger):
"""
Separate class used for monitoring health of litellm-adjacent services (redis/postgres).
@@ -38,6 +54,37 @@ class ServiceLogging(CustomLogger):
if "prometheus_system" in litellm.service_callback:
self.prometheusServicesLogger = PrometheusServicesLogger()
+ def _resolve_otel_service_logger(self, callback: Any) -> Optional[Any]:
+ """Resolve the OTel logger (legacy or V2) to emit a service span on.
+
+ Returns the logger instance whose ``async_service_*_hook`` should fire for
+ this ``callback``, or ``None`` when ``callback`` is not an OTel callback.
+
+ The V2 ``OpenTelemetryV2`` logger is a plain ``CustomLogger`` and is NOT a
+ subclass of the legacy ``OpenTelemetry``, so the legacy ``isinstance``
+ check alone misses it — which is why redis/postgres service spans never
+ showed up under ``LITELLM_OTEL_V2``. Match both the legacy and V2 types,
+ whether the callback is the logger instance itself or the ``"otel"`` string
+ (which routes to the proxy's registered ``open_telemetry_logger``).
+ """
+ otel_v2_cls = _get_otel_v2_class()
+
+ def _is_otel_logger(obj: Any) -> bool:
+ if isinstance(obj, OpenTelemetry):
+ return True
+ return otel_v2_cls is not None and isinstance(obj, otel_v2_cls)
+
+ if _is_otel_logger(callback):
+ return callback
+ if callback == "otel":
+ from litellm.proxy.proxy_server import open_telemetry_logger
+
+ if open_telemetry_logger is not None and _is_otel_logger(
+ open_telemetry_logger
+ ):
+ return open_telemetry_logger
+ return None
+
def service_success_hook(
self,
service: ServiceTypes,
@@ -129,6 +176,13 @@ class ServiceLogging(CustomLogger):
event_metadata=event_metadata,
)
+ # OTel loggers already fired this event. ``service_callback`` can hold more
+ # than one reference that resolves to the *same* logger — the ``"otel"``
+ # string AND the registered instance both map to ``open_telemetry_logger``
+ # (the V2 logger self-registers its instance even when the string is
+ # present, unlike V1). Without this guard each such reference emits its own
+ # span, so a single DB call shows up as duplicate ``postgres ...`` spans.
+ emitted_otel_logger_ids: set = set()
for callback in litellm.service_callback:
if callback == "prometheus_system":
await self.init_prometheus_services_logger_if_none()
@@ -144,19 +198,18 @@ class ServiceLogging(CustomLogger):
end_time=end_time,
event_metadata=event_metadata,
)
- elif callback == "otel" or isinstance(callback, OpenTelemetry):
- _otel_logger_to_use: Optional[OpenTelemetry] = None
- if isinstance(callback, OpenTelemetry):
- _otel_logger_to_use = callback
- else:
- from litellm.proxy.proxy_server import open_telemetry_logger
-
- if open_telemetry_logger is not None and isinstance(
- open_telemetry_logger, OpenTelemetry
- ):
- _otel_logger_to_use = open_telemetry_logger
-
- if _otel_logger_to_use is not None and parent_otel_span is not None:
+ else:
+ _otel_logger_to_use = self._resolve_otel_service_logger(callback)
+ # No ``parent_otel_span is not None`` gate: a background service
+ # call (no request on the stack) has no parent, and dropping it
+ # here is what hid those calls from traces entirely. The OTel
+ # logger decides what to do with a missing parent — legacy V1
+ # no-ops, V2 emits a root span (and skips metrics-only pings).
+ if (
+ _otel_logger_to_use is not None
+ and id(_otel_logger_to_use) not in emitted_otel_logger_ids
+ ):
+ emitted_otel_logger_ids.add(id(_otel_logger_to_use))
await _otel_logger_to_use.async_service_success_hook(
payload=payload,
parent_otel_span=parent_otel_span,
@@ -238,6 +291,9 @@ class ServiceLogging(CustomLogger):
event_metadata=event_metadata,
)
+ # Dedupe OTel loggers per event — see ``async_service_success_hook`` for why
+ # the same logger can be referenced twice in ``service_callback``.
+ emitted_otel_logger_ids: set = set()
for callback in litellm.service_callback:
if callback == "prometheus_system":
await self.init_prometheus_services_logger_if_none()
@@ -255,22 +311,19 @@ class ServiceLogging(CustomLogger):
end_time=end_time,
event_metadata=event_metadata,
)
- elif callback == "otel" or isinstance(callback, OpenTelemetry):
- _otel_logger_to_use: Optional[OpenTelemetry] = None
- if isinstance(callback, OpenTelemetry):
- _otel_logger_to_use = callback
- else:
- from litellm.proxy.proxy_server import open_telemetry_logger
-
- if open_telemetry_logger is not None and isinstance(
- open_telemetry_logger, OpenTelemetry
- ):
- _otel_logger_to_use = open_telemetry_logger
+ else:
+ _otel_logger_to_use = self._resolve_otel_service_logger(callback)
if not isinstance(error, str):
error = str(error)
- if _otel_logger_to_use is not None and parent_otel_span is not None:
+ # See the success hook: no parent gate, so background failures
+ # are traced too. V1 no-ops without a parent; V2 emits a root.
+ if (
+ _otel_logger_to_use is not None
+ and id(_otel_logger_to_use) not in emitted_otel_logger_ids
+ ):
+ emitted_otel_logger_ids.add(id(_otel_logger_to_use))
await _otel_logger_to_use.async_service_failure_hook(
payload=payload,
error=error,
diff --git a/litellm/constants.py b/litellm/constants.py
index f72528eb170..ae98b37d6e6 100644
--- a/litellm/constants.py
+++ b/litellm/constants.py
@@ -1409,6 +1409,13 @@ LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME = "litellm_internal_jobs"
# Prometheus metrics, audit trails, or any other downstream consumer.
LITELLM_PROXY_MASTER_KEY_ALIAS = "litellm_proxy_master_key"
+# Marker placed in ``model_call_details`` on a synthetic ``Logging`` object that
+# records a proxy-gate error (auth/rate-limit rejection) for a request that never
+# reached an upstream provider. Tracing callbacks key off it to avoid fabricating
+# an LLM-call span for a call that did not happen. See
+# ``ProxyLogging._handle_logging_proxy_only_error``.
+LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL = "litellm_no_upstream_llm_call"
+
# Key Rotation Constants
LITELLM_KEY_ROTATION_ENABLED = os.getenv("LITELLM_KEY_ROTATION_ENABLED", "false")
LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS = int(
diff --git a/litellm/integrations/otel/README.md b/litellm/integrations/otel/README.md
new file mode 100644
index 00000000000..99b3ecea162
--- /dev/null
+++ b/litellm/integrations/otel/README.md
@@ -0,0 +1,258 @@
+# OpenTelemetry instrumentation
+
+This package produces OpenTelemetry traces for LiteLLM. It is enabled by the
+`LITELLM_OTEL_V2` environment variable (`is_otel_v2_enabled()` in
+[`config.py`](./model/config.py)); when unset, nothing in this package runs.
+
+## What gets traced
+
+A traced proxy request produces one trace with two kinds of spans:
+
+```
+SERVER span "POST /v1/chat/completions" ← FastAPI instrumentation
+├── INTERNAL span "auth /v1/chat/completions" ← auth phase ┐
+│ ├── CLIENT span "postgres get_key_object" ← datastore call │
+│ └── CLIENT span "postgres get_team_membership" │
+├── INTERNAL span "execute_guardrail …" ← guardrail │ this package
+├── CLIENT span "chat gpt-4o" ← LLM call │
+└── CLIENT span "batch_write_to_db …" ← spend write ┘
+```
+
+The gen-ai spans are siblings under the server span. In particular the guardrail
+span is a sibling of the LLM call, not a child of it: pre/during/post-call
+guardrail hooks are part of the request lifecycle (a pre-call guardrail runs
+before the LLM call even starts), so they belong directly under the server span,
+alongside the LLM call.
+
+Request-level spans (LLM call, guardrail) parent to the server span via an
+**explicit anchor** — `context.set_request_root_span` captures the server span
+once at request entry, and `resolve_request_span_context` reads it — rather than
+to whatever span is momentarily active. Ambient-only parenting was wrong at two
+boundaries: inside the live `auth` phase span the active span is `auth` (so the
+span would nest under auth), and a pass-through request closes its span from a
+detached `asyncio.create_task` where the server span is no longer active (so the
+span orphaned into its own trace). The anchor — a contextvar inherited by those
+child tasks — gives a stable parent in both cases. DB/service spans keep ambient
+parenting so an auth DB lookup still nests under `auth`.
+
+**Which service calls become spans (`spans.span_role_for_service`).** LiteLLM's
+service-logging layer instruments many internal functions, but only some are
+traceable units of work:
+
+- **`DB_CALL` (CLIENT)** — outbound datastore calls (redis, postgres,
+ `batch_write_to_db`), carrying `db.system.name` / `db.operation.name` semconv.
+- **`SERVICE` (INTERNAL)** — genuine internal work worth a span (background
+ budget/reset jobs, pod-lock manager).
+- **metrics-only (no span)** — `self` (the `track_llm_api_timing` wrapper, which
+ duplicates the LLM-call span), `router` (duplicates the request), and
+ `proxy_pre_call` (a guardrail's real span is `execute_guardrail …`). These
+ still feed Prometheus/Datadog through their own hooks; they just never enter
+ the trace. `auth` is also excluded here because it gets a **live phase span**
+ instead (see below).
+
+Spans are named `"{service} {call_type}"` (e.g. `"redis set"`) so repeated calls
+to one service stay distinguishable. Like every other span they parent to the
+**ambient** context, falling back to the threaded `litellm_parent_otel_span` only
+when ambient has no live span; a background job with neither starts its own root
+trace. Caller-supplied `event_metadata` is **sanitized** before it reaches a span
+(primitives only, no live objects, no secrets/headers, bounded) — see
+`payloads.sanitize_event_metadata`.
+
+**Live phase spans.** `auth` is wrapped in a real, active span
+(`logger.phase_span`) for the duration of authentication, so the DB lookups it
+triggers nest **under** it instead of flattening onto the server span. Identity
+Baggage (team/key/user) is seeded once the key resolves, so every post-auth span
+inherits it; auth-internal DB lookups that run before the key is known stay
+unlabeled, which is correct.
+
+**Status.** On success a span's status is left `UNSET` (the semconv default,
+matching the FastAPI server span); only a genuine error sets `ERROR`.
+
+- **Server spans** (one per HTTP route) are created by the
+ `opentelemetry-instrumentation-fastapi` package. It stamps `http.*` attributes
+ and extracts inbound `traceparent` headers. This package does **not** create
+ or modify server spans — request routes never touch spans.
+- **Gen-AI spans** (LLM calls, guardrails, internal service calls) are created
+ by this package from LiteLLM's logging callbacks. Request-level spans parent to
+ the server span via the captured anchor; DB/service spans parent to the active
+ span (ambient) so they nest under the request phase that triggered them.
+
+Both kinds share a single `TracerProvider`, so they belong to the same trace
+and export through the same configured exporters. FastAPI middleware can only be
+added before the app starts serving, so the app is instrumented at
+import time **without** a provider — it binds to the OTel global
+`ProxyTracerProvider`. Once config (and the callbacks) is loaded, the proxy
+publishes the chosen logger's `TracerProvider` as the global via
+`trace.set_tracer_provider(...)`, and the server spans delegate to it. When a
+preset callback (`arize`, `langfuse_otel`, …) is configured, its provider
+becomes the global, so server spans export to that backend too.
+
+## How a request flows
+
+1. **App creation** (`proxy_server` import): when the gate is on,
+ `mount.instrument_fastapi_app(app)` calls `FastAPIInstrumentor.instrument_app`
+ with no provider (the middleware stack is frozen once the app serves, so this
+ can't wait for startup). It binds to the OTel global `ProxyTracerProvider`. Noisy
+ non-LLM routes are excluded by default (`mount._DEFAULT_EXCLUDED_ROUTES`): health
+ checks (`/health*`), the Prometheus scrape (`/metrics`), and static UI/docs assets
+ (`/litellm-asset-prefix`, `/_next`, `/ui`, `/swagger`, `/docs`, `/redoc`,
+ `/openapi.json`, favicons, `/.well-known`) — so load-balancer polling, metric
+ scrapes, and asset fetches don't flood traces. Entries are substring-matched, so
+ `/metrics` also drops the `/model/metrics` admin-analytics spans. Set
+ `OTEL_PYTHON_FASTAPI_EXCLUDED_URLS` to override the whole set (e.g. `""` to trace
+ everything, or your own comma-separated path list).
+2. **Startup** (`proxy_server.proxy_startup_event`): after the config (and
+ callbacks) is loaded, the already-registered preset `OpenTelemetryV2` logger
+ is reused — or a generic one reading `OTEL_*` envs is built when no preset is
+ configured — and its `TracerProvider` is published as the OTel global with
+ `trace.set_tracer_provider(...)`. The proxy tracer then delegates to it, so
+ server spans and gen-ai spans share one provider and the same trace.
+3. **Request**: the FastAPI instrumentation starts the server span and makes it
+ the active context for the request task. The proxy's first call into the V2
+ logger (`create_litellm_proxy_request_started_span`, at the auth boundary)
+ **captures it as the request anchor** (`set_request_root_span`), so every later
+ request-level span has a stable explicit parent regardless of what is active
+ when it emits.
+4. **LLM call span (born at the boundary)**: `OpenTelemetryV2.log_pre_api_call`
+ runs synchronously in the request task, just before the upstream call, and
+ **opens** the LLM-call span there, parented to the anchored server span
+ (`resolve_request_span_context`). The open span is held in a bounded cache keyed
+ by `litellm_call_id` (a primitive the callback kwargs carry at both `pre_call`
+ and close), so no live `Span` ever travels through a `litellm_params` metadata
+ dict. For the boundary hook to fire at all, the logger is registered into
+ `litellm.input_callback` — the list `Logging.pre_call` iterates. The async
+ success/failure callback later
+ **closes** it: it builds an `LLMCallSpanData` from the typed
+ `standard_logging_object` (token usage and cost are computed only by then),
+ stamps the attributes, sets status, and ends the span. The sync callback is a
+ no-op (closing is async-only). When `pre_call` runs off the request task — a
+ sync-only provider driven through a thread pool, where contextvars (and so the
+ anchor) don't follow — no parent is visible there, so creation is **deferred**
+ to the async callback, whose worker context was copied from the request task at
+ enqueue and so still carries the anchor. **Pass-through** endpoints call
+ `logging_obj.pre_call` in the request task too, then close from a detached
+ `asyncio.create_task`; the anchor (not the by-then-inactive server span) keeps
+ their LLM-call span in the request's trace. `pre_call` is litellm's generic
+ "log the attempt" hook, so it also fires for synthetic proxy-gate error logs
+ (auth/rate-limit rejections); those carry `LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL`
+ and are skipped, so a request rejected before reaching a provider never produces
+ a phantom CLIENT span.
+5. **Guardrails / services**: the post-call and service hooks emit guardrail and
+ service spans the same way — typed data → engine → span. Service spans
+ (Redis/Postgres) are dispatched by `litellm/_service_logger.py`, which
+ recognizes the V2 `OpenTelemetryV2` logger (a plain `CustomLogger`, not a
+ subclass of the legacy `OpenTelemetry`). It hands every service call to the
+ logger — including calls with no parent span — and the V2 adapter decides the
+ role (`DB_CALL` vs `SERVICE`), the parent (ambient → threaded → root), and
+ whether the call is a traceable operation or a metrics-only ping. Guardrail
+ span data is built from the typed, provider-agnostic
+ `StandardLoggingGuardrailInformation` — no single provider's field shape is
+ assumed.
+6. **Export**: each span ends and is handed to the provider's span processors,
+ which export to the configured backends (OTLP, console, in-memory, …).
+
+## Components
+
+### Sources of truth (`model/`, no OpenTelemetry import)
+
+These define the shape of a span without depending on the OTel SDK, so they can
+be imported anywhere. They live in [`model/`](./model) and form a closed set —
+nothing here imports outside it:
+
+- [`semconv.py`](./model/semconv.py) — attribute-key constants (`gen_ai.*`, `http.*`,
+ `litellm.*`), the GenAI operation/provider enums, and the functions that map
+ LiteLLM provider/call-type strings onto convention values.
+- [`spans.py`](./model/spans.py) — the span registry: every span role, its OTel span
+ kind, its place in the hierarchy, and its name builder.
+- [`payloads.py`](./model/payloads.py) — frozen dataclasses (`LLMCallSpanData`,
+ `GuardrailSpanData`, `ServiceSpanData`, …) built from heterogeneous logging
+ payloads via `from_*` classmethods.
+- [`config.py`](./model/config.py) — `OpenTelemetryV2Config`, a pydantic-settings
+ model that reads `OTEL_*` / `LITELLM_OTEL_*` env vars, plus the feature gate.
+ `capture_span_content` gates whether prompt/response bodies may be written as
+ span attributes; it defaults **off** (`no_content`). The Baggage allowlists are
+ configurable, not hard-coded: set `LITELLM_OTEL_BAGGAGE_PROMOTED_KEYS` /
+ `LITELLM_OTEL_BAGGAGE_METADATA_KEYS` (comma-separated) as env vars, or
+ `baggage_promoted_keys` / `baggage_metadata_keys` (YAML lists) under
+ `callback_settings.otel` in `config.yaml` — the latter reach the config through
+ the logger's constructor kwargs.
+- [`baggage.py`](./model/baggage.py) — the single definition of which request-identity
+ values are promoted into Baggage (so child spans inherit them) and under which
+ attribute keys.
+- [`utils.py`](./model/utils.py) — value coercion, JSON serialization, and
+ extractor-table application, shared across the package.
+
+### Engine
+
+- [`emitter.py`](./emitter.py) — `SpanEmitter.emit(role, data)`: dedupe → start
+ the span → run the mapper chain to stamp attributes → set status → end. It
+ owns no attribute keys. The dedupe set (which coalesces the sync+async firing
+ of one request) is a bounded LRU so it can't grow without limit.
+- [`mappers/`](./mappers) — each mapper turns typed span data into a flat
+ `{attribute key: value}` dict. They compose: listing several mapper names in
+ the config layers multiple attribute vocabularies onto the same span.
+ - `genai` — the canonical OpenTelemetry GenAI vocabulary, always present.
+ - `legacy` — an additional vocabulary using the older semconv-ai / Traceloop
+ attribute key names, for backends that read those.
+ - `openinference`, `langfuse`, `weave`, `langtrace` — vendor vocabularies.
+ - `resolve_mappers(names)` turns config names into mapper instances.
+
+### Plumbing (`plumbing/`)
+
+The OTel-SDK wiring. Everything here imports only `model/` and each other; it
+lives in [`plumbing/`](./plumbing):
+
+- [`providers.py`](./plumbing/providers.py) — builds the `TracerProvider`, its exporters
+ (from `ExporterSpec`s), and the span processor that copies allowlisted Baggage
+ entries onto every span. `register_exporter_factory(kind, factory)` lets a
+ preset contribute a custom exporter `kind` (e.g. one that fetches an auth
+ token lazily) without coupling this module to any vendor.
+- [`context.py`](./plumbing/context.py) — trace-context and Baggage read/write helpers.
+- [`routing.py`](./plumbing/routing.py) — `TenantTracerCache`: when a request carries
+ team/key-scoped vendor credentials, route its spans through a credential-keyed
+ `TracerProvider` so one logger serves many tenants. The cache is a bounded LRU
+ that flushes + shuts down evicted providers, since the key derives from
+ request-supplied credentials and must not grow (or leak threads) without limit.
+- [`metrics.py`](./plumbing/metrics.py) — GenAI client metric instruments.
+
+### Adapter
+
+- [`logger.py`](./logger.py) — `OpenTelemetryV2`, a `CustomLogger` that
+ translates LiteLLM's logging callbacks into typed span data and hands them to
+ the engine. The LLM-call span is opened at the `log_pre_api_call` boundary
+ (parented to the live server span via ambient context) and closed at the async
+ success/failure callback; the open span is held in a bounded cache keyed by
+ `litellm_call_id`, never threaded through a metadata dict. The logger registers
+ itself into `litellm.input_callback` so `Logging.pre_call` fires the boundary
+ hook.
+- [`mount.py`](./mount.py) — `instrument_fastapi_app(app)`, the single call site
+ that attaches `opentelemetry-instrumentation-fastapi` for SERVER spans. It owns
+ the health-check exclusion default (`OTEL_PYTHON_FASTAPI_EXCLUDED_URLS`) and the
+ passthrough span-naming hook (`PASSTHROUGH_PREFIXES`) so `proxy_server` carries
+ no OTel detail. A safe no-op when the gate is off or the instrumentation package
+ is absent; must be called at app-creation time (the middleware stack freezes
+ once the app serves).
+
+### Presets
+
+- [`presets/`](./presets) — each preset reads one integration's env vars and
+ returns an `OpenTelemetryV2Config` (exporter destination + mapper vocabularies
+ + resource attributes). `PRESET_BY_CALLBACK` maps a callback name (`"arize"`,
+ `"langfuse_otel"`, …) to its preset. Integrations that support team/key-scoped
+ credentials also provide a per-request OTLP header builder
+ (`DYNAMIC_HEADERS_BY_CALLBACK`). Presets do **no** network I/O at build time:
+ AgentOps, for example, mints its JWT lazily inside a custom exporter on the
+ first export (in the `BatchSpanProcessor` worker thread), never on the event
+ loop.
+
+## Extending
+
+- **A new attribute vocabulary for a backend**: add a mapper in `mappers/`
+ (a class with a `map(data) -> AttributeMap` method, typically built from
+ `key -> extractor` tables) and register it in `mappers/__init__._MAPPER_BY_NAME`.
+- **A new integration**: add a preset in `presets/` that returns an
+ `OpenTelemetryV2Config`, and register it in `presets/__init__.PRESET_BY_CALLBACK`.
+ If it supports dynamic credentials, add a header builder to
+ `DYNAMIC_HEADERS_BY_CALLBACK`.
+- **A new span kind**: add a role to `spans.py` (registry entry + name builder),
+ a payload dataclass in `payloads.py`, and a branch in the relevant mapper(s).
diff --git a/litellm/integrations/otel/__init__.py b/litellm/integrations/otel/__init__.py
new file mode 100644
index 00000000000..42a84a85fbd
--- /dev/null
+++ b/litellm/integrations/otel/__init__.py
@@ -0,0 +1,102 @@
+"""Typed, semconv-aligned OpenTelemetry instrumentation for LiteLLM.
+
+The three sources of truth — attribute keys (:mod:`semconv`), the span and
+hierarchy registry (:mod:`spans`), and the typed span-data inputs
+(:mod:`payloads`) — plus :mod:`config` are exported here and are free of any
+``opentelemetry`` import. The engine layer (``emitter``, ``providers``,
+``context``, ``metrics``) and the ``CustomLogger`` adapter (``logger``) are
+reached via their submodule paths so that importing this package never
+requires the OTel SDK.
+
+The ``LITELLM_OTEL_V2`` env var gates whether the factory in
+``litellm_core_utils.litellm_logging`` constructs the ``OpenTelemetryV2``
+class (from :mod:`logger`).
+"""
+
+from litellm.integrations.otel.model.config import (
+ OTEL_V2_ENV,
+ OpenTelemetryV2Config,
+ is_otel_v2_enabled,
+)
+from litellm.integrations.otel.model.baggage import (
+ BAGGAGE_PROMOTED_KEYS,
+ DEFAULT_BAGGAGE_METADATA_KEYS,
+ promoted_baggage,
+)
+from litellm.integrations.otel.model.metadata import (
+ RequestContext,
+ RequestIdentity,
+)
+from litellm.integrations.otel.model.payloads import (
+ GuardrailSpanData,
+ LLMCallSpanData,
+ LLMRequestParams,
+ LLMUsage,
+ ProxyRequestSpanData,
+ ServerInfo,
+ ServiceSpanData,
+ SpanError,
+)
+from litellm.integrations.otel.model.semconv import (
+ DB,
+ Error,
+ GenAI,
+ GenAIOperation,
+ GenAIProvider,
+ HTTP,
+ LiteLLM,
+ Metric,
+ Server,
+ resolve_operation,
+ resolve_provider,
+)
+from litellm.integrations.otel.model.spans import (
+ SPAN_REGISTRY,
+ LiteLLMSpanKind,
+ SpanRole,
+ SpanSpec,
+ db_system,
+ span_role_for_service,
+ validate_registry,
+)
+
+__all__ = [
+ # config
+ "OTEL_V2_ENV",
+ "OpenTelemetryV2Config",
+ "is_otel_v2_enabled",
+ # semconv
+ "BAGGAGE_PROMOTED_KEYS",
+ "DB",
+ "DEFAULT_BAGGAGE_METADATA_KEYS",
+ "Error",
+ "GenAI",
+ "GenAIOperation",
+ "GenAIProvider",
+ "HTTP",
+ "LiteLLM",
+ "Metric",
+ "Server",
+ "resolve_operation",
+ "resolve_provider",
+ # spans
+ "SPAN_REGISTRY",
+ "LiteLLMSpanKind",
+ "SpanRole",
+ "SpanSpec",
+ "db_system",
+ "span_role_for_service",
+ "validate_registry",
+ # payloads
+ "GuardrailSpanData",
+ "LLMCallSpanData",
+ "LLMRequestParams",
+ "LLMUsage",
+ "ProxyRequestSpanData",
+ "RequestContext",
+ "RequestIdentity",
+ "ServerInfo",
+ "ServiceSpanData",
+ "SpanError",
+ "promoted_baggage",
+]
diff --git a/litellm/integrations/otel/emitter.py b/litellm/integrations/otel/emitter.py
new file mode 100644
index 00000000000..cae6514efdf
--- /dev/null
+++ b/litellm/integrations/otel/emitter.py
@@ -0,0 +1,175 @@
+"""The span engine: dedup, start, run the mapper chain, set status, end."""
+
+from collections import OrderedDict
+from typing import Callable, Sequence
+
+from opentelemetry.context import Context
+from opentelemetry.trace import Span, Tracer
+from opentelemetry.trace.status import Status, StatusCode
+
+from litellm.integrations.otel.model.config import OpenTelemetryV2Config
+from litellm.integrations.otel.mappers import resolve_mappers
+from litellm.integrations.otel.mappers.base import AttributeMapper, SpanData
+from litellm.integrations.otel.model.payloads import (
+ GuardrailSpanData,
+ LLMCallSpanData,
+ ServiceSpanData,
+)
+from litellm.integrations.otel.plumbing.providers import to_otel_span_kind
+from litellm.integrations.otel.model.semconv import Error
+from litellm.integrations.otel.model.spans import (
+ SPAN_REGISTRY,
+ SpanRole,
+ guardrail_span_name,
+ llm_call_span_name,
+ service_span_name,
+)
+
+# Roles emit() knows how to name and emit. PROXY_REQUEST and the management
+# routes are SERVER spans owned by the mounted FastAPI instrumentor, so they
+# have no builder here.
+_NAME_BUILDERS: dict[SpanRole, Callable[..., str]] = {
+ SpanRole.LLM_CALL: llm_call_span_name,
+ SpanRole.GUARDRAIL: guardrail_span_name,
+ # DB_CALL and SERVICE are both built from ServiceSpanData; they differ only in
+ # span kind (CLIENT vs INTERNAL) and attribute vocabulary, not in naming.
+ SpanRole.DB_CALL: service_span_name,
+ SpanRole.SERVICE: service_span_name,
+}
+
+# Cap on the dedup cache. It only needs to coalesce the sync+async firing window
+# of a single in-flight request, so a bounded LRU keeps memory flat on a
+# long-running proxy while still covering every concurrently-open call.
+_DEDUP_CACHE_MAX = 10_000
+
+
+class SpanEmitter:
+ def __init__(
+ self,
+ tracer: Tracer,
+ config: OpenTelemetryV2Config,
+ mappers: Sequence[AttributeMapper] | None = None,
+ ) -> None:
+ self._tracer = tracer
+ self._config = config
+ # The mapper chain is the sole source of span attributes. When not
+ # passed in, resolve it from the config so there's one source of truth.
+ self._mappers: list[AttributeMapper] = (
+ list(mappers)
+ if mappers is not None
+ else resolve_mappers(config.mapper_names)
+ )
+ # Bounded LRU (ordered by insertion / most-recent touch). Storing keys
+ # only — the value is unused — so it behaves like a capped set.
+ self._emitted: "OrderedDict[tuple[str, SpanRole], None]" = OrderedDict()
+
+ # -- low-level helpers --------------------------------------------------- #
+
+ def start_span(
+ self,
+ role: SpanRole,
+ name: str,
+ parent_context: Context | None = None,
+ start_time_ns: int | None = None,
+ *,
+ tracer: Tracer | None = None,
+ ) -> Span:
+ """Start a span for ``role`` without dedup or attribute mapping.
+
+ For callers that own and manage their own span lifecycle. ``tracer``
+ overrides the bound tracer for this span only, used for per-request
+ multi-tenant credential routing.
+ """
+ return (tracer or self._tracer).start_span(
+ name,
+ context=parent_context,
+ kind=to_otel_span_kind(SPAN_REGISTRY[role].kind),
+ start_time=start_time_ns,
+ )
+
+ def _seen(self, dedup_key: str | None, role: SpanRole) -> bool:
+ """Return True once a ``(dedup_key, role)`` pair has been emitted.
+
+ Guards against emitting the same span twice when a streaming call
+ fires both a sync and an async logging callback.
+ """
+ if not dedup_key:
+ return False
+ marker = (dedup_key, role)
+ if marker in self._emitted:
+ self._emitted.move_to_end(marker)
+ return True
+ self._emitted[marker] = None
+ if len(self._emitted) > _DEDUP_CACHE_MAX:
+ self._emitted.popitem(last=False) # evict least-recently-used
+ return False
+
+ # -- the engine ---------------------------------------------------------- #
+
+ def emit(
+ self,
+ role: SpanRole,
+ data: SpanData,
+ parent_context: Context | None = None,
+ *,
+ start_time_ns: int | None = None,
+ end_time_ns: int | None = None,
+ tracer: Tracer | None = None,
+ ) -> Span | None:
+ """Emit one complete span: dedup, start, map attributes, status, end.
+
+ Return the span, or ``None`` if it was deduplicated away. ``tracer``
+ overrides the bound tracer for this span, used for per-request routing.
+ """
+ # Only LLM-call spans carry a dedup key; LLM-call and service spans
+ # carry an ``error`` field. ``isinstance`` narrows the type for mypy and
+ # keeps the engine free of duck-typed attribute reads.
+ dedup_key = data.identity.call_id if isinstance(data, LLMCallSpanData) else None
+ if self._seen(dedup_key, role):
+ return None
+ span = self.start_span(
+ role,
+ _NAME_BUILDERS[role](data),
+ parent_context=parent_context,
+ start_time_ns=start_time_ns,
+ tracer=tracer,
+ )
+ self.finish_span(role, span, data, end_time_ns=end_time_ns)
+ return span
+
+ def finish_span(
+ self,
+ role: SpanRole,
+ span: Span,
+ data: SpanData,
+ *,
+ end_time_ns: int | None = None,
+ ) -> None:
+ """Stamp attributes + status on an already-started ``span`` and end it.
+
+ The counterpart to :meth:`start_span` for callers that own a span's
+ lifecycle — the LLM-call span is opened at the request's ``pre_call``
+ boundary (so it parents to the live server span via real ambient context,
+ never a span threaded through a metadata dict) and closed here once the
+ typed payload is available. The span name is (re)built from the now-known
+ data, since the boundary opener only has a provisional name.
+ """
+ span.update_name(_NAME_BUILDERS[role](data))
+ for mapper in self._mappers:
+ for key, value in mapper.map(data).items():
+ span.set_attribute(key, value)
+ error = (
+ data.error
+ if isinstance(data, (LLMCallSpanData, ServiceSpanData, GuardrailSpanData))
+ else None
+ )
+ if error and (error.error_type or error.message):
+ span.set_attribute(Error.TYPE, error.error_type or "error")
+ span.set_status(
+ Status(StatusCode.ERROR, error.message or error.error_type or "error")
+ )
+ # On success leave the status UNSET (the semconv default) rather than
+ # forcing OK — that matches the FastAPI server span and avoids implying a
+ # span-level health signal litellm doesn't actually evaluate. Only a
+ # genuine error sets a status.
+ span.end(end_time=end_time_ns)
diff --git a/litellm/integrations/otel/logger.py b/litellm/integrations/otel/logger.py
new file mode 100644
index 00000000000..007b41df0a3
--- /dev/null
+++ b/litellm/integrations/otel/logger.py
@@ -0,0 +1,495 @@
+"""``CustomLogger`` adapter on the OpenTelemetry span engine."""
+
+from collections import OrderedDict
+from contextlib import contextmanager
+from datetime import datetime
+from typing import TYPE_CHECKING, Any, Iterator, Mapping, cast
+
+from opentelemetry.context import attach, get_current
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.trace import Span, Tracer, get_current_span, use_span
+
+import litellm
+from litellm.integrations.custom_logger import CustomLogger
+from litellm.integrations.otel.model.baggage import promoted_baggage
+from litellm.integrations.otel.model.config import OpenTelemetryV2Config
+from litellm.integrations.otel.plumbing.context import (
+ is_recordable_span,
+ resolve_parent_context,
+ resolve_request_span_context,
+ set_request_baggage,
+ set_request_root_span,
+)
+from litellm.integrations.otel.emitter import SpanEmitter
+from litellm.integrations.otel.mappers import resolve_mappers
+from litellm.integrations.otel.model.metadata import (
+ LLMCallEvent,
+ RequestIdentity,
+ guardrail_entries_from_request_data,
+ model_from_request_data,
+)
+from litellm.integrations.otel.model.payloads import (
+ GuardrailSpanData,
+ LLMCallSpanData,
+ ServiceSpanData,
+ SpanError,
+)
+from litellm.integrations.otel.plumbing.providers import (
+ build_tracer_provider,
+ get_tracer,
+)
+from litellm.integrations.otel.plumbing.routing import TenantTracerCache
+from litellm.integrations.otel.model.spans import SpanRole, span_role_for_service
+from litellm.integrations.otel.model.utils import to_ns
+
+if TYPE_CHECKING:
+ from litellm.types.utils import StandardLoggingGuardrailInformation
+
+LITELLM_TRACER_NAME = "litellm"
+
+# Any callback whose class belongs to one of these modules is "the OTel
+# callback" for proxy-global-registration purposes.
+_OTEL_MODULES = (
+ "litellm.integrations.otel",
+ "litellm.integrations.opentelemetry",
+)
+
+
+# Cap on the open-call carrier map. A span opened at ``pre_call`` that never
+# reaches a success/failure callback (e.g. a stream that only fires stream
+# events) would otherwise linger; bounding the map evicts the oldest so memory
+# stays flat on a long-running proxy while covering every concurrent in-flight
+# call.
+_OPEN_CALLS_MAX = 10_000
+
+
+class _LLMCallSpan:
+ """The state carried from the ``pre_call`` boundary to span close.
+
+ ``span`` is the live span when it could be opened at the boundary (the server
+ span was ambient), or ``None`` when creation was deferred because no ambient
+ parent was visible — in which case the async callback creates it against its
+ own (worker-copied) ambient context using ``start_time_ns``. The presence of
+ a carrier for a call at all is the proof that ``pre_call`` ran, i.e. that an
+ upstream call was actually attempted.
+ """
+
+ __slots__ = ("span", "start_time_ns")
+
+ def __init__(self, span: "Span | None", start_time_ns: int | None) -> None:
+ self.span = span
+ self.start_time_ns = start_time_ns
+
+
+class OpenTelemetryV2(CustomLogger):
+ """The ``CustomLogger`` for OpenTelemetry."""
+
+ def __init__(
+ self,
+ config: OpenTelemetryV2Config | None = None,
+ callback_name: str | None = None,
+ tracer_provider: TracerProvider | None = None,
+ logger_provider: Any | None = None, # reserved for OTel logs
+ meter_provider: Any | None = None, # reserved for metrics
+ **kwargs: Any,
+ ) -> None:
+ super().__init__(**kwargs)
+ self.config: OpenTelemetryV2Config = config or OpenTelemetryV2Config(**kwargs)
+ self.callback_name = callback_name
+ self._tracer_provider: TracerProvider = (
+ tracer_provider
+ if tracer_provider is not None
+ else build_tracer_provider(self.config)
+ )
+ self.tracer: Tracer = get_tracer(self._tracer_provider, LITELLM_TRACER_NAME)
+ self._emitter = SpanEmitter(
+ self.tracer, self.config, mappers=resolve_mappers(self.config.mapper_names)
+ )
+ self._tenant_tracers = TenantTracerCache(
+ self.config, callback_name, LITELLM_TRACER_NAME
+ )
+ self._open_llm_calls: "OrderedDict[str, _LLMCallSpan]" = OrderedDict()
+ self._init_otel_logger_on_litellm_proxy()
+
+ # ====================================================================== #
+ # Proxy global registration
+ # ====================================================================== #
+
+ def _register_in_callback_list(self, callbacks: list) -> None:
+ already_otel = any(
+ cb.__class__.__module__.startswith(_OTEL_MODULES)
+ for cb in callbacks
+ if hasattr(cb, "__class__")
+ )
+ if not already_otel:
+ callbacks.append(self)
+
+ def _init_otel_logger_on_litellm_proxy(self) -> None:
+ try:
+ from litellm.proxy import proxy_server
+ except Exception:
+ return
+ try:
+ self._register_in_callback_list(litellm.service_callback)
+ self._register_in_callback_list(litellm.input_callback)
+ self._register_in_callback_list(litellm._async_success_callback)
+ self._register_in_callback_list(litellm._async_failure_callback)
+ except Exception:
+ pass
+ if getattr(proxy_server, "open_telemetry_logger", None) is None:
+ setattr(proxy_server, "open_telemetry_logger", self)
+
+ # ====================================================================== #
+ # LLM-call callbacks — the span is opened at the ``pre_call`` boundary and
+ # closed here. See ``log_pre_api_call``.
+ # ====================================================================== #
+
+ def log_pre_api_call(self, model, messages, kwargs):
+ """Open the LLM-call span at the call boundary.
+
+ Runs synchronously inside the request task, before the upstream call —
+ the one place where the live server span is genuinely the ambient OTel
+ context — so the span parents to it natively, with no span threaded
+ through a metadata dict. The open span is stashed on the per-request
+ ``LiteLLMLoggingObj`` (a typed object) and closed in the async callback.
+
+ When no recordable parent is visible (``pre_call`` was driven from a thread
+ pool for a sync-only provider, where contextvars — and so the anchor —
+ don't follow), creation is deferred: only the start time is recorded, and
+ the async callback — whose worker context was copied from the request task
+ and so still carries the anchor — creates the span then.
+
+ Synthetic proxy-gate error logs (auth/rate-limit rejections) also fire this
+ hook but never made an upstream call; they are tagged and skipped so no
+ phantom LLM-call span is produced.
+ """
+ call = LLMCallEvent.from_dict(kwargs)
+ if call.is_no_upstream_call:
+ return
+ call_id = call.call_id
+ if call_id is None:
+ return
+ # Idempotent: a retried call may re-enter ``pre_call`` with the same
+ # call id; keep the first span so its start time is the true one.
+ if call_id in self._open_llm_calls:
+ return
+ start_time_ns = to_ns(datetime.now())
+ span: Span | None = None
+ # Parent to the request's anchored root span (stable across the request),
+ # falling back to ambient on the SDK path. Open the span live only when
+ # that resolves to a recordable parent; otherwise defer to the close
+ # callback (the thread-pool case, where the anchor isn't visible here).
+ parent_context = resolve_request_span_context()
+ if is_recordable_span(get_current_span(parent_context)):
+ span = self._emitter.start_span(
+ SpanRole.LLM_CALL,
+ call.provisional_span_name,
+ parent_context=parent_context,
+ start_time_ns=start_time_ns,
+ tracer=self._tenant_tracers.tracer_for(
+ self.tracer, call.dynamic_params
+ ),
+ )
+ self._open_llm_calls[call_id] = _LLMCallSpan(
+ span=span, start_time_ns=start_time_ns
+ )
+ # Evict the oldest open call if the map is over budget. A call that opens
+ # but never closes (a stream that only fires stream events) would linger
+ # otherwise; the evicted span is simply dropped (never exported).
+ if len(self._open_llm_calls) > _OPEN_CALLS_MAX:
+ self._open_llm_calls.popitem(last=False)
+
+ async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+ self._close_llm_call(kwargs, start_time, end_time)
+
+ async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
+ self._close_llm_call(kwargs, start_time, end_time)
+
+ def _close_llm_call(
+ self,
+ kwargs: Mapping[str, Any],
+ start_time: datetime | float | None,
+ end_time: datetime | float | None,
+ ) -> Span | None:
+ """Finish the LLM-call span opened at ``pre_call`` (or create it deferred).
+
+ No carrier for this call id means ``pre_call`` never ran — the request was
+ rejected at the gate or blocked by a pre-call guardrail before any upstream
+ call — so there is nothing to record and no phantom span.
+ """
+ call = LLMCallEvent.from_dict(kwargs)
+ call_id = call.call_id
+ # ``pop`` is the dedup: this method runs from both the success and failure
+ # paths, and whichever fires first removes the carrier and closes the span.
+ carrier = self._open_llm_calls.pop(call_id, None) if call_id else None
+ if carrier is None:
+ return None
+ payload = call.payload
+ if payload is None:
+ if carrier.span is not None:
+ # Opened at the boundary but the payload never materialized — end
+ # it (named provisionally) so it isn't leaked as an open span.
+ carrier.span.end(end_time=to_ns(end_time))
+ return None
+ data = LLMCallSpanData.from_standard_logging_payload(
+ payload, capture_content=self.config.capture_span_content
+ )
+ end_time_ns = to_ns(end_time)
+ if carrier.span is not None:
+ # Born at the boundary: stamp attributes from the typed payload, set
+ # status, and end it. Its parent (the server span) was captured at
+ # creation from real ambient context.
+ self._emitter.finish_span(
+ SpanRole.LLM_CALL, carrier.span, data, end_time_ns=end_time_ns
+ )
+ return carrier.span
+ # Deferred: ``pre_call`` saw no recordable parent, so create the span now.
+ # The worker copied the request task's context, which carries the anchored
+ # root span — parent to it (ambient fallback on the SDK path). Seed identity
+ # Baggage so the span — and the SDK path, which has none — is labeled
+ # consistently.
+ parent_ctx = resolve_request_span_context()
+ bag = promoted_baggage(
+ data.identity,
+ data.request_model,
+ promoted_keys=tuple(self.config.baggage_promoted_keys),
+ metadata_keys=tuple(self.config.baggage_metadata_keys),
+ )
+ if bag:
+ parent_ctx = set_request_baggage(bag, context=parent_ctx)
+ return self._emitter.emit(
+ SpanRole.LLM_CALL,
+ data,
+ parent_context=parent_ctx,
+ start_time_ns=carrier.start_time_ns,
+ end_time_ns=end_time_ns,
+ tracer=self._tenant_tracers.tracer_for(self.tracer, call.dynamic_params),
+ )
+
+ # ====================================================================== #
+ # Service hooks
+ # ====================================================================== #
+
+ async def async_service_success_hook(
+ self,
+ payload: Any,
+ parent_otel_span: Span | None = None,
+ start_time: datetime | float | None = None,
+ end_time: datetime | float | None = None,
+ event_metadata: dict | None = None,
+ ) -> None:
+ self._emit_service(
+ payload,
+ parent_otel_span=parent_otel_span,
+ start_time=start_time,
+ end_time=end_time,
+ event_metadata=event_metadata,
+ error_override=None,
+ )
+
+ async def async_service_failure_hook(
+ self,
+ payload: Any,
+ error: str | None = "",
+ parent_otel_span: Span | None = None,
+ start_time: datetime | float | None = None,
+ end_time: datetime | float | None = None,
+ event_metadata: dict | None = None,
+ ) -> None:
+ self._emit_service(
+ payload,
+ parent_otel_span=parent_otel_span,
+ start_time=start_time,
+ end_time=end_time,
+ event_metadata=event_metadata,
+ error_override=error or "error",
+ )
+
+ def _emit_service(
+ self,
+ payload: Any,
+ *,
+ parent_otel_span: Span | None,
+ start_time: datetime | float | None,
+ end_time: datetime | float | None,
+ event_metadata: dict | None,
+ error_override: str | None,
+ ) -> Span | None:
+ data = ServiceSpanData.from_payload(payload, event_metadata=event_metadata)
+ # Decide whether this service call is a span at all, and of what kind.
+ # ``None`` means metrics-only (framework instrumentation that duplicates a
+ # gen-AI span — ``self``/``router``/``proxy_pre_call`` — or ``auth``, which
+ # gets a live phase span instead). Those still feed Prometheus/Datadog via
+ # their own hooks; they just never enter the trace.
+ role = span_role_for_service(data.service_name)
+ if role is None:
+ return None
+ # A metrics-only ping with neither timing nor a parent (in-memory queue
+ # gauges) is not a traceable operation; a span for it would be a
+ # zero-duration root with no context, so skip it. Real background work
+ # (budget/reset jobs, spend flush) passes start/end times and still emits
+ # as a root; anything with a parent emits regardless.
+ if (
+ error_override is None
+ and start_time is None
+ and end_time is None
+ and parent_otel_span is None
+ ):
+ return None
+ if error_override is not None and data.error is None:
+ data = ServiceSpanData(
+ service_name=data.service_name,
+ call_type=data.call_type,
+ error=SpanError(message=error_override),
+ event_metadata=data.event_metadata,
+ )
+ # Parent like every other span: ambient context first (so identity Baggage
+ # rides along and the call nests under whatever request phase is active —
+ # e.g. a DB lookup under the live ``auth`` span), falling back to the
+ # server span the proxy threaded as ``parent_otel_span``. A background
+ # service call has neither, so it starts its own root trace.
+ parent_context = resolve_parent_context(threaded=parent_otel_span)
+ return self._emitter.emit(
+ role,
+ data,
+ parent_context=parent_context,
+ start_time_ns=to_ns(start_time),
+ end_time_ns=to_ns(end_time),
+ )
+
+ # ====================================================================== #
+ # async_post_call_* hooks — emit guardrail spans. The server span's status
+ # / errors are the FastAPI instrumentor's job, so we don't touch it here.
+ # ====================================================================== #
+
+ def seed_request_identity(self, user_api_key_dict: Any, model: Any = None) -> None:
+ """Attach request-identity Baggage to the current context + server span.
+
+ Seeding identity into Baggage makes **every** span emitted afterwards for
+ this request — LLM call, guardrail, DB call — inherit it via
+ ``LiteLLMBaggageSpanProcessor``. Called once at the auth boundary (as soon
+ as the key resolves) so post-auth spans are labeled consistently; the
+ Baggage rides the request task's contextvar from there on. Auth-internal
+ DB lookups that run before the key is known stay unlabeled — identity
+ isn't determined yet, which is correct.
+ """
+ try:
+ identity = RequestIdentity.from_user_api_key_auth(user_api_key_dict)
+ bag = promoted_baggage(
+ identity,
+ model,
+ promoted_keys=tuple(self.config.baggage_promoted_keys),
+ metadata_keys=tuple(self.config.baggage_metadata_keys),
+ )
+ if bag:
+ # Attach (no detach): the contextvar is scoped to this request's
+ # asyncio task and is reclaimed when the task ends.
+ attach(set_request_baggage(bag, context=get_current()))
+ # The server span was started by the instrumentor before this ran,
+ # so the Baggage processor (which only fires at span start) won't
+ # backfill it — stamp identity on it directly.
+ server_span = get_current_span()
+ if is_recordable_span(server_span):
+ # Re-capture the anchor here too: this runs post-auth with the
+ # server span active and covers entrypoints that bypass
+ # ``create_litellm_proxy_request_started_span`` (e.g. the SDK
+ # path's ``async_pre_call_hook``). Idempotent.
+ set_request_root_span(server_span)
+ for key, value in bag.items():
+ server_span.set_attribute(key, value)
+ except Exception:
+ pass
+
+ @contextmanager
+ def start_phase_span(self, name: str) -> "Iterator[Span]":
+ span = self._emitter.start_span(SpanRole.SERVICE, name)
+ with use_span(span, end_on_exit=True):
+ yield span
+
+ async def async_pre_call_hook(
+ self,
+ user_api_key_dict: Any,
+ cache: Any,
+ data: dict,
+ call_type: Any,
+ ) -> dict:
+ self.seed_request_identity(
+ user_api_key_dict,
+ model=model_from_request_data(data),
+ )
+ return data
+
+ async def async_post_call_success_hook(
+ self,
+ data: Mapping[str, Any],
+ user_api_key_dict: Any,
+ response: Any,
+ ) -> Any:
+ self._emit_guardrail_spans(data)
+ return response
+
+ async def async_post_call_failure_hook(
+ self,
+ request_data: Mapping[str, Any],
+ original_exception: BaseException | None,
+ user_api_key_dict: Any,
+ traceback_str: str | None = None,
+ ) -> None:
+ self._emit_guardrail_spans(request_data)
+
+ def _emit_guardrail_spans(self, request_data: Mapping[str, Any]) -> None:
+ # A guardrail is a sibling of the LLM call under the request's root span,
+ # so parent it to the explicit anchor — not the active span, which on the
+ # failure path can be the live ``auth`` phase span (post-call failure hooks
+ # run from inside it on an auth rejection). Emit with the guardrail's actual
+ # execution window so a pre_call guardrail is placed before the LLM call
+ # rather than at post-call emission time.
+ guardrails = guardrail_entries_from_request_data(request_data)
+ if not guardrails:
+ return
+ parent_ctx = resolve_request_span_context()
+ for entry in guardrails:
+ data = GuardrailSpanData.from_logging_entry(
+ cast("StandardLoggingGuardrailInformation", entry)
+ )
+ self._emitter.emit(
+ SpanRole.GUARDRAIL,
+ data,
+ parent_context=parent_ctx,
+ start_time_ns=to_ns(data.start_time),
+ end_time_ns=to_ns(data.end_time),
+ )
+
+ def create_litellm_proxy_request_started_span(
+ self, start_time: datetime, headers: Mapping[str, str] | None
+ ) -> Span | None:
+ span = get_current_span()
+ if not is_recordable_span(span):
+ return None
+ set_request_root_span(span)
+ return span
+
+
+def _registered_v2_logger() -> "OpenTelemetryV2 | None":
+ try:
+ from litellm.proxy import proxy_server
+ except Exception:
+ return None
+ logger = getattr(proxy_server, "open_telemetry_logger", None)
+ return logger if isinstance(logger, OpenTelemetryV2) else None
+
+
+def seed_request_identity(user_api_key_dict: Any, model: Any = None) -> None:
+ logger = _registered_v2_logger()
+ if logger is not None:
+ logger.seed_request_identity(user_api_key_dict, model=model)
+
+
+@contextmanager
+def phase_span(name: str) -> "Iterator[Span | None]":
+ logger = _registered_v2_logger()
+ if logger is None:
+ yield None
+ return
+ with logger.start_phase_span(name) as span:
+ yield span
diff --git a/litellm/integrations/otel/mappers/__init__.py b/litellm/integrations/otel/mappers/__init__.py
new file mode 100644
index 00000000000..012e63f1bee
--- /dev/null
+++ b/litellm/integrations/otel/mappers/__init__.py
@@ -0,0 +1,58 @@
+"""Attribute mappers: pure ``LLMCallSpanData -> {attribute key: value}`` functions.
+
+Composition over inheritance: vocabularies layer onto the same span. Listing
+``["genai", "openinference"]`` in ``config.mapper_names`` makes every span
+carry both the canonical ``gen_ai.*`` keys and the OpenInference (Arize +
+Phoenix) keys. Add ``"langfuse"`` and it works for all three backends at once.
+"""
+
+from typing import Callable, Iterable
+
+from litellm.integrations.otel.mappers.base import (
+ AttributeMap,
+ AttributeMapper,
+ AttrValue,
+)
+from litellm.integrations.otel.mappers.genai import GenAIMapper
+from litellm.integrations.otel.mappers.langfuse import LangfuseMapper
+from litellm.integrations.otel.mappers.langtrace import LangtraceMapper
+from litellm.integrations.otel.mappers.legacy import LegacyMapper
+from litellm.integrations.otel.mappers.openinference import OpenInferenceMapper
+from litellm.integrations.otel.mappers.weave import WeaveMapper
+
+# Registry keyed by ``config.mapper_names`` entries.
+_MAPPER_BY_NAME: dict[str, Callable[[], AttributeMapper]] = {
+ "genai": GenAIMapper,
+ "legacy": LegacyMapper,
+ "openinference": OpenInferenceMapper,
+ "langfuse": LangfuseMapper,
+ "weave": WeaveMapper,
+ "langtrace": LangtraceMapper,
+}
+
+
+def resolve_mappers(names: Iterable[str]) -> list[AttributeMapper]:
+ """Resolve mapper names to instances. Unknown names raise ``ValueError``."""
+ out: list[AttributeMapper] = []
+ for name in names:
+ factory = _MAPPER_BY_NAME.get(name)
+ if factory is None:
+ raise ValueError(
+ f"unknown mapper name {name!r}; known: " f"{sorted(_MAPPER_BY_NAME)}"
+ )
+ out.append(factory())
+ return out
+
+
+__all__ = [
+ "AttributeMap",
+ "AttributeMapper",
+ "AttrValue",
+ "GenAIMapper",
+ "LangfuseMapper",
+ "LangtraceMapper",
+ "LegacyMapper",
+ "OpenInferenceMapper",
+ "WeaveMapper",
+ "resolve_mappers",
+]
diff --git a/litellm/integrations/otel/mappers/base.py b/litellm/integrations/otel/mappers/base.py
new file mode 100644
index 00000000000..e8fb5af9797
--- /dev/null
+++ b/litellm/integrations/otel/mappers/base.py
@@ -0,0 +1,36 @@
+"""Mapper protocol and attribute value types."""
+
+from typing import Sequence
+
+from typing_extensions import Protocol, runtime_checkable
+
+from litellm.integrations.otel.model.payloads import (
+ GuardrailSpanData,
+ LLMCallSpanData,
+ ServiceSpanData,
+)
+
+AttrScalar = str | bool | int | float
+# Mirrors ``opentelemetry.util.types.AttributeValue`` (homogeneous sequences)
+# without importing the SDK, so mappers stay OTel-free.
+AttrValue = (
+ AttrScalar | Sequence[str] | Sequence[bool] | Sequence[int] | Sequence[float]
+)
+AttributeMap = dict[str, AttrValue]
+
+# The closed set of span-data types the engine routes through the mapper chain.
+# Server spans (PROXY_REQUEST + management routes) belong to the mounted FastAPI
+# instrumentor, not the mapper chain.
+SpanData = LLMCallSpanData | GuardrailSpanData | ServiceSpanData
+
+
+@runtime_checkable
+class AttributeMapper(Protocol):
+ """Maps a typed span input to a flat dict of OTel span attributes.
+
+ One method per mapper, dispatched internally on the ``data`` type. The
+ engine calls this uniformly for every span kind — mappers that don't speak
+ a given type return ``{}``. This is why the engine contains no attribute keys.
+ """
+
+ def map(self, data: SpanData) -> AttributeMap: ...
diff --git a/litellm/integrations/otel/mappers/genai.py b/litellm/integrations/otel/mappers/genai.py
new file mode 100644
index 00000000000..57fa51ea1fb
--- /dev/null
+++ b/litellm/integrations/otel/mappers/genai.py
@@ -0,0 +1,137 @@
+"""Canonical OpenTelemetry GenAI semantic-convention mapper (always active).
+
+Owns the attribute schema for every span kind the engine emits — LLM call,
+guardrail, and service — so the engine itself never references attribute keys.
+
+Each span kind declares its schema as a flat ``attribute key -> extractor``
+table: one lambda per mapping operation, applied against the typed span data.
+"""
+
+from typing import Callable
+
+from litellm.integrations.otel.mappers.base import AttributeMap, AttrValue, SpanData
+from litellm.integrations.otel.mappers.utils import collect, drop_none
+from litellm.integrations.otel.model.payloads import (
+ GuardrailSpanData,
+ LLMCallSpanData,
+ ServiceSpanData,
+ ToolDefinition,
+)
+from litellm.integrations.otel.model.semconv import DB, Error, GenAI, LiteLLM, Server
+from litellm.integrations.otel.model.spans import db_system
+
+
+class GenAIMapper:
+
+ _LLM_CALL_ATTRS: dict[str, Callable[[LLMCallSpanData], AttrValue | None]] = {
+ GenAI.OPERATION_NAME: lambda d: d.operation.value,
+ GenAI.PROVIDER_NAME: lambda d: d.provider or None,
+ GenAI.REQUEST_MODEL: lambda d: d.request_model or None,
+ GenAI.REQUEST_TEMPERATURE: lambda d: d.request_params.temperature,
+ GenAI.REQUEST_TOP_P: lambda d: d.request_params.top_p,
+ GenAI.REQUEST_TOP_K: lambda d: d.request_params.top_k,
+ GenAI.REQUEST_MAX_TOKENS: lambda d: d.request_params.max_tokens,
+ GenAI.REQUEST_FREQUENCY_PENALTY: lambda d: d.request_params.frequency_penalty,
+ GenAI.REQUEST_PRESENCE_PENALTY: lambda d: d.request_params.presence_penalty,
+ GenAI.REQUEST_STOP_SEQUENCES: lambda d: (
+ list(d.request_params.stop_sequences)
+ if d.request_params.stop_sequences
+ else None
+ ),
+ GenAI.REQUEST_SEED: lambda d: d.request_params.seed,
+ GenAI.RESPONSE_MODEL: lambda d: d.response_model,
+ GenAI.RESPONSE_ID: lambda d: d.response_id,
+ GenAI.RESPONSE_FINISH_REASONS: lambda d: (
+ list(d.finish_reasons) if d.finish_reasons else None
+ ),
+ GenAI.USAGE_INPUT_TOKENS: lambda d: d.usage.input_tokens,
+ GenAI.USAGE_OUTPUT_TOKENS: lambda d: d.usage.output_tokens,
+ Error.TYPE: lambda d: d.error.error_type if d.error else None,
+ Server.ADDRESS: lambda d: d.server.address if d.server else None,
+ Server.PORT: lambda d: d.server.port if d.server else None,
+ LiteLLM.CALL_ID: lambda d: d.identity.call_id or None,
+ # The provider/underlying model is only known once routing has picked a
+ # deployment, so it can't ride identity Baggage (seeded at auth, before
+ # routing) onto the boundary-born LLM span — stamp it directly here.
+ LiteLLM.PROVIDER_MODEL: lambda d: d.identity.provider_model or None,
+ f"{LiteLLM.COST_PREFIX}total": lambda d: d.response_cost,
+ LiteLLM.REQUEST_STREAMING: lambda d: d.is_streaming,
+ }
+
+ _TOOL_ATTRS: dict[str, Callable[[ToolDefinition], AttrValue | None]] = {
+ "name": lambda t: t.name,
+ "description": lambda t: t.description or None,
+ "parameters": lambda t: t.parameters_json or None,
+ }
+
+ _GUARDRAIL_ATTRS: dict[str, Callable[[GuardrailSpanData], AttrValue | None]] = {
+ LiteLLM.GUARDRAIL_NAME: lambda d: d.guardrail_name,
+ LiteLLM.GUARDRAIL_MODE: lambda d: d.mode,
+ LiteLLM.GUARDRAIL_STATUS: lambda d: d.status,
+ LiteLLM.GUARDRAIL_PROVIDER: lambda d: d.provider,
+ LiteLLM.GUARDRAIL_ACTION: lambda d: d.action,
+ LiteLLM.GUARDRAIL_RESPONSE: lambda d: d.response_json,
+ LiteLLM.GUARDRAIL_VIOLATION_CATEGORIES: lambda d: (
+ list(d.violation_categories) if d.violation_categories else None
+ ),
+ LiteLLM.GUARDRAIL_CONFIDENCE_SCORE: lambda d: d.confidence_score,
+ LiteLLM.GUARDRAIL_RISK_SCORE: lambda d: d.risk_score,
+ LiteLLM.GUARDRAIL_MASKED_ENTITY_COUNT: lambda d: d.masked_entity_count,
+ LiteLLM.GUARDRAIL_DURATION: lambda d: d.duration,
+ LiteLLM.GUARDRAIL_ID: lambda d: d.guardrail_id,
+ LiteLLM.GUARDRAIL_POLICY_TEMPLATE: lambda d: d.policy_template,
+ LiteLLM.GUARDRAIL_DETECTION_METHOD: lambda d: d.detection_method,
+ }
+
+ _SERVICE_ATTRS: dict[str, Callable[[ServiceSpanData], AttrValue | None]] = {
+ LiteLLM.SERVICE_NAME: lambda d: d.service_name,
+ LiteLLM.SERVICE_CALL_TYPE: lambda d: d.call_type,
+ }
+
+ def map(self, data: SpanData) -> AttributeMap:
+ match data:
+ case LLMCallSpanData():
+ return self._llm_call(data)
+ case GuardrailSpanData():
+ return self._guardrail(data)
+ case ServiceSpanData():
+ return self._service(data)
+ case _:
+ return {}
+
+ @classmethod
+ def _llm_call(cls, data: LLMCallSpanData) -> AttributeMap:
+ attrs = collect(cls._LLM_CALL_ATTRS, data)
+ attrs.update(
+ drop_none(
+ {
+ f"gen_ai.tool.{idx}.{suffix}": extract(tool)
+ for idx, tool in enumerate(data.tools)
+ for suffix, extract in cls._TOOL_ATTRS.items()
+ }
+ )
+ )
+ return attrs
+
+ @classmethod
+ def _guardrail(cls, data: GuardrailSpanData) -> AttributeMap:
+ return collect(cls._GUARDRAIL_ATTRS, data)
+
+ @classmethod
+ def _service(cls, data: ServiceSpanData) -> AttributeMap:
+ attrs = collect(cls._SERVICE_ATTRS, data)
+ # An outbound datastore call (DB_CALL / CLIENT span) also carries db.*
+ # semconv. Internal services (router, budget jobs, …) have no db.system,
+ # so they get only the litellm.service.* keys above.
+ system = db_system(data.service_name)
+ if system is not None:
+ attrs[DB.SYSTEM_NAME] = system
+ if data.call_type:
+ attrs[DB.OPERATION_NAME] = data.call_type
+ attrs.update(
+ {
+ f"{LiteLLM.METADATA_PREFIX}{key}": value
+ for key, value in data.event_metadata.items()
+ }
+ )
+ return attrs
diff --git a/litellm/integrations/otel/mappers/langfuse.py b/litellm/integrations/otel/mappers/langfuse.py
new file mode 100644
index 00000000000..14c9fd01d05
--- /dev/null
+++ b/litellm/integrations/otel/mappers/langfuse.py
@@ -0,0 +1,84 @@
+"""Langfuse OTLP attribute mapper.
+
+Langfuse ingests OTLP spans and reads from its own vendor namespace
+(``langfuse.observation.*``, ``langfuse.trace.*``). Compose this mapper after
+``GenAIMapper`` to send canonical + Langfuse-flavored spans simultaneously.
+
+Every attribute is declared as a ``key -> extractor`` table entry (one callable
+per mapping operation): ``_LLM_CALL_ATTRS`` for scalars and ``_BLOB_ATTRS`` for
+the JSON-serialized payloads. ``_llm_call`` just applies both tables.
+"""
+
+import json
+from typing import Callable
+
+from litellm.integrations.otel.mappers.base import AttributeMap, AttrValue, SpanData
+from litellm.integrations.otel.mappers.utils import (
+ collect,
+ json_if,
+ output_messages,
+ serialize_messages,
+)
+from litellm.integrations.otel.model.payloads import (
+ LLMCallSpanData,
+ LLMRequestParams,
+ LLMUsage,
+)
+
+
+class LangfuseMapper:
+
+ _LLM_CALL_ATTRS: dict[str, Callable[[LLMCallSpanData], AttrValue | None]] = {
+ "langfuse.observation.type": lambda d: "generation",
+ "langfuse.observation.model.name": lambda d: d.request_model or None,
+ "langfuse.observation.metadata.provider": lambda d: d.provider or None,
+ "langfuse.observation.id": lambda d: d.identity.call_id or None,
+ "langfuse.trace.metadata.team_id": lambda d: d.identity.team_id or None,
+ "langfuse.trace.metadata.team_alias": lambda d: d.identity.team_alias or None,
+ }
+
+ # Sub-tables folded into their respective JSON blobs.
+ _MODEL_PARAMS: dict[str, Callable[[LLMRequestParams], AttrValue | None]] = {
+ "temperature": lambda rp: rp.temperature,
+ "top_p": lambda rp: rp.top_p,
+ "max_tokens": lambda rp: rp.max_tokens,
+ "frequency_penalty": lambda rp: rp.frequency_penalty,
+ "presence_penalty": lambda rp: rp.presence_penalty,
+ "seed": lambda rp: rp.seed,
+ }
+ _USAGE_FIELDS: dict[str, Callable[[LLMUsage], AttrValue | None]] = {
+ "input": lambda u: u.input_tokens,
+ "output": lambda u: u.output_tokens,
+ "total": lambda u: u.total_tokens,
+ }
+
+ # JSON-payload attributes: each builder returns the serialized blob or None.
+ _BLOB_ATTRS: dict[str, Callable[[LLMCallSpanData], AttrValue | None]] = {
+ "langfuse.observation.model.parameters": lambda d: json_if(
+ collect(LangfuseMapper._MODEL_PARAMS, d.request_params)
+ ),
+ "langfuse.observation.input": lambda d: serialize_messages(d.messages_in),
+ "langfuse.observation.output": lambda d: serialize_messages(output_messages(d)),
+ "langfuse.observation.usage_details": lambda d: json_if(
+ collect(LangfuseMapper._USAGE_FIELDS, d.usage)
+ ),
+ "langfuse.observation.cost_details": lambda d: (
+ json.dumps({"total": d.response_cost})
+ if d.response_cost is not None
+ else None
+ ),
+ }
+
+ def map(self, data: SpanData) -> AttributeMap:
+ match data:
+ case LLMCallSpanData():
+ return self._llm_call(data)
+ case _:
+ return {}
+
+ @classmethod
+ def _llm_call(cls, data: LLMCallSpanData) -> AttributeMap:
+ return {
+ **collect(cls._LLM_CALL_ATTRS, data),
+ **collect(cls._BLOB_ATTRS, data),
+ }
diff --git a/litellm/integrations/otel/mappers/langtrace.py b/litellm/integrations/otel/mappers/langtrace.py
new file mode 100644
index 00000000000..7c0f30e57dd
--- /dev/null
+++ b/litellm/integrations/otel/mappers/langtrace.py
@@ -0,0 +1,64 @@
+"""Langtrace attribute mapper.
+
+Produces Langtrace's attribute vocabulary so a span can be ingested by a
+Langtrace backend. Compose it alongside other mappers like any other
+vocabulary.
+
+Scalar attributes are declared as a flat ``key -> extractor`` table (one lambda
+per mapping operation); the prompt/completion blobs are serialized as a tail.
+"""
+
+from typing import Callable
+
+from litellm.integrations.otel.mappers.base import AttributeMap, AttrValue, SpanData
+from litellm.integrations.otel.mappers.utils import (
+ collect,
+ json_or_none,
+ output_messages,
+)
+from litellm.integrations.otel.model.payloads import LLMCallSpanData
+
+
+class LangtraceMapper:
+
+ _LLM_CALL_ATTRS: dict[str, Callable[[LLMCallSpanData], AttrValue | None]] = {
+ "gen_ai.operation.name": lambda d: "chat",
+ "langtrace.service.name": lambda d: d.provider or None,
+ "llm.model": lambda d: d.request_model or None,
+ "gen_ai.response.model": lambda d: d.response_model or None,
+ "gen_ai.response_id": lambda d: d.response_id or None,
+ "gen_ai.system_fingerprint": lambda d: d.system_fingerprint or None,
+ "llm.temperature": lambda d: d.request_params.temperature,
+ "llm.top_p": lambda d: d.request_params.top_p,
+ "llm.top_k": lambda d: d.request_params.top_k,
+ "llm.max_tokens": lambda d: d.request_params.max_tokens,
+ "llm.frequency_penalty": lambda d: d.request_params.frequency_penalty,
+ "llm.presence_penalty": lambda d: d.request_params.presence_penalty,
+ "llm.stream": lambda d: d.is_streaming,
+ "llm.token.counts.prompt": lambda d: d.usage.input_tokens,
+ "llm.token.counts.completion": lambda d: d.usage.output_tokens,
+ "llm.token.counts.total": lambda d: d.usage.total_tokens,
+ }
+
+ _BLOB_ATTRS: dict[str, Callable[[LLMCallSpanData], AttrValue | None]] = {
+ "llm.prompts": lambda d: (
+ json_or_none(list(d.messages_in)) if d.messages_in else None
+ ),
+ "llm.completions": lambda d: (
+ json_or_none(output_messages(d)) if d.choices_out else None
+ ),
+ }
+
+ def map(self, data: SpanData) -> AttributeMap:
+ match data:
+ case LLMCallSpanData():
+ return self._llm_call(data)
+ case _:
+ return {}
+
+ @classmethod
+ def _llm_call(cls, data: LLMCallSpanData) -> AttributeMap:
+ return {
+ **collect(cls._LLM_CALL_ATTRS, data),
+ **collect(cls._BLOB_ATTRS, data),
+ }
diff --git a/litellm/integrations/otel/mappers/legacy.py b/litellm/integrations/otel/mappers/legacy.py
new file mode 100644
index 00000000000..20ffe8b0dd8
--- /dev/null
+++ b/litellm/integrations/otel/mappers/legacy.py
@@ -0,0 +1,97 @@
+"""Mapper for the older semantic-convention attribute vocabulary.
+
+Emits attributes under the semconv-ai / Traceloop key names (e.g.
+``gen_ai.system``, ``gen_ai.usage.prompt_tokens``, ``llm.is_streaming``) plus a
+few bare, unprefixed service keys (``service``, ``call_type``, ``error``), for
+backends that consume those names.
+
+Like ``GenAIMapper``, each span kind declares its schema as a flat
+``attribute key -> extractor`` table: one lambda per mapping operation.
+"""
+
+from typing import Callable, Final
+
+from litellm.integrations.otel.mappers.base import AttributeMap, AttrValue, SpanData
+from litellm.integrations.otel.mappers.utils import collect, drop_none
+from litellm.integrations.otel.model.payloads import (
+ LLMCallSpanData,
+ ServiceSpanData,
+ ToolDefinition,
+)
+
+# Attribute keys in the semconv-ai / Traceloop vocabulary.
+_LEGACY_SYSTEM: Final = "gen_ai.system"
+_LEGACY_PROMPT_TOKENS: Final = "gen_ai.usage.prompt_tokens"
+_LEGACY_COMPLETION_TOKENS: Final = "gen_ai.usage.completion_tokens"
+_LEGACY_TOTAL_TOKENS: Final = "gen_ai.usage.total_tokens"
+_LEGACY_IS_STREAMING: Final = "llm.is_streaming"
+_LEGACY_TOP_K: Final = "llm.top_k"
+_LEGACY_FREQUENCY_PENALTY: Final = "llm.frequency_penalty"
+_LEGACY_PRESENCE_PENALTY: Final = "llm.presence_penalty"
+_LEGACY_STOP_SEQUENCES: Final = "llm.chat.stop_sequences"
+_LEGACY_SERVICE: Final = "service"
+_LEGACY_CALL_TYPE: Final = "call_type"
+_LEGACY_ERROR: Final = "error"
+
+
+class LegacyMapper:
+ """Emits LLM-call and service attributes under the older key names."""
+
+ _LLM_CALL_ATTRS: dict[str, Callable[[LLMCallSpanData], AttrValue | None]] = {
+ _LEGACY_SYSTEM: lambda d: d.provider or None,
+ _LEGACY_PROMPT_TOKENS: lambda d: d.usage.input_tokens,
+ _LEGACY_COMPLETION_TOKENS: lambda d: d.usage.output_tokens,
+ _LEGACY_TOTAL_TOKENS: lambda d: d.usage.total_tokens,
+ _LEGACY_IS_STREAMING: lambda d: d.is_streaming,
+ _LEGACY_TOP_K: lambda d: d.request_params.top_k,
+ _LEGACY_FREQUENCY_PENALTY: lambda d: d.request_params.frequency_penalty,
+ _LEGACY_PRESENCE_PENALTY: lambda d: d.request_params.presence_penalty,
+ _LEGACY_STOP_SEQUENCES: lambda d: (
+ list(d.request_params.stop_sequences)
+ if d.request_params.stop_sequences
+ else None
+ ),
+ }
+
+ _TOOL_ATTRS: dict[str, Callable[[ToolDefinition], AttrValue | None]] = {
+ "name": lambda t: t.name,
+ "description": lambda t: t.description or None,
+ "parameters": lambda t: t.parameters_json or None,
+ }
+
+ _SERVICE_ATTRS: dict[str, Callable[[ServiceSpanData], AttrValue | None]] = {
+ _LEGACY_SERVICE: lambda d: d.service_name,
+ _LEGACY_CALL_TYPE: lambda d: d.call_type,
+ _LEGACY_ERROR: lambda d: (
+ d.error.message if d.error is not None and d.error.message else None
+ ),
+ }
+
+ def map(self, data: SpanData) -> AttributeMap:
+ match data:
+ case LLMCallSpanData():
+ return self._llm_call(data)
+ case ServiceSpanData():
+ return self._service(data)
+ case _:
+ return {}
+
+ @classmethod
+ def _llm_call(cls, data: LLMCallSpanData) -> AttributeMap:
+ attrs = collect(cls._LLM_CALL_ATTRS, data)
+ attrs.update(
+ drop_none(
+ {
+ f"llm.request.functions.{idx}.{suffix}": extract(tool)
+ for idx, tool in enumerate(data.tools)
+ for suffix, extract in cls._TOOL_ATTRS.items()
+ }
+ )
+ )
+ return attrs
+
+ @classmethod
+ def _service(cls, data: ServiceSpanData) -> AttributeMap:
+ attrs = collect(cls._SERVICE_ATTRS, data)
+ attrs.update(dict(data.event_metadata))
+ return attrs
diff --git a/litellm/integrations/otel/mappers/openinference.py b/litellm/integrations/otel/mappers/openinference.py
new file mode 100644
index 00000000000..d8195cbe03d
--- /dev/null
+++ b/litellm/integrations/otel/mappers/openinference.py
@@ -0,0 +1,128 @@
+"""OpenInference attribute mapper (Arize + Arize-Phoenix shared vocabulary).
+
+Spec: https://github.com/Arize-ai/openinference/tree/main/spec — the standard
+both Arize and Phoenix consume. Composing this mapper after ``GenAIMapper``
+gives the same span both vocabularies, so a single trace lights up Arize +
+Phoenix + any other OpenInference-aware backend simultaneously.
+"""
+
+import json
+from typing import Callable, Sequence
+
+from litellm.integrations.otel.mappers.base import AttributeMap, AttrValue, SpanData
+from litellm.integrations.otel.mappers.utils import (
+ collect,
+ drop_none,
+ json_if,
+ message_content,
+ output_messages,
+)
+from litellm.integrations.otel.model.payloads import (
+ LLMCallSpanData,
+ LLMRequestParams,
+ ToolDefinition,
+)
+
+
+class OpenInferenceMapper:
+ """Emits OpenInference attributes for LLM_CALL spans.
+
+ Key families (per the OpenInference spec):
+ - ``openinference.span.kind`` — discriminator (``"LLM"`` here)
+ - ``llm.model_name`` / ``llm.provider`` / ``llm.invocation_parameters``
+ - ``llm.input_messages.{i}.message.role`` / ``...content``
+ - ``llm.output_messages.{i}.message.role`` / ``...content``
+ - ``llm.token_count.prompt`` / ``...completion`` / ``...total``
+ - ``input.value`` / ``output.value`` — JSON-serialized request / response
+ """
+
+ _LLM_CALL_ATTRS: dict[str, Callable[[LLMCallSpanData], AttrValue | None]] = {
+ "openinference.span.kind": lambda d: "LLM",
+ "llm.model_name": lambda d: d.request_model or None,
+ "llm.provider": lambda d: d.provider or None,
+ "llm.token_count.prompt": lambda d: d.usage.input_tokens,
+ "llm.token_count.completion": lambda d: d.usage.output_tokens,
+ "llm.token_count.total": lambda d: d.usage.total_tokens,
+ }
+
+ # Folded into the ``llm.invocation_parameters`` JSON blob.
+ _INVOCATION_PARAMS: dict[str, Callable[[LLMRequestParams], AttrValue | None]] = {
+ "temperature": lambda rp: rp.temperature,
+ "top_p": lambda rp: rp.top_p,
+ "top_k": lambda rp: rp.top_k,
+ "max_tokens": lambda rp: rp.max_tokens,
+ "frequency_penalty": lambda rp: rp.frequency_penalty,
+ "presence_penalty": lambda rp: rp.presence_penalty,
+ "seed": lambda rp: rp.seed,
+ }
+
+ # Per-tool extractors, keyed by the ``llm.tools.{idx}.*`` suffix.
+ _TOOL_ATTRS: dict[str, Callable[[ToolDefinition], AttrValue | None]] = {
+ "tool.name": lambda t: t.name,
+ "tool.description": lambda t: t.description or None,
+ "tool.json_schema": lambda t: t.parameters_json or None,
+ }
+
+ # JSON-payload attributes: each builder returns the serialized blob or None.
+ _BLOB_ATTRS: dict[str, Callable[[LLMCallSpanData], AttrValue | None]] = {
+ "llm.invocation_parameters": lambda d: json_if(
+ collect(OpenInferenceMapper._INVOCATION_PARAMS, d.request_params)
+ ),
+ }
+
+ def map(self, data: SpanData) -> AttributeMap:
+ match data:
+ case LLMCallSpanData():
+ return self._llm_call(data)
+ case _:
+ return {}
+
+ @classmethod
+ def _llm_call(cls, data: LLMCallSpanData) -> AttributeMap:
+ return {
+ **collect(cls._LLM_CALL_ATTRS, data),
+ **collect(cls._BLOB_ATTRS, data),
+ **cls._messages("llm.input_messages", "input.value", data.messages_in),
+ **cls._messages(
+ "llm.output_messages", "output.value", output_messages(data)
+ ),
+ **cls._tools(data),
+ }
+
+ @staticmethod
+ def _messages(
+ prefix: str, value_key: str, messages: Sequence[object]
+ ) -> AttributeMap:
+ """Per-message ``{prefix}.{idx}.message.*`` keys + the ``value_key`` blob."""
+ parsed = [
+ (m.get("role") if isinstance(m, dict) else None, message_content(m))
+ for m in messages
+ ]
+ attrs = drop_none(
+ {
+ key: value
+ for idx, (role, content) in enumerate(parsed)
+ for key, value in (
+ (
+ f"{prefix}.{idx}.message.role",
+ role if isinstance(role, str) else None,
+ ),
+ (f"{prefix}.{idx}.message.content", content),
+ )
+ }
+ )
+ if parsed:
+ attrs[value_key] = json.dumps(
+ [{"role": role, "content": content} for role, content in parsed]
+ )
+ return attrs
+
+ @classmethod
+ def _tools(cls, data: LLMCallSpanData) -> AttributeMap:
+ return drop_none(
+ {
+ f"llm.tools.{idx}.{suffix}": extract(tool)
+ for idx, tool in enumerate(data.tools)
+ for suffix, extract in cls._TOOL_ATTRS.items()
+ }
+ )
diff --git a/litellm/integrations/otel/mappers/utils.py b/litellm/integrations/otel/mappers/utils.py
new file mode 100644
index 00000000000..6228fc8bbe7
--- /dev/null
+++ b/litellm/integrations/otel/mappers/utils.py
@@ -0,0 +1,76 @@
+"""Shared helpers for the attribute mappers.
+
+Small, mapper-agnostic utilities — JSON serialization, message extraction, and
+extractor-table application — pulled out of the individual mapper modules so
+they live in one place.
+"""
+
+import json
+from typing import Callable, Mapping, Sequence
+
+from litellm.integrations.otel.mappers.base import AttributeMap, AttrValue
+from litellm.integrations.otel.model.payloads import LLMCallSpanData
+
+
+def drop_none(values: Mapping[str, AttrValue | None]) -> AttributeMap:
+ """Return ``values`` with ``None``-valued entries removed."""
+ return {k: v for k, v in values.items() if v is not None}
+
+
+def collect(table: Mapping[str, Callable], source: object) -> AttributeMap:
+ """Apply an extractor table to ``source``, dropping ``None`` results."""
+ return drop_none({key: extract(source) for key, extract in table.items()})
+
+
+def json_if(payload: Mapping[str, object]) -> str | None:
+ """JSON-serialize ``payload`` only when it's non-empty; else ``None``."""
+ return json.dumps(payload) if payload else None
+
+
+def json_or_none(value: object) -> str | None:
+ """JSON-serialize ``value`` (falling back to ``str``); ``None`` on failure."""
+ try:
+ return json.dumps(value, default=str)
+ except Exception:
+ return None
+
+
+def stringify_message(message: object) -> str | None:
+ """JSON-serialize a chat message dict; ``None`` if not a dict or on failure."""
+ if not isinstance(message, dict):
+ return None
+ try:
+ return json.dumps(message, default=str)
+ except Exception:
+ return None
+
+
+def serialize_messages(messages: Sequence[object]) -> str | None:
+ """Round-trip a sequence of message dicts through ``stringify_message``."""
+ serialized = [
+ json.loads(s) for s in (stringify_message(m) for m in messages) if s is not None
+ ]
+ return json.dumps(serialized) if serialized else None
+
+
+def message_content(message: object) -> str | None:
+ """Extract the textual ``content`` from a chat message dict."""
+ if not isinstance(message, dict):
+ return None
+ content = message.get("content")
+ if isinstance(content, str):
+ return content
+ if isinstance(content, list):
+ # multimodal: concatenate text parts only
+ parts = [
+ part.get("text", "")
+ for part in content
+ if isinstance(part, dict) and part.get("type") == "text"
+ ]
+ return "".join(p for p in parts if isinstance(p, str)) or None
+ return None
+
+
+def output_messages(data: LLMCallSpanData) -> list:
+ """The ``message`` payload of each response choice."""
+ return [c.get("message") for c in data.choices_out if isinstance(c, dict)]
diff --git a/litellm/integrations/otel/mappers/weave.py b/litellm/integrations/otel/mappers/weave.py
new file mode 100644
index 00000000000..54b07299271
--- /dev/null
+++ b/litellm/integrations/otel/mappers/weave.py
@@ -0,0 +1,48 @@
+"""Weave (W&B) attribute mapper.
+
+Weave consumes OpenInference + a small set of Weave-specific keys (display
+name, thread id, output value). This mapper layers the latter on top of
+OpenInference's vocabulary — compose ``["genai", "openinference", "weave"]``
+to feed a Weave backend.
+"""
+
+from typing import Callable
+
+from litellm.integrations.otel.mappers.base import AttributeMap, AttrValue, SpanData
+from litellm.integrations.otel.mappers.utils import collect, json_or_none
+from litellm.integrations.otel.model.payloads import LLMCallSpanData
+
+
+class WeaveMapper:
+ """Maps ``LLMCallSpanData`` to Weave's vendor attributes."""
+
+ _LLM_CALL_ATTRS: dict[str, Callable[[LLMCallSpanData], AttrValue | None]] = {
+ # ``display_name`` has the form ``"{operation} {model}"``. The span
+ # name already covers that, but Weave reads this attribute too.
+ "weave.display_name": lambda d: (
+ f"{d.operation.value} {d.request_model}" if d.request_model else None
+ ),
+ "weave.call_id": lambda d: d.identity.call_id or None,
+ }
+
+ # JSON-payload attributes: each builder returns the serialized blob or None.
+ _BLOB_ATTRS: dict[str, Callable[[LLMCallSpanData], AttrValue | None]] = {
+ # Weave treats the response choices as the "output" payload.
+ "weave.output": lambda d: (
+ json_or_none(list(d.choices_out)) if d.choices_out else None
+ ),
+ }
+
+ def map(self, data: SpanData) -> AttributeMap:
+ match data:
+ case LLMCallSpanData():
+ return self._llm_call(data)
+ case _:
+ return {}
+
+ @classmethod
+ def _llm_call(cls, data: LLMCallSpanData) -> AttributeMap:
+ return {
+ **collect(cls._LLM_CALL_ATTRS, data),
+ **collect(cls._BLOB_ATTRS, data),
+ }
diff --git a/litellm/integrations/otel/model/__init__.py b/litellm/integrations/otel/model/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/litellm/integrations/otel/model/baggage.py b/litellm/integrations/otel/model/baggage.py
new file mode 100644
index 00000000000..67dd64e3914
--- /dev/null
+++ b/litellm/integrations/otel/model/baggage.py
@@ -0,0 +1,76 @@
+"""Baggage promotion: request-identity values carried across child spans.
+
+A bounded set of identity values is written into OpenTelemetry Baggage on the
+LLM-call span so that child spans (guardrail, service) inherit them.
+``providers.LiteLLMBaggageSpanProcessor`` reads Baggage at span start and stamps
+the allowlisted keys onto every span.
+
+This module is the single place baggage is defined: ``_PROMOTABLE`` maps each
+promotable attribute key to how its value is read, and the two ``*_KEYS``
+defaults select what is promoted unless the config overrides them.
+"""
+
+from collections.abc import Callable
+from typing import Final
+
+from litellm.integrations.otel.model.metadata import RequestIdentity
+from litellm.integrations.otel.model.semconv import GenAI, LiteLLM
+
+# Attribute key -> value extractor over (identity, request_model). The single
+# definition of what may be promoted and under which key.
+_PROMOTABLE: Final[dict[str, Callable[[RequestIdentity, str | None], str | None]]] = {
+ LiteLLM.TEAM_ID: lambda identity, model: identity.team_id,
+ LiteLLM.TEAM_ALIAS: lambda identity, model: identity.team_alias,
+ LiteLLM.TEAM_METADATA: lambda identity, model: identity.team_metadata,
+ LiteLLM.KEY_HASH: lambda identity, model: identity.key_hash,
+ LiteLLM.END_USER: lambda identity, model: identity.end_user,
+ GenAI.REQUEST_MODEL: lambda identity, model: model,
+ LiteLLM.PROVIDER_MODEL: lambda identity, model: identity.provider_model,
+}
+
+# Keys promoted by default (a subset of ``_PROMOTABLE``). ``END_USER`` is
+# promotable but off by default — it identifies an individual user, so stamping
+# it onto every span is opt-in via ``config.baggage_promoted_keys``.
+BAGGAGE_PROMOTED_KEYS: Final[tuple[str, ...]] = (
+ LiteLLM.TEAM_ID,
+ LiteLLM.TEAM_ALIAS,
+ LiteLLM.TEAM_METADATA,
+ LiteLLM.KEY_HASH,
+ GenAI.REQUEST_MODEL,
+ LiteLLM.PROVIDER_MODEL,
+)
+
+# Metadata sub-keys eligible for promotion under the ``litellm.metadata.*``
+# namespace. The full metadata blob is never promoted; only this allowlist is.
+DEFAULT_BAGGAGE_METADATA_KEYS: Final[tuple[str, ...]] = (
+ "user_api_key_org_id",
+ "user_api_key_user_id",
+ "user_api_key_alias",
+ "user_api_key_end_user_id",
+ "requester_ip_address",
+)
+
+
+def promoted_baggage(
+ identity: RequestIdentity,
+ request_model: str | None,
+ promoted_keys: tuple[str, ...],
+ metadata_keys: tuple[str, ...] = DEFAULT_BAGGAGE_METADATA_KEYS,
+) -> dict[str, str]:
+ """Identity values to write into Baggage, filtered to ``promoted_keys``.
+
+ ``promoted_keys`` selects from ``_PROMOTABLE``; ``metadata_keys`` selects
+ sub-keys of ``identity.metadata`` to promote under ``litellm.metadata.*``.
+ Empty values are dropped.
+ """
+ out: dict[str, str] = {}
+ for key, extract in _PROMOTABLE.items():
+ if key in promoted_keys:
+ value = extract(identity, request_model)
+ if value:
+ out[key] = value
+ for meta_key in metadata_keys:
+ value = identity.metadata.get(meta_key)
+ if value:
+ out[f"{LiteLLM.METADATA_PREFIX}{meta_key}"] = value
+ return out
diff --git a/litellm/integrations/otel/model/config.py b/litellm/integrations/otel/model/config.py
new file mode 100644
index 00000000000..f78e1515ea1
--- /dev/null
+++ b/litellm/integrations/otel/model/config.py
@@ -0,0 +1,236 @@
+"""Typed configuration for the OpenTelemetry instrumentation."""
+
+from typing import Any, List
+
+from pydantic import AliasChoices, BaseModel, Field, field_validator, model_validator
+from pydantic_settings import BaseSettings, NoDecode, SettingsConfigDict
+from typing_extensions import Annotated
+
+from litellm.integrations.otel.model.baggage import (
+ BAGGAGE_PROMOTED_KEYS,
+ DEFAULT_BAGGAGE_METADATA_KEYS,
+)
+
+#: Master feature-flag env var. The logger is inert until this is truthy.
+OTEL_V2_ENV = "LITELLM_OTEL_V2"
+
+
+class CaptureMessageContent(str):
+ NO_CONTENT = "no_content"
+ SPAN_ONLY = "span_only"
+ EVENT_ONLY = "event_only"
+ SPAN_AND_EVENT = "span_and_event"
+
+
+class _OTelV2Flag(BaseSettings):
+ model_config = SettingsConfigDict(extra="ignore")
+
+ enabled: bool = Field(default=False, validation_alias=AliasChoices(OTEL_V2_ENV))
+
+
+def is_otel_v2_enabled() -> bool:
+ return _OTelV2Flag().enabled
+
+
+class ExporterSpec(BaseModel):
+ """One span-export destination.
+
+ The shared ``TracerProvider`` attaches one ``SpanProcessor`` per spec, so
+ listing several specs sends every span to all of them at once (e.g. Arize +
+ Phoenix + your own Honeycomb).
+ """
+
+ model_config = {"extra": "forbid"}
+
+ kind: str = Field(
+ default="console",
+ description="console | in_memory | otlp_http | otlp_grpc | ",
+ )
+ endpoint: str | None = None
+ headers: str | None = None
+ options: dict[str, str] | None = Field(
+ default=None,
+ description=(
+ "Factory-specific configuration for a custom exporter ``kind`` "
+ "registered via ``providers.register_exporter_factory`` (e.g. an "
+ "API key a lazy-auth exporter fetches a token with). Ignored by the "
+ "built-in console/in_memory/otlp exporters."
+ ),
+ )
+ use_simple_processor: bool | None = Field(
+ default=None,
+ description=(
+ "Force SimpleSpanProcessor regardless of exporter kind. Default: "
+ "auto (Simple for console/in_memory, Batch otherwise)."
+ ),
+ )
+
+
+class OpenTelemetryV2Config(BaseSettings):
+ model_config = SettingsConfigDict(populate_by_name=True, extra="ignore")
+
+ # ----- single-destination shorthand, read from standard OTEL_* envs ----- #
+ exporter: str = Field(
+ default="console",
+ validation_alias=AliasChoices("OTEL_EXPORTER", "OTEL_EXPORTER_OTLP_PROTOCOL"),
+ description=(
+ "Exporter kind for the single-destination shorthand. The model "
+ "validator folds this (with ``endpoint`` / ``headers``) into a "
+ "one-entry ``exporters`` list when ``exporters`` is empty; set "
+ "``exporters`` directly for multiple destinations."
+ ),
+ )
+ endpoint: str | None = Field(
+ default=None,
+ validation_alias=AliasChoices("OTEL_ENDPOINT", "OTEL_EXPORTER_OTLP_ENDPOINT"),
+ )
+ headers: str | None = Field(
+ default=None,
+ validation_alias=AliasChoices("OTEL_HEADERS", "OTEL_EXPORTER_OTLP_HEADERS"),
+ )
+ service_name: str = Field(
+ default="litellm", validation_alias=AliasChoices("OTEL_SERVICE_NAME")
+ )
+ deployment_environment: str | None = Field(
+ default=None, validation_alias=AliasChoices("OTEL_ENVIRONMENT_NAME")
+ )
+
+ enable_metrics: bool = Field(
+ default=False,
+ validation_alias=AliasChoices("LITELLM_OTEL_INTEGRATION_ENABLE_METRICS"),
+ )
+ enable_events: bool = Field(
+ default=False,
+ validation_alias=AliasChoices("LITELLM_OTEL_INTEGRATION_ENABLE_EVENTS"),
+ )
+ capture_message_content: str = Field(
+ default=CaptureMessageContent.NO_CONTENT,
+ validation_alias=AliasChoices(
+ "OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT"
+ ),
+ )
+ legacy_compat: bool = Field(
+ default=True, validation_alias=AliasChoices("LITELLM_OTEL_LEGACY_COMPAT")
+ )
+
+ # ----- explicit multi-destination / vocabulary configuration ------------ #
+
+ exporters: list[ExporterSpec] = Field(
+ default_factory=list,
+ description=(
+ "One destination per spec. The shared TracerProvider attaches a "
+ "SpanProcessor per entry. When empty, the model validator folds "
+ "the ``exporter`` / ``endpoint`` / ``headers`` shorthand into a "
+ "single spec so there is always at least one destination."
+ ),
+ )
+
+ mapper_names: Annotated[List[str], NoDecode] = Field(
+ default_factory=lambda: ["genai"],
+ description=(
+ "Ordered attribute vocabularies to emit. ``genai`` is the "
+ "canonical OTel GenAI vocabulary and is always placed first. "
+ "Vendor names: ``openinference`` (Arize + Phoenix), ``langfuse``, "
+ "``weave``, ``langtrace``."
+ ),
+ )
+
+ resource_attributes: dict[str, str] = Field(
+ default_factory=dict,
+ description=(
+ "Extra Resource attributes beyond ``service.name`` and "
+ "``deployment.environment`` (e.g. integration-specific markers)."
+ ),
+ )
+
+ baggage_promoted_keys: Annotated[List[str], NoDecode] = Field(
+ default_factory=lambda: list(BAGGAGE_PROMOTED_KEYS),
+ validation_alias=AliasChoices(
+ "baggage_promoted_keys", "LITELLM_OTEL_BAGGAGE_PROMOTED_KEYS"
+ ),
+ description=(
+ "Identity attribute keys written into Baggage and stamped on every "
+ "child span (e.g. ``litellm.team.id``). Configure via the "
+ "``LITELLM_OTEL_BAGGAGE_PROMOTED_KEYS`` env var (comma-separated) or "
+ "``callback_settings.otel.baggage_promoted_keys`` in config.yaml (a "
+ "YAML list)."
+ ),
+ )
+ baggage_metadata_keys: Annotated[List[str], NoDecode] = Field(
+ default_factory=lambda: list(DEFAULT_BAGGAGE_METADATA_KEYS),
+ validation_alias=AliasChoices(
+ "baggage_metadata_keys", "LITELLM_OTEL_BAGGAGE_METADATA_KEYS"
+ ),
+ description=(
+ "Metadata sub-keys promoted under the ``litellm.metadata.*`` "
+ "namespace. Configure via the ``LITELLM_OTEL_BAGGAGE_METADATA_KEYS`` "
+ "env var (comma-separated) or "
+ "``callback_settings.otel.baggage_metadata_keys`` in config.yaml."
+ ),
+ )
+
+ @field_validator(
+ "baggage_promoted_keys",
+ "baggage_metadata_keys",
+ "mapper_names",
+ mode="before",
+ )
+ @classmethod
+ def _split_csv(cls, value: Any) -> Any:
+ """Accept a comma-separated string for list fields.
+
+ Env vars are strings, but these fields are lists. Pydantic-settings would
+ otherwise require JSON for a list env var; splitting on commas here lets
+ an operator write ``LITELLM_OTEL_BAGGAGE_PROMOTED_KEYS=litellm.team.id,litellm.api_key.hash``.
+ YAML lists (from ``callback_settings.otel.*``) and real lists pass through
+ unchanged.
+ """
+ if isinstance(value, str):
+ return [item.strip() for item in value.split(",") if item.strip()]
+ return value
+
+ @model_validator(mode="after")
+ def _normalize(self) -> "OpenTelemetryV2Config":
+ # An endpoint with the default exporter kind implies OTLP/HTTP.
+ if self.endpoint and self.exporter == "console":
+ self.exporter = "otlp_http"
+ # When no explicit destinations are given, fold the single-destination
+ # shorthand into one spec so the provider always has a destination.
+ if not self.exporters:
+ self.exporters = [
+ ExporterSpec(
+ kind=self.exporter,
+ endpoint=self.endpoint,
+ headers=self.headers,
+ )
+ ]
+ # Ensure ``genai`` is always present and first.
+ names = list(self.mapper_names)
+ if "genai" in names:
+ names = ["genai"] + [n for n in names if n != "genai"]
+ else:
+ names = ["genai"] + names
+ # When enabled, also emit attribute keys under their semconv-ai /
+ # Traceloop names via the ``legacy`` mapper. Append it at the tail so
+ # the canonical ``genai`` keys win on any conflict.
+ if self.legacy_compat and "legacy" not in names:
+ names.append("legacy")
+ self.mapper_names = names
+ return self
+
+ @property
+ def capture_span_content(self) -> bool:
+ """Whether prompt/response content may be stamped as span attributes.
+
+ Defaults off (``no_content``): an operator must opt in before message
+ bodies leave the process, so a user request can never force its prompt
+ or completion into the configured backend while capture is disabled.
+ """
+ return self.capture_message_content in (
+ CaptureMessageContent.SPAN_ONLY,
+ CaptureMessageContent.SPAN_AND_EVENT,
+ )
+
+ @classmethod
+ def from_env(cls) -> "OpenTelemetryV2Config":
+ return cls()
diff --git a/litellm/integrations/otel/model/metadata.py b/litellm/integrations/otel/model/metadata.py
new file mode 100644
index 00000000000..f0ea0a608c6
--- /dev/null
+++ b/litellm/integrations/otel/model/metadata.py
@@ -0,0 +1,315 @@
+"""The single translation layer between a request's metadata and the spans.
+
+Every relevant field litellm exposes about a request — the user-facing model,
+the model actually dispatched to the provider, the deployment, and the caller's
+identity (team, key, end-user) — is parsed **once**, here, out of the
+``StandardLoggingPayload`` (or a ``UserAPIKeyAuth`` at the auth boundary). Span
+data, baggage promotion, and the mappers then read these typed fields instead of
+each digging into the raw ``metadata`` / ``hidden_params`` dicts.
+
+Two models live here because a request's identity is known *before* its model
+resolution is:
+
+* :class:`RequestIdentity` — team / key / end-user, seeded into Baggage at the
+ auth boundary (``from_user_api_key_auth``), before routing has picked a
+ deployment. ``provider_model`` is therefore absent from that early seed and is
+ only filled in from the payload once the call closes.
+* :class:`RequestContext` — the full picture available at close: the resolved
+ request vs. provider model split, plus the response model, model group, model
+ id, and api base, wrapping the :class:`RequestIdentity`.
+
+The request-vs-provider model split is the subtle part. On the proxy a caller
+asks for a *model group* (e.g. ``gpt-4o``) that routes to a concrete deployment
+(e.g. ``azure/my-deployment``); the two are distinct and both worth recording.
+``StandardLoggingPayload`` exposes them as:
+
+* ``model_group`` — the user-facing name the caller requested.
+* ``model`` — already reconstructed (see ``reconstruct_model_name``) to the name
+ litellm dispatched to the provider (the deployment, provider-prefixed).
+* ``hidden_params.litellm_model_name`` — a secondary source for the dispatched
+ model (populated only on some call paths, e.g. files).
+
+So ``gen_ai.request.model`` is the *group* (falling back to the call model on the
+SDK path, which has no group), and ``litellm.provider.model`` is the *dispatched*
+model. They coincide on the SDK path, which is correct.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Mapping, cast
+
+from litellm.constants import LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL
+from litellm.integrations.otel.model.semconv import resolve_operation
+from litellm.integrations.otel.model.utils import as_str
+
+if TYPE_CHECKING:
+ from litellm.types.utils import StandardLoggingPayload
+
+
+@dataclass(frozen=True)
+class RequestIdentity:
+ call_id: str | None = None
+ team_id: str | None = None
+ team_alias: str | None = None
+ # The team's free-form metadata dict, JSON-serialized (empty/missing -> None).
+ team_metadata: str | None = None
+ key_hash: str | None = None
+ end_user: str | None = None
+ # The model litellm dispatched to the provider. Only known once the call
+ # completes (routing has picked a deployment), so it's absent from the
+ # auth-time seed and filled only from the payload.
+ provider_model: str | None = None
+ metadata: Mapping[str, str] = field(default_factory=dict)
+
+ @classmethod
+ def from_payload(cls, payload: "StandardLoggingPayload") -> "RequestIdentity":
+ """Parse caller identity out of a closed request's payload metadata.
+
+ ``provider_model`` is resolved here too (see :func:`resolve_provider_model`)
+ so the identity carried into Baggage labels every span with the dispatched
+ model, not just the user-facing one.
+ """
+ raw_meta = cast(Mapping[str, object], payload.get("metadata") or {})
+ metadata = {
+ key: str(value)
+ for key, value in raw_meta.items()
+ if isinstance(value, (str, bool, int, float))
+ }
+ return cls(
+ call_id=as_str(payload.get("litellm_call_id")) or as_str(payload.get("id")),
+ # StandardLoggingMetadata's canonical key is ``user_api_key_team_id``;
+ # the bare ``team_id`` is a legacy alias and is often empty, so prefer
+ # the canonical key and fall back to the alias.
+ team_id=as_str(raw_meta.get("user_api_key_team_id"))
+ or as_str(raw_meta.get("team_id")),
+ team_alias=as_str(raw_meta.get("user_api_key_team_alias"))
+ or as_str(raw_meta.get("team_alias")),
+ team_metadata=_team_metadata_json(
+ raw_meta.get("user_api_key_team_metadata")
+ ),
+ key_hash=as_str(raw_meta.get("user_api_key_hash")),
+ end_user=as_str(payload.get("end_user"))
+ or as_str(raw_meta.get("user_api_key_end_user_id")),
+ provider_model=resolve_provider_model(payload),
+ metadata=metadata,
+ )
+
+ @classmethod
+ def from_user_api_key_auth(cls, auth: object) -> "RequestIdentity":
+ """Identity from a ``UserAPIKeyAuth`` (duck-typed to keep this module
+ free of a proxy import).
+
+ Used in the pre-call hook to seed Baggage early — before any LLM,
+ guardrail, or service span is created — so the whole request's spans
+ inherit identity, not just the LLM-call span. Metadata sub-keys use the
+ ``user_api_key_*`` names that ``baggage.DEFAULT_BAGGAGE_METADATA_KEYS``
+ promotes.
+ """
+ get = lambda name: getattr(auth, name, None) # noqa: E731
+ metadata = {
+ meta_key: str(value)
+ for meta_key, attr in (
+ ("user_api_key_user_id", "user_id"),
+ ("user_api_key_org_id", "org_id"),
+ ("user_api_key_alias", "key_alias"),
+ ("user_api_key_end_user_id", "end_user_id"),
+ )
+ if (value := get(attr))
+ }
+ return cls(
+ team_id=as_str(get("team_id")),
+ team_alias=as_str(get("team_alias")),
+ team_metadata=_team_metadata_json(get("team_metadata")),
+ key_hash=as_str(get("api_key")),
+ end_user=as_str(get("end_user_id")),
+ # ``provider_model`` is unknown at the auth boundary — routing hasn't
+ # picked a deployment yet — so it's only populated from the payload.
+ metadata=metadata,
+ )
+
+
+@dataclass(frozen=True)
+class RequestContext:
+ """The fully-resolved view of a closed request, parsed once from the payload.
+
+ ``request_model`` is the user-facing requested model and ``provider_model``
+ (on :attr:`identity`) is the model litellm dispatched to the provider; the two
+ differ on the proxy (group vs. deployment) and coincide on the SDK path.
+ """
+
+ request_model: str
+ response_model: str | None
+ model_group: str | None
+ model_id: str | None
+ api_base: str | None
+ identity: RequestIdentity
+
+ @property
+ def provider_model(self) -> str | None:
+ """The dispatched-model name, carried on the identity for Baggage."""
+ return self.identity.provider_model
+
+ @classmethod
+ def from_standard_logging_payload(
+ cls, payload: "StandardLoggingPayload"
+ ) -> "RequestContext":
+ raw_meta = cast(Mapping[str, object], payload.get("metadata") or {})
+ hidden = cast(Mapping[str, object], payload.get("hidden_params") or {})
+ raw_response = payload.get("response")
+ response = cast(
+ Mapping[str, object], raw_response if isinstance(raw_response, dict) else {}
+ )
+ model_group = as_str(payload.get("model_group")) or as_str(
+ raw_meta.get("model_group")
+ )
+ return cls(
+ # The user asked for the group; fall back to the call model on the SDK
+ # path, which has no group. Empty string (never None) so the span name
+ # builder and the mapper see a plain string.
+ request_model=model_group or as_str(payload.get("model")) or "",
+ response_model=as_str(response.get("model")),
+ model_group=model_group,
+ model_id=as_str(payload.get("model_id"))
+ or _model_info_id(raw_meta.get("model_info")),
+ api_base=as_str(payload.get("api_base")) or as_str(hidden.get("api_base")),
+ identity=RequestIdentity.from_payload(payload),
+ )
+
+
+# --- live-callback kwargs parsing ------------------------------------------- #
+#
+# The model and helpers below parse the *live* callback ``kwargs`` god object (and
+# the raw pre/post-call ``data`` dicts) — the untyped request state that reaches a
+# ``CustomLogger`` before, or instead of, a ``StandardLoggingPayload``. They live
+# here, with the payload/auth parsers, so every read out of a request's raw dicts
+# is in one place rather than scattered across the ``CustomLogger``.
+
+
+@dataclass(frozen=True)
+class LLMCallEvent:
+ """The typed view of the live callback ``kwargs`` (``model_call_details``).
+
+ litellm hands every callback an untyped ``kwargs`` god object. The fields the
+ OTel logger needs out of it are parsed **once**, here, so the ``CustomLogger``
+ reads typed attributes instead of digging into the dict at each boundary.
+ """
+
+ # The ``litellm_call_id`` correlating ``pre_call`` with the close callback.
+ # Present in ``model_call_details`` at ``pre_call`` and in both the kwargs and
+ # the ``standard_logging_object`` at success/failure, so it's a stable key for
+ # the open-call carrier — no back-reference to the logging object required (the
+ # object isn't reachable from the callback kwargs at ``pre_call`` time).
+ call_id: str | None
+ # The ``StandardLoggingPayload`` carried on a success/failure callback; ``None``
+ # at ``pre_call``, or when the call closed before any payload materialized (so
+ # there is nothing to stamp on the span).
+ payload: "StandardLoggingPayload | None"
+ # The ``standard_callback_dynamic_params`` routing the call to a per-tenant
+ # tracer (its own exporter/endpoint), or ``None`` when the call isn't scoped.
+ dynamic_params: Any
+ # True for synthetic proxy-gate logs (auth / rate-limit rejections): they fire
+ # the ``pre_call`` hook but never made an upstream call, so they get no span.
+ is_no_upstream_call: bool
+ # A best-effort ``"{operation} {model}"`` name known at ``pre_call`` time. The
+ # span is renamed from the typed payload at close (``finish_span``); this only
+ # needs to be reasonable for a span that never gets closed (a leak).
+ provisional_span_name: str
+
+ @classmethod
+ def from_dict(cls, kwargs: Mapping[str, Any]) -> "LLMCallEvent":
+ raw_payload = kwargs.get("standard_logging_object")
+ payload = cast("StandardLoggingPayload", raw_payload) if raw_payload else None
+ operation = resolve_operation(as_str(kwargs.get("call_type")))
+ model = as_str(kwargs.get("model")) or ""
+ return cls(
+ call_id=_call_id(payload, kwargs),
+ payload=payload,
+ dynamic_params=kwargs.get("standard_callback_dynamic_params"),
+ is_no_upstream_call=bool(kwargs.get(LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL)),
+ provisional_span_name=f"{operation.value} {model}".strip(),
+ )
+
+
+def _call_id(
+ payload: "StandardLoggingPayload | None", kwargs: Mapping[str, Any]
+) -> str | None:
+ """The call id from the payload (when closed) or the bare kwargs (at pre_call)."""
+ if payload is not None:
+ call_id = as_str(payload.get("litellm_call_id")) or as_str(payload.get("id"))
+ if call_id:
+ return call_id
+ return as_str(kwargs.get("litellm_call_id"))
+
+
+def model_from_request_data(data: object) -> str | None:
+ """The user-facing ``model`` from a pre-call ``data`` dict (``None`` if absent).
+
+ Read at the auth boundary to label early Baggage before routing has resolved
+ a deployment; ``data`` is duck-typed since it arrives untyped from the proxy.
+ """
+ if isinstance(data, Mapping):
+ return as_str(data.get("model"))
+ return None
+
+
+def guardrail_entries_from_request_data(
+ request_data: Mapping[str, Any],
+) -> list[dict]:
+ """The guardrail-information dicts buried in ``metadata`` of a post-call dict.
+
+ ``standard_logging_guardrail_information`` is stored as either a single dict
+ or a list of them; normalize to a list of dicts (dropping non-dict noise) so
+ the caller just iterates. Empty list when none are present.
+ """
+ metadata = request_data.get("metadata")
+ if not isinstance(metadata, Mapping):
+ return []
+ info = metadata.get("standard_logging_guardrail_information")
+ if isinstance(info, Mapping):
+ return [cast(dict, info)]
+ if isinstance(info, list):
+ return [entry for entry in info if isinstance(entry, dict)]
+ return []
+
+
+def resolve_provider_model(payload: "StandardLoggingPayload") -> str | None:
+ """The model litellm dispatched to the provider, from the payload.
+
+ Prefers the explicit ``hidden_params.litellm_model_name`` (set on call paths
+ that know it, e.g. files), then the top-level ``model`` — which
+ ``reconstruct_model_name`` has already resolved to the deployment's
+ provider-prefixed name. Returns ``None`` only when neither is present.
+ """
+ raw_meta = cast(Mapping[str, object], payload.get("metadata") or {})
+ hidden = cast(Mapping[str, object], payload.get("hidden_params") or {})
+ return (
+ # ``deployment`` survives only on paths that don't strip it from metadata;
+ # harmless (and most precise) to prefer it when present.
+ as_str(raw_meta.get("deployment"))
+ or as_str(hidden.get("litellm_model_name"))
+ or as_str(payload.get("model"))
+ )
+
+
+def _model_info_id(model_info: object) -> str | None:
+ """The deployment id from a ``metadata.model_info`` sub-dict, if present."""
+ if isinstance(model_info, Mapping):
+ return as_str(model_info.get("id"))
+ return None
+
+
+def _team_metadata_json(value: object) -> str | None:
+ """JSON-serialize a team's metadata dict for a single Baggage value.
+
+ Returns ``None`` for a missing, non-dict, or empty mapping so the empty case
+ is dropped rather than promoting a useless ``"{}"``. Keys are sorted for a
+ stable, diff-friendly serialization.
+ """
+ if not isinstance(value, Mapping) or not value:
+ return None
+ try:
+ return json.dumps(value, default=str, sort_keys=True)
+ except Exception:
+ return None
diff --git a/litellm/integrations/otel/model/payloads.py b/litellm/integrations/otel/model/payloads.py
new file mode 100644
index 00000000000..65b50d0fc12
--- /dev/null
+++ b/litellm/integrations/otel/model/payloads.py
@@ -0,0 +1,468 @@
+"""Typed span-data inputs: frozen dataclasses the engine and mappers consume."""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import TYPE_CHECKING, ClassVar, Mapping, cast
+from urllib.parse import urlsplit
+
+from litellm.integrations.otel.model.metadata import (
+ RequestContext,
+ RequestIdentity,
+)
+from litellm.integrations.otel.model.semconv import (
+ GenAIOperation,
+ resolve_operation,
+ resolve_provider,
+)
+from litellm.integrations.otel.model.utils import (
+ as_bool,
+ as_float,
+ as_int,
+ as_str,
+ as_str_tuple,
+)
+
+# ``RequestIdentity`` and the request-metadata translation now live in
+# :mod:`metadata`; re-exported here so existing ``model.payloads`` imports keep
+# resolving it.
+__all__ = [
+ "RequestContext",
+ "RequestIdentity",
+ "GuardrailSpanData",
+ "LLMCallSpanData",
+ "LLMRequestParams",
+ "LLMUsage",
+ "ProxyRequestSpanData",
+ "ServerInfo",
+ "ServiceSpanData",
+ "SpanError",
+ "ToolDefinition",
+]
+
+if TYPE_CHECKING:
+ from litellm.types.services import ServiceLoggerPayload
+ from litellm.types.utils import (
+ StandardLoggingGuardrailInformation,
+ StandardLoggingPayload,
+ )
+
+
+# --- typed sub-structures ---------------------------------------------------- #
+
+
+@dataclass(frozen=True)
+class LLMRequestParams:
+ temperature: float | None = None
+ top_p: float | None = None
+ top_k: int | None = None
+ max_tokens: int | None = None
+ frequency_penalty: float | None = None
+ presence_penalty: float | None = None
+ stop_sequences: tuple[str, ...] | None = None
+ seed: int | None = None
+
+ @classmethod
+ def from_model_parameters(cls, params: Mapping[str, object]) -> "LLMRequestParams":
+ max_tokens = as_int(params.get("max_tokens"))
+ if max_tokens is None:
+ max_tokens = as_int(params.get("max_completion_tokens"))
+ return cls(
+ temperature=as_float(params.get("temperature")),
+ top_p=as_float(params.get("top_p")),
+ top_k=as_int(params.get("top_k")),
+ max_tokens=max_tokens,
+ frequency_penalty=as_float(params.get("frequency_penalty")),
+ presence_penalty=as_float(params.get("presence_penalty")),
+ stop_sequences=as_str_tuple(params.get("stop")),
+ seed=as_int(params.get("seed")),
+ )
+
+
+@dataclass(frozen=True)
+class LLMUsage:
+ input_tokens: int | None = None
+ output_tokens: int | None = None
+ total_tokens: int | None = None
+
+
+@dataclass(frozen=True)
+class SpanError:
+ error_type: str | None = None
+ message: str | None = None
+
+
+@dataclass(frozen=True)
+class ServerInfo:
+ address: str | None = None
+ port: int | None = None
+
+ @classmethod
+ def from_api_base(cls, api_base: str | None) -> ServerInfo | None:
+ if not api_base:
+ return None
+ parsed = urlsplit(api_base if "://" in api_base else f"//{api_base}")
+ if not parsed.hostname:
+ return None
+ return cls(address=parsed.hostname, port=parsed.port)
+
+
+@dataclass(frozen=True)
+class GuardrailSpanData:
+ guardrail_name: str
+ mode: str | None = None
+ status: str | None = None
+ masked_entity_count: int | None = None
+ provider: str | None = None
+ action: str | None = None
+ # The guardrail verdict / provider response (e.g. the moderation result),
+ # JSON-serialized. This is the detail that belongs on the guardrail span.
+ response_json: str | None = None
+ violation_categories: tuple[str, ...] = ()
+ confidence_score: float | None = None
+ risk_score: float | None = None
+ duration: float | None = None
+ # Actual execution window (epoch seconds) from the logging entry, so the span
+ # is placed when the guardrail really ran — a pre_call guardrail before the
+ # LLM call — rather than at post-call emission time.
+ start_time: float | None = None
+ end_time: float | None = None
+ # Provider-agnostic configuration/detection metadata (see
+ # ``StandardLoggingGuardrailInformation``). Present for any guardrail that
+ # populates them, not just one provider's shape.
+ guardrail_id: str | None = None
+ policy_template: str | None = None
+ detection_method: str | None = None
+ # Set when the guardrail intervened/blocked or failed, so the emitter marks
+ # the span ERROR — a blocking guardrail is an error outcome for that span.
+ error: SpanError | None = None
+
+ # Guardrail statuses that mean the guardrail did not pass the request through.
+ _ERROR_STATUSES: ClassVar[frozenset[str]] = frozenset(
+ {"guardrail_intervened", "guardrail_failed_to_respond"}
+ )
+
+ @classmethod
+ def from_logging_entry(
+ cls, entry: "StandardLoggingGuardrailInformation"
+ ) -> "GuardrailSpanData":
+ """Build from one ``standard_logging_guardrail_information`` entry.
+
+ Reads the canonical, provider-agnostic ``StandardLoggingGuardrailInformation``
+ keys only — no guessing at a single provider's field names. Values that are
+ typed as enums or lists (e.g. ``guardrail_mode``) are normalized to a
+ stable string rather than assumed to already be plain strings.
+ """
+ get = cast(Mapping[str, object], entry).get
+ status = as_str(get("guardrail_status"))
+ response = get("guardrail_response")
+ error = (
+ SpanError(error_type=status, message=as_str(get("guardrail_action")))
+ if status in cls._ERROR_STATUSES
+ else None
+ )
+ return cls(
+ guardrail_name=as_str(get("guardrail_name")) or "guardrail",
+ mode=_guardrail_mode_str(get("guardrail_mode")),
+ status=status,
+ masked_entity_count=_total_masked_entities(get("masked_entity_count")),
+ provider=as_str(get("guardrail_provider")),
+ action=as_str(get("guardrail_action")),
+ response_json=_json_or_none(response) if response is not None else None,
+ violation_categories=as_str_tuple(get("violation_categories")) or (),
+ confidence_score=as_float(get("confidence_score")),
+ risk_score=as_float(get("risk_score")),
+ duration=as_float(get("duration")),
+ start_time=as_float(get("start_time")),
+ end_time=as_float(get("end_time")),
+ guardrail_id=as_str(get("guardrail_id")),
+ policy_template=as_str(get("policy_template")),
+ detection_method=as_str(get("detection_method")),
+ error=error,
+ )
+
+
+@dataclass(frozen=True)
+class ServiceSpanData:
+ service_name: str
+ call_type: str | None = None
+ error: SpanError | None = None
+ # Caller-supplied attributes to stamp on the service span, passed through
+ # from ``async_service_*_hook(event_metadata=...)``. The mapper owns how
+ # these are namespaced: the canonical vocabulary uses ``litellm.metadata.*``
+ # keys, the semconv-ai / Traceloop vocabulary uses the bare key names.
+ event_metadata: Mapping[str, str] = field(default_factory=dict)
+
+ @classmethod
+ def from_payload(
+ cls,
+ payload: "ServiceLoggerPayload",
+ event_metadata: Mapping[str, object] | None = None,
+ ) -> "ServiceSpanData":
+ # ``payload.service`` is a ``ServiceTypes(str, Enum)`` and ``error`` is
+ # ``Optional[str]`` on the Pydantic model — no defensive reads needed.
+ # ``event_metadata`` is sanitized: the legacy service decorators pass raw
+ # call-site data (live objects, full request metadata, response headers),
+ # none of which belongs on a span.
+ return cls(
+ service_name=payload.service.value,
+ call_type=payload.call_type,
+ error=SpanError(message=payload.error) if payload.error else None,
+ event_metadata=sanitize_event_metadata(event_metadata),
+ )
+
+
+@dataclass(frozen=True)
+class ProxyRequestSpanData:
+ http_method: str
+ route: str
+ url_path: str | None = None
+ status_code: int | None = None
+ identity: RequestIdentity | None = None
+
+
+# --- the primary LLM-call model ---------------------------------------------- #
+
+
+@dataclass(frozen=True)
+class ToolDefinition:
+ """A single function/tool declared on a chat-completion request."""
+
+ name: str
+ description: str | None = None
+ parameters_json: str | None = (
+ None # JSON-serialized schema (str so it's an AttrValue)
+ )
+
+
+@dataclass(frozen=True)
+class LLMCallSpanData:
+ operation: GenAIOperation
+ provider: str
+ request_model: str
+ response_model: str | None
+ response_id: str | None
+ request_params: LLMRequestParams
+ usage: LLMUsage
+ finish_reasons: tuple[str, ...]
+ error: SpanError | None
+ response_cost: float | None
+ server: ServerInfo | None
+ identity: RequestIdentity
+ is_streaming: bool | None = None
+ tools: tuple[ToolDefinition, ...] = ()
+ # Raw messages and response, needed by vendor mappers (OpenInference,
+ # Langfuse, Weave) that stamp message-level attributes. ``messages_in`` is
+ # the request payload; ``choices_out`` mirrors ``response.choices`` from
+ # the StandardLoggingPayload. Both are tuples of immutable mappings so the
+ # dataclass stays hashable and frozen.
+ messages_in: tuple[Mapping[str, object], ...] = ()
+ choices_out: tuple[Mapping[str, object], ...] = ()
+ system_fingerprint: str | None = None
+
+ @classmethod
+ def from_standard_logging_payload(
+ cls, payload: "StandardLoggingPayload", capture_content: bool = False
+ ) -> "LLMCallSpanData":
+ params = cast(Mapping[str, object], payload.get("model_parameters") or {})
+ # The single parse of the request's metadata — the request-vs-provider
+ # model split, the response model, api base, and identity all come from
+ # here rather than being re-derived from the raw payload dicts.
+ context = RequestContext.from_standard_logging_payload(payload)
+ # Normalize ``response`` to a dict once so the content/id reads below are a
+ # plain ``.get`` — no repeated ``isinstance`` guards.
+ raw_response = payload.get("response")
+ response = cast(
+ Mapping[str, object], raw_response if isinstance(raw_response, dict) else {}
+ )
+ choices_out = _dicts(response.get("choices"))
+ # ``finish_reasons`` is metadata, not content, so derive it from
+ # ``choices_out`` before gating. The raw message/choice bodies are only
+ # retained when content capture is enabled (see ``capture_span_content``);
+ # otherwise the content-bearing mappers receive empty sequences and emit
+ # no prompt/response text.
+ finish_reasons = _finish_reasons(choices_out)
+ return cls(
+ operation=resolve_operation(as_str(payload.get("call_type"))),
+ provider=resolve_provider(as_str(payload.get("custom_llm_provider"))),
+ request_model=context.request_model,
+ response_model=context.response_model,
+ response_id=as_str(response.get("id")),
+ request_params=LLMRequestParams.from_model_parameters(params),
+ usage=LLMUsage(
+ input_tokens=as_int(payload.get("prompt_tokens")),
+ output_tokens=as_int(payload.get("completion_tokens")),
+ total_tokens=as_int(payload.get("total_tokens")),
+ ),
+ finish_reasons=finish_reasons,
+ error=_parse_error(payload),
+ response_cost=as_float(payload.get("response_cost")),
+ server=ServerInfo.from_api_base(context.api_base),
+ identity=context.identity,
+ is_streaming=as_bool(payload.get("stream")),
+ tools=_extract_tools(params),
+ messages_in=_dicts(payload.get("messages")) if capture_content else (),
+ choices_out=choices_out if capture_content else (),
+ system_fingerprint=as_str(response.get("system_fingerprint")),
+ )
+
+
+# --- service event_metadata sanitization ------------------------------------ #
+
+# Substrings (case-insensitive) of keys that must never reach a span: secrets,
+# tokens, and raw request/response dumps the legacy service decorators pass.
+_SENSITIVE_METADATA_SUBSTRINGS: tuple[str, ...] = (
+ "api_key",
+ "token",
+ "secret",
+ "password",
+ "cookie",
+ "authorization",
+ "header",
+ "hidden_params",
+)
+# Keys that carry raw call-site internals — live objects, full kwargs/args. The
+# operation name is already the span's ``call_type``, so ``function_name`` is
+# redundant.
+_DROP_METADATA_KEYS: frozenset = frozenset(
+ {"function_kwargs", "function_args", "function_name"}
+)
+_MAX_METADATA_VALUE_LEN = 1024
+_MAX_METADATA_ITEMS = 32
+
+
+def sanitize_event_metadata(
+ event_metadata: Mapping[str, object] | None,
+) -> dict[str, str]:
+ """Reduce caller-supplied ``event_metadata`` to span-safe string attributes.
+
+ Keeps only primitive values (str/int/float/bool) under non-sensitive keys —
+ never ``repr()``-ing objects, dicts, or lists, never stamping secrets/headers,
+ and bounding the count and per-value length. This is the single chokepoint:
+ both the GenAI and legacy mappers read the cleaned result.
+ """
+ if not event_metadata:
+ return {}
+ clean: dict[str, str] = {}
+ for key, value in event_metadata.items():
+ if len(clean) >= _MAX_METADATA_ITEMS:
+ break
+ if not isinstance(key, str) or key in _DROP_METADATA_KEYS:
+ continue
+ lowered = key.lower()
+ if any(token in lowered for token in _SENSITIVE_METADATA_SUBSTRINGS):
+ continue
+ # ``bool`` is a subclass of ``int``, so it's covered. Non-primitive values
+ # (objects, dicts, lists) are dropped rather than stringified.
+ if isinstance(value, (str, int, float)):
+ clean[key] = str(value)[:_MAX_METADATA_VALUE_LEN]
+ return clean
+
+
+def _json_or_none(value: object) -> str | None:
+ """JSON-serialize ``value`` (already-string values pass through). ``None`` on failure."""
+ if isinstance(value, str):
+ return value
+ try:
+ return json.dumps(value, default=str)
+ except Exception:
+ return None
+
+
+def _guardrail_mode_str(value: object) -> str | None:
+ """Normalize ``guardrail_mode`` to a stable string.
+
+ ``guardrail_mode`` is typed as a ``GuardrailEventHooks`` enum, a list of them,
+ or a ``GuardrailMode`` — not a plain string. Emit the enum *value* (e.g.
+ ``"pre_call"``) rather than ``str(enum)`` (``"GuardrailEventHooks.pre_call"``),
+ and join a list of modes so a guardrail that runs at multiple hooks is
+ represented faithfully.
+ """
+ if value is None:
+ return None
+ if isinstance(value, (list, tuple)):
+ parts: list[str] = []
+ for item in value:
+ if item is None:
+ continue
+ part = as_str(item.value) if isinstance(item, Enum) else as_str(item)
+ if part:
+ parts.append(part)
+ return ",".join(parts) or None
+ if isinstance(value, Enum):
+ return as_str(value.value)
+ return as_str(value)
+
+
+def _total_masked_entities(value: object) -> int | None:
+ """``masked_entity_count`` is a ``{entity_type: count}`` map — sum to a total."""
+ if isinstance(value, Mapping):
+ total = sum(v for v in value.values() if isinstance(v, int))
+ return total or None
+ return as_int(value)
+
+
+def _dicts(value: object) -> tuple[Mapping[str, object], ...]:
+ """The dict items of ``value`` (when it's a list), as a tuple. Else empty."""
+ if not isinstance(value, list):
+ return ()
+ return tuple(item for item in value if isinstance(item, dict))
+
+
+def _finish_reasons(choices: tuple[Mapping[str, object], ...]) -> tuple[str, ...]:
+ """Non-empty ``finish_reason`` of each response choice."""
+ return tuple(r for c in choices if (r := as_str(c.get("finish_reason"))))
+
+
+def _parse_error(payload: "StandardLoggingPayload") -> SpanError | None:
+ """A ``SpanError`` for a failed request, or ``None`` on success."""
+ if payload.get("status") != "failure":
+ return None
+ info = cast(Mapping[str, object], payload.get("error_information") or {})
+ return SpanError(
+ error_type=as_str(info.get("error_class")) or as_str(info.get("error_code")),
+ message=as_str(info.get("error_message")) or as_str(payload.get("error_str")),
+ )
+
+
+def _tool_from_entry(entry: object) -> ToolDefinition | None:
+ """One ``tools``/``functions`` entry → ``ToolDefinition``, or ``None`` if unusable."""
+ if not isinstance(entry, dict):
+ return None
+ fn = entry.get("function") if "function" in entry else entry
+ if not isinstance(fn, dict):
+ return None
+ name = as_str(fn.get("name"))
+ if not name:
+ return None
+ params = fn.get("parameters")
+ parameters_json: str | None = None
+ if params is not None:
+ try:
+ parameters_json = json.dumps(params, default=str)
+ except Exception:
+ parameters_json = None
+ return ToolDefinition(
+ name=name,
+ description=as_str(fn.get("description")),
+ parameters_json=parameters_json,
+ )
+
+
+def _extract_tools(
+ model_parameters: Mapping[str, object],
+) -> tuple[ToolDefinition, ...]:
+ """Pull declared tools from request params (OpenAI / Anthropic shape).
+
+ Accepts the chat-completion ``tools=[{"type":"function", "function":
+ {...}}, ...]`` shape, and falls back to the ``functions=[...]`` shape.
+ Returns an empty tuple when neither is present.
+ """
+ raw_tools = model_parameters.get("tools")
+ if not isinstance(raw_tools, list):
+ raw_tools = model_parameters.get("functions") # ``functions`` shape
+ if not isinstance(raw_tools, list):
+ return ()
+ return tuple(t for entry in raw_tools if (t := _tool_from_entry(entry)) is not None)
diff --git a/litellm/integrations/otel/model/semconv.py b/litellm/integrations/otel/model/semconv.py
new file mode 100644
index 00000000000..1c6c30eda0d
--- /dev/null
+++ b/litellm/integrations/otel/model/semconv.py
@@ -0,0 +1,201 @@
+"""
+Keys follow the OpenTelemetry GenAI semantic conventions (experimental). Anything
+without a semconv equivalent lives under the ``litellm.*`` vendor namespace.
+"""
+
+from enum import Enum
+from typing import Final
+
+
+class GenAIOperation(str, Enum):
+ """Values for ``gen_ai.operation.name``."""
+
+ CHAT = "chat"
+ TEXT_COMPLETION = "text_completion"
+ EMBEDDINGS = "embeddings"
+ GENERATE_CONTENT = "generate_content"
+ CREATE_AGENT = "create_agent" # reserved for future agent spans
+ INVOKE_AGENT = "invoke_agent" # reserved for future agent spans
+ EXECUTE_TOOL = "execute_tool" # reserved for future tool spans
+
+
+class GenAIProvider(str, Enum):
+ """Common values for the ``gen_ai.provider.name`` attribute."""
+
+ OPENAI = "openai"
+ ANTHROPIC = "anthropic"
+ AWS_BEDROCK = "aws.bedrock"
+ AZURE_AI_OPENAI = "azure.ai.openai"
+ AZURE_AI_INFERENCE = "azure.ai.inference"
+ GCP_GEMINI = "gcp.gemini"
+ GCP_VERTEX_AI = "gcp.vertex_ai"
+ COHERE = "cohere"
+ MISTRAL_AI = "mistral_ai"
+ DEEPSEEK = "deepseek"
+ GROQ = "groq"
+ PERPLEXITY = "perplexity"
+ X_AI = "x_ai"
+ IBM_WATSONX_AI = "ibm.watsonx.ai"
+
+
+class GenAI:
+ """Canonical OTel GenAI span-attribute keys."""
+
+ # request
+ OPERATION_NAME: Final = "gen_ai.operation.name"
+ PROVIDER_NAME: Final = "gen_ai.provider.name"
+ REQUEST_MODEL: Final = "gen_ai.request.model"
+ REQUEST_TEMPERATURE: Final = "gen_ai.request.temperature"
+ REQUEST_TOP_P: Final = "gen_ai.request.top_p"
+ REQUEST_TOP_K: Final = "gen_ai.request.top_k"
+ REQUEST_MAX_TOKENS: Final = "gen_ai.request.max_tokens"
+ REQUEST_FREQUENCY_PENALTY: Final = "gen_ai.request.frequency_penalty"
+ REQUEST_PRESENCE_PENALTY: Final = "gen_ai.request.presence_penalty"
+ REQUEST_STOP_SEQUENCES: Final = "gen_ai.request.stop_sequences"
+ REQUEST_SEED: Final = "gen_ai.request.seed"
+ REQUEST_CHOICE_COUNT: Final = "gen_ai.request.choice.count"
+ REQUEST_ENCODING_FORMATS: Final = "gen_ai.request.encoding_formats"
+ # response
+ RESPONSE_ID: Final = "gen_ai.response.id"
+ RESPONSE_MODEL: Final = "gen_ai.response.model"
+ RESPONSE_FINISH_REASONS: Final = "gen_ai.response.finish_reasons"
+ # usage
+ USAGE_INPUT_TOKENS: Final = "gen_ai.usage.input_tokens"
+ USAGE_OUTPUT_TOKENS: Final = "gen_ai.usage.output_tokens"
+ # content (opt-in, gated by capture mode)
+ INPUT_MESSAGES: Final = "gen_ai.input.messages"
+ OUTPUT_MESSAGES: Final = "gen_ai.output.messages"
+ SYSTEM_INSTRUCTIONS: Final = "gen_ai.system_instructions"
+ OUTPUT_TYPE: Final = "gen_ai.output.type"
+ CONVERSATION_ID: Final = "gen_ai.conversation.id"
+ # agent / tool (reserved)
+ AGENT_ID: Final = "gen_ai.agent.id"
+ AGENT_NAME: Final = "gen_ai.agent.name"
+ TOOL_NAME: Final = "gen_ai.tool.name"
+ TOOL_CALL_ID: Final = "gen_ai.tool.call.id"
+
+
+class Error:
+ TYPE: Final = "error.type"
+
+
+class Server:
+ ADDRESS: Final = "server.address"
+ PORT: Final = "server.port"
+
+
+class DB:
+ """Database / cache client-span keys (OTel ``db.*`` semconv).
+
+ Stamped on ``DB_CALL`` spans (redis / postgres), which are CLIENT spans for
+ outbound datastore calls — not on the INTERNAL ``SERVICE`` spans.
+ """
+
+ SYSTEM_NAME: Final = "db.system.name"
+ OPERATION_NAME: Final = "db.operation.name"
+
+
+class HTTP:
+ """HTTP server-span keys. Belong on the SERVER span only (never promoted)."""
+
+ REQUEST_METHOD: Final = "http.request.method"
+ ROUTE: Final = "http.route"
+ RESPONSE_STATUS_CODE: Final = "http.response.status_code"
+ URL_PATH: Final = "url.path"
+
+
+class LiteLLM:
+ """Vendor-extension keys (no semconv equivalent). Always ``litellm.*``."""
+
+ CALL_ID: Final = "litellm.call_id"
+ COST_PREFIX: Final = "litellm.cost."
+ METADATA_PREFIX: Final = "litellm.metadata."
+ TEAM_ID: Final = "litellm.team.id"
+ TEAM_ALIAS: Final = "litellm.team.alias"
+ # The team's free-form metadata dict, JSON-serialized into a single value.
+ TEAM_METADATA: Final = "litellm.team.metadata"
+ KEY_HASH: Final = "litellm.api_key.hash"
+ END_USER: Final = "litellm.end_user.id"
+ # The model string litellm actually sent to the provider (the deployment's
+ # ``litellm_params.model``), distinct from the user-facing ``gen_ai.request.model``.
+ PROVIDER_MODEL: Final = "litellm.provider.model"
+ REQUEST_STREAMING: Final = "litellm.request.streaming"
+ GUARDRAIL_NAME: Final = "litellm.guardrail.name"
+ GUARDRAIL_MODE: Final = "litellm.guardrail.mode"
+ GUARDRAIL_STATUS: Final = "litellm.guardrail.status"
+ GUARDRAIL_PROVIDER: Final = "litellm.guardrail.provider"
+ GUARDRAIL_ACTION: Final = "litellm.guardrail.action"
+ GUARDRAIL_RESPONSE: Final = "litellm.guardrail.response"
+ GUARDRAIL_VIOLATION_CATEGORIES: Final = "litellm.guardrail.violation_categories"
+ GUARDRAIL_CONFIDENCE_SCORE: Final = "litellm.guardrail.confidence_score"
+ GUARDRAIL_RISK_SCORE: Final = "litellm.guardrail.risk_score"
+ GUARDRAIL_MASKED_ENTITY_COUNT: Final = "litellm.guardrail.masked_entity_count"
+ GUARDRAIL_DURATION: Final = "litellm.guardrail.duration"
+ GUARDRAIL_ID: Final = "litellm.guardrail.id"
+ GUARDRAIL_POLICY_TEMPLATE: Final = "litellm.guardrail.policy_template"
+ GUARDRAIL_DETECTION_METHOD: Final = "litellm.guardrail.detection_method"
+ SERVICE_NAME: Final = "litellm.service.name"
+ SERVICE_CALL_TYPE: Final = "litellm.service.call_type"
+ PREPROCESSING_MS: Final = "litellm.preprocessing.duration_ms"
+
+
+class Metric:
+ """GenAI metric instrument names."""
+
+ TOKEN_USAGE: Final = "gen_ai.client.token.usage"
+ OPERATION_DURATION: Final = "gen_ai.client.operation.duration"
+
+
+# litellm ``custom_llm_provider`` -> ``gen_ai.provider.name`` value.
+_PROVIDER_BY_LITELLM: dict[str, GenAIProvider] = {
+ "openai": GenAIProvider.OPENAI,
+ "text-completion-openai": GenAIProvider.OPENAI,
+ "azure": GenAIProvider.AZURE_AI_OPENAI,
+ "azure_ai": GenAIProvider.AZURE_AI_INFERENCE,
+ "anthropic": GenAIProvider.ANTHROPIC,
+ "bedrock": GenAIProvider.AWS_BEDROCK,
+ "bedrock_converse": GenAIProvider.AWS_BEDROCK,
+ "vertex_ai": GenAIProvider.GCP_VERTEX_AI,
+ "vertex_ai_beta": GenAIProvider.GCP_VERTEX_AI,
+ "gemini": GenAIProvider.GCP_GEMINI,
+ "cohere": GenAIProvider.COHERE,
+ "cohere_chat": GenAIProvider.COHERE,
+ "mistral": GenAIProvider.MISTRAL_AI,
+ "deepseek": GenAIProvider.DEEPSEEK,
+ "groq": GenAIProvider.GROQ,
+ "perplexity": GenAIProvider.PERPLEXITY,
+ "xai": GenAIProvider.X_AI,
+ "watsonx": GenAIProvider.IBM_WATSONX_AI,
+}
+
+# litellm ``call_type`` -> ``gen_ai.operation.name``.
+_OPERATION_BY_CALL_TYPE: dict[str, GenAIOperation] = {
+ "completion": GenAIOperation.CHAT,
+ "acompletion": GenAIOperation.CHAT,
+ "completion_with_retries": GenAIOperation.CHAT,
+ "text_completion": GenAIOperation.TEXT_COMPLETION,
+ "atext_completion": GenAIOperation.TEXT_COMPLETION,
+ "embedding": GenAIOperation.EMBEDDINGS,
+ "aembedding": GenAIOperation.EMBEDDINGS,
+ "responses": GenAIOperation.CHAT,
+ "aresponses": GenAIOperation.CHAT,
+}
+
+
+def resolve_provider(custom_llm_provider: str | None) -> str:
+ """Map a litellm provider string to a ``gen_ai.provider.name`` value.
+
+ Unknown providers pass through verbatim — the convention explicitly allows
+ provider-specific values, so an unmapped name is still valid.
+ """
+ if not custom_llm_provider:
+ return ""
+ mapped = _PROVIDER_BY_LITELLM.get(custom_llm_provider.lower())
+ return mapped.value if mapped is not None else custom_llm_provider
+
+
+def resolve_operation(call_type: str | None) -> GenAIOperation:
+ """Map a litellm ``call_type`` to a ``gen_ai.operation.name`` value."""
+ if not call_type:
+ return GenAIOperation.CHAT
+ return _OPERATION_BY_CALL_TYPE.get(call_type.lower(), GenAIOperation.CHAT)
diff --git a/litellm/integrations/otel/model/spans.py b/litellm/integrations/otel/model/spans.py
new file mode 100644
index 00000000000..e4876f4ee58
--- /dev/null
+++ b/litellm/integrations/otel/model/spans.py
@@ -0,0 +1,203 @@
+"""
+This module declares every span the instrumentation can emit and the hierarchy.
+
+Span-name patterns live here as typed builder functions.
+
+Canonical hierarchy::
+
+ PROXY_REQUEST (SERVER, root) # owned by the FastAPI instrumentor
+ ├── SERVICE (INTERNAL) # auth phase span (live; see logger.phase_span)
+ │ └── DB_CALL (CLIENT) # its key/user/team lookups nest here
+ ├── GUARDRAIL (INTERNAL) # request-lifecycle hook, sibling of LLM_CALL
+ ├── LLM_CALL (CLIENT)
+ └── DB_CALL (CLIENT) # e.g. the spend-log write
+
+Guardrails parent to PROXY_REQUEST, not LLM_CALL: pre/during/post-call guardrail
+hooks are orchestrated by the request lifecycle (a pre-call guardrail runs
+before the LLM call even starts), so a guardrail is a sibling of the LLM call,
+not a child of it. The emitter parents every span to the ambient OTel context
+(the active server span), which matches this.
+
+Not every service call becomes a span — :func:`span_role_for_service` decides:
+
+- ``DB_CALL`` (CLIENT) — outbound datastores (redis, postgres,
+ ``batch_write_to_db``), carrying ``db.*`` semconv.
+- ``SERVICE`` (INTERNAL) — genuine internal work worth a span (background
+ budget/reset jobs, pod-lock manager).
+- ``None`` (metrics-only) — framework instrumentation that duplicates a gen-AI
+ span (``self`` = the ``track_llm_api_timing`` wrapper, ``router``,
+ ``proxy_pre_call``) or ``auth`` (which gets a live phase span instead). These
+ still feed Prometheus/Datadog; they just never enter the trace.
+
+``DB_CALL`` and ``SERVICE`` are built from the same ``ServiceSpanData``; only the
+role (hence span kind and attribute vocabulary) differs. A service call can fire
+outside any request (a background job), in which case it parents to no server
+span and starts its own root trace rather than being dropped.
+
+Management/admin endpoints are ordinary FastAPI routes — their SERVER spans are
+owned by the instrumentor too, so they don't appear as a role here.
+"""
+
+from dataclasses import dataclass
+from enum import Enum
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+ from litellm.integrations.otel.model.payloads import (
+ GuardrailSpanData,
+ LLMCallSpanData,
+ ProxyRequestSpanData,
+ ServiceSpanData,
+ )
+
+
+class SpanRole(str, Enum):
+ PROXY_REQUEST = "proxy_request"
+ LLM_CALL = "llm_call"
+ GUARDRAIL = "guardrail"
+ DB_CALL = "db_call"
+ SERVICE = "service"
+
+
+class LiteLLMSpanKind(str, Enum):
+ SERVER = "server"
+ CLIENT = "client"
+ INTERNAL = "internal"
+ PRODUCER = "producer"
+ CONSUMER = "consumer"
+
+
+@dataclass(frozen=True)
+class SpanSpec:
+ role: SpanRole
+ kind: LiteLLMSpanKind
+ parent: SpanRole | None
+
+
+SPAN_REGISTRY: dict[SpanRole, SpanSpec] = {
+ SpanRole.PROXY_REQUEST: SpanSpec(
+ SpanRole.PROXY_REQUEST, LiteLLMSpanKind.SERVER, parent=None
+ ),
+ SpanRole.LLM_CALL: SpanSpec(
+ SpanRole.LLM_CALL, LiteLLMSpanKind.CLIENT, parent=SpanRole.PROXY_REQUEST
+ ),
+ SpanRole.GUARDRAIL: SpanSpec(
+ SpanRole.GUARDRAIL, LiteLLMSpanKind.INTERNAL, parent=SpanRole.PROXY_REQUEST
+ ),
+ SpanRole.DB_CALL: SpanSpec(
+ SpanRole.DB_CALL, LiteLLMSpanKind.CLIENT, parent=SpanRole.PROXY_REQUEST
+ ),
+ SpanRole.SERVICE: SpanSpec(
+ SpanRole.SERVICE, LiteLLMSpanKind.INTERNAL, parent=SpanRole.PROXY_REQUEST
+ ),
+}
+
+
+# ``ServiceTypes`` value -> ``db.system.name``. These are outbound datastore
+# calls and become CLIENT ``DB_CALL`` spans; ``redis_``-prefixed names cover the
+# redis-backed spend queues. Any service not mapped here is litellm-internal work
+# and stays an INTERNAL ``SERVICE`` span. This table is the single source of
+# datastore knowledge — both the role classifier and the mapper read it.
+_DB_SYSTEM_BY_SERVICE: dict[str, str] = {
+ "redis": "redis",
+ "postgres": "postgresql",
+ "batch_write_to_db": "postgresql",
+}
+
+
+def db_system(service_name: str) -> str | None:
+ """The ``db.system.name`` for a datastore service, else ``None``.
+
+ ``None`` means the service is not an outbound datastore call. Redis-backed
+ spend queues (``redis_*``) map to ``redis``.
+ """
+ if service_name in _DB_SYSTEM_BY_SERVICE:
+ return _DB_SYSTEM_BY_SERVICE[service_name]
+ if service_name.startswith("redis_"):
+ return "redis"
+ return None
+
+
+# ``ServiceTypes`` values that are NOT emitted as spans — they are framework
+# instrumentation that either duplicates a gen-AI span or has a better home as a
+# Prometheus/Datadog metric. They still flow to those metric backends via their
+# own hooks; the v2 logger just does not put them in the trace:
+#
+# - ``self`` — ``track_llm_api_timing`` wraps the LLM call; the
+# ``chat {model}`` CLIENT span already represents it.
+# - ``router`` — wraps the whole request; duplicates the server span.
+# - ``proxy_pre_call`` — per-callback pre-call timing; a guardrail's real span
+# is ``execute_guardrail {name}``.
+# - ``auth`` — emitted instead as a live phase span (see
+# ``logger.phase_span``) so its DB lookups nest under it,
+# not as a flat post-hoc service span.
+_METRICS_ONLY_SERVICES: frozenset[str] = frozenset(
+ {"self", "router", "proxy_pre_call", "auth"}
+)
+
+
+def span_role_for_service(service_name: str) -> SpanRole | None:
+ """The span role for a service call, or ``None`` when it must not be a span.
+
+ ``DB_CALL`` for outbound datastores, ``SERVICE`` for genuine internal work
+ worth a span (background jobs), and ``None`` for framework instrumentation
+ that duplicates a gen-AI span or belongs in metrics only
+ (see ``_METRICS_ONLY_SERVICES``).
+ """
+ if service_name in _METRICS_ONLY_SERVICES:
+ return None
+ return SpanRole.DB_CALL if db_system(service_name) is not None else SpanRole.SERVICE
+
+
+# --- span name builders (the naming convention, per role) ------------------- #
+
+
+# The name the FastAPI instrumentor gives the root server span. V2 never creates
+# this span (the instrumentor owns it), but it anchors request-level spans to it
+# and tests assert against it by name, so the literal lives here with the rest of
+# the span vocabulary rather than being duplicated at each call site.
+LITELLM_PROXY_REQUEST_SPAN_NAME = "Received Proxy Server Request"
+
+
+def llm_call_span_name(data: "LLMCallSpanData") -> str:
+ """``"{operation} {model}"`` e.g. ``"chat gpt-4o"`` (GenAI semconv)."""
+ model = data.request_model or ""
+ return f"{data.operation.value} {model}".strip()
+
+
+def proxy_request_span_name(data: "ProxyRequestSpanData") -> str:
+ """``"{method} {route}"`` (HTTP semconv)."""
+ return f"{data.http_method} {data.route}".strip()
+
+
+def guardrail_span_name(data: "GuardrailSpanData") -> str:
+ return f"execute_guardrail {data.guardrail_name}".strip()
+
+
+def service_span_name(data: "ServiceSpanData") -> str:
+ """``"{service} {call_type}"`` e.g. ``"redis set"`` — service name alone when
+ no call type is known, so identically-named calls stay distinguishable."""
+ return f"{data.service_name} {data.call_type or ''}".strip()
+
+
+def root_roles() -> list[SpanRole]:
+ """Roles that start a new trace (no in-process parent)."""
+ return [role for role, spec in SPAN_REGISTRY.items() if spec.parent is None]
+
+
+def child_roles(parent: SpanRole) -> list[SpanRole]:
+ return [role for role, spec in SPAN_REGISTRY.items() if spec.parent == parent]
+
+
+def validate_registry(
+ registry: dict[SpanRole, SpanSpec] | None = None,
+) -> None:
+ reg = registry if registry is not None else SPAN_REGISTRY
+ for role, spec in reg.items():
+ if spec.role is not role:
+ raise ValueError(f"SPAN_REGISTRY[{role}] has mismatched role {spec.role}")
+ if spec.parent is not None and spec.parent not in reg:
+ raise ValueError(f"span role {role} declares unknown parent {spec.parent}")
+ missing = [role for role in SpanRole if role not in reg]
+ if missing:
+ raise ValueError(f"SPAN_REGISTRY is missing roles: {missing}")
diff --git a/litellm/integrations/otel/model/utils.py b/litellm/integrations/otel/model/utils.py
new file mode 100644
index 00000000000..f37afc97879
--- /dev/null
+++ b/litellm/integrations/otel/model/utils.py
@@ -0,0 +1,103 @@
+"""Shared, OpenTelemetry-free helpers for the otel integration.
+
+Generic value coercion (for reading heterogeneous logging-payload dicts), time
+conversion, and header parsing — pulled out of the individual modules so they
+live in one place. Deliberately free of any ``opentelemetry`` import so the
+OTel-free sources of truth (payloads, semconv, spans, config) can use it too.
+"""
+
+from datetime import datetime
+
+
+def as_str(value: object) -> str | None:
+ if value is None:
+ return None
+ if isinstance(value, str):
+ return value
+ return str(value)
+
+
+def as_int(value: object) -> int | None:
+ if isinstance(value, bool):
+ return int(value)
+ if isinstance(value, int):
+ return value
+ if isinstance(value, float):
+ return int(value)
+ if isinstance(value, str):
+ try:
+ return int(value)
+ except ValueError:
+ return None
+ return None
+
+
+def as_float(value: object) -> float | None:
+ if isinstance(value, bool):
+ return float(value)
+ if isinstance(value, (int, float)):
+ return float(value)
+ if isinstance(value, str):
+ try:
+ return float(value)
+ except ValueError:
+ return None
+ return None
+
+
+def as_bool(value: object) -> bool | None:
+ if value is None:
+ return None
+ if isinstance(value, bool):
+ return value
+ return bool(value)
+
+
+def as_str_tuple(value: object) -> tuple[str, ...] | None:
+ if value is None:
+ return None
+ if isinstance(value, str):
+ return (value,)
+ if isinstance(value, (list, tuple)):
+ return tuple(str(v) for v in value)
+ return None
+
+
+def to_ns(value: datetime | float | int | None) -> int | None:
+ """Coerce a datetime / epoch value to integer nanoseconds."""
+ if value is None:
+ return None
+ if isinstance(value, datetime):
+ return int(value.timestamp() * 1e9)
+ if isinstance(value, (int, float)) and not isinstance(value, bool):
+ return int(float(value) * 1e9)
+ return None
+
+
+def to_seconds(value: datetime | float | int | str | None) -> float | None:
+ """Coerce a datetime / epoch / formatted-string value to epoch seconds."""
+ if value is None:
+ return None
+ if isinstance(value, datetime):
+ return value.timestamp()
+ if isinstance(value, (int, float)) and not isinstance(value, bool):
+ return float(value)
+ if isinstance(value, str):
+ for fmt in ("%Y-%m-%d %H:%M:%S.%f", "%Y-%m-%d %H:%M:%S"):
+ try:
+ return datetime.strptime(value, fmt).timestamp()
+ except ValueError:
+ continue
+ return None
+
+
+def parse_headers(raw: str | None) -> dict[str, str]:
+ """Parse an OTLP ``"k=v,k=v"`` header string into a dict."""
+ headers: dict[str, str] = {}
+ if not raw:
+ return headers
+ for pair in raw.split(","):
+ if "=" in pair:
+ key, _, value = pair.partition("=")
+ headers[key.strip()] = value.strip()
+ return headers
diff --git a/litellm/integrations/otel/mount.py b/litellm/integrations/otel/mount.py
new file mode 100644
index 00000000000..ebcf4aa35af
--- /dev/null
+++ b/litellm/integrations/otel/mount.py
@@ -0,0 +1,130 @@
+"""FastAPI server-span instrumentation — the proxy mounts this at app creation.
+
+``opentelemetry-instrumentation-fastapi`` creates the SERVER span for each HTTP
+route and extracts inbound ``traceparent`` headers. This module owns the one call
+site that attaches it to the proxy app, plus the passthrough span-naming hook, so
+``proxy_server`` stays free of OTel details.
+
+The ``FastAPIInstrumentor`` import is kept lazy (inside :func:`instrument_fastapi_app`,
+after the gate check) so importing this module never requires the optional
+``opentelemetry-instrumentation-fastapi`` package and pulls in nothing OTel-related
+when the feature gate is off.
+"""
+
+import os
+from typing import Any
+
+from litellm._logging import verbose_logger
+from litellm.integrations.otel.model.config import is_otel_v2_enabled
+
+# Routes excluded from server-span tracing by default: high-frequency pollers and
+# static UI/docs assets, none of which are LLM traffic. Entries are substring-matched
+# against the request path (unanchored, so they survive a ``server_root_path`` prefix
+# and each entry also covers everything beneath it — e.g. ``/health`` covers
+# ``/health/readiness``). Operators override the whole set via the standard
+# ``OTEL_PYTHON_FASTAPI_EXCLUDED_URLS`` env var (set "" to trace everything).
+_DEFAULT_EXCLUDED_ROUTES = (
+ "/health", # load-balancer liveness/readiness polling
+ "/metrics", # Prometheus scrape (also drops the /model/metrics admin analytics)
+ "/litellm-asset-prefix", # hashed UI asset bundles
+ "/_next", # Next.js static JS/CSS chunks (root-level mount)
+ "/ui", # admin UI single-page app
+ "/swagger", # static Swagger UI assets
+ "/docs", # FastAPI Swagger docs page
+ "/redoc", # FastAPI ReDoc docs page
+ "/openapi.json", # OpenAPI schema
+ "favicon", # /favicon.ico + /get_favicon
+ "/.well-known", # UI config discovery
+)
+_DEFAULT_EXCLUDED_URLS = ",".join(_DEFAULT_EXCLUDED_ROUTES)
+
+# Passthrough routes are catch-alls (e.g. "/openai/{endpoint:path}"), so the
+# default OTel server-span name "{method} {route}" collapses every upstream
+# endpoint into "POST /openai/{endpoint:path}". The hook below renames those spans
+# to the real request path so each endpoint is distinguishable. Non-catch-all
+# routes keep their low-cardinality template name.
+PASSTHROUGH_PREFIXES = frozenset(
+ {
+ "openai",
+ "openai_passthrough",
+ "anthropic",
+ "azure",
+ "azure_ai",
+ "bedrock",
+ "cohere",
+ "cursor",
+ "gemini",
+ "mistral",
+ "vllm",
+ "vertex_ai",
+ "vertex-ai",
+ "assemblyai",
+ "eu.assemblyai",
+ "milvus",
+ }
+)
+
+
+def _passthrough_span_name_hook(span: Any, scope: dict) -> None:
+ """FastAPI ``server_request_hook``: give passthrough server spans a useful name.
+
+ The instrumentation matches the route at span creation, so both the span name
+ and ``http.route`` are set to the catch-all template (``/openai/{endpoint:path}``)
+ before this hook runs. Rewrite both to the real request path so each upstream
+ endpoint is distinguishable. (The ASGI ``http receive``/``http send`` sub-spans
+ can't be renamed from here — their name is captured at creation — so they are
+ dropped via ``exclude_spans`` at instrumentation time.)
+ """
+ try:
+ if span is None or not span.is_recording():
+ return
+ path = scope.get("path") or ""
+ method = scope.get("method") or ""
+ first_segment = path.lstrip("/").split("/", 1)[0]
+ if first_segment in PASSTHROUGH_PREFIXES:
+ span.update_name(f"{method} {path}".strip())
+ span.set_attribute("http.route", path)
+ except Exception:
+ pass
+
+
+def instrument_fastapi_app(app: Any) -> None:
+ """Attach OTel server-span instrumentation to the proxy FastAPI app.
+
+ Safe no-op when the V2 gate is off or ``opentelemetry-instrumentation-fastapi``
+ is unavailable. This MUST be called at app-creation time — once the lifespan
+ runs, the middleware stack is frozen and ``instrument_app`` raises "Cannot add
+ middleware after an application has started".
+
+ No ``TracerProvider`` is passed, so the instrumentation binds to the OTel global
+ ``ProxyTracerProvider``; the proxy publishes the real provider as the global
+ after config load (see ``proxy_startup_event``), and the proxy delegates to it.
+ That way server spans and gen-ai spans share one provider and the same trace.
+ """
+ try:
+ if not is_otel_v2_enabled():
+ return
+
+ # Lazy: only the V2-enabled path needs the optional
+ # ``opentelemetry-instrumentation-fastapi`` package, which is not part of the
+ # base ``litellm[proxy]`` install. Importing it at module top would make
+ # ``proxy_server``'s unconditional ``import`` of this module crash when the
+ # package is absent, even with the gate off.
+ from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
+
+ excluded_urls = (
+ os.environ.get("OTEL_PYTHON_FASTAPI_EXCLUDED_URLS")
+ if "OTEL_PYTHON_FASTAPI_EXCLUDED_URLS" in os.environ
+ else _DEFAULT_EXCLUDED_URLS
+ )
+ FastAPIInstrumentor.instrument_app(
+ app,
+ excluded_urls=excluded_urls,
+ server_request_hook=_passthrough_span_name_hook,
+ # Drop the ASGI "http receive"/"http send" lifecycle sub-spans: they
+ # are low-value noise and (for passthrough) carry the catch-all route
+ # template in their name, which can't be rewritten from a hook.
+ exclude_spans=["receive", "send"],
+ )
+ except Exception as e:
+ verbose_logger.debug("Skipping OTel V2 FastAPI instrumentation: %s", e)
diff --git a/litellm/integrations/otel/plumbing/__init__.py b/litellm/integrations/otel/plumbing/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/litellm/integrations/otel/plumbing/context.py b/litellm/integrations/otel/plumbing/context.py
new file mode 100644
index 00000000000..64790da814b
--- /dev/null
+++ b/litellm/integrations/otel/plumbing/context.py
@@ -0,0 +1,127 @@
+"""Trace-context + Baggage helpers."""
+
+from contextvars import ContextVar
+from typing import Mapping
+
+from opentelemetry import baggage
+from opentelemetry.context import Context, get_current
+from opentelemetry.trace import Span, get_current_span, set_span_in_context
+from opentelemetry.trace.propagation.tracecontext import (
+ TraceContextTextMapPropagator,
+)
+
+_PROPAGATOR = TraceContextTextMapPropagator()
+
+# The request's root span — the FastAPI-owned SERVER span — captured ONCE when the
+# proxy first resolves it, so request-level spans (the LLM call, guardrails) can
+# parent to it EXPLICITLY instead of to whatever span happens to be active at the
+# instant they are emitted. Ambient-only parenting (``get_current_span()``) is
+# wrong at two boundaries:
+# * inside the ``auth`` phase span the active span is the auth span, so an LLM /
+# guardrail span emitted there would nest under auth instead of being its
+# sibling; and
+# * in a detached success task (pass-through logs success from a fire-and-forget
+# ``asyncio.create_task``) the server span may not be active at all, orphaning
+# the span into a brand-new trace.
+# A ``ContextVar`` (not a request attribute) so it rides the request task's context
+# and is inherited by ``asyncio.create_task`` children — i.e. the async logging
+# callbacks that close the span. It is never reset: the contextvar dies with the
+# request task, so there is nothing to leak.
+_request_root_span: "ContextVar[Span | None]" = ContextVar(
+ "litellm_otel_request_root_span", default=None
+)
+
+
+def set_request_root_span(span: Span) -> None:
+ """Anchor the request's root (server) span for explicit child parenting.
+
+ No-ops for a non-recordable span so a bad capture can never replace a good one
+ with a phantom parent. Idempotent — the proxy captures the same server span at
+ more than one entry point.
+ """
+ if is_recordable_span(span):
+ _request_root_span.set(span)
+
+
+def request_root_span() -> "Span | None":
+ """The anchored request root span, or ``None`` outside a proxy request."""
+ span = _request_root_span.get()
+ return span if is_recordable_span(span) else None
+
+
+def set_request_baggage(
+ values: Mapping[str, str], context: Context | None = None
+) -> Context:
+ """Return a context with ``values`` written into Baggage."""
+ ctx = context
+ for key, value in values.items():
+ ctx = baggage.set_baggage(key, value, context=ctx)
+ return ctx if ctx is not None else (context or get_current())
+
+
+def get_baggage_attributes(context: Context | None = None) -> dict[str, str]:
+ """All Baggage entries on ``context`` as strings."""
+ return {key: str(value) for key, value in baggage.get_all(context).items()}
+
+
+def context_from_span(span: Span, context: Context | None = None) -> Context:
+ """A context with ``span`` as the active span (for explicit parenting)."""
+ return set_span_in_context(span, context=context)
+
+
+def resolve_parent_context(threaded: Span | None = None) -> Context:
+ """The context a child span should parent under.
+
+ Ambient-first: parent to the active OTel context (the server span, restored
+ by the logging worker or active in the request task), falling back to a span
+ passed explicitly (``threaded``) only when the ambient context has no
+ recordable span — e.g. a background service call with no request on the
+ stack. When neither is recordable the ambient context is returned unchanged,
+ so the span starts a new root trace.
+
+ Only service/DB spans pass ``threaded`` (the ``parent_otel_span`` handed to
+ the service hook). Request-level spans — the LLM call and guardrails — are
+ created where the server span is genuinely ambient, so they never need it.
+ """
+ ctx = get_current()
+ if is_recordable_span(threaded) and not is_recordable_span(get_current_span(ctx)):
+ ctx = context_from_span(threaded, context=ctx) # type: ignore[arg-type]
+ return ctx
+
+
+def resolve_request_span_context() -> Context:
+ """The parent context for a request-level span (the LLM call, a guardrail).
+
+ These are direct children of the request's root server span — siblings of the
+ ``auth`` phase span and of each other, never nested under whatever span is
+ momentarily active. So prefer the explicitly anchored root span; fall back to
+ ambient context only when there is no anchor (the SDK / no-proxy path), where
+ the span legitimately starts its own root trace.
+
+ Unlike :func:`resolve_parent_context` (used by DB/service spans, which DO want
+ to nest under the active phase span, e.g. an auth DB lookup under ``auth``),
+ this never returns the active span when an anchor exists.
+ """
+ root = request_root_span()
+ if root is not None:
+ return context_from_span(root)
+ return get_current()
+
+
+def is_recordable_span(obj: object) -> bool:
+ """True if ``obj`` is a live span with a valid context (safe to parent under)."""
+ if not isinstance(obj, Span):
+ return False
+ try:
+ ctx = obj.get_span_context()
+ except Exception:
+ return False
+ return ctx is not None and ctx.is_valid
+
+
+def extract_traceparent(headers: Mapping[str, str]) -> Context | None:
+ """Extract a remote parent context from incoming HTTP headers, if present."""
+ if not any(key.lower() == "traceparent" for key in headers):
+ return None
+ carrier = {str(key).lower(): value for key, value in headers.items()}
+ return _PROPAGATOR.extract(carrier)
diff --git a/litellm/integrations/otel/plumbing/metrics.py b/litellm/integrations/otel/plumbing/metrics.py
new file mode 100644
index 00000000000..edd120f91e6
--- /dev/null
+++ b/litellm/integrations/otel/plumbing/metrics.py
@@ -0,0 +1,28 @@
+"""GenAI client metrics (token usage + operation duration histograms)."""
+
+from dataclasses import dataclass
+
+from opentelemetry.metrics import Histogram, Meter
+
+from litellm.integrations.otel.model.semconv import Metric
+
+
+@dataclass(frozen=True)
+class GenAIMetrics:
+ token_usage: Histogram
+ operation_duration: Histogram
+
+
+def create_genai_metrics(meter: Meter) -> GenAIMetrics:
+ return GenAIMetrics(
+ token_usage=meter.create_histogram(
+ name=Metric.TOKEN_USAGE,
+ unit="{token}",
+ description="Number of tokens used per GenAI request.",
+ ),
+ operation_duration=meter.create_histogram(
+ name=Metric.OPERATION_DURATION,
+ unit="s",
+ description="GenAI operation duration.",
+ ),
+ )
diff --git a/litellm/integrations/otel/plumbing/providers.py b/litellm/integrations/otel/plumbing/providers.py
new file mode 100644
index 00000000000..40a0e41b905
--- /dev/null
+++ b/litellm/integrations/otel/plumbing/providers.py
@@ -0,0 +1,220 @@
+"""Provider / exporter factory + the Baggage span processor."""
+
+from typing import Callable, Iterable
+
+from opentelemetry import baggage
+from opentelemetry.context import Context
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace import ReadableSpan, SpanProcessor, TracerProvider
+from opentelemetry.sdk.trace.export import (
+ BatchSpanProcessor,
+ ConsoleSpanExporter,
+ SimpleSpanProcessor,
+ SpanExporter,
+)
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import (
+ InMemorySpanExporter,
+)
+from opentelemetry.trace import Span, SpanKind, Tracer
+
+from litellm.integrations.otel.model.config import ExporterSpec, OpenTelemetryV2Config
+from litellm.integrations.otel.model.semconv import LiteLLM
+from litellm.integrations.otel.model.spans import LiteLLMSpanKind
+
+# Re-exported so ``providers.parse_headers`` remains a stable entry point.
+from litellm.integrations.otel.model.utils import parse_headers as parse_headers
+
+_SPAN_KIND_BY_ROLE_KIND: dict[LiteLLMSpanKind, SpanKind] = {
+ LiteLLMSpanKind.SERVER: SpanKind.SERVER,
+ LiteLLMSpanKind.CLIENT: SpanKind.CLIENT,
+ LiteLLMSpanKind.INTERNAL: SpanKind.INTERNAL,
+ LiteLLMSpanKind.PRODUCER: SpanKind.PRODUCER,
+ LiteLLMSpanKind.CONSUMER: SpanKind.CONSUMER,
+}
+
+
+def to_otel_span_kind(kind: LiteLLMSpanKind) -> SpanKind:
+ return _SPAN_KIND_BY_ROLE_KIND[kind]
+
+
+# Custom exporter factories keyed by ``ExporterSpec.kind``. A preset registers
+# one here when its destination needs construction logic the built-in kinds
+# can't express — e.g. an exporter that fetches an auth token lazily on its
+# first export (off the event loop) instead of blocking at config-build time.
+# Keeping the registry here lets this module stay vendor-agnostic: the factory
+# lives with the integration that needs it.
+_EXPORTER_FACTORIES: dict[str, Callable[[ExporterSpec], SpanExporter]] = {}
+
+
+def register_exporter_factory(
+ kind: str, factory: Callable[[ExporterSpec], SpanExporter]
+) -> None:
+ """Register a custom exporter ``factory`` for the exporter ``kind``."""
+ _EXPORTER_FACTORIES[kind.lower()] = factory
+
+
+class LiteLLMBaggageSpanProcessor(SpanProcessor):
+ """Stamps an allowlisted set of Baggage entries onto every span at start."""
+
+ def __init__(
+ self,
+ allowed_keys: Iterable[str],
+ allowed_prefixes: tuple[str, ...] = (LiteLLM.METADATA_PREFIX,),
+ ) -> None:
+ self._allowed_keys = frozenset(allowed_keys)
+ self._allowed_prefixes = tuple(allowed_prefixes)
+
+ def _is_allowed(self, key: str) -> bool:
+ return key in self._allowed_keys or any(
+ key.startswith(prefix) for prefix in self._allowed_prefixes
+ )
+
+ def on_start(self, span: Span, parent_context: Context | None = None) -> None:
+ for key, value in baggage.get_all(parent_context).items():
+ if self._is_allowed(key) and isinstance(value, (str, bool, int, float)):
+ span.set_attribute(key, value)
+
+ def on_end(self, span: ReadableSpan) -> None: # noqa: D401 - no-op
+ return None
+
+ def shutdown(self) -> None:
+ return None
+
+ def force_flush(self, timeout_millis: int = 30000) -> bool:
+ return True
+
+
+def _otlp_traces_endpoint(endpoint: str | None) -> str | None:
+ """Point an OTLP/HTTP base endpoint at the ``/v1/traces`` signal path.
+
+ ``OTEL_EXPORTER_OTLP_ENDPOINT`` is a base URL (e.g. ``http://host:4318``).
+ The OTLP/HTTP exporter only appends the ``/v1/traces`` path when it reads
+ that env var itself; when an endpoint is passed explicitly it is used
+ verbatim, so a base URL would POST to the root and the collector returns
+ 404. Append the signal path here (leaving an already-correct path intact).
+ """
+ if not endpoint:
+ return endpoint
+ endpoint = endpoint.rstrip("/")
+ # Splunk Observability uses ``/v2/trace/otlp``; never rewrite it.
+ if endpoint.endswith("/v1/traces") or "/v2/trace/otlp" in endpoint:
+ return endpoint
+ for other_signal in ("/v1/logs", "/v1/metrics"):
+ if endpoint.endswith(other_signal):
+ return endpoint[: -len(other_signal)] + "/v1/traces"
+ return endpoint + "/v1/traces"
+
+
+def _exporter_from_spec(spec: ExporterSpec) -> SpanExporter:
+ kind = (spec.kind or "console").lower()
+ factory = _EXPORTER_FACTORIES.get(kind)
+ if factory is not None:
+ return factory(spec)
+ if kind in ("in_memory", "inmemory", "memory"):
+ return InMemorySpanExporter()
+ if kind in ("otlp_http", "http", "http/protobuf", "http/json"):
+ from opentelemetry.exporter.otlp.proto.http.trace_exporter import (
+ OTLPSpanExporter as HTTPExporter,
+ )
+
+ return HTTPExporter(
+ endpoint=_otlp_traces_endpoint(spec.endpoint),
+ headers=parse_headers(spec.headers),
+ )
+ if kind in ("otlp_grpc", "grpc"):
+ from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import (
+ OTLPSpanExporter as GRPCExporter,
+ )
+
+ return GRPCExporter(endpoint=spec.endpoint, headers=parse_headers(spec.headers))
+ return ConsoleSpanExporter()
+
+
+def _processor_for(exporter: SpanExporter, use_simple: bool | None) -> SpanProcessor:
+ """Pick a Simple or Batch span processor for ``exporter``.
+
+ When ``use_simple`` is unset, default to Simple for console and in-memory
+ exporters (spans export synchronously, which tests rely on) and Batch for
+ everything else (the right export semantics for production).
+ """
+ if use_simple is None:
+ use_simple = isinstance(exporter, (ConsoleSpanExporter, InMemorySpanExporter))
+ return SimpleSpanProcessor(exporter) if use_simple else BatchSpanProcessor(exporter)
+
+
+def build_span_exporter(config: OpenTelemetryV2Config) -> SpanExporter:
+ """Build a single exporter from the top-level config fields.
+
+ Convenience for the common single-exporter case (and for tests): reads the
+ ``exporter`` / ``endpoint`` / ``headers`` fields. To configure multiple
+ exporters, populate ``config.exporters`` directly.
+ """
+ return _exporter_from_spec(
+ ExporterSpec(
+ kind=config.exporter, endpoint=config.endpoint, headers=config.headers
+ )
+ )
+
+
+def build_resource(config: OpenTelemetryV2Config) -> Resource:
+ attributes: dict[str, str] = {"service.name": config.service_name}
+ if config.deployment_environment:
+ attributes["deployment.environment"] = config.deployment_environment
+ attributes.update(config.resource_attributes)
+ return Resource.create(attributes)
+
+
+def build_tracer_provider(
+ config: OpenTelemetryV2Config,
+ exporter: SpanExporter | None = None,
+ baggage_processor: SpanProcessor | None = None,
+ use_simple_processor: bool | None = None,
+) -> TracerProvider:
+ """Build the shared :class:`TracerProvider`.
+
+ Attach the Baggage processor first (so identity attributes land on each
+ span before any export decision), then add one ``SpanProcessor`` per
+ ``config.exporters`` entry — this is what fans spans out to multiple
+ backends. ``exporter`` and ``use_simple_processor`` are explicit overrides:
+ pass a single exporter to attach exactly that one (used by tests).
+ """
+ provider = TracerProvider(resource=build_resource(config))
+ if baggage_processor is None:
+ baggage_processor = LiteLLMBaggageSpanProcessor(
+ allowed_keys=config.baggage_promoted_keys
+ )
+ provider.add_span_processor(baggage_processor)
+
+ if exporter is not None:
+ provider.add_span_processor(_processor_for(exporter, use_simple_processor))
+ return provider
+
+ # ``config._normalize`` guarantees at least one spec (it folds the top-level
+ # ``exporter``/``endpoint``/``headers`` fields in when ``exporters`` is empty).
+ for spec in config.exporters:
+ exp = _exporter_from_spec(spec)
+ provider.add_span_processor(
+ _processor_for(
+ exp,
+ (
+ spec.use_simple_processor
+ if spec.use_simple_processor is not None
+ else use_simple_processor
+ ),
+ )
+ )
+ return provider
+
+
+def get_tracer(provider: TracerProvider, name: str = "litellm") -> Tracer:
+ return provider.get_tracer(name)
+
+
+def in_memory_provider(
+ config: OpenTelemetryV2Config | None = None,
+) -> tuple[TracerProvider, InMemorySpanExporter]:
+ """Convenience for tests: a provider exporting to an in-memory buffer."""
+ cfg = config or OpenTelemetryV2Config(exporter="in_memory")
+ exporter = InMemorySpanExporter()
+ provider = build_tracer_provider(cfg, exporter=exporter)
+ return provider, exporter
diff --git a/litellm/integrations/otel/plumbing/routing.py b/litellm/integrations/otel/plumbing/routing.py
new file mode 100644
index 00000000000..4d0943a263a
--- /dev/null
+++ b/litellm/integrations/otel/plumbing/routing.py
@@ -0,0 +1,101 @@
+"""Per-request multi-tenant tracer routing.
+
+When a request carries team/key vendor credentials in
+``standard_callback_dynamic_params``, its spans must export through a
+``TracerProvider`` whose OTLP headers carry those credentials.
+``TenantTracerCache`` builds and caches one provider per distinct credential
+set, and otherwise hands back the logger's default tracer. This lets a single
+logger fan requests out to many tenants without needing a logger per tenant.
+"""
+
+from collections import OrderedDict
+from typing import Any, Mapping
+
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.trace import Tracer
+
+from litellm._logging import verbose_logger
+from litellm.integrations.otel.model.config import OpenTelemetryV2Config
+from litellm.integrations.otel.presets import dynamic_otlp_headers
+from litellm.integrations.otel.plumbing.providers import (
+ build_tracer_provider,
+ get_tracer,
+)
+
+# Exporter kinds that ignore headers — never rewritten with dynamic credentials.
+_NON_OTLP_KINDS = ("console", "in_memory", "inmemory", "memory")
+
+# Cap on distinct credential-scoped providers held at once. ``dynamic_params``
+# can be populated from request metadata, so an unbounded cache lets a caller
+# spawn one ``TracerProvider`` (plus its ``BatchSpanProcessor`` background
+# thread) per unique credential set and exhaust the proxy. The LRU bound keeps
+# the working set of active tenants resident while flushing and shutting down
+# evicted providers so their threads are reclaimed.
+_MAX_CACHED_PROVIDERS = 256
+
+
+def _shutdown_provider(provider: TracerProvider) -> None:
+ """Flush + stop an evicted provider's processors (reclaims their threads).
+
+ ``TracerProvider.shutdown`` force-flushes each ``SpanProcessor`` before
+ stopping it, so any spans already handed to a ``BatchSpanProcessor`` are
+ exported rather than dropped. Best-effort: a shutdown failure must not break
+ the request that triggered the eviction.
+ """
+ try:
+ provider.shutdown()
+ except Exception as e: # pragma: no cover - defensive
+ verbose_logger.debug("OTel V2: error shutting down evicted provider: %s", e)
+
+
+class TenantTracerCache:
+ """Credential-scoped ``TracerProvider`` cache keyed by the dynamic headers."""
+
+ def __init__(
+ self,
+ config: OpenTelemetryV2Config,
+ callback_name: str | None,
+ tracer_name: str,
+ ) -> None:
+ self._config = config
+ self._callback_name = callback_name
+ self._tracer_name = tracer_name
+ self._providers: "OrderedDict[tuple[tuple[str, str], ...], TracerProvider]" = (
+ OrderedDict()
+ )
+
+ def tracer_for(self, default: Tracer, dynamic_params: Any) -> Tracer:
+ """Return the tracer for this request.
+
+ Use ``default`` unless the request's dynamic credentials require a
+ credential-scoped tracer, in which case build (or reuse) one. The cache
+ is a bounded LRU: the least-recently-used provider is flushed and shut
+ down on overflow so its exporter threads don't accumulate.
+ """
+ headers = dynamic_otlp_headers(self._callback_name, dynamic_params)
+ if not headers:
+ return default
+ cache_key = tuple(sorted(headers.items()))
+ provider = self._providers.get(cache_key)
+ if provider is not None:
+ self._providers.move_to_end(cache_key)
+ else:
+ provider = build_tracer_provider(self._config_with_headers(headers))
+ self._providers[cache_key] = provider
+ if len(self._providers) > _MAX_CACHED_PROVIDERS:
+ _, evicted = self._providers.popitem(last=False)
+ _shutdown_provider(evicted)
+ return get_tracer(provider, self._tracer_name)
+
+ def _config_with_headers(self, headers: Mapping[str, str]) -> OpenTelemetryV2Config:
+ """Clone the config, replacing OTLP exporter headers with ``headers``."""
+ header_str = ",".join(f"{key}={value}" for key, value in headers.items())
+ exporters = [
+ (
+ spec
+ if spec.kind.lower() in _NON_OTLP_KINDS
+ else spec.model_copy(update={"headers": header_str})
+ )
+ for spec in self._config.exporters
+ ]
+ return self._config.model_copy(update={"exporters": exporters})
diff --git a/litellm/integrations/otel/presets/__init__.py b/litellm/integrations/otel/presets/__init__.py
new file mode 100644
index 00000000000..c69d257ab52
--- /dev/null
+++ b/litellm/integrations/otel/presets/__init__.py
@@ -0,0 +1,78 @@
+"""Integration presets — each one returns an :class:`OpenTelemetryV2Config`.
+
+A preset is a callable that reads an integration's env vars and returns an
+``OpenTelemetryV2Config`` describing the exporter destination, the mapper
+vocabularies to apply, and any resource attributes. ``PRESET_BY_CALLBACK``
+maps a callback name (``"arize"``, ``"langfuse_otel"``, ...) to its preset so
+the factory in ``litellm_logging`` can resolve a name and build a single
+``OpenTelemetryV2`` instance from the result.
+"""
+
+from typing import Callable
+
+from litellm.integrations.otel.presets.agentops import agentops_preset
+from litellm.integrations.otel.presets.arize import arize_dynamic_headers, arize_preset
+from litellm.integrations.otel.presets.base import Preset
+from litellm.integrations.otel.presets.langfuse import (
+ langfuse_dynamic_headers,
+ langfuse_preset,
+)
+from litellm.integrations.otel.presets.langtrace import langtrace_preset
+from litellm.integrations.otel.presets.levo import levo_preset
+from litellm.integrations.otel.presets.phoenix import phoenix_preset
+from litellm.integrations.otel.presets.weave import weave_dynamic_headers, weave_preset
+from litellm.types.utils import StandardCallbackDynamicParams
+
+#: Callback name → preset. The ``Preset`` annotation makes mypy verify every
+#: registered value matches the preset interface.
+PRESET_BY_CALLBACK: dict[str, Preset] = {
+ "agentops": agentops_preset,
+ "arize": arize_preset,
+ "arize_phoenix": phoenix_preset,
+ "langfuse_otel": langfuse_preset,
+ "langtrace": langtrace_preset,
+ "levo": levo_preset,
+ "weave_otel": weave_preset,
+}
+
+#: Callback name → per-request OTLP header builder (team/key multi-tenant
+#: routing). Only integrations that support dynamic credentials appear here —
+#: Arize-Phoenix/Langtrace/Levo/AgentOps don't, so they use the logger's
+#: default tracer.
+DYNAMIC_HEADERS_BY_CALLBACK: dict[
+ str, Callable[[StandardCallbackDynamicParams], dict[str, str]]
+] = {
+ "arize": arize_dynamic_headers,
+ "langfuse_otel": langfuse_dynamic_headers,
+ "weave_otel": weave_dynamic_headers,
+}
+
+
+def dynamic_otlp_headers(
+ callback_name: str | None,
+ dynamic_params: StandardCallbackDynamicParams | None,
+) -> dict[str, str] | None:
+ """Per-request OTLP headers for ``callback_name``, or ``None`` if N/A.
+
+ ``None`` means "no per-request routing" — the caller uses its default tracer.
+ """
+ builder = DYNAMIC_HEADERS_BY_CALLBACK.get(callback_name or "")
+ if builder is None or not dynamic_params:
+ return None
+ headers = builder(dynamic_params)
+ return headers or None
+
+
+__all__ = [
+ "PRESET_BY_CALLBACK",
+ "DYNAMIC_HEADERS_BY_CALLBACK",
+ "Preset",
+ "dynamic_otlp_headers",
+ "agentops_preset",
+ "arize_preset",
+ "langfuse_preset",
+ "langtrace_preset",
+ "levo_preset",
+ "phoenix_preset",
+ "weave_preset",
+]
diff --git a/litellm/integrations/otel/presets/agentops.py b/litellm/integrations/otel/presets/agentops.py
new file mode 100644
index 00000000000..5a12818fd99
--- /dev/null
+++ b/litellm/integrations/otel/presets/agentops.py
@@ -0,0 +1,139 @@
+"""AgentOps preset — OTLP/HTTP to AgentOps' endpoint with a lazily-fetched JWT.
+
+AgentOps authenticates with a short-lived JWT minted from the API key. Fetching
+it is blocking network I/O, so it must never run on the event loop: callback
+construction (where presets are built) can run inside the proxy's async startup
+or, in the SDK, on the first request. Instead of fetching at config-build time,
+this preset registers a custom exporter (``kind="agentops"``) that mints the JWT
+**on its first export** — which the ``BatchSpanProcessor`` runs in its own
+worker thread, off any event loop — and caches it for the process lifetime.
+"""
+
+from typing import Any
+
+import httpx
+from pydantic import Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+from litellm._logging import verbose_logger
+from litellm.integrations.otel.model.config import ExporterSpec, OpenTelemetryV2Config
+from litellm.integrations.otel.plumbing.providers import register_exporter_factory
+
+_AGENTOPS_ENDPOINT = "https://otlp.agentops.cloud/v1/traces"
+_AGENTOPS_AUTH_ENDPOINT = "https://api.agentops.ai/v3/auth/token"
+_AGENTOPS_EXPORTER_KIND = "agentops"
+
+
+class _AgentOpsSettings(BaseSettings):
+ model_config = SettingsConfigDict(case_sensitive=False, extra="ignore")
+
+ api_key: str | None = Field(default=None, validation_alias="AGENTOPS_API_KEY")
+ service_name: str = Field(
+ default="agentops", validation_alias="AGENTOPS_SERVICE_NAME"
+ )
+ environment: str | None = Field(
+ default=None, validation_alias="AGENTOPS_ENVIRONMENT"
+ )
+
+
+def agentops_preset(
+ *,
+ config_overrides: OpenTelemetryV2Config | None = None,
+) -> OpenTelemetryV2Config:
+ """Build the AgentOps config without any network I/O.
+
+ The ``agentops`` exporter mints (and caches) the JWT lazily on its first
+ export, so this stays non-blocking. ``project.id`` is therefore not a
+ resource attribute — it is encoded in the JWT, which AgentOps uses to route
+ the trace to the right project.
+ """
+ settings = _AgentOpsSettings()
+ base = config_overrides or OpenTelemetryV2Config()
+ return base.model_copy(
+ update={
+ "exporters": [
+ *base.exporters,
+ ExporterSpec(
+ kind=_AGENTOPS_EXPORTER_KIND,
+ endpoint=_AGENTOPS_ENDPOINT,
+ options=(
+ {"api_key": settings.api_key} if settings.api_key else None
+ ),
+ ),
+ ],
+ "resource_attributes": {
+ **base.resource_attributes,
+ "service.name": settings.service_name,
+ "telemetry.sdk.name": "agentops",
+ **(
+ {"deployment.environment": settings.environment}
+ if settings.environment
+ else {}
+ ),
+ },
+ }
+ )
+
+
+def _build_agentops_exporter(spec: ExporterSpec) -> Any:
+ """Factory for the ``agentops`` exporter kind: a lazy-auth OTLP/HTTP exporter."""
+ from opentelemetry.exporter.otlp.proto.http.trace_exporter import (
+ OTLPSpanExporter,
+ )
+
+ class _LazyAuthAgentOpsExporter(OTLPSpanExporter):
+ """OTLP/HTTP exporter that mints the AgentOps JWT on its first export.
+
+ ``export`` runs in the ``BatchSpanProcessor`` worker thread, so the
+ blocking token fetch never touches an event loop. The result is cached
+ after the first attempt (success or failure) so it runs at most once.
+ """
+
+ def __init__(self, *, endpoint: str | None, api_key: str | None) -> None:
+ super().__init__(endpoint=endpoint)
+ self._agentops_api_key = api_key
+ self._auth_resolved = False
+
+ def _ensure_authenticated(self) -> None:
+ if self._auth_resolved:
+ return
+ self._auth_resolved = True
+ if not self._agentops_api_key:
+ return
+ try:
+ token = _fetch_agentops_jwt(self._agentops_api_key).get("token")
+ if token:
+ # ``_session`` is the requests.Session the base exporter
+ # POSTs through; updating its Authorization header is how the
+ # minted JWT reaches every subsequent export.
+ self._session.headers["Authorization"] = f"Bearer {token}"
+ except Exception as e:
+ verbose_logger.debug("AgentOps JWT fetch failed: %s", e)
+
+ def export(self, spans: Any) -> Any:
+ self._ensure_authenticated()
+ return super().export(spans)
+
+ options = spec.options or {}
+ return _LazyAuthAgentOpsExporter(
+ endpoint=spec.endpoint, api_key=options.get("api_key")
+ )
+
+
+def _fetch_agentops_jwt(api_key: str) -> dict[str, Any]:
+ # Own a short-lived client rather than ``_get_httpx_client()``: that returns
+ # a process-wide cached ``HTTPHandler`` whose connection pool is shared by
+ # every caller, so closing it here would break concurrent/subsequent
+ # requests. This one-shot auth call gets its own client to close.
+ with httpx.Client(timeout=10) as client:
+ response = client.post(
+ url=_AGENTOPS_AUTH_ENDPOINT,
+ headers={"Content-Type": "application/json", "Connection": "keep-alive"},
+ json={"api_key": api_key},
+ )
+ if response.status_code != 200:
+ raise RuntimeError(f"Failed to fetch AgentOps token: {response.text}")
+ return response.json()
+
+
+register_exporter_factory(_AGENTOPS_EXPORTER_KIND, _build_agentops_exporter)
diff --git a/litellm/integrations/otel/presets/arize.py b/litellm/integrations/otel/presets/arize.py
new file mode 100644
index 00000000000..4df15125f5a
--- /dev/null
+++ b/litellm/integrations/otel/presets/arize.py
@@ -0,0 +1,75 @@
+"""Arize preset — OTLP exporter to Arize + OpenInference vocabulary."""
+
+from pydantic import Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+from litellm.integrations.arize.arize import ArizeLogger as _V1ArizeLogger
+from litellm.integrations.otel.model.config import ExporterSpec, OpenTelemetryV2Config
+from litellm.integrations.otel.presets.utils import ensure_mappers
+from litellm.types.utils import StandardCallbackDynamicParams
+
+
+class _ArizeSettings(BaseSettings):
+ model_config = SettingsConfigDict(case_sensitive=False, extra="ignore")
+
+ # Standard OTLP headers env var, used as the fallback when no Arize
+ # credentials are configured.
+ otlp_traces_headers: str | None = Field(
+ default=None, validation_alias="OTEL_EXPORTER_OTLP_TRACES_HEADERS"
+ )
+
+
+def arize_preset(
+ *,
+ config_overrides: OpenTelemetryV2Config | None = None,
+) -> OpenTelemetryV2Config:
+ arize_cfg = _V1ArizeLogger.get_arize_config()
+ headers = _arize_headers(arize_cfg)
+ base = config_overrides or OpenTelemetryV2Config()
+ return base.model_copy(
+ update={
+ "exporters": [
+ *base.exporters,
+ ExporterSpec(
+ kind=arize_cfg.protocol or "otlp_grpc",
+ endpoint=arize_cfg.endpoint or "https://otlp.arize.com/v1",
+ headers=headers,
+ ),
+ ],
+ "mapper_names": ensure_mappers(base.mapper_names, "openinference"),
+ "resource_attributes": {
+ **base.resource_attributes,
+ **(
+ {"model_id": arize_cfg.project_name}
+ if arize_cfg.project_name
+ else {}
+ ),
+ },
+ }
+ )
+
+
+def _arize_headers(arize_cfg) -> str | None:
+ pieces = []
+ if arize_cfg.space_id or arize_cfg.space_key:
+ pieces.append(f"space_id={arize_cfg.space_id or arize_cfg.space_key}")
+ if arize_cfg.api_key:
+ pieces.append(f"api_key={arize_cfg.api_key}")
+ if not pieces:
+ # Fall back to the standard OTLP headers env var when no Arize
+ # credentials are configured.
+ return _ArizeSettings().otlp_traces_headers
+ return ",".join(pieces)
+
+
+def arize_dynamic_headers(params: StandardCallbackDynamicParams) -> dict[str, str]:
+ """Per-request Arize OTLP headers from team/key dynamic params."""
+ headers: dict[str, str] = {}
+ # ``arize_space_key`` is the suggested param and wins over ``arize_space_id``.
+ space = params.get("arize_space_key") or params.get("arize_space_id")
+ if space:
+ headers["arize-space-id"] = space
+ api_key = params.get("arize_api_key")
+ if api_key:
+ headers["api_key"] = api_key
+ return headers
diff --git a/litellm/integrations/otel/presets/base.py b/litellm/integrations/otel/presets/base.py
new file mode 100644
index 00000000000..b50908e7652
--- /dev/null
+++ b/litellm/integrations/otel/presets/base.py
@@ -0,0 +1,25 @@
+"""Preset interface.
+
+A preset is a callable that reads its integration's env vars and produces an
+:class:`OpenTelemetryV2Config` (exporter list + mapper-name list + resource
+attributes). This ``Protocol`` pins that contract so ``PRESET_BY_CALLBACK`` and
+the factory in ``litellm_logging`` are type-checked structurally against it,
+matching the ``AttributeMapper`` protocol the mappers use.
+"""
+
+from typing import Protocol, runtime_checkable
+
+from litellm.integrations.otel.model.config import OpenTelemetryV2Config
+
+
+@runtime_checkable
+class Preset(Protocol):
+ """Reads an integration's env config and returns an ``OpenTelemetryV2Config``.
+
+ ``config_overrides`` lets one preset layer onto another's config (or onto
+ test-supplied defaults); the factory calls presets with no arguments.
+ """
+
+ def __call__(
+ self, *, config_overrides: OpenTelemetryV2Config | None = None
+ ) -> OpenTelemetryV2Config: ...
diff --git a/litellm/integrations/otel/presets/langfuse.py b/litellm/integrations/otel/presets/langfuse.py
new file mode 100644
index 00000000000..011545384b9
--- /dev/null
+++ b/litellm/integrations/otel/presets/langfuse.py
@@ -0,0 +1,43 @@
+"""Langfuse-OTEL preset."""
+
+from litellm.integrations.langfuse.langfuse_otel import (
+ LangfuseOtelLogger as _V1Langfuse,
+)
+from litellm.integrations.otel.model.config import ExporterSpec, OpenTelemetryV2Config
+from litellm.integrations.otel.presets.utils import ensure_mappers
+from litellm.types.utils import StandardCallbackDynamicParams
+
+
+def langfuse_preset(
+ *,
+ config_overrides: OpenTelemetryV2Config | None = None,
+) -> OpenTelemetryV2Config:
+ cfg = _V1Langfuse.get_langfuse_otel_config()
+ kind = cfg.exporter if isinstance(cfg.exporter, str) else "otlp_http"
+ base = config_overrides or OpenTelemetryV2Config()
+ return base.model_copy(
+ update={
+ "exporters": [
+ *base.exporters,
+ ExporterSpec(
+ kind=kind,
+ endpoint=cfg.endpoint,
+ headers=cfg.headers,
+ ),
+ ],
+ "mapper_names": ensure_mappers(base.mapper_names, "langfuse"),
+ }
+ )
+
+
+def langfuse_dynamic_headers(params: StandardCallbackDynamicParams) -> dict[str, str]:
+ """Per-request Langfuse OTLP headers from team/key dynamic params."""
+ public_key = params.get("langfuse_public_key")
+ secret_key = params.get("langfuse_secret_key")
+ if public_key and secret_key:
+ return {
+ "Authorization": _V1Langfuse._get_langfuse_authorization_header(
+ public_key=public_key, secret_key=secret_key
+ )
+ }
+ return {}
diff --git a/litellm/integrations/otel/presets/langtrace.py b/litellm/integrations/otel/presets/langtrace.py
new file mode 100644
index 00000000000..acdbaf870d3
--- /dev/null
+++ b/litellm/integrations/otel/presets/langtrace.py
@@ -0,0 +1,22 @@
+"""Langtrace preset — Langtrace consumes generic OTLP + a vendor mapper."""
+
+from litellm.integrations.otel.model.config import OpenTelemetryV2Config
+from litellm.integrations.otel.presets.utils import ensure_mappers
+
+
+def langtrace_preset(
+ *,
+ config_overrides: OpenTelemetryV2Config | None = None,
+) -> OpenTelemetryV2Config:
+ """Compose the Langtrace mapper on top of the customer's OTLP destination.
+
+ Unlike Arize / Phoenix / Langfuse, Langtrace doesn't ship its own endpoint
+ — users point their existing OTLP collector at Langtrace and just
+ need the vendor attribute schema applied to outgoing spans.
+ """
+ base = config_overrides or OpenTelemetryV2Config()
+ return base.model_copy(
+ update={
+ "mapper_names": ensure_mappers(base.mapper_names, "langtrace"),
+ }
+ )
diff --git a/litellm/integrations/otel/presets/levo.py b/litellm/integrations/otel/presets/levo.py
new file mode 100644
index 00000000000..4c4cba982a4
--- /dev/null
+++ b/litellm/integrations/otel/presets/levo.py
@@ -0,0 +1,24 @@
+"""Levo preset — OTLP/HTTP to a Levo collector with org+workspace headers."""
+
+from litellm.integrations.levo.levo import LevoLogger as _V1Levo
+from litellm.integrations.otel.model.config import ExporterSpec, OpenTelemetryV2Config
+
+
+def levo_preset(
+ *,
+ config_overrides: OpenTelemetryV2Config | None = None,
+) -> OpenTelemetryV2Config:
+ cfg = _V1Levo.get_levo_config()
+ base = config_overrides or OpenTelemetryV2Config()
+ return base.model_copy(
+ update={
+ "exporters": [
+ *base.exporters,
+ ExporterSpec(
+ kind="otlp_http",
+ endpoint=cfg.endpoint,
+ headers=cfg.otlp_auth_headers,
+ ),
+ ],
+ }
+ )
diff --git a/litellm/integrations/otel/presets/phoenix.py b/litellm/integrations/otel/presets/phoenix.py
new file mode 100644
index 00000000000..4c2b165ffca
--- /dev/null
+++ b/litellm/integrations/otel/presets/phoenix.py
@@ -0,0 +1,48 @@
+"""Arize-Phoenix preset."""
+
+from pydantic import AliasChoices, Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+from litellm.integrations.arize.arize_phoenix import (
+ ArizePhoenixLogger as _V1Phoenix,
+)
+from litellm.integrations.otel.model.config import ExporterSpec, OpenTelemetryV2Config
+from litellm.integrations.otel.presets.utils import ensure_mappers
+
+
+class _PhoenixSettings(BaseSettings):
+ model_config = SettingsConfigDict(case_sensitive=False, extra="ignore")
+
+ project_name: str = Field(
+ default="default",
+ validation_alias=AliasChoices(
+ "PHOENIX_PROJECT_NAME", "PHOENIX_COLLECTOR_PROJECT_NAME"
+ ),
+ )
+
+
+def phoenix_preset(
+ *,
+ config_overrides: OpenTelemetryV2Config | None = None,
+) -> OpenTelemetryV2Config:
+ cfg = _V1Phoenix.get_arize_phoenix_config()
+ headers = cfg.otlp_auth_headers if hasattr(cfg, "otlp_auth_headers") else None
+ project_name = _PhoenixSettings().project_name
+ base = config_overrides or OpenTelemetryV2Config()
+ return base.model_copy(
+ update={
+ "exporters": [
+ *base.exporters,
+ ExporterSpec(
+ kind=cfg.protocol if hasattr(cfg, "protocol") else "otlp_http",
+ endpoint=cfg.endpoint,
+ headers=headers,
+ ),
+ ],
+ "mapper_names": ensure_mappers(base.mapper_names, "openinference"),
+ "resource_attributes": {
+ **base.resource_attributes,
+ "openinference.project.name": project_name,
+ },
+ }
+ )
diff --git a/litellm/integrations/otel/presets/utils.py b/litellm/integrations/otel/presets/utils.py
new file mode 100644
index 00000000000..fdf8184441d
--- /dev/null
+++ b/litellm/integrations/otel/presets/utils.py
@@ -0,0 +1,16 @@
+"""Shared helpers for the integration presets."""
+
+from typing import Iterable
+
+
+def ensure_mappers(mapper_names: Iterable[str], *names: str) -> list[str]:
+ """Return ``mapper_names`` with each of ``names`` appended if not already present.
+
+ Order is preserved and duplicates are skipped, so composing several presets
+ (or re-applying one) never double-adds a vocabulary.
+ """
+ result = list(mapper_names)
+ for name in names:
+ if name not in result:
+ result.append(name)
+ return result
diff --git a/litellm/integrations/otel/presets/weave.py b/litellm/integrations/otel/presets/weave.py
new file mode 100644
index 00000000000..9fc03c84a6d
--- /dev/null
+++ b/litellm/integrations/otel/presets/weave.py
@@ -0,0 +1,43 @@
+"""Weave (W&B) preset."""
+
+from litellm.integrations.otel.model.config import ExporterSpec, OpenTelemetryV2Config
+from litellm.integrations.otel.presets.utils import ensure_mappers
+from litellm.integrations.weave.weave_otel import (
+ _get_weave_authorization_header,
+ get_weave_otel_config,
+)
+from litellm.types.utils import StandardCallbackDynamicParams
+
+
+def weave_preset(
+ *,
+ config_overrides: OpenTelemetryV2Config | None = None,
+) -> OpenTelemetryV2Config:
+ weave_cfg = get_weave_otel_config()
+ base = config_overrides or OpenTelemetryV2Config()
+ return base.model_copy(
+ update={
+ "exporters": [
+ *base.exporters,
+ ExporterSpec(
+ kind=weave_cfg.protocol or "otlp_http",
+ endpoint=weave_cfg.endpoint,
+ headers=weave_cfg.otlp_auth_headers,
+ ),
+ ],
+ # Weave consumes OpenInference + a small Weave-specific overlay.
+ "mapper_names": ensure_mappers(base.mapper_names, "openinference", "weave"),
+ }
+ )
+
+
+def weave_dynamic_headers(params: StandardCallbackDynamicParams) -> dict[str, str]:
+ """Per-request Weave OTLP headers from team/key dynamic params."""
+ headers: dict[str, str] = {}
+ api_key = params.get("wandb_api_key")
+ if api_key:
+ headers["Authorization"] = _get_weave_authorization_header(api_key=api_key)
+ project_id = params.get("weave_project_id")
+ if project_id:
+ headers["project_id"] = project_id
+ return headers
diff --git a/litellm/integrations/otel/runtime.py b/litellm/integrations/otel/runtime.py
new file mode 100644
index 00000000000..ac3b991c971
--- /dev/null
+++ b/litellm/integrations/otel/runtime.py
@@ -0,0 +1,38 @@
+"""SDK-free entrypoints for proxy-core call sites (auth, …).
+
+Proxy code may run without the OpenTelemetry SDK installed, so it must not import
+``litellm.integrations.otel.logger`` (which imports the SDK at module scope) at
+module load. These wrappers import it lazily and no-op when the SDK is absent or
+V2 is not the active logger — so a call site can wrap a request phase or seed
+identity unconditionally.
+"""
+
+from contextlib import contextmanager
+from typing import Any, Iterator
+
+
+@contextmanager
+def phase_span(name: str) -> "Iterator[Any]":
+ """Run a request phase inside a live active span so its DB/service calls nest.
+
+ Yields ``None`` (a plain no-op) when the OTel SDK is unavailable or V2 is not
+ the active logger.
+ """
+ try:
+ from litellm.integrations.otel.logger import phase_span as _phase_span
+ except Exception:
+ yield None
+ return
+ with _phase_span(name) as span:
+ yield span
+
+
+def seed_request_identity(user_api_key_dict: Any, model: Any = None) -> None:
+ """Seed request-identity Baggage at the auth boundary (no-op without V2)."""
+ try:
+ from litellm.integrations.otel.logger import (
+ seed_request_identity as _seed_request_identity,
+ )
+ except Exception:
+ return
+ _seed_request_identity(user_api_key_dict, model=model)
diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py
index 606d28585bc..c127b3873a7 100644
--- a/litellm/litellm_core_utils/litellm_logging.py
+++ b/litellm/litellm_core_utils/litellm_logging.py
@@ -3792,6 +3792,9 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
try:
custom_logger_init_args = custom_logger_init_args or {}
if logging_integration == "agentops": # Add AgentOps initialization
+ _v2 = _maybe_construct_otel_v2("agentops", _in_memory_loggers)
+ if _v2 is not None:
+ return _v2 # type: ignore
for callback in _in_memory_loggers:
if isinstance(callback, AgentOps):
return callback # type: ignore
@@ -3944,6 +3947,9 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
_in_memory_loggers.append(_opik_logger)
return _opik_logger # type: ignore
elif logging_integration == "arize":
+ _v2 = _maybe_construct_otel_v2("arize", _in_memory_loggers)
+ if _v2 is not None:
+ return _v2 # type: ignore
from litellm.integrations.opentelemetry import (
OpenTelemetry,
OpenTelemetryConfig,
@@ -3973,6 +3979,9 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
_in_memory_loggers.append(_arize_otel_logger)
return _arize_otel_logger # type: ignore
elif logging_integration == "arize_phoenix":
+ _v2 = _maybe_construct_otel_v2("arize_phoenix", _in_memory_loggers)
+ if _v2 is not None:
+ return _v2 # type: ignore
from litellm.integrations.opentelemetry import (
OpenTelemetry,
OpenTelemetryConfig,
@@ -4003,6 +4012,9 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
_in_memory_loggers.append(_arize_phoenix_otel_logger)
return _arize_phoenix_otel_logger # type: ignore
elif logging_integration == "levo":
+ _v2 = _maybe_construct_otel_v2("levo", _in_memory_loggers)
+ if _v2 is not None:
+ return _v2 # type: ignore
from litellm.integrations.levo.levo import LevoLogger
from litellm.integrations.opentelemetry import (
OpenTelemetry,
@@ -4028,6 +4040,28 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
_in_memory_loggers.append(_levo_otel_logger)
return _levo_otel_logger # type: ignore
elif logging_integration == "otel":
+ # Gate the new typed V2 adapter behind LITELLM_OTEL_V2. When off,
+ # the legacy 3,227-line god-class is used unchanged. The two are
+ # never registered simultaneously — the dedup loop below treats
+ # any module under ``litellm.integrations.otel`` or
+ # ``litellm.integrations.opentelemetry`` as "the OTel callback".
+ from litellm.integrations.otel.model.config import is_otel_v2_enabled
+
+ if is_otel_v2_enabled():
+ from litellm.integrations.otel.logger import OpenTelemetryV2
+
+ for callback in _in_memory_loggers:
+ if type(callback) is OpenTelemetryV2:
+ return callback # type: ignore
+ otel_logger_v2 = OpenTelemetryV2(
+ **_get_custom_logger_settings_from_proxy_server(
+ callback_name=logging_integration
+ )
+ )
+ _in_memory_loggers.append(otel_logger_v2)
+ _maybe_auto_initialize_arize_phoenix(_in_memory_loggers)
+ return otel_logger_v2 # type: ignore
+
from litellm.integrations.opentelemetry import OpenTelemetry
for callback in _in_memory_loggers:
@@ -4166,6 +4200,9 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
elif logging_integration == "langtrace":
if "LANGTRACE_API_KEY" not in os.environ:
raise ValueError("LANGTRACE_API_KEY not found in environment variables")
+ _v2 = _maybe_construct_otel_v2("langtrace", _in_memory_loggers)
+ if _v2 is not None:
+ return _v2 # type: ignore
from litellm.integrations.opentelemetry import (
OpenTelemetry,
@@ -4206,6 +4243,9 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
_in_memory_loggers.append(langfuse_logger)
return langfuse_logger # type: ignore
elif logging_integration == "langfuse_otel":
+ _v2 = _maybe_construct_otel_v2("langfuse_otel", _in_memory_loggers)
+ if _v2 is not None:
+ return _v2 # type: ignore
from litellm.integrations.langfuse.langfuse_otel import LangfuseOtelLogger
for callback in _in_memory_loggers:
@@ -4222,6 +4262,9 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
_in_memory_loggers.append(_otel_logger)
return _otel_logger # type: ignore
elif logging_integration == "weave_otel":
+ _v2 = _maybe_construct_otel_v2("weave_otel", _in_memory_loggers)
+ if _v2 is not None:
+ return _v2 # type: ignore
from litellm.integrations.opentelemetry import OpenTelemetryConfig
from litellm.integrations.weave.weave_otel import (
WeaveOtelLogger,
@@ -4370,6 +4413,42 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
return None
+def _maybe_construct_otel_v2(
+ callback_name: str, _in_memory_loggers: list
+) -> Optional[Any]:
+ """If ``LITELLM_OTEL_V2`` is on, build (or reuse) a single ``OpenTelemetryV2``
+ instance configured via the preset for ``callback_name``.
+
+ Returns ``None`` when V2 is off OR when there's no preset registered for
+ ``callback_name`` — callers should then fall through to the legacy path.
+ """
+ from litellm.integrations.otel.model.config import is_otel_v2_enabled
+
+ if not is_otel_v2_enabled():
+ return None
+ from litellm.integrations.otel.logger import OpenTelemetryV2
+ from litellm.integrations.otel.presets import PRESET_BY_CALLBACK
+
+ preset_fn = PRESET_BY_CALLBACK.get(callback_name)
+ if preset_fn is None:
+ return None
+ for callback in _in_memory_loggers:
+ if (
+ isinstance(callback, OpenTelemetryV2)
+ and getattr(callback, "callback_name", None) == callback_name
+ ):
+ return callback
+ try:
+ config = preset_fn()
+ except Exception:
+ # If env vars are missing or the preset raises, defer to the legacy path
+ # so customers get the same error story they had before V2 landed.
+ return None
+ v2_logger = OpenTelemetryV2(config=config, callback_name=callback_name)
+ _in_memory_loggers.append(v2_logger)
+ return v2_logger
+
+
def _maybe_auto_initialize_arize_phoenix(_in_memory_loggers: list) -> None:
"""
Auto-initialize ArizePhoenixLogger when Phoenix env vars are detected.
diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py
index 813b9826b37..828b719e299 100644
--- a/litellm/proxy/auth/user_api_key_auth.py
+++ b/litellm/proxy/auth/user_api_key_auth.py
@@ -21,6 +21,8 @@ from fastapi.security.api_key import APIKeyHeader
import litellm
from litellm._logging import verbose_logger, verbose_proxy_logger
from litellm._service_logger import ServiceLogging
+from litellm.integrations.otel.model.config import is_otel_v2_enabled
+from litellm.integrations.otel.runtime import phase_span, seed_request_identity
from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS
from litellm.litellm_core_utils.dd_tracing import tracer
from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value
@@ -694,11 +696,17 @@ def _ensure_parent_otel_span_on_request_state(request: Request) -> None:
start_time=start_time,
headers=_safe_get_request_headers(request),
)
- open_telemetry_logger.set_proxy_request_route_attributes(
- parent_otel_span,
- url_path=get_request_route(request=request),
- http_route=get_request_route_template(request),
+ # Under V2 the FastAPI instrumentor stamps http.route / url.path on the server
+ # span; only the legacy logger needs these set explicitly.
+ set_route_attrs = getattr(
+ open_telemetry_logger, "set_proxy_request_route_attributes", None
)
+ if not is_otel_v2_enabled() and set_route_attrs is not None:
+ set_route_attrs(
+ parent_otel_span,
+ url_path=get_request_route(request=request),
+ http_route=get_request_route_template(request),
+ )
request.state.parent_otel_span = parent_otel_span
@@ -2187,73 +2195,85 @@ async def user_api_key_auth(
route: str = get_request_route(request=request)
## CHECK IF ROUTE IS ALLOWED
- user_api_key_auth_obj = await _user_api_key_auth_builder(
- request=request,
- api_key=api_key,
- azure_api_key_header=azure_api_key_header,
- anthropic_api_key_header=anthropic_api_key_header,
- google_ai_studio_api_key_header=google_ai_studio_api_key_header,
- azure_apim_header=azure_apim_header,
- request_data=request_data,
- custom_litellm_key_header=custom_litellm_key_header,
- )
- user_api_key_auth_obj.budget_reservation = None
-
- ## ENSURE DISABLE ROUTE WORKS ACROSS ALL USER AUTH FLOWS ##
- RouteChecks.should_call_route(
- route=route, valid_token=user_api_key_auth_obj, request=request
- )
-
- # Single authorization point. Builder paths MUST NOT call common_checks.
- # Route through the same exception handler the builder uses so
- # authorization failures (ProxyException, or plain Exception from
- # admin-only-route / model-access / budget checks) surface as
- # ProxyException consistently with pre-refactor behavior.
- try:
- await _run_centralized_common_checks(
- user_api_key_auth_obj=user_api_key_auth_obj,
+ # Run the whole auth phase inside a live ``auth`` span so the DB lookups it
+ # triggers (key/user/team object reads) nest under it instead of flattening
+ # onto the server span. No-op when OTel V2 isn't active.
+ with phase_span(f"auth {route}"):
+ user_api_key_auth_obj = await _user_api_key_auth_builder(
request=request,
- request_data=request_data,
- route=route,
- )
- except Exception as e:
- return await UserAPIKeyAuthExceptionHandler._handle_authentication_error(
- e=e,
- request=request,
- request_data=request_data,
- route=route,
- parent_otel_span=user_api_key_auth_obj.parent_otel_span,
api_key=api_key,
+ azure_api_key_header=azure_api_key_header,
+ anthropic_api_key_header=anthropic_api_key_header,
+ google_ai_studio_api_key_header=google_ai_studio_api_key_header,
+ azure_apim_header=azure_apim_header,
+ request_data=request_data,
+ custom_litellm_key_header=custom_litellm_key_header,
+ )
+ user_api_key_auth_obj.budget_reservation = None
+
+ ## ENSURE DISABLE ROUTE WORKS ACROSS ALL USER AUTH FLOWS ##
+ RouteChecks.should_call_route(
+ route=route, valid_token=user_api_key_auth_obj, request=request
)
- # Defense-in-depth: ``_user_api_key_auth_builder`` has multiple early-return
- # paths (no master key, /user/auth route, JWT short-circuits) that bypass
- # the end-user resolution block. If those paths produced an auth obj
- # without an ``end_user_id`` set, fall back to extracting from the request
- # body so spend logs are still attributed correctly. Validation honours
- # ``litellm.validate_end_user_id_in_db``.
- if user_api_key_auth_obj.end_user_id is None:
- from litellm.proxy.proxy_server import (
- prisma_client,
- proxy_logging_obj,
- user_api_key_cache,
- )
-
- raw_end_user_id = get_end_user_id_from_request_body(
- request_data, _safe_get_request_headers(request)
- )
- if raw_end_user_id is not None:
- resolved_end_user_id = await resolve_and_validate_end_user_id(
- raw_end_user_id=raw_end_user_id,
- prisma_client=prisma_client,
- user_api_key_cache=user_api_key_cache,
- parent_otel_span=user_api_key_auth_obj.parent_otel_span,
- proxy_logging_obj=proxy_logging_obj,
+ # Single authorization point. Builder paths MUST NOT call common_checks.
+ # Route through the same exception handler the builder uses so
+ # authorization failures (ProxyException, or plain Exception from
+ # admin-only-route / model-access / budget checks) surface as
+ # ProxyException consistently with pre-refactor behavior.
+ try:
+ await _run_centralized_common_checks(
+ user_api_key_auth_obj=user_api_key_auth_obj,
+ request=request,
+ request_data=request_data,
route=route,
)
- if resolved_end_user_id is not None:
- user_api_key_auth_obj.end_user_id = resolved_end_user_id
+ except Exception as e:
+ return await UserAPIKeyAuthExceptionHandler._handle_authentication_error(
+ e=e,
+ request=request,
+ request_data=request_data,
+ route=route,
+ parent_otel_span=user_api_key_auth_obj.parent_otel_span,
+ api_key=api_key,
+ )
+ # Defense-in-depth: ``_user_api_key_auth_builder`` has multiple early-return
+ # paths (no master key, /user/auth route, JWT short-circuits) that bypass
+ # the end-user resolution block. If those paths produced an auth obj
+ # without an ``end_user_id`` set, fall back to extracting from the request
+ # body so spend logs are still attributed correctly. Validation honours
+ # ``litellm.validate_end_user_id_in_db``.
+ if user_api_key_auth_obj.end_user_id is None:
+ from litellm.proxy.proxy_server import (
+ prisma_client,
+ proxy_logging_obj,
+ user_api_key_cache,
+ )
+
+ raw_end_user_id = get_end_user_id_from_request_body(
+ request_data, _safe_get_request_headers(request)
+ )
+ if raw_end_user_id is not None:
+ resolved_end_user_id = await resolve_and_validate_end_user_id(
+ raw_end_user_id=raw_end_user_id,
+ prisma_client=prisma_client,
+ user_api_key_cache=user_api_key_cache,
+ parent_otel_span=user_api_key_auth_obj.parent_otel_span,
+ proxy_logging_obj=proxy_logging_obj,
+ route=route,
+ )
+ if resolved_end_user_id is not None:
+ user_api_key_auth_obj.end_user_id = resolved_end_user_id
+
+ # Identity is now resolved. Seed it AFTER the auth span closes so the Baggage
+ # persists on the request task (detaching the span's context token inside the
+ # ``with`` would unwind a Baggage attach made within it) and every post-auth
+ # span — pre-call, LLM call, guardrail, spend write — inherits team/key/user.
+ seed_request_identity(
+ user_api_key_auth_obj,
+ model=request_data.get("model") if isinstance(request_data, dict) else None,
+ )
user_api_key_auth_obj.request_route = normalize_request_route(route)
return user_api_key_auth_obj
diff --git a/litellm/proxy/db/log_db_metrics.py b/litellm/proxy/db/log_db_metrics.py
index 5c795155324..eb4961062df 100644
--- a/litellm/proxy/db/log_db_metrics.py
+++ b/litellm/proxy/db/log_db_metrics.py
@@ -7,13 +7,21 @@ ServiceLogger() then sends DB logs to Prometheus, OTEL, Datadog etc
import asyncio
from datetime import datetime
from functools import wraps
-from typing import Callable, Dict, Tuple
+from typing import Callable, Dict, Optional, Tuple
from litellm._service_logger import ServiceTypes
-from litellm.litellm_core_utils.core_helpers import (
- _get_parent_otel_span_from_kwargs,
- get_litellm_metadata_from_kwargs,
-)
+from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs
+
+
+def _safe_db_event_metadata(kwargs: Dict) -> Optional[Dict[str, str]]:
+ """Minimal, non-sensitive ``event_metadata`` for a DB service log.
+
+ The raw ``kwargs``/``args`` carry live objects (Prisma client, OTel spans)
+ and secrets (tokens), none of which belongs on a span — so we surface only
+ the table name when present. Everything else is dropped.
+ """
+ table_name = kwargs.get("table_name")
+ return {"table_name": table_name} if isinstance(table_name, str) else None
def log_db_metrics(func):
@@ -52,11 +60,7 @@ def log_db_metrics(func):
duration=(end_time - start_time).total_seconds(),
start_time=start_time,
end_time=end_time,
- event_metadata={
- "function_name": func.__name__,
- "function_kwargs": kwargs,
- "function_args": args,
- },
+ event_metadata=_safe_db_event_metadata(kwargs),
)
)
elif (
@@ -71,8 +75,9 @@ def log_db_metrics(func):
kwargs=passed_kwargs
)
if parent_otel_span is not None:
- metadata = get_litellm_metadata_from_kwargs(kwargs=passed_kwargs)
-
+ # No metadata dump: identity rides on Baggage, and the full
+ # request metadata (auth blob, response headers, tokens) must
+ # not land on a span.
asyncio.create_task(
proxy_logging_obj.service_logging_obj.async_service_success_hook(
service=ServiceTypes.BATCH_WRITE_TO_DB,
@@ -81,7 +86,7 @@ def log_db_metrics(func):
duration=0.0,
start_time=start_time,
end_time=end_time,
- event_metadata=metadata,
+ event_metadata=None,
)
)
# end of logging to otel
@@ -134,9 +139,5 @@ async def _handle_logging_db_exception(
duration=(end_time - start_time).total_seconds(),
start_time=start_time,
end_time=end_time,
- event_metadata={
- "function_name": func.__name__,
- "function_kwargs": kwargs,
- "function_args": args,
- },
+ event_metadata=_safe_db_event_metadata(kwargs),
)
diff --git a/litellm/proxy/management_helpers/utils.py b/litellm/proxy/management_helpers/utils.py
index 888ccc78188..287dc38a1fb 100644
--- a/litellm/proxy/management_helpers/utils.py
+++ b/litellm/proxy/management_helpers/utils.py
@@ -8,6 +8,7 @@ from fastapi import HTTPException, Request
import litellm
from litellm._logging import verbose_logger
+from litellm.integrations.otel.model.config import is_otel_v2_enabled
from litellm._uuid import uuid
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
from litellm.proxy._types import ( # key request types; user request types; team request types; customer request types
@@ -458,6 +459,12 @@ async def _emit_management_endpoint_otel_span(
if open_telemetry_logger is None:
return
+ # Under V2 OTel, management endpoints are ordinary FastAPI routes already
+ # spanned by the mounted instrumentor — there is no management hook to fire, so
+ # skip the payload build entirely. The legacy logger still needs the hook.
+ if is_otel_v2_enabled():
+ return
+
http_request: Optional[Request] = kwargs.get("http_request")
if http_request is not None:
# Inline import — auth_utils participates in a proxy import cycle.
diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
index 67af176c1eb..5aa0f6cb184 100644
--- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
+++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
@@ -801,9 +801,14 @@ async def pass_through_request( # noqa: PLR0915
)
## LOGGING OBJECT ## - initialize before pre_call_hook so guardrails can access it
+ # Surface the requested model (when the body carries one) so logging/spans
+ # read e.g. ``chat gpt-4o`` instead of ``chat unknown``.
+ passthrough_model = (
+ _parsed_body.get("model") if isinstance(_parsed_body, dict) else None
+ ) or "unknown"
start_time = datetime.now()
logging_obj = Logging(
- model="unknown",
+ model=passthrough_model,
messages=[{"role": "user", "content": safe_dumps(_parsed_body)}],
stream=False,
call_type="pass_through_endpoint",
@@ -855,7 +860,7 @@ async def pass_through_request( # noqa: PLR0915
# done for supporting 'parallel_request_limiter.py' with pass-through endpoints
logging_obj.update_environment_variables(
- model="unknown",
+ model=passthrough_model,
user="unknown",
optional_params={},
litellm_params=kwargs["litellm_params"],
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 8fbe6d97dbc..2a085077434 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -825,6 +825,37 @@ async def proxy_startup_event(app: FastAPI): # noqa: PLR0915
if isinstance(worker_config, dict):
await initialize(**worker_config)
+ ## V2 OTEL: now that config (and therefore the callbacks) is loaded, publish
+ ## the chosen V2 logger's TracerProvider as the OTel global. The FastAPI
+ ## instrumentation mounted at app-creation binds to the global provider, so
+ ## this is what makes server spans and gen-ai spans share one provider and
+ ## land in the same trace. Prefer an already-registered preset logger
+ ## (arize, langfuse, …) so server spans export to that backend too; otherwise
+ ## build a generic one from OTEL_* envs. ``set_tracer_provider`` only takes
+ ## effect once, so the first configured logger wins.
+ try:
+ from litellm.integrations.otel.model.config import is_otel_v2_enabled
+
+ if is_otel_v2_enabled():
+ from opentelemetry import trace as _otel_trace
+
+ from litellm.integrations.otel.logger import OpenTelemetryV2
+
+ _otel_v2_logger = (
+ next(
+ (
+ cb
+ for cb in litellm.service_callback
+ if isinstance(cb, OpenTelemetryV2)
+ ),
+ None,
+ )
+ or OpenTelemetryV2()
+ )
+ _otel_trace.set_tracer_provider(_otel_v2_logger._tracer_provider)
+ except Exception as e:
+ verbose_proxy_logger.debug("Skipping OTel V2 provider setup: %s", e)
+
# check if DATABASE_URL in environment - load from there
if prisma_client is None:
_db_url: Optional[str] = get_secret("DATABASE_URL", None) # type: ignore
@@ -1074,6 +1105,16 @@ app = FastAPI(
strict_content_type=False,
)
+## V2 OTEL: instrument the FastAPI app for server spans (gated by
+## LITELLM_OTEL_V2). This MUST run at app-creation time — once the lifespan runs,
+## the middleware stack is frozen and ``instrument_app`` raises "Cannot add
+## middleware after an application has started". See
+## ``litellm.integrations.otel.mount`` for the full rationale; the call is a safe
+## no-op when the gate is off or the instrumentation package is unavailable.
+from litellm.integrations.otel.mount import instrument_fastapi_app
+
+instrument_fastapi_app(app)
+
vertex_live_passthrough_vertex_base = VertexBase()
@@ -1238,6 +1279,17 @@ def _close_dangling_otel_server_span(request: Request, status_code: int) -> None
return
if open_telemetry_logger is None:
return
+ # Under OTel V2 the FastAPI instrumentor owns the server span (parent_otel_span
+ # is that same span), and it records the error + ends it itself. Ending it here
+ # would end it early — losing the http.* attributes the instrumentor stamps on
+ # completion — and double-end it. Leave it to the instrumentor.
+ try:
+ from litellm.integrations.otel.model.config import is_otel_v2_enabled
+
+ if is_otel_v2_enabled():
+ return
+ except Exception:
+ pass
try:
from opentelemetry.trace import Status, StatusCode
diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py
index 14f7f411e41..0e72f47e224 100644
--- a/litellm/proxy/utils.py
+++ b/litellm/proxy/utils.py
@@ -30,7 +30,11 @@ from typing import (
)
from litellm import _custom_logger_compatible_callbacks_literal
-from litellm.constants import DEFAULT_MODEL_CREATED_AT_TIME, MAX_TEAM_LIST_LIMIT
+from litellm.constants import (
+ DEFAULT_MODEL_CREATED_AT_TIME,
+ LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL,
+ MAX_TEAM_LIST_LIMIT,
+)
from litellm.proxy._types import (
DB_CONNECTION_ERROR_TYPES,
CommonProxyErrors,
@@ -557,6 +561,26 @@ class ProxyLogging:
for idx, initialized_callback in string_callbacks_to_replace.items():
litellm.callbacks[idx] = initialized_callback
+ # Fan ``litellm.callbacks`` (the "all events" registry) out into the
+ # success/failure event lists eagerly, at startup. ``completion()`` does
+ # this lazily in ``function_setup`` on the first call, but request paths
+ # that build their own logging object and never run ``function_setup`` —
+ # notably pass-through endpoints — read ``litellm._async_success_callback``
+ # directly. Without this, a config-registered logger (e.g. ``otel``) is
+ # invisible to pass-through traffic until some other request warms the
+ # global lists. The manager dedupes, so this is idempotent with
+ # ``function_setup``.
+ for callback in litellm.callbacks:
+ if isinstance(callback, CustomLogger):
+ litellm.logging_callback_manager.add_litellm_success_callback(callback)
+ litellm.logging_callback_manager.add_litellm_failure_callback(callback)
+ litellm.logging_callback_manager.add_litellm_async_success_callback(
+ callback
+ )
+ litellm.logging_callback_manager.add_litellm_async_failure_callback(
+ callback
+ )
+
async def update_request_status(
self, litellm_call_id: str, status: Literal["success", "fail"]
):
@@ -2171,6 +2195,15 @@ class ProxyLogging:
# async_post_call_failure_hook — skip pre_call and failure handlers.
if litellm_logging_obj.call_type == CallTypes.pass_through.value:
return
+ # This is a proxy-gate error (auth/rate-limit) for a request that never
+ # reached a provider. ``pre_call`` below still fires every callback's
+ # input hook so the failure is logged — but tracing callbacks must not
+ # fabricate an LLM-call span for a call that did not happen (and, since
+ # this runs inside the live ``auth`` phase span, would otherwise nest it
+ # under auth). The marker tells them to skip span creation.
+ litellm_logging_obj.model_call_details[
+ LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL
+ ] = True
litellm_logging_obj.pre_call(
input=input,
api_key="",
diff --git a/pyproject.toml b/pyproject.toml
index 6e84afad17e..2e371630fd4 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -119,6 +119,7 @@ proxy-runtime = [
"opentelemetry-api==1.28.0",
"opentelemetry-sdk==1.28.0",
"opentelemetry-exporter-otlp==1.28.0",
+ "opentelemetry-instrumentation-fastapi==0.49b0",
"ddtrace>=2.19.0,<3.0",
"sentry-sdk>=2.21.0,<3.0",
"mangum>=0.17.0,<1.0",
@@ -160,6 +161,7 @@ dev = [
"opentelemetry-api==1.28.0",
"opentelemetry-sdk==1.28.0",
"opentelemetry-exporter-otlp==1.28.0",
+ "opentelemetry-instrumentation-fastapi==0.49b0",
"langfuse==2.59.7",
"fastapi-offline==1.7.6",
"fakeredis==2.34.1",
@@ -178,6 +180,7 @@ proxy-dev = [
"opentelemetry-api==1.28.0",
"opentelemetry-sdk==1.28.0",
"opentelemetry-exporter-otlp==1.28.0",
+ "opentelemetry-instrumentation-fastapi==0.49b0",
"azure-identity==1.25.2",
"a2a-sdk==0.3.24",
]
diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_baggage.py b/tests/test_litellm/integrations/otel/test_otel_v2_baggage.py
new file mode 100644
index 00000000000..fd4a7e141e5
--- /dev/null
+++ b/tests/test_litellm/integrations/otel/test_otel_v2_baggage.py
@@ -0,0 +1,172 @@
+"""Tests for Baggage-based promotion of request-scoped attributes onto every span,
+and the two antipattern boundaries: http.* is never promoted, and the full
+metadata blob is never promoted (only the bounded allowlist)."""
+
+import pytest
+
+pytest.importorskip("opentelemetry")
+
+from litellm.integrations.otel import ( # noqa: E402
+ GenAI,
+ HTTP,
+ LiteLLM,
+ OpenTelemetryV2Config,
+ promoted_baggage,
+)
+from litellm.integrations.otel.plumbing import context as ctx_mod # noqa: E402
+from litellm.integrations.otel.plumbing import providers # noqa: E402
+from litellm.integrations.otel.emitter import SpanEmitter # noqa: E402
+from litellm.integrations.otel.model.payloads import ( # noqa: E402
+ GuardrailSpanData,
+ LLMCallSpanData,
+ ServiceSpanData,
+)
+from litellm.integrations.otel.model.baggage import BAGGAGE_PROMOTED_KEYS # noqa: E402
+from litellm.integrations.otel.model.spans import SpanRole # noqa: E402
+
+
+def _payload():
+ return {
+ "call_type": "acompletion",
+ "custom_llm_provider": "openai",
+ "model": "gpt-4o",
+ "prompt_tokens": 1,
+ "completion_tokens": 1,
+ "total_tokens": 2,
+ "metadata": {
+ "team_id": "t1",
+ "team_alias": "team one",
+ "user_api_key_hash": "hsh",
+ "user_api_key_org_id": "org1",
+ "user_api_key_team_metadata": {"tier": "gold", "cost_center": "42"},
+ "private_note": "do-not-promote",
+ },
+ "status": "success",
+ "litellm_call_id": "call_1",
+ "hidden_params": {"litellm_model_name": "azure/my-deployment"},
+ }
+
+
+def _engine_and_exporter(config=None):
+ cfg = config or OpenTelemetryV2Config(exporter="in_memory")
+ provider, exporter = providers.in_memory_provider(cfg)
+ tracer = providers.get_tracer(provider, "litellm-baggage-test")
+ return SpanEmitter(tracer, cfg), exporter
+
+
+def test_identity_promoted_onto_every_span():
+ engine, exporter = _engine_and_exporter()
+ data = LLMCallSpanData.from_standard_logging_payload(_payload())
+ bag = promoted_baggage(data.identity, data.request_model, BAGGAGE_PROMOTED_KEYS)
+ ctx = ctx_mod.set_request_baggage(bag)
+
+ root = engine.start_span(SpanRole.PROXY_REQUEST, "POST /chat/completions", ctx)
+ root_ctx = ctx_mod.context_from_span(root, ctx)
+ engine.emit(SpanRole.LLM_CALL, data, parent_context=root_ctx)
+ engine.emit(
+ SpanRole.GUARDRAIL, GuardrailSpanData("presidio", status="success"), root_ctx
+ )
+ engine.emit(SpanRole.SERVICE, ServiceSpanData("redis", call_type="set"), root_ctx)
+ root.end()
+
+ spans = exporter.get_finished_spans()
+ assert len(spans) == 4
+ for span in spans:
+ assert span.attributes.get(LiteLLM.TEAM_ID) == "t1"
+ assert span.attributes.get(LiteLLM.TEAM_ALIAS) == "team one"
+ assert span.attributes.get(GenAI.REQUEST_MODEL) == "gpt-4o"
+
+
+def test_team_metadata_and_provider_model_promoted():
+ """The team's metadata dict (JSON) and the provider/underlying model name are
+ promoted onto every span, alongside the user-facing ``gen_ai.request.model``."""
+ import json
+
+ engine, exporter = _engine_and_exporter()
+ data = LLMCallSpanData.from_standard_logging_payload(_payload())
+ bag = promoted_baggage(data.identity, data.request_model, BAGGAGE_PROMOTED_KEYS)
+ ctx = ctx_mod.set_request_baggage(bag)
+ engine.emit(SpanRole.SERVICE, ServiceSpanData("redis", call_type="set"), ctx)
+ (span,) = exporter.get_finished_spans()
+
+ # team metadata: the whole dict, JSON-serialized into one value
+ assert json.loads(span.attributes[LiteLLM.TEAM_METADATA]) == {
+ "tier": "gold",
+ "cost_center": "42",
+ }
+ # provider model is distinct from the user-facing request model
+ assert span.attributes.get(LiteLLM.PROVIDER_MODEL) == "azure/my-deployment"
+ assert span.attributes.get(GenAI.REQUEST_MODEL) == "gpt-4o"
+
+
+def test_empty_team_metadata_is_dropped():
+ """An absent/empty team_metadata dict must not promote a useless ``"{}"``."""
+ payload = _payload()
+ payload["metadata"]["user_api_key_team_metadata"] = {}
+ payload["hidden_params"] = {}
+ data = LLMCallSpanData.from_standard_logging_payload(payload)
+ assert data.identity.team_metadata is None
+ # With no explicit dispatched-model source (hidden_params emptied), the
+ # provider model falls back to the call model — so it's present, not dropped.
+ assert data.identity.provider_model == "gpt-4o"
+ bag = promoted_baggage(data.identity, data.request_model, BAGGAGE_PROMOTED_KEYS)
+ assert LiteLLM.TEAM_METADATA not in bag
+ assert bag[LiteLLM.PROVIDER_MODEL] == "gpt-4o"
+
+
+def test_allowlisted_metadata_subkey_promoted_blob_excluded():
+ engine, exporter = _engine_and_exporter()
+ data = LLMCallSpanData.from_standard_logging_payload(_payload())
+ bag = promoted_baggage(data.identity, data.request_model, BAGGAGE_PROMOTED_KEYS)
+ ctx = ctx_mod.set_request_baggage(bag)
+ engine.emit(SpanRole.SERVICE, ServiceSpanData("redis", call_type="set"), ctx)
+ (span,) = exporter.get_finished_spans()
+ # allowlisted metadata sub-key is promoted
+ assert (
+ span.attributes.get(f"{LiteLLM.METADATA_PREFIX}user_api_key_org_id") == "org1"
+ )
+ # non-allowlisted metadata is NOT promoted (no full-blob dumping)
+ assert all("private_note" not in k for k in span.attributes)
+
+
+def test_http_attributes_never_promoted():
+ """Even if http.* is present in baggage, the processor must not stamp it on
+ child spans (it belongs on the SERVER span only)."""
+ engine, exporter = _engine_and_exporter()
+ ctx = ctx_mod.set_request_baggage(
+ {
+ LiteLLM.TEAM_ID: "t1",
+ HTTP.ROUTE: "/chat/completions",
+ HTTP.REQUEST_METHOD: "POST",
+ }
+ )
+ engine.emit(SpanRole.SERVICE, ServiceSpanData("redis", call_type="set"), ctx)
+ (span,) = exporter.get_finished_spans()
+ assert span.attributes.get(LiteLLM.TEAM_ID) == "t1"
+ assert HTTP.ROUTE not in span.attributes
+ assert HTTP.REQUEST_METHOD not in span.attributes
+
+
+def test_arbitrary_upstream_baggage_not_promoted():
+ engine, exporter = _engine_and_exporter()
+ ctx = ctx_mod.set_request_baggage(
+ {LiteLLM.TEAM_ID: "t1", "some.upstream.key": "leak"}
+ )
+ engine.emit(SpanRole.SERVICE, ServiceSpanData("redis", call_type="set"), ctx)
+ (span,) = exporter.get_finished_spans()
+ assert span.attributes.get(LiteLLM.TEAM_ID) == "t1"
+ assert "some.upstream.key" not in span.attributes
+
+
+def test_baggage_processor_allowlist_can_be_widened():
+ cfg = OpenTelemetryV2Config(
+ exporter="in_memory",
+ baggage_promoted_keys=[LiteLLM.TEAM_ID, "custom.key"],
+ )
+ engine, exporter = _engine_and_exporter(cfg)
+ ctx = ctx_mod.set_request_baggage({"custom.key": "v", LiteLLM.TEAM_ALIAS: "ta"})
+ engine.emit(SpanRole.SERVICE, ServiceSpanData("redis"), ctx)
+ (span,) = exporter.get_finished_spans()
+ assert span.attributes.get("custom.key") == "v"
+ # team_alias not in this config's allowlist -> not promoted
+ assert LiteLLM.TEAM_ALIAS not in span.attributes
diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_components.py b/tests/test_litellm/integrations/otel/test_otel_v2_components.py
new file mode 100644
index 00000000000..86d84bd8100
--- /dev/null
+++ b/tests/test_litellm/integrations/otel/test_otel_v2_components.py
@@ -0,0 +1,463 @@
+"""Coverage for the engine-layer components: providers/exporters, context +
+baggage helpers, metrics, the typed coercion helpers, mapper branches, span-name
+builders, and the registry validator's failure paths. Needs the OTel SDK."""
+
+import pytest
+
+pytest.importorskip("opentelemetry")
+
+from opentelemetry.sdk.metrics import MeterProvider # noqa: E402
+from opentelemetry.sdk.metrics.export import InMemoryMetricReader # noqa: E402
+from opentelemetry.sdk.trace.export import ( # noqa: E402
+ BatchSpanProcessor,
+ ConsoleSpanExporter,
+ SimpleSpanProcessor,
+)
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( # noqa: E402
+ InMemorySpanExporter,
+)
+from opentelemetry.trace import SpanKind # noqa: E402
+
+from litellm.integrations.otel.plumbing import context as ctx_mod # noqa: E402
+from litellm.integrations.otel.plumbing import providers # noqa: E402
+from litellm.integrations.otel.model.config import OpenTelemetryV2Config # noqa: E402
+from litellm.integrations.otel.mappers.genai import GenAIMapper # noqa: E402
+from litellm.integrations.otel.mappers.legacy import LegacyMapper # noqa: E402
+from litellm.integrations.otel.plumbing.metrics import (
+ create_genai_metrics,
+) # noqa: E402
+from litellm.integrations.otel.model.payloads import ( # noqa: E402
+ GuardrailSpanData,
+ LLMCallSpanData,
+ LLMRequestParams,
+ LLMUsage,
+ ProxyRequestSpanData,
+ RequestIdentity,
+ ServerInfo,
+ ServiceSpanData,
+ SpanError,
+)
+from litellm.integrations.otel.model.semconv import GenAI, GenAIOperation
+from litellm.integrations.otel.model.spans import ( # noqa: E402
+ SPAN_REGISTRY,
+ LiteLLMSpanKind,
+ SpanRole,
+ SpanSpec,
+ db_system,
+ guardrail_span_name,
+ proxy_request_span_name,
+ service_span_name,
+ span_role_for_service,
+ validate_registry,
+)
+from litellm.integrations.otel.model.utils import ( # noqa: E402
+ as_bool,
+ as_float,
+ as_int,
+ as_str,
+ as_str_tuple,
+)
+
+# --- typed coercion helpers ------------------------------------------------- #
+
+
+def test_as_str():
+ assert as_str(None) is None
+ assert as_str("x") == "x"
+ assert as_str(5) == "5"
+
+
+def test_as_int():
+ assert as_int(True) == 1
+ assert as_int(3) == 3
+ assert as_int(3.9) == 3
+ assert as_int("7") == 7
+ assert as_int("nope") is None
+ assert as_int(None) is None
+
+
+def test_as_float():
+ assert as_float(True) == 1.0
+ assert as_float(2) == 2.0
+ assert as_float("1.5") == 1.5
+ assert as_float("nope") is None
+ assert as_float(None) is None
+
+
+def test_as_bool():
+ assert as_bool(None) is None
+ assert as_bool(True) is True
+ assert as_bool(1) is True
+ assert as_bool(0) is False
+
+
+def test_as_str_tuple():
+ assert as_str_tuple(None) is None
+ assert as_str_tuple("a") == ("a",)
+ assert as_str_tuple(["a", 2]) == ("a", "2")
+ assert as_str_tuple(123) is None
+
+
+def test_request_params_max_completion_tokens_fallback():
+ params = LLMRequestParams.from_model_parameters({"max_completion_tokens": 99})
+ assert params.max_tokens == 99
+
+
+def test_server_info_from_api_base():
+ assert ServerInfo.from_api_base(None) is None
+ assert ServerInfo.from_api_base("api.host.com:8080") == ServerInfo(
+ "api.host.com", 8080
+ )
+ assert ServerInfo.from_api_base("https://h.com/v1") == ServerInfo("h.com", None)
+ # scheme present but empty netloc -> no hostname
+ assert ServerInfo.from_api_base("http:///v1") is None
+
+
+def test_service_span_data_from_payload():
+ class _Service:
+ value = "redis"
+
+ class _Payload:
+ service = _Service()
+ call_type = "async_set_cache"
+ error = None
+
+ data = ServiceSpanData.from_payload(_Payload())
+ assert data.service_name == "redis"
+ assert data.call_type == "async_set_cache"
+ assert data.error is None
+
+ class _FailPayload:
+ service = _Service()
+ call_type = "async_set_cache"
+ error = "boom"
+
+ failed = ServiceSpanData.from_payload(_FailPayload())
+ assert failed.error is not None
+ assert failed.error.message == "boom"
+
+
+# --- span name builders ----------------------------------------------------- #
+
+
+def test_name_builders():
+ assert (
+ proxy_request_span_name(ProxyRequestSpanData("POST", "/chat/completions"))
+ == "POST /chat/completions"
+ )
+ # "{service} {call_type}" so same-service calls stay distinguishable; the
+ # service name alone when there's no call type.
+ assert service_span_name(ServiceSpanData("redis", call_type="set")) == "redis set"
+ assert service_span_name(ServiceSpanData("redis")) == "redis"
+ assert (
+ guardrail_span_name(GuardrailSpanData("presidio"))
+ == "execute_guardrail presidio"
+ )
+
+
+# --- registry validator failure paths --------------------------------------- #
+
+
+def test_validate_registry_detects_role_mismatch():
+ bad = {SpanRole.LLM_CALL: SpanSpec(SpanRole.SERVICE, LiteLLMSpanKind.CLIENT, None)}
+ with pytest.raises(ValueError, match="mismatched role"):
+ validate_registry(bad)
+
+
+def test_validate_registry_detects_unknown_parent():
+ bad = {
+ SpanRole.LLM_CALL: SpanSpec(
+ SpanRole.LLM_CALL, LiteLLMSpanKind.CLIENT, parent=SpanRole.PROXY_REQUEST
+ )
+ }
+ with pytest.raises(ValueError, match="unknown parent"):
+ validate_registry(bad)
+
+
+def test_validate_registry_detects_missing_roles():
+ partial = {
+ SpanRole.PROXY_REQUEST: SPAN_REGISTRY[SpanRole.PROXY_REQUEST],
+ }
+ with pytest.raises(ValueError, match="missing roles"):
+ validate_registry(partial)
+
+
+# --- mappers (full branch coverage) ----------------------------------------- #
+
+
+def _full_llm_call():
+ return LLMCallSpanData(
+ operation=GenAIOperation.CHAT,
+ provider="openai",
+ request_model="gpt-4o",
+ response_model="gpt-4o-2024",
+ response_id="resp_1",
+ request_params=LLMRequestParams(
+ temperature=0.7,
+ top_p=0.9,
+ top_k=40,
+ max_tokens=256,
+ frequency_penalty=0.1,
+ presence_penalty=0.2,
+ stop_sequences=("STOP",),
+ seed=42,
+ ),
+ usage=LLMUsage(input_tokens=10, output_tokens=5, total_tokens=15),
+ finish_reasons=("stop",),
+ error=None,
+ response_cost=0.002,
+ server=ServerInfo("api.openai.com", 443),
+ identity=RequestIdentity(call_id="c1"),
+ is_streaming=True,
+ )
+
+
+def test_genai_mapper_all_request_params():
+ attrs = GenAIMapper().map(_full_llm_call())
+ assert attrs[GenAI.REQUEST_TOP_P] == 0.9
+ assert attrs[GenAI.REQUEST_TOP_K] == 40
+ assert attrs[GenAI.REQUEST_MAX_TOKENS] == 256
+ assert attrs[GenAI.REQUEST_FREQUENCY_PENALTY] == 0.1
+ assert attrs[GenAI.REQUEST_PRESENCE_PENALTY] == 0.2
+ assert attrs[GenAI.REQUEST_STOP_SEQUENCES] == ["STOP"]
+ assert attrs[GenAI.REQUEST_SEED] == 42
+ assert attrs["server.port"] == 443
+
+
+def test_genai_mapper_guardrail_and_service():
+ from litellm.integrations.otel.model.semconv import LiteLLM
+
+ g = GenAIMapper().map(GuardrailSpanData("presidio", mode="pre"))
+ assert g[LiteLLM.GUARDRAIL_NAME] == "presidio"
+ assert g[LiteLLM.GUARDRAIL_MODE] == "pre"
+
+ # A datastore service (redis) also gets db.* semconv.
+ s = GenAIMapper().map(ServiceSpanData("redis", call_type="set"))
+ assert s[LiteLLM.SERVICE_NAME] == "redis"
+ assert s[LiteLLM.SERVICE_CALL_TYPE] == "set"
+ assert s["db.system.name"] == "redis"
+ assert s["db.operation.name"] == "set"
+
+ # An internal service (router) gets no db.* keys.
+ internal = GenAIMapper().map(ServiceSpanData("router", call_type="acompletion"))
+ assert internal[LiteLLM.SERVICE_NAME] == "router"
+ assert "db.system.name" not in internal
+
+
+def test_legacy_mapper_all_request_params():
+ attrs = LegacyMapper().map(_full_llm_call())
+ assert attrs["llm.top_k"] == 40
+ assert attrs["llm.frequency_penalty"] == 0.1
+ assert attrs["llm.presence_penalty"] == 0.2
+ assert attrs["llm.chat.stop_sequences"] == ["STOP"]
+ assert attrs["gen_ai.usage.total_tokens"] == 15
+
+
+def test_legacy_mapper_covers_service_with_v1_bare_keys():
+ """Service spans dual-emit V1's bare ``service``/``call_type``/``error`` keys."""
+ attrs = LegacyMapper().map(
+ ServiceSpanData("redis", call_type="set", event_metadata={"k": "v"}),
+ )
+ assert attrs["service"] == "redis"
+ assert attrs["call_type"] == "set"
+ assert attrs["k"] == "v" # event_metadata is stamped bare (V1 behavior)
+
+
+def test_legacy_mapper_skips_guardrail_role():
+ """Guardrail spans never had a V1 vocabulary; legacy mapper returns ``{}``."""
+ assert LegacyMapper().map(GuardrailSpanData("presidio")) == {}
+
+
+# --- metrics ---------------------------------------------------------------- #
+
+
+def test_create_genai_metrics_records():
+ reader = InMemoryMetricReader()
+ meter = MeterProvider(metric_readers=[reader]).get_meter("test")
+ metrics = create_genai_metrics(meter)
+ metrics.token_usage.record(10, {"x": "y"})
+ metrics.operation_duration.record(0.5, {"x": "y"})
+ data = reader.get_metrics_data()
+ assert data is not None
+
+
+# --- context + baggage helpers ---------------------------------------------- #
+
+
+def test_extract_traceparent():
+ valid = {"traceparent": "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"}
+ assert ctx_mod.extract_traceparent(valid) is not None
+ assert ctx_mod.extract_traceparent({"x": "y"}) is None
+
+
+def test_set_request_baggage_empty_returns_context():
+ assert ctx_mod.set_request_baggage({}) is not None
+
+
+def test_get_baggage_attributes_roundtrip():
+ ctx = ctx_mod.set_request_baggage({"litellm.team.id": "t1"})
+ assert ctx_mod.get_baggage_attributes(ctx)["litellm.team.id"] == "t1"
+
+
+# --- providers -------------------------------------------------------------- #
+
+
+def test_to_otel_span_kind_covers_all():
+ assert providers.to_otel_span_kind(LiteLLMSpanKind.SERVER) is SpanKind.SERVER
+ assert providers.to_otel_span_kind(LiteLLMSpanKind.CLIENT) is SpanKind.CLIENT
+ assert providers.to_otel_span_kind(LiteLLMSpanKind.INTERNAL) is SpanKind.INTERNAL
+ assert providers.to_otel_span_kind(LiteLLMSpanKind.PRODUCER) is SpanKind.PRODUCER
+ assert providers.to_otel_span_kind(LiteLLMSpanKind.CONSUMER) is SpanKind.CONSUMER
+
+
+def test_parse_headers():
+ assert providers.parse_headers(None) == {}
+ assert providers.parse_headers("a=1,b=2") == {"a": "1", "b": "2"}
+ assert providers.parse_headers("no-equals") == {}
+
+
+def test_otlp_traces_endpoint_normalization():
+ norm = providers._otlp_traces_endpoint
+ # A base endpoint gets the signal path appended (the common OTLP env shape).
+ assert norm("http://collector:4318") == "http://collector:4318/v1/traces"
+ assert norm("http://collector:4318/") == "http://collector:4318/v1/traces"
+ # An already-correct path is left intact.
+ assert norm("http://collector:4318/v1/traces") == "http://collector:4318/v1/traces"
+ # Another signal's path is rewritten to traces.
+ assert norm("http://collector:4318/v1/logs") == "http://collector:4318/v1/traces"
+ # Splunk's path is preserved; None passes through.
+ assert (
+ norm("https://x.splunk.com/v2/trace/otlp")
+ == "https://x.splunk.com/v2/trace/otlp"
+ )
+ assert norm(None) is None
+
+
+def test_build_span_exporter_variants():
+ assert isinstance(
+ providers.build_span_exporter(OpenTelemetryV2Config(exporter="console")),
+ ConsoleSpanExporter,
+ )
+ assert isinstance(
+ providers.build_span_exporter(OpenTelemetryV2Config(exporter="in_memory")),
+ InMemorySpanExporter,
+ )
+ assert isinstance(
+ providers.build_span_exporter(OpenTelemetryV2Config(exporter="unknown")),
+ ConsoleSpanExporter,
+ )
+ http_exporter = providers.build_span_exporter(
+ OpenTelemetryV2Config(exporter="otlp_http", endpoint="http://h:4318")
+ )
+ assert "OTLPSpanExporter" in type(http_exporter).__name__
+ grpc_exporter = providers.build_span_exporter(
+ OpenTelemetryV2Config(exporter="otlp_grpc", endpoint="http://h:4317")
+ )
+ assert "OTLPSpanExporter" in type(grpc_exporter).__name__
+
+
+def test_build_resource_includes_deployment_environment():
+ resource = providers.build_resource(
+ OpenTelemetryV2Config(service_name="svc", deployment_environment="prod")
+ )
+ assert resource.attributes["service.name"] == "svc"
+ assert resource.attributes["deployment.environment"] == "prod"
+
+
+def test_build_tracer_provider_processor_selection():
+ cfg = OpenTelemetryV2Config(exporter="in_memory")
+ simple = providers.build_tracer_provider(cfg, exporter=InMemorySpanExporter())
+ batch = providers.build_tracer_provider(
+ cfg, exporter=ConsoleSpanExporter(), use_simple_processor=False
+ )
+ # both build without error; assert the requested processor type was used
+ simple_procs = simple._active_span_processor._span_processors
+ batch_procs = batch._active_span_processor._span_processors
+ assert any(isinstance(p, SimpleSpanProcessor) for p in simple_procs)
+ assert any(isinstance(p, BatchSpanProcessor) for p in batch_procs)
+
+
+def test_baggage_processor_lifecycle_noops():
+ proc = providers.LiteLLMBaggageSpanProcessor(allowed_keys=["litellm.team.id"])
+ # no-op lifecycle hooks must not raise
+ assert proc.on_end(None) is None # type: ignore[arg-type]
+ assert proc.shutdown() is None
+ assert proc.force_flush() is True
+
+
+def test_emitter_without_call_id_is_not_deduped():
+ from litellm.integrations.otel.emitter import SpanEmitter
+
+ cfg = OpenTelemetryV2Config(exporter="in_memory")
+ provider, exporter = providers.in_memory_provider(cfg)
+ engine = SpanEmitter(providers.get_tracer(provider, "t"), cfg)
+ data = LLMCallSpanData(
+ operation=GenAIOperation.CHAT,
+ provider="openai",
+ request_model="gpt-4o",
+ response_model=None,
+ response_id=None,
+ request_params=LLMRequestParams(),
+ usage=LLMUsage(),
+ finish_reasons=(),
+ error=SpanError(error_type="X", message=None),
+ response_cost=None,
+ server=None,
+ identity=RequestIdentity(call_id=None),
+ )
+ engine.emit(SpanRole.LLM_CALL, data)
+ engine.emit(SpanRole.LLM_CALL, data) # no call_id -> not deduped
+ assert len(exporter.get_finished_spans()) == 2
+
+
+# --- service taxonomy: which calls become spans, and of what kind ----------- #
+
+
+def test_span_role_for_service_classifies_datastores_internal_and_metrics_only():
+ # Outbound datastores -> DB_CALL (CLIENT), with a db.system.
+ for name in (
+ "redis",
+ "postgres",
+ "batch_write_to_db",
+ "redis_daily_spend_update_queue",
+ ):
+ assert span_role_for_service(name) is SpanRole.DB_CALL
+ assert db_system(name) is not None
+ # Genuine internal work worth a span -> SERVICE (INTERNAL).
+ assert span_role_for_service("reset_budget_job") is SpanRole.SERVICE
+ assert db_system("reset_budget_job") is None
+ # Framework instrumentation that duplicates a gen-AI span (or gets a live
+ # phase span) -> None: never emitted as a service span.
+ for name in ("self", "router", "proxy_pre_call", "auth"):
+ assert span_role_for_service(name) is None
+
+
+# --- event_metadata sanitization -------------------------------------------- #
+
+
+def test_sanitize_event_metadata_drops_objects_dumps_and_secrets():
+ from litellm.integrations.otel.model.payloads import sanitize_event_metadata
+
+ clean = sanitize_event_metadata(
+ {
+ "table_name": "combined_view", # safe primitive -> kept
+ "count": 3, # primitive -> kept (stringified)
+ "function_kwargs": {"prisma_client": object()}, # denylisted key
+ "function_args": (1, 2), # denylisted key
+ "user_api_key_auth": "blob", # 'auth' substring -> dropped
+ "api_key": "sk-secret", # 'api_key' substring -> dropped
+ "set-cookie": "x", # 'cookie' substring -> dropped
+ "hidden_params": "headers...", # denylisted substring
+ "obj": object(), # non-primitive value -> dropped
+ "nested": {"x": 1}, # non-primitive value -> dropped
+ }
+ )
+ assert clean == {"table_name": "combined_view", "count": "3"}
+
+
+def test_sanitize_event_metadata_caps_value_length_and_handles_none():
+ from litellm.integrations.otel.model.payloads import sanitize_event_metadata
+
+ assert sanitize_event_metadata(None) == {}
+ big = sanitize_event_metadata({"k": "v" * 5000})
+ assert len(big["k"]) == 1024
diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_config_baggage_parenting_guardrails.py b/tests/test_litellm/integrations/otel/test_otel_v2_config_baggage_parenting_guardrails.py
new file mode 100644
index 00000000000..dcaff3c911a
--- /dev/null
+++ b/tests/test_litellm/integrations/otel/test_otel_v2_config_baggage_parenting_guardrails.py
@@ -0,0 +1,237 @@
+"""Behavior of three V2 OTel instrumentation areas:
+
+1. Baggage allowlists are configurable via env vars and config.yaml
+ (``callback_settings.otel.*``), not just hard-coded.
+2. Pass-through LLM-call spans nest under the proxy server span because they are
+ opened at the ``pre_call`` boundary in the request task (where the server span
+ is ambient) — no span threaded through metadata.
+3. Guardrail span data is built from the typed
+ ``StandardLoggingGuardrailInformation`` shape (provider-agnostic), not from
+ one provider's assumed field names.
+"""
+
+import asyncio
+
+import pytest
+
+pytest.importorskip("opentelemetry")
+
+from opentelemetry import trace # noqa: E402
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( # noqa: E402
+ InMemorySpanExporter,
+)
+
+from litellm.integrations.otel import LiteLLM, OpenTelemetryV2Config # noqa: E402
+from litellm.integrations.otel.plumbing import providers # noqa: E402
+from litellm.integrations.otel.model.baggage import ( # noqa: E402
+ BAGGAGE_PROMOTED_KEYS,
+ DEFAULT_BAGGAGE_METADATA_KEYS,
+)
+from litellm.integrations.otel.logger import OpenTelemetryV2 # noqa: E402
+from litellm.integrations.otel.model.payloads import GuardrailSpanData # noqa: E402
+from litellm.integrations.otel.model.spans import ( # noqa: E402
+ LITELLM_PROXY_REQUEST_SPAN_NAME,
+ SpanRole,
+)
+
+# --------------------------------------------------------------------------- #
+# Area 1 — baggage allowlists configurable
+# --------------------------------------------------------------------------- #
+
+
+def test_baggage_keys_default_when_unset():
+ cfg = OpenTelemetryV2Config()
+ assert cfg.baggage_promoted_keys == list(BAGGAGE_PROMOTED_KEYS)
+ assert cfg.baggage_metadata_keys == list(DEFAULT_BAGGAGE_METADATA_KEYS)
+
+
+def test_baggage_promoted_keys_from_env_csv(monkeypatch):
+ monkeypatch.setenv(
+ "LITELLM_OTEL_BAGGAGE_PROMOTED_KEYS",
+ f"{LiteLLM.TEAM_ID}, {LiteLLM.KEY_HASH}",
+ )
+ monkeypatch.setenv(
+ "LITELLM_OTEL_BAGGAGE_METADATA_KEYS",
+ "user_api_key_user_id,requester_ip_address",
+ )
+ cfg = OpenTelemetryV2Config()
+ # Whitespace around comma-separated entries is trimmed.
+ assert cfg.baggage_promoted_keys == [LiteLLM.TEAM_ID, LiteLLM.KEY_HASH]
+ assert cfg.baggage_metadata_keys == [
+ "user_api_key_user_id",
+ "requester_ip_address",
+ ]
+
+
+def test_baggage_keys_from_config_yaml_kwargs():
+ """``callback_settings.otel.*`` reaches the config through the logger kwargs."""
+ logger = OpenTelemetryV2(
+ baggage_promoted_keys=[LiteLLM.TEAM_ALIAS],
+ baggage_metadata_keys=["user_api_key_alias"],
+ )
+ assert logger.config.baggage_promoted_keys == [LiteLLM.TEAM_ALIAS]
+ assert logger.config.baggage_metadata_keys == ["user_api_key_alias"]
+
+
+def test_baggage_processor_allowlist_uses_config_keys():
+ cfg = OpenTelemetryV2Config(
+ exporter="in_memory", baggage_promoted_keys=[LiteLLM.TEAM_ID]
+ )
+ provider, exporter = providers.in_memory_provider(cfg)
+ from litellm.integrations.otel.plumbing import context as ctx_mod
+ from litellm.integrations.otel.emitter import SpanEmitter
+ from litellm.integrations.otel.model.payloads import ServiceSpanData
+
+ engine = SpanEmitter(providers.get_tracer(provider, "t"), cfg)
+ ctx = ctx_mod.set_request_baggage({LiteLLM.TEAM_ID: "t1", LiteLLM.TEAM_ALIAS: "ta"})
+ engine.emit(SpanRole.SERVICE, ServiceSpanData("redis"), ctx)
+ (span,) = exporter.get_finished_spans()
+ assert span.attributes.get(LiteLLM.TEAM_ID) == "t1"
+ assert LiteLLM.TEAM_ALIAS not in span.attributes # not in this allowlist
+
+
+# --------------------------------------------------------------------------- #
+# Area 2 — pass-through LLM span parents to the ambient server span
+# --------------------------------------------------------------------------- #
+
+
+def _logger():
+ cfg = OpenTelemetryV2Config(exporter="in_memory")
+ exporter = InMemorySpanExporter()
+ tracer_provider = providers.build_tracer_provider(cfg, exporter=exporter)
+ return OpenTelemetryV2(config=cfg, tracer_provider=tracer_provider), exporter
+
+
+def _payload():
+ return {
+ "call_type": "pass_through_endpoint",
+ "custom_llm_provider": "openai",
+ "model": "gpt-4o",
+ "prompt_tokens": 1,
+ "completion_tokens": 1,
+ "total_tokens": 2,
+ "status": "success",
+ "litellm_call_id": "call_pt",
+ "metadata": {},
+ "hidden_params": {},
+ }
+
+
+def test_passthrough_llm_span_parents_to_ambient_server_span():
+ """Pass-through calls ``logging_obj.pre_call`` in the request task, where the
+ server span is the ambient context — so the LLM-call span is opened there and
+ parents to it natively, with no ``litellm_parent_otel_span`` threading. The
+ later (possibly detached) success callback only closes the already-parented
+ span, so it never becomes a separate root trace."""
+ logger, exporter = _logger()
+ server = logger._emitter.start_span(
+ SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME
+ )
+ kwargs = {
+ "standard_logging_object": _payload(),
+ "litellm_params": {"metadata": {}},
+ }
+ # pre_call runs in the request task (server span ambient); success closes it.
+ with trace.use_span(server, end_on_exit=False):
+ logger.log_pre_api_call(model="gpt-4o", messages=[], kwargs=kwargs)
+ asyncio.run(logger.async_log_success_event(kwargs, None, None, None))
+ server.end()
+
+ by_name = {s.name: s for s in exporter.get_finished_spans()}
+ llm_span = by_name["chat gpt-4o"]
+ assert llm_span.parent is not None
+ assert llm_span.parent.span_id == server.get_span_context().span_id
+
+
+def test_llm_span_unaffected_by_phase_span_active_at_close():
+ """The LLM-call span's parent is captured at the ``pre_call`` boundary (under
+ the server span), so a phase span (e.g. ``auth``) that happens to be ambient
+ when the *close* callback fires can't re-parent it. This is the structural
+ successor to the old auth-failure-401 case where the LLM log nested under
+ ``auth``: the span is now born after auth, parented to the request root."""
+ logger, exporter = _logger()
+ server = logger._emitter.start_span(
+ SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME
+ )
+ kwargs = {
+ "standard_logging_object": _payload(),
+ "litellm_params": {"metadata": {}},
+ }
+ with trace.use_span(server, end_on_exit=False):
+ logger.log_pre_api_call(model="gpt-4o", messages=[], kwargs=kwargs)
+ # A phase span is ambient when the close callback fires — must not re-parent.
+ phase = logger._emitter.start_span(SpanRole.SERVICE, "auth /v1/chat/completions")
+ with trace.use_span(phase, end_on_exit=False):
+ asyncio.run(logger.async_log_success_event(kwargs, None, None, None))
+ phase.end()
+ server.end()
+ by_name = {s.name: s for s in exporter.get_finished_spans()}
+ llm_span = by_name["chat gpt-4o"]
+ assert llm_span.parent.span_id == server.get_span_context().span_id
+
+
+# --------------------------------------------------------------------------- #
+# Area 3 — typed, provider-agnostic guardrail span data
+# --------------------------------------------------------------------------- #
+
+
+def test_guardrail_mode_enum_normalized_to_value():
+ from litellm.types.guardrails import GuardrailEventHooks
+
+ d = GuardrailSpanData.from_logging_entry(
+ {
+ "guardrail_name": "bedrock-guardrail",
+ "guardrail_mode": GuardrailEventHooks.pre_call,
+ "guardrail_status": "success",
+ }
+ )
+ # The enum *value* ("pre_call"), not "GuardrailEventHooks.pre_call".
+ assert d.mode == "pre_call"
+
+
+def test_guardrail_mode_list_of_enums_joined():
+ from litellm.types.guardrails import GuardrailEventHooks
+
+ d = GuardrailSpanData.from_logging_entry(
+ {
+ "guardrail_name": "g",
+ "guardrail_mode": [
+ GuardrailEventHooks.pre_call,
+ GuardrailEventHooks.post_call,
+ ],
+ "guardrail_status": "success",
+ }
+ )
+ assert d.mode == "pre_call,post_call"
+
+
+def test_guardrail_typed_metadata_fields_mapped_to_span():
+ from litellm.integrations.otel.mappers.genai import GenAIMapper
+
+ d = GuardrailSpanData.from_logging_entry(
+ {
+ "guardrail_name": "eu-pii",
+ "guardrail_status": "success",
+ "guardrail_id": "gd-eu-pii-001",
+ "policy_template": "EU AI Act Article 5",
+ "detection_method": "presidio",
+ }
+ )
+ assert d.guardrail_id == "gd-eu-pii-001"
+ assert d.policy_template == "EU AI Act Article 5"
+ assert d.detection_method == "presidio"
+ attrs = GenAIMapper().map(d)
+ assert attrs[LiteLLM.GUARDRAIL_ID] == "gd-eu-pii-001"
+ assert attrs[LiteLLM.GUARDRAIL_POLICY_TEMPLATE] == "EU AI Act Article 5"
+ assert attrs[LiteLLM.GUARDRAIL_DETECTION_METHOD] == "presidio"
+
+
+def test_guardrail_ignores_non_canonical_provider_keys():
+ """Only canonical ``StandardLoggingGuardrailInformation`` keys are read; a
+ provider's ad-hoc bare ``name``/``status``/``mode`` keys are not assumed."""
+ d = GuardrailSpanData.from_logging_entry(
+ {"name": "bare", "status": "blocked", "mode": "pre"} # type: ignore[typeddict-unknown-key]
+ )
+ assert d.guardrail_name == "guardrail" # fell back to the default
+ assert d.status is None
+ assert d.mode is None
diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_dynamic.py b/tests/test_litellm/integrations/otel/test_otel_v2_dynamic.py
new file mode 100644
index 00000000000..1150c2c51c3
--- /dev/null
+++ b/tests/test_litellm/integrations/otel/test_otel_v2_dynamic.py
@@ -0,0 +1,131 @@
+"""Per-request multi-tenant credential routing (V1 parity)."""
+
+import os
+import sys
+
+sys.path.insert(0, os.path.abspath("../../../.."))
+
+from opentelemetry.trace import NoOpTracer
+
+from litellm.integrations.otel.model.config import ExporterSpec, OpenTelemetryV2Config
+from litellm.integrations.otel.presets import dynamic_otlp_headers
+from litellm.integrations.otel.plumbing.routing import TenantTracerCache
+
+
+def _cache(callback_name, exporters=None):
+ cfg = OpenTelemetryV2Config(exporters=exporters or [ExporterSpec(kind="in_memory")])
+ return TenantTracerCache(cfg, callback_name, "litellm")
+
+
+# --- header builders mirror the V1 construct_dynamic_otel_headers overrides --- #
+
+
+def test_arize_dynamic_headers():
+ headers = dynamic_otlp_headers(
+ "arize", {"arize_space_id": "S", "arize_api_key": "K"}
+ )
+ assert headers == {"arize-space-id": "S", "api_key": "K"}
+
+
+def test_arize_space_key_overrides_space_id():
+ headers = dynamic_otlp_headers(
+ "arize", {"arize_space_id": "S", "arize_space_key": "SK"}
+ )
+ assert headers == {"arize-space-id": "SK"}
+
+
+def test_langfuse_dynamic_headers_need_both_keys():
+ assert dynamic_otlp_headers("langfuse_otel", {"langfuse_public_key": "pk"}) is None
+ headers = dynamic_otlp_headers(
+ "langfuse_otel", {"langfuse_public_key": "pk", "langfuse_secret_key": "sk"}
+ )
+ assert headers is not None and "Authorization" in headers
+
+
+def test_weave_dynamic_headers():
+ headers = dynamic_otlp_headers(
+ "weave_otel", {"wandb_api_key": "w", "weave_project_id": "p"}
+ )
+ assert headers is not None
+ assert "Authorization" in headers and headers["project_id"] == "p"
+
+
+def test_non_participating_callbacks_have_no_routing():
+ # Phoenix subclasses the base in V1 (no override) → no dynamic routing.
+ assert dynamic_otlp_headers("arize_phoenix", {"arize_api_key": "K"}) is None
+ assert dynamic_otlp_headers("langtrace", {"arize_api_key": "K"}) is None
+ assert dynamic_otlp_headers(None, {"arize_api_key": "K"}) is None
+
+
+def test_no_dynamic_params_is_no_routing():
+ assert dynamic_otlp_headers("arize", None) is None
+ assert dynamic_otlp_headers("arize", {}) is None
+
+
+# --- TenantTracerCache routes + caches a TracerProvider per credential set --- #
+
+
+def test_provider_cached_per_credential_set():
+ cache = _cache("arize")
+ default = NoOpTracer()
+ creds_a = {"arize_space_id": "S", "arize_api_key": "K"}
+ creds_b = {"arize_space_id": "S2", "arize_api_key": "K2"}
+
+ cache.tracer_for(default, creds_a)
+ cache.tracer_for(default, creds_a) # same set → reuse, no new provider
+ assert len(cache._providers) == 1
+ cache.tracer_for(default, creds_b) # new set → new provider
+ assert len(cache._providers) == 2
+
+
+def test_provider_cache_is_bounded_and_evicts_lru(monkeypatch):
+ # The cache key derives from request-supplied dynamic credentials, so it
+ # must be bounded — an unbounded cache lets a caller spawn one provider (and
+ # its background exporter thread) per unique credential set. On overflow the
+ # least-recently-used provider is evicted and shut down.
+ from litellm.integrations.otel.plumbing import routing as routing_mod
+
+ monkeypatch.setattr(routing_mod, "_MAX_CACHED_PROVIDERS", 2)
+ shut_down = []
+ monkeypatch.setattr(
+ routing_mod, "_shutdown_provider", lambda p: shut_down.append(p)
+ )
+
+ cache = _cache("arize")
+ default = NoOpTracer()
+
+ def creds(space):
+ return {"arize_space_id": space, "arize_api_key": "K"}
+
+ cache.tracer_for(default, creds("1"))
+ cache.tracer_for(default, creds("2"))
+ cache.tracer_for(default, creds("1")) # touch "1" → "2" is now LRU
+ cache.tracer_for(default, creds("3")) # overflow → evict "2"
+
+ assert len(cache._providers) == 2
+ assert len(shut_down) == 1 # exactly the evicted provider was shut down
+
+
+def test_no_dynamic_params_uses_default_tracer():
+ cache = _cache("arize")
+ default = NoOpTracer()
+ assert cache.tracer_for(default, {}) is default
+ assert cache._providers == {}
+
+
+def test_non_participating_callback_uses_default_tracer():
+ cache = _cache("arize_phoenix")
+ default = NoOpTracer()
+ assert cache.tracer_for(default, {"arize_api_key": "K"}) is default
+ assert cache._providers == {}
+
+
+def test_dynamic_headers_applied_to_otlp_exporter_only():
+ cache = _cache(
+ "arize",
+ exporters=[ExporterSpec(kind="otlp_http"), ExporterSpec(kind="in_memory")],
+ )
+ new_cfg = cache._config_with_headers({"arize-space-id": "S", "api_key": "K"})
+ otlp, in_mem = new_cfg.exporters
+ assert otlp.headers == "arize-space-id=S,api_key=K"
+ assert in_mem.headers is None # console/in_memory left untouched
diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py b/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py
new file mode 100644
index 00000000000..2dbedda1ab6
--- /dev/null
+++ b/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py
@@ -0,0 +1,229 @@
+"""Golden tests for the OTel v2 engine: span shape, kinds, semconv attributes,
+legacy dual-emit, hierarchy, error status, and idempotency. Needs the OTel SDK."""
+
+import pytest
+
+pytest.importorskip("opentelemetry")
+
+from opentelemetry.trace import SpanKind # noqa: E402
+from opentelemetry.trace.status import StatusCode # noqa: E402
+
+from litellm.integrations.otel import ( # noqa: E402
+ GenAI,
+ LiteLLM,
+ OpenTelemetryV2Config,
+)
+from litellm.integrations.otel.plumbing import context as ctx_mod # noqa: E402
+from litellm.integrations.otel.plumbing import providers # noqa: E402
+from litellm.integrations.otel.emitter import SpanEmitter # noqa: E402
+from litellm.integrations.otel.model.payloads import ( # noqa: E402
+ GuardrailSpanData,
+ LLMCallSpanData,
+ ServiceSpanData,
+)
+from litellm.integrations.otel.model.spans import SPAN_REGISTRY, SpanRole # noqa: E402
+
+
+def _payload(**overrides):
+ payload = {
+ "call_type": "acompletion",
+ "custom_llm_provider": "openai",
+ "model": "gpt-4o",
+ "prompt_tokens": 10,
+ "completion_tokens": 5,
+ "total_tokens": 15,
+ "stream": False,
+ "model_parameters": {"temperature": 0.7, "max_tokens": 256, "top_k": 40},
+ "response": {
+ "id": "resp_1",
+ "model": "gpt-4o-2024",
+ "choices": [{"finish_reason": "stop"}],
+ },
+ "metadata": {"team_id": "t1", "team_alias": "team one"},
+ "api_base": "https://api.openai.com:443/v1",
+ "status": "success",
+ "litellm_call_id": "call_1",
+ "response_cost": 0.002,
+ "hidden_params": {},
+ }
+ payload.update(overrides)
+ return payload
+
+
+def _engine(legacy_compat=True):
+ cfg = OpenTelemetryV2Config(exporter="in_memory", legacy_compat=legacy_compat)
+ provider, exporter = providers.in_memory_provider(cfg)
+ tracer = providers.get_tracer(provider, "litellm-test")
+ return SpanEmitter(tracer, cfg), exporter
+
+
+def test_llm_call_span_golden():
+ engine, exporter = _engine()
+ data = LLMCallSpanData.from_standard_logging_payload(_payload())
+ engine.emit(SpanRole.LLM_CALL, data)
+ (span,) = exporter.get_finished_spans()
+ assert span.name == "chat gpt-4o"
+ assert span.kind is SpanKind.CLIENT
+ a = span.attributes
+ assert a[GenAI.OPERATION_NAME] == "chat"
+ assert a[GenAI.PROVIDER_NAME] == "openai"
+ assert a[GenAI.REQUEST_MODEL] == "gpt-4o"
+ assert a[GenAI.RESPONSE_MODEL] == "gpt-4o-2024"
+ assert a[GenAI.RESPONSE_ID] == "resp_1"
+ assert a[GenAI.USAGE_INPUT_TOKENS] == 10
+ assert a[GenAI.USAGE_OUTPUT_TOKENS] == 5
+ assert a[GenAI.RESPONSE_FINISH_REASONS] == ("stop",)
+ assert a[GenAI.REQUEST_TEMPERATURE] == 0.7
+ assert a["server.address"] == "api.openai.com"
+ assert a[LiteLLM.CALL_ID] == "call_1"
+ assert a["litellm.cost.total"] == 0.002
+ # Success leaves status UNSET (semconv default), not forced OK.
+ assert span.status.status_code is StatusCode.UNSET
+
+
+def test_legacy_dual_emit_on():
+ engine, exporter = _engine(legacy_compat=True)
+ engine.emit(
+ SpanRole.LLM_CALL, LLMCallSpanData.from_standard_logging_payload(_payload())
+ )
+ (span,) = exporter.get_finished_spans()
+ # canonical AND legacy keys are both present
+ assert span.attributes[GenAI.USAGE_OUTPUT_TOKENS] == 5
+ assert span.attributes["gen_ai.usage.completion_tokens"] == 5
+ assert span.attributes["gen_ai.system"] == "openai"
+
+
+def test_legacy_dual_emit_off():
+ engine, exporter = _engine(legacy_compat=False)
+ engine.emit(
+ SpanRole.LLM_CALL, LLMCallSpanData.from_standard_logging_payload(_payload())
+ )
+ (span,) = exporter.get_finished_spans()
+ # canonical present, legacy absent
+ assert span.attributes[GenAI.USAGE_OUTPUT_TOKENS] == 5
+ assert "gen_ai.usage.completion_tokens" not in span.attributes
+ assert "gen_ai.system" not in span.attributes
+
+
+def test_error_span_sets_status_and_error_type():
+ engine, exporter = _engine()
+ payload = _payload(
+ status="failure",
+ error_information={"error_class": "RateLimitError", "error_message": "429"},
+ )
+ engine.emit(
+ SpanRole.LLM_CALL, LLMCallSpanData.from_standard_logging_payload(payload)
+ )
+ (span,) = exporter.get_finished_spans()
+ assert span.status.status_code is StatusCode.ERROR
+ assert span.attributes["error.type"] == "RateLimitError"
+
+
+def test_hierarchy_and_kinds_match_registry():
+ engine, exporter = _engine()
+ data = LLMCallSpanData.from_standard_logging_payload(_payload())
+ root = engine.start_span(SpanRole.PROXY_REQUEST, "POST /chat/completions")
+ root_ctx = ctx_mod.context_from_span(root)
+ engine.emit(SpanRole.LLM_CALL, data, parent_context=root_ctx)
+ engine.emit(
+ SpanRole.GUARDRAIL, GuardrailSpanData("presidio", status="success"), root_ctx
+ )
+ # An outbound datastore call (DB_CALL) and an internal service call differ in
+ # span kind; both are named "{service} {call_type}".
+ engine.emit(SpanRole.DB_CALL, ServiceSpanData("redis", call_type="set"), root_ctx)
+ engine.emit(
+ SpanRole.SERVICE, ServiceSpanData("router", call_type="acompletion"), root_ctx
+ )
+ root.end()
+
+ by_name = {s.name: s for s in exporter.get_finished_spans()}
+ root_id = root.get_span_context().span_id
+ assert by_name["chat gpt-4o"].parent.span_id == root_id
+ assert by_name["execute_guardrail presidio"].parent.span_id == root_id
+ assert by_name["redis set"].parent.span_id == root_id
+ assert by_name["router acompletion"].parent.span_id == root_id
+ # kinds come straight from the registry
+ assert by_name["chat gpt-4o"].kind is SpanKind.CLIENT
+ assert by_name["execute_guardrail presidio"].kind is SpanKind.INTERNAL
+ assert by_name["redis set"].kind is SpanKind.CLIENT
+ assert by_name["router acompletion"].kind is SpanKind.INTERNAL
+ assert by_name["POST /chat/completions"].kind is SpanKind.SERVER
+
+
+def test_idempotent_dual_fire():
+ engine, exporter = _engine()
+ data = LLMCallSpanData.from_standard_logging_payload(_payload())
+ first = engine.emit(SpanRole.LLM_CALL, data)
+ second = engine.emit(SpanRole.LLM_CALL, data) # same call_id -> deduped
+ assert first is not None
+ assert second is None
+ assert len(exporter.get_finished_spans()) == 1
+
+
+def test_dedup_cache_is_bounded(monkeypatch):
+ """The dedup cache only needs to coalesce one request's sync+async fire, so
+ it is a bounded LRU — every unique call_id must not accumulate forever on a
+ long-running proxy."""
+ from litellm.integrations.otel import emitter as emitter_mod
+
+ monkeypatch.setattr(emitter_mod, "_DEDUP_CACHE_MAX", 3)
+ engine, _ = _engine()
+ for i in range(10):
+ engine.emit(
+ SpanRole.LLM_CALL,
+ LLMCallSpanData.from_standard_logging_payload(
+ _payload(litellm_call_id=f"call_{i}")
+ ),
+ )
+ assert len(engine._emitted) <= 3
+
+
+def test_service_error_span():
+ from litellm.integrations.otel.model.payloads import SpanError
+
+ engine, exporter = _engine()
+ engine.emit(
+ SpanRole.SERVICE,
+ ServiceSpanData(
+ "postgres", call_type="query", error=SpanError("DBError", "boom")
+ ),
+ )
+ (span,) = exporter.get_finished_spans()
+ assert span.status.status_code is StatusCode.ERROR
+ assert span.attributes["error.type"] == "DBError"
+ assert span.attributes[LiteLLM.SERVICE_NAME] == "postgres"
+
+
+def test_guardrail_block_span_is_error_and_carries_verdict():
+ engine, exporter = _engine()
+ data = GuardrailSpanData.from_logging_entry(
+ {
+ "guardrail_name": "openai-moderation",
+ "guardrail_mode": "pre_call",
+ "guardrail_status": "guardrail_intervened",
+ "guardrail_provider": "openai",
+ "guardrail_response": {"violated_categories": ["violence"]},
+ "masked_entity_count": {"EMAIL": 2},
+ }
+ )
+ engine.emit(SpanRole.GUARDRAIL, data)
+ (span,) = exporter.get_finished_spans()
+ assert span.status.status_code is StatusCode.ERROR # intervention → ERROR
+ a = span.attributes
+ assert a[LiteLLM.GUARDRAIL_STATUS] == "guardrail_intervened"
+ assert a[LiteLLM.GUARDRAIL_PROVIDER] == "openai"
+ assert "violence" in a[LiteLLM.GUARDRAIL_RESPONSE] # the verdict rides the span
+ assert a[LiteLLM.GUARDRAIL_MASKED_ENTITY_COUNT] == 2
+
+
+def test_guardrail_success_span_is_unset():
+ """On success the status is left UNSET (semconv default) — not forced OK."""
+ engine, exporter = _engine()
+ engine.emit(
+ SpanRole.GUARDRAIL,
+ GuardrailSpanData.from_logging_entry(
+ {"guardrail_name": "g", "guardrail_status": "success"}
+ ),
+ )
+ (span,) = exporter.get_finished_spans()
+ assert span.status.status_code is StatusCode.UNSET
diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_logger.py b/tests/test_litellm/integrations/otel/test_otel_v2_logger.py
new file mode 100644
index 00000000000..c1cde55c58e
--- /dev/null
+++ b/tests/test_litellm/integrations/otel/test_otel_v2_logger.py
@@ -0,0 +1,1122 @@
+"""Tests for the V2 ``OpenTelemetryV2`` CustomLogger adapter.
+
+Exercises the callback surface the existing call sites use: the LLM-call span
+opened at the ``pre_call`` boundary and closed at async success/failure, service
+hooks, proxy SERVER span lifecycle (start + setters), parent-context resolution
+(ambient context), and Baggage promotion onto child spans.
+"""
+
+import asyncio
+import contextlib
+from datetime import datetime, timezone
+
+import pytest
+
+pytest.importorskip("opentelemetry")
+
+from opentelemetry import trace # noqa: E402
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( # noqa: E402
+ InMemorySpanExporter,
+)
+from opentelemetry.trace import SpanKind # noqa: E402
+from opentelemetry.trace.status import StatusCode # noqa: E402
+
+from litellm.integrations.otel import ( # noqa: E402
+ GenAI,
+ LiteLLM,
+ OpenTelemetryV2Config,
+)
+from litellm.integrations.otel.plumbing import providers # noqa: E402
+from litellm.integrations.otel.plumbing.context import (
+ set_request_root_span,
+) # noqa: E402
+from litellm.integrations.otel.logger import OpenTelemetryV2 # noqa: E402
+from litellm.integrations.otel.model.spans import ( # noqa: E402
+ LITELLM_PROXY_REQUEST_SPAN_NAME,
+ SpanRole,
+)
+from litellm.integrations.otel.model.utils import to_ns, to_seconds # noqa: E402
+
+# --------------------------------------------------------------------------- #
+# Fixtures
+# --------------------------------------------------------------------------- #
+
+
+@pytest.fixture(autouse=True)
+def _reset_request_root_span():
+ """Clear the request-root-span anchor around every test.
+
+ In production each request runs in its own asyncio task whose context is a
+ fresh copy, so the anchor never leaks between requests. The test process
+ shares one context, so reset it explicitly to keep tests order-independent.
+ """
+ from litellm.integrations.otel.plumbing import context as _otel_context
+
+ _otel_context._request_root_span.set(None)
+ yield
+ _otel_context._request_root_span.set(None)
+
+
+def _payload(**overrides):
+ payload = {
+ "call_type": "acompletion",
+ "custom_llm_provider": "openai",
+ "model": "gpt-4o",
+ "prompt_tokens": 10,
+ "completion_tokens": 5,
+ "total_tokens": 15,
+ "stream": False,
+ "model_parameters": {"temperature": 0.7, "max_tokens": 256},
+ "response": {
+ "id": "resp_1",
+ "model": "gpt-4o-2024",
+ "choices": [{"finish_reason": "stop"}],
+ },
+ "metadata": {
+ "team_id": "t1",
+ "team_alias": "team one",
+ "user_api_key_hash": "hsh",
+ },
+ "api_base": "https://api.openai.com:443/v1",
+ "status": "success",
+ "litellm_call_id": "call_1",
+ "response_cost": 0.002,
+ "hidden_params": {},
+ }
+ payload.update(overrides)
+ return payload
+
+
+def _kwargs(payload=None):
+ return {
+ # ``litellm_call_id`` (here carried inside the payload) correlates the
+ # pre_call boundary with the close callback — the carrier is keyed by it.
+ "standard_logging_object": payload if payload is not None else _payload(),
+ "litellm_params": {"metadata": {}},
+ }
+
+
+def _logger(legacy_compat=True):
+ cfg = OpenTelemetryV2Config(exporter="in_memory", legacy_compat=legacy_compat)
+ exporter = InMemorySpanExporter()
+ tracer_provider = providers.build_tracer_provider(cfg, exporter=exporter)
+ return OpenTelemetryV2(config=cfg, tracer_provider=tracer_provider), exporter
+
+
+def _emit_llm(logger, kwargs=None, *, ambient=None, fail=False):
+ """Drive the real boundary flow: open at ``pre_call`` then close at the async
+ callback. ``ambient``, if given, is the span that is the active OTel context
+ while ``pre_call`` runs (the server span) so the LLM span parents to it."""
+ if kwargs is None:
+ kwargs = _kwargs()
+ payload = kwargs.get("standard_logging_object") or {}
+ with (
+ trace.use_span(ambient, end_on_exit=False)
+ if ambient is not None
+ else contextlib.nullcontext()
+ ):
+ logger.log_pre_api_call(model=payload.get("model"), messages=[], kwargs=kwargs)
+ hook = logger.async_log_failure_event if fail else logger.async_log_success_event
+ asyncio.run(hook(kwargs, None, None, None))
+ return kwargs
+
+
+# --------------------------------------------------------------------------- #
+# Time helpers
+# --------------------------------------------------------------------------- #
+
+
+def test_to_ns_handles_datetime_and_float():
+ dt = datetime(2026, 5, 26, 12, 0, 0, tzinfo=timezone.utc)
+ assert to_ns(dt) == int(dt.timestamp() * 1e9)
+ assert to_ns(1.5) == 1_500_000_000
+ assert to_ns(None) is None
+ assert to_ns(True) is None # bool is rejected — not a real epoch value
+
+
+def test_to_seconds_parses_string_formats():
+ assert to_seconds("2026-05-26 12:00:00.123") is not None
+ assert to_seconds("2026-05-26 12:00:00") is not None
+ assert to_seconds("nonsense") is None
+ assert to_seconds(None) is None
+ assert to_seconds(1.5) == 1.5
+
+
+# --------------------------------------------------------------------------- #
+# LLM-call callbacks
+# --------------------------------------------------------------------------- #
+
+
+def test_async_log_success_event_emits_llm_call_span():
+ logger, exporter = _logger()
+ _emit_llm(logger)
+ (span,) = exporter.get_finished_spans()
+ assert span.name == "chat gpt-4o"
+ assert span.kind is SpanKind.CLIENT
+ assert span.attributes[GenAI.OPERATION_NAME] == "chat"
+ assert span.attributes[GenAI.REQUEST_MODEL] == "gpt-4o"
+ assert span.attributes[LiteLLM.CALL_ID] == "call_1"
+ # Success leaves status UNSET (semconv default), not forced OK.
+ assert span.status.status_code is StatusCode.UNSET
+
+
+def test_async_log_failure_event_marks_error_status():
+ logger, exporter = _logger()
+ payload = _payload(
+ status="failure",
+ error_information={"error_class": "RateLimitError", "error_message": "429"},
+ )
+ _emit_llm(logger, _kwargs(payload=payload), fail=True)
+ (span,) = exporter.get_finished_spans()
+ assert span.status.status_code is StatusCode.ERROR
+ assert span.attributes["error.type"] == "RateLimitError"
+
+
+def test_sync_log_event_is_noop():
+ """V2 closes the span async-only; the sync callback runs out-of-context, so
+ it no-ops (the span stays open on the carrier until the async callback)."""
+ logger, exporter = _logger()
+ kwargs = _kwargs()
+ logger.log_pre_api_call(model="gpt-4o", messages=[], kwargs=kwargs)
+ logger.log_success_event(kwargs, None, None, None)
+ logger.log_failure_event(kwargs, None, None, None)
+ assert exporter.get_finished_spans() == ()
+
+
+def test_missing_standard_logging_object_is_noop():
+ """No carrier (``pre_call`` never ran) → the callback emits nothing."""
+ logger, exporter = _logger()
+ asyncio.run(
+ logger.async_log_success_event({"litellm_params": {}}, None, None, None)
+ )
+ assert exporter.get_finished_spans() == ()
+
+
+def test_no_span_when_pre_call_never_ran():
+ """A request rejected before the upstream call — at the auth/budget gate, or
+ blocked by a pre-call guardrail — never reaches ``pre_call``, so there is no
+ carrier and the failure log produces no phantom CLIENT span. This replaces the
+ old post-hoc heuristics: "did pre_call run?" is the only signal needed."""
+ logger, exporter = _logger()
+ payload = _payload(
+ status="failure",
+ error_information={"error_class": "ProxyException", "error_code": "401"},
+ )
+ # No log_pre_api_call: the call never started.
+ asyncio.run(
+ logger.async_log_failure_event(_kwargs(payload=payload), None, None, None)
+ )
+ assert exporter.get_finished_spans() == () # no phantom LLM span
+
+
+def test_real_llm_failure_still_emitted():
+ """A genuine LLM failure: ``pre_call`` ran (the call was attempted), so the
+ CLIENT span is opened at the boundary and closed ERROR."""
+ logger, exporter = _logger()
+ payload = _payload(
+ status="failure",
+ error_information={"error_class": "RateLimitError", "error_code": "429"},
+ )
+ _emit_llm(logger, _kwargs(payload=payload), fail=True)
+ (span,) = exporter.get_finished_spans()
+ assert span.name == "chat gpt-4o"
+ assert span.status.status_code is StatusCode.ERROR
+
+
+def test_idempotent_on_repeat_callback():
+ """The carrier is the dedup: once the async callback closes the span and
+ clears the carrier, a second callback firing emits nothing."""
+ logger, exporter = _logger()
+ kwargs = _kwargs()
+ logger.log_pre_api_call(model="gpt-4o", messages=[], kwargs=kwargs)
+ asyncio.run(logger.async_log_success_event(kwargs, None, None, None))
+ asyncio.run(logger.async_log_success_event(kwargs, None, None, None))
+ assert len(exporter.get_finished_spans()) == 1
+
+
+def test_pre_call_idempotent_keeps_first_span():
+ """A retried call may re-enter ``pre_call`` with the same call id; the first
+ span (with the true start time) is kept, not replaced."""
+ logger, _ = _logger()
+ kwargs = _kwargs()
+ server = logger._emitter.start_span(
+ SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME
+ )
+ with trace.use_span(server, end_on_exit=False):
+ logger.log_pre_api_call(model="gpt-4o", messages=[], kwargs=kwargs)
+ first = logger._open_llm_calls["call_1"]
+ logger.log_pre_api_call(model="gpt-4o", messages=[], kwargs=kwargs)
+ second = logger._open_llm_calls["call_1"]
+ server.end()
+ assert first is second # not overwritten
+
+
+# --------------------------------------------------------------------------- #
+# Parent resolution — ambient context at the boundary (no metadata threading)
+# --------------------------------------------------------------------------- #
+
+
+def test_llm_span_parents_to_ambient_server_span():
+ """The span is opened at ``pre_call`` while the server span is the active
+ context, so it nests under it natively (no ``litellm_parent_otel_span``)."""
+ logger, exporter = _logger()
+ server = logger._emitter.start_span(
+ SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME
+ )
+ _emit_llm(logger, ambient=server)
+ server.end()
+ by_name = {s.name: s for s in exporter.get_finished_spans()}
+ llm_span = by_name["chat gpt-4o"]
+ assert llm_span.parent is not None
+ assert llm_span.parent.span_id == server.get_span_context().span_id
+
+
+def test_llm_span_is_root_without_ambient_server_span():
+ """No server span at ``pre_call`` → creation is deferred and the span is a
+ root of its own trace (the SDK / no-proxy path)."""
+ logger, exporter = _logger()
+ _emit_llm(logger)
+ (span,) = exporter.get_finished_spans()
+ assert span.parent is None # standalone (no proxy server span) → root
+
+
+# --------------------------------------------------------------------------- #
+# Explicit request-root-span anchor — request-level spans (LLM call, guardrail)
+# parent to the captured server span, NOT to whatever span is momentarily
+# active. Regression cover for the two ambient-only failure modes:
+# * auth: the LLM/guardrail span must not nest under the live ``auth`` span;
+# * pass-through: the span must not orphan when closed off the request task.
+# --------------------------------------------------------------------------- #
+
+
+def test_llm_span_anchors_to_root_even_inside_active_phase_span():
+ """Bug 1: a synthetic error log can fire ``pre_call`` while the ``auth`` phase
+ span is the *active* context. The LLM span must still parent to the request
+ root (the server span), never to the auth span it happens to be nested in."""
+ logger, exporter = _logger()
+ server = logger._emitter.start_span(
+ SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME
+ )
+ set_request_root_span(server)
+ kwargs = _kwargs()
+ # ``auth`` phase span is the active span when pre_call + close run.
+ with trace.use_span(server, end_on_exit=False):
+ with logger.start_phase_span("auth /chat/completions"):
+ logger.log_pre_api_call(model="gpt-4o", messages=[], kwargs=kwargs)
+ asyncio.run(logger.async_log_success_event(kwargs, None, None, None))
+ server.end()
+ by_name = {s.name: s for s in exporter.get_finished_spans()}
+ llm_span = by_name["chat gpt-4o"]
+ auth_span = by_name["auth /chat/completions"]
+ # Parented to the server root, NOT the auth span it was emitted inside.
+ assert llm_span.parent.span_id == server.get_span_context().span_id
+ assert llm_span.parent.span_id != auth_span.get_span_context().span_id
+
+
+def test_live_llm_span_anchors_to_root_with_no_active_span():
+ """Bug 2 (pass-through), live path: even with no span active at ``pre_call``,
+ the anchor is a recordable parent, so the span opens live under the server root
+ instead of orphaning — and the detached close just ends it, in the right
+ trace."""
+ logger, exporter = _logger()
+ server = logger._emitter.start_span(
+ SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME
+ )
+ set_request_root_span(server)
+ kwargs = _kwargs()
+ logger.log_pre_api_call(model="gpt-4o", messages=[], kwargs=kwargs)
+ assert logger._open_llm_calls["call_1"].span is not None # live, via anchor
+ asyncio.run(logger.async_log_success_event(kwargs, None, None, None))
+ server.end()
+ by_name = {s.name: s for s in exporter.get_finished_spans()}
+ llm_span = by_name["chat gpt-4o"]
+ assert llm_span.parent.span_id == server.get_span_context().span_id
+ assert llm_span.context.trace_id == server.get_span_context().trace_id
+
+
+def test_deferred_llm_span_reads_anchor_at_close():
+ """Bug 2, deferred path: when the anchor isn't visible at ``pre_call`` (a
+ sync-only provider's thread-pool call) the span defers; the close — back on the
+ request task, anchor visible — must parent it to the root, not orphan it."""
+ logger, exporter = _logger()
+ server = logger._emitter.start_span(
+ SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME
+ )
+ kwargs = _kwargs()
+ # pre_call with NO anchor and no active span → deferred.
+ logger.log_pre_api_call(model="gpt-4o", messages=[], kwargs=kwargs)
+ assert logger._open_llm_calls["call_1"].span is None # deferred
+ # Anchor becomes visible at close (worker copied the request task's context).
+ set_request_root_span(server)
+ asyncio.run(logger.async_log_success_event(kwargs, None, None, None))
+ server.end()
+ by_name = {s.name: s for s in exporter.get_finished_spans()}
+ llm_span = by_name["chat gpt-4o"]
+ assert llm_span.parent.span_id == server.get_span_context().span_id
+ assert llm_span.context.trace_id == server.get_span_context().trace_id
+
+
+def test_synthetic_error_log_produces_no_llm_span():
+ """Bug 1 root cause: a proxy-gate error log (auth/rate-limit) fires ``pre_call``
+ for a request that never reached a provider. Tagged with
+ ``LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL``, it must open no carrier and emit no
+ LLM-call span — even though the failure callback also fires."""
+ from litellm.constants import LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL
+
+ logger, exporter = _logger()
+ server = logger._emitter.start_span(
+ SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME
+ )
+ set_request_root_span(server)
+ payload = _payload(
+ status="failure",
+ error_information={"error_class": "ProxyException", "error_code": "401"},
+ )
+ kwargs = _kwargs(payload=payload)
+ kwargs[LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL] = True
+ with trace.use_span(server, end_on_exit=False):
+ with logger.start_phase_span("auth /chat/completions"):
+ logger.log_pre_api_call(model="gpt-4o", messages=[], kwargs=kwargs)
+ assert "call_1" not in logger._open_llm_calls # no carrier opened
+ asyncio.run(logger.async_log_failure_event(kwargs, None, None, None))
+ server.end()
+ names = {s.name for s in exporter.get_finished_spans()}
+ assert "chat gpt-4o" not in names # no phantom LLM span
+ assert "auth /chat/completions" in names # auth span itself still recorded
+
+
+def test_create_request_started_span_captures_anchor():
+ """``create_litellm_proxy_request_started_span`` doubles as the anchor capture
+ point: the active server span becomes the request root for later spans."""
+ from litellm.integrations.otel.plumbing.context import request_root_span
+
+ logger, _ = _logger()
+ server = logger._emitter.start_span(
+ SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME
+ )
+ with trace.use_span(server, end_on_exit=False):
+ returned = logger.create_litellm_proxy_request_started_span(
+ start_time=datetime.now(), headers=None
+ )
+ server.end()
+ assert returned.get_span_context().span_id == server.get_span_context().span_id
+ assert (
+ request_root_span().get_span_context().span_id
+ == server.get_span_context().span_id
+ )
+
+
+def test_guardrail_span_anchors_to_root_inside_active_phase_span():
+ """A guardrail emitted from a failure hook that runs inside the live ``auth``
+ span must still be a sibling of the LLM call under the request root, not a
+ child of auth."""
+ logger, exporter = _logger()
+ server = logger._emitter.start_span(
+ SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME
+ )
+ set_request_root_span(server)
+ request_data = {
+ "metadata": {
+ "standard_logging_guardrail_information": {
+ "guardrail_name": "my_guard",
+ "guardrail_status": "success",
+ }
+ }
+ }
+ with trace.use_span(server, end_on_exit=False):
+ with logger.start_phase_span("auth /chat/completions"):
+ logger._emit_guardrail_spans(request_data)
+ server.end()
+ by_name = {s.name: s for s in exporter.get_finished_spans()}
+ guard = by_name["execute_guardrail my_guard"]
+ auth_span = by_name["auth /chat/completions"]
+ assert guard.parent.span_id == server.get_span_context().span_id
+ assert guard.parent.span_id != auth_span.get_span_context().span_id
+
+
+def test_real_logging_pre_call_opens_span_end_to_end():
+ """Regression guard: a real ``LiteLLMLoggingObj.pre_call`` must fire
+ ``log_pre_api_call`` on the V2 logger (via ``litellm.input_callback``), so the
+ boundary span is opened and then closed by the success callback. If the logger
+ is not wired into ``input_callback``, no span is produced at all."""
+ import litellm
+ from litellm.litellm_core_utils.litellm_logging import Logging
+
+ logger, exporter = _logger()
+ # Register exactly this logger as the (only) input callback pre_call iterates.
+ monkeypatch = pytest.MonkeyPatch()
+ monkeypatch.setattr(litellm, "input_callback", [logger], raising=False)
+ try:
+ logging_obj = Logging(
+ model="gpt-4o",
+ messages=[{"role": "user", "content": "hi"}],
+ stream=False,
+ call_type="acompletion",
+ start_time=datetime.now(),
+ litellm_call_id="call_e2e",
+ function_id="fn",
+ )
+ # The wrapper always runs this before pre_call — it's what seeds
+ # ``litellm_params`` and ``litellm_call_id`` into ``model_call_details``
+ # (the call id is how the close callback correlates back to this span).
+ logging_obj.update_environment_variables(
+ litellm_params={"metadata": {}},
+ optional_params={},
+ model="gpt-4o",
+ )
+ # pre_call fires log_pre_api_call → opens the boundary span on the obj.
+ logging_obj.pre_call(input="hi", api_key="sk-test")
+ # The success callback closes it, reading the typed payload.
+ logging_obj.model_call_details["standard_logging_object"] = _payload(
+ litellm_call_id="call_e2e"
+ )
+ asyncio.run(
+ logger.async_log_success_event(
+ logging_obj.model_call_details, None, None, None
+ )
+ )
+ finally:
+ monkeypatch.undo()
+ (span,) = exporter.get_finished_spans()
+ assert span.name == "chat gpt-4o"
+
+
+def test_deferred_span_parents_to_ambient_at_close():
+ """When ``pre_call`` runs off the request task (a sync-only provider driven
+ through a thread pool, where contextvars don't follow), no ambient parent is
+ visible there, so span creation is deferred. The async callback — whose worker
+ context was copied from the request task and so still carries the server span —
+ then creates it parented to that server span, not as an orphan root."""
+ logger, exporter = _logger()
+ kwargs = _kwargs()
+ # pre_call with NO ambient span (the thread-pool case) → deferred.
+ logger.log_pre_api_call(model="gpt-4o", messages=[], kwargs=kwargs)
+ server = logger._emitter.start_span(
+ SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME
+ )
+ # The close callback runs with the (worker-copied) server span ambient.
+ with trace.use_span(server, end_on_exit=False):
+ asyncio.run(logger.async_log_success_event(kwargs, None, None, None))
+ server.end()
+ by_name = {s.name: s for s in exporter.get_finished_spans()}
+ llm_span = by_name["chat gpt-4o"]
+ assert llm_span.parent.span_id == server.get_span_context().span_id
+
+
+# Inbound ``traceparent`` propagation is now the FastAPI instrumentor's job
+# (see proxy_server's startup mount + ``test_otel_v2_mount``), not the logger's.
+
+
+# --------------------------------------------------------------------------- #
+# Baggage promotion (LLM call writes identity into baggage so child spans
+# inherit team/key/model attrs).
+# --------------------------------------------------------------------------- #
+
+
+def test_baggage_identity_promoted_onto_llm_call():
+ """On the deferred (SDK / no-proxy) path the callback seeds identity Baggage
+ from the payload so the span is still labeled with team/key. (On the proxy
+ boundary path identity rides in from auth-seeded ambient Baggage instead.)"""
+ logger, exporter = _logger()
+ _emit_llm(logger)
+ (span,) = exporter.get_finished_spans()
+ assert span.attributes[LiteLLM.TEAM_ID] == "t1"
+ assert span.attributes[LiteLLM.TEAM_ALIAS] == "team one"
+ assert span.attributes[GenAI.REQUEST_MODEL] == "gpt-4o"
+
+
+class _Auth:
+ """Stub matching the ``UserAPIKeyAuth`` fields the logger reads."""
+
+ team_id = "t1"
+ team_alias = "team one"
+ team_metadata = {"tier": "gold", "cost_center": "42"}
+ api_key = "hash1"
+ user_id = "u1"
+ org_id = None
+ key_alias = "k1"
+ end_user_id = None
+
+
+def test_provider_model_and_team_metadata_on_real_boundary_flow():
+ """End-to-end on the proxy boundary path (the gap a pure-emitter test misses):
+
+ - ``litellm.team.metadata`` is known at auth, so it rides identity Baggage
+ seeded there onto EVERY span (server + LLM call).
+ - ``litellm.provider.model`` is only known once routing picks a deployment
+ (in the payload at close), AFTER the auth seed and AFTER the boundary span
+ starts — so it can't ride Baggage. It's stamped directly on the LLM-call
+ span by the mapper, and is absent from the server span (which starts first).
+ """
+ import json
+
+ logger, exporter = _logger()
+ server = logger._emitter.start_span(
+ SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME
+ )
+ payload = _payload(
+ hidden_params={"litellm_model_name": "azure/my-deployment"},
+ metadata={
+ "user_api_key_team_id": "t1",
+ "user_api_key_team_alias": "team one",
+ "user_api_key_hash": "hash1",
+ "user_api_key_team_metadata": {"tier": "gold", "cost_center": "42"},
+ },
+ )
+ kwargs = _kwargs(payload=payload)
+ with trace.use_span(server, end_on_exit=False):
+ # auth boundary: seed identity (provider model unknown here)
+ logger.seed_request_identity(_Auth(), model="gpt-4o")
+ # pre_call boundary opens the LLM span; success closes it from the payload
+ logger.log_pre_api_call(model="gpt-4o", messages=[], kwargs=kwargs)
+ asyncio.run(logger.async_log_success_event(kwargs, None, None, None))
+ server.end()
+
+ spans = {s.name: s for s in exporter.get_finished_spans()}
+ llm = spans["chat gpt-4o"]
+ srv = spans[LITELLM_PROXY_REQUEST_SPAN_NAME]
+ # provider model: on the LLM call span, NOT the server span
+ assert llm.attributes[LiteLLM.PROVIDER_MODEL] == "azure/my-deployment"
+ assert LiteLLM.PROVIDER_MODEL not in srv.attributes
+ # team metadata: on every span, JSON-serialized
+ expected = {"tier": "gold", "cost_center": "42"}
+ assert json.loads(llm.attributes[LiteLLM.TEAM_METADATA]) == expected
+ assert json.loads(srv.attributes[LiteLLM.TEAM_METADATA]) == expected
+
+
+def test_pre_call_hook_seeds_baggage_onto_server_and_child_spans():
+ """The pre-call hook seeds identity Baggage in the request context so the
+ server span (stamped directly) AND later child spans (service here, via the
+ Baggage processor) carry identity — not just the LLM-call span."""
+ logger, exporter = _logger()
+ server = logger._emitter.start_span(
+ SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME
+ )
+
+ async def _flow():
+ # pre-call seeds baggage + stamps the active server span
+ await logger.async_pre_call_hook(
+ _Auth(), None, {"model": "gpt-4o"}, "completion"
+ )
+ # a later service call (same task) must inherit the identity
+ await logger.async_service_success_hook(
+ payload=_ServicePayload("redis", "set"), parent_otel_span=server
+ )
+
+ with trace.use_span(server, end_on_exit=False):
+ asyncio.run(_flow())
+ server.end()
+
+ spans = {s.name: s for s in exporter.get_finished_spans()}
+ redis = spans["redis set"]
+ assert redis.attributes[LiteLLM.TEAM_ID] == "t1"
+ assert redis.attributes[LiteLLM.KEY_HASH] == "hash1"
+ assert redis.attributes[f"{LiteLLM.METADATA_PREFIX}user_api_key_user_id"] == "u1"
+ srv = spans[LITELLM_PROXY_REQUEST_SPAN_NAME]
+ assert (
+ srv.attributes[LiteLLM.TEAM_ID] == "t1"
+ ) # stamped directly on the server span
+ assert srv.attributes[f"{LiteLLM.METADATA_PREFIX}user_api_key_user_id"] == "u1"
+
+
+# --------------------------------------------------------------------------- #
+# Service hooks (Phase 3)
+# --------------------------------------------------------------------------- #
+
+
+class _Service:
+ """Stub matching ``ServiceTypes(str, Enum)``."""
+
+ def __init__(self, value):
+ self.value = value
+
+
+class _ServicePayload:
+ def __init__(self, service="redis", call_type="set", error=None):
+ self.service = _Service(service)
+ self.call_type = call_type
+ self.error = error
+
+
+def _service_parent(logger):
+ """Helper: a live PROXY_REQUEST span to parent service spans under."""
+ return logger._emitter.start_span(
+ SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME
+ )
+
+
+def test_async_service_success_hook_emits_service_span():
+ logger, exporter = _logger()
+ parent = _service_parent(logger)
+ try:
+ asyncio.run(
+ logger.async_service_success_hook(
+ payload=_ServicePayload("redis", "set"),
+ parent_otel_span=parent,
+ event_metadata={"key1": "val1"},
+ )
+ )
+ finally:
+ parent.end()
+ by_name = {s.name: s for s in exporter.get_finished_spans()}
+ # Name disambiguates calls to the same service; redis is an outbound
+ # datastore call, so it's a CLIENT span with db.* semconv.
+ span = by_name["redis set"]
+ assert span.kind is SpanKind.CLIENT
+ assert span.attributes["db.system.name"] == "redis"
+ assert span.attributes["db.operation.name"] == "set"
+ assert span.attributes[LiteLLM.SERVICE_NAME] == "redis"
+ assert span.attributes[LiteLLM.SERVICE_CALL_TYPE] == "set"
+ # Canonical (V2) namespaced metadata key
+ assert span.attributes[f"{LiteLLM.METADATA_PREFIX}key1"] == "val1"
+ # V1 bare key (legacy dual-emit)
+ assert span.attributes["key1"] == "val1"
+ assert span.attributes["service"] == "redis" # V1 bare key
+ assert span.attributes["call_type"] == "set" # V1 bare key
+ # Success leaves status UNSET (semconv default), not forced OK.
+ assert span.status.status_code is StatusCode.UNSET
+
+
+def test_async_service_failure_hook_marks_error_status():
+ logger, exporter = _logger()
+ parent = _service_parent(logger)
+ try:
+ asyncio.run(
+ logger.async_service_failure_hook(
+ payload=_ServicePayload("postgres", "query"),
+ error="boom",
+ parent_otel_span=parent,
+ )
+ )
+ finally:
+ parent.end()
+ by_name = {s.name: s for s in exporter.get_finished_spans()}
+ span = by_name["postgres query"]
+ assert span.kind is SpanKind.CLIENT
+ assert span.attributes["db.system.name"] == "postgresql"
+ assert span.status.status_code is StatusCode.ERROR
+ # Without an explicit error_type from the payload, V2 stamps the fallback.
+ assert span.attributes["error.type"] == "error"
+ assert span.attributes[LiteLLM.SERVICE_NAME] == "postgres"
+
+
+def test_async_service_failure_hook_preserves_payload_error_over_override():
+ """When the payload itself carries an error, that takes precedence over the override."""
+ logger, exporter = _logger()
+ parent = _service_parent(logger)
+ try:
+ asyncio.run(
+ logger.async_service_failure_hook(
+ payload=_ServicePayload("postgres", "query", error="db-down"),
+ error="override-only-used-when-payload-clean",
+ parent_otel_span=parent,
+ )
+ )
+ finally:
+ parent.end()
+ by_name = {s.name: s for s in exporter.get_finished_spans()}
+ span = by_name["postgres query"]
+ assert span.status.status_code is StatusCode.ERROR
+ assert "db-down" in (span.status.description or "")
+
+
+def test_metrics_only_ping_without_timing_or_parent_is_noop():
+ """A success with no timing and no parent is a prometheus-only ping (the
+ per-request ``self`` latency hook, in-memory queue gauges) — not a traceable
+ operation, so no span is emitted."""
+ logger, exporter = _logger()
+ asyncio.run(
+ logger.async_service_success_hook(
+ payload=_ServicePayload(), parent_otel_span=None
+ )
+ )
+ assert exporter.get_finished_spans() == ()
+
+
+def test_background_service_call_with_timing_emits_root_span():
+ """A background datastore call (no request → no parent) but with real timing
+ still emits — as its own root trace — instead of being dropped."""
+ logger, exporter = _logger()
+ asyncio.run(
+ logger.async_service_success_hook(
+ payload=_ServicePayload("postgres", "query"),
+ parent_otel_span=None,
+ start_time=1.0,
+ end_time=2.0,
+ )
+ )
+ spans = exporter.get_finished_spans()
+ assert [s.name for s in spans] == ["postgres query"]
+ # No parent → it's a root span of its own trace.
+ assert spans[0].parent is None
+ assert spans[0].kind is SpanKind.CLIENT
+
+
+def test_internal_service_call_is_internal_kind_without_db_attrs():
+ """A genuine internal service (background job) is an INTERNAL span, no db.*."""
+ logger, exporter = _logger()
+ asyncio.run(
+ logger.async_service_success_hook(
+ payload=_ServicePayload("reset_budget_job", "reset_budget"),
+ parent_otel_span=None,
+ start_time=1.0,
+ end_time=2.0,
+ )
+ )
+ span = exporter.get_finished_spans()[0]
+ assert span.name == "reset_budget_job reset_budget"
+ assert span.kind is SpanKind.INTERNAL
+ assert "db.system.name" not in span.attributes
+ assert span.attributes[LiteLLM.SERVICE_NAME] == "reset_budget_job"
+
+
+def test_metrics_only_services_emit_no_span():
+ """self / router / proxy_pre_call / auth duplicate gen-AI spans or get a live
+ phase span — they are metrics-only and must not produce a service span."""
+ for service in ("self", "router", "proxy_pre_call", "auth"):
+ logger, exporter = _logger()
+ asyncio.run(
+ logger.async_service_success_hook(
+ payload=_ServicePayload(service, "x"),
+ parent_otel_span=None,
+ start_time=1.0,
+ end_time=2.0,
+ )
+ )
+ assert exporter.get_finished_spans() == (), f"{service} should emit no span"
+
+
+def test_service_span_inherits_parent_when_provided():
+ logger, exporter = _logger()
+ parent = logger._emitter.start_span(
+ SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME
+ )
+ try:
+ asyncio.run(
+ logger.async_service_success_hook(
+ payload=_ServicePayload(), parent_otel_span=parent
+ )
+ )
+ finally:
+ parent.end()
+ by_name = {s.name: s for s in exporter.get_finished_spans()}
+ assert (
+ by_name["redis set"].parent.span_id
+ == by_name[LITELLM_PROXY_REQUEST_SPAN_NAME].get_span_context().span_id
+ )
+
+
+def test_service_span_prefers_ambient_context_over_threaded_parent():
+ """Service/DB spans parent to the active (ambient) span when there is one, so
+ they nest under whatever phase is active (e.g. a DB lookup under the live
+ ``auth`` span). The threaded ``parent_otel_span`` is only a fallback for when
+ ambient has no live span (a background service call)."""
+ logger, exporter = _logger()
+ ambient = logger._emitter.start_span(SpanRole.LLM_CALL, "chat gpt-4o")
+ threaded = logger._emitter.start_span(
+ SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME
+ )
+ try:
+ with trace.use_span(ambient, end_on_exit=False):
+ asyncio.run(
+ logger.async_service_success_hook(
+ payload=_ServicePayload("redis", "get"),
+ parent_otel_span=threaded,
+ )
+ )
+ finally:
+ ambient.end()
+ threaded.end()
+ by_name = {s.name: s for s in exporter.get_finished_spans()}
+ assert by_name["redis get"].parent.span_id == ambient.get_span_context().span_id
+
+
+# --------------------------------------------------------------------------- #
+# Proxy SERVER span lifecycle
+# --------------------------------------------------------------------------- #
+
+
+def test_create_proxy_request_started_span_returns_ambient_span():
+ """V2 doesn't create a server span (the instrumentor does), but it returns
+ the active server span so the proxy can thread it as the service-span parent
+ — service logging only fires the OTel hook when that parent is non-None."""
+ logger, exporter = _logger()
+ # No ambient recordable span → None (and creates nothing).
+ assert (
+ logger.create_litellm_proxy_request_started_span(
+ start_time=datetime.now(timezone.utc), headers={"traceparent": "x"}
+ )
+ is None
+ )
+ assert exporter.get_finished_spans() == ()
+ # With an active server span, return it (do NOT create a new one).
+ server = logger._emitter.start_span(
+ SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME
+ )
+ with trace.use_span(server, end_on_exit=False):
+ got = logger.create_litellm_proxy_request_started_span(
+ start_time=datetime.now(timezone.utc), headers=None
+ )
+ server.end()
+ assert got is server
+
+
+# --------------------------------------------------------------------------- #
+# Constructor / proxy global guard
+# --------------------------------------------------------------------------- #
+
+
+def test_constructor_accepts_v1_compatible_kwargs():
+ """Mirrors V1's positional shape — config / callback_name / providers / **kwargs."""
+ cfg = OpenTelemetryV2Config(exporter="in_memory")
+ tp = providers.build_tracer_provider(cfg)
+ logger = OpenTelemetryV2(
+ config=cfg,
+ callback_name="otel",
+ tracer_provider=tp,
+ logger_provider=None,
+ meter_provider=None,
+ turn_off_message_logging=True,
+ )
+ assert logger.callback_name == "otel"
+ assert logger.turn_off_message_logging is True
+ assert logger.tracer is not None
+
+
+def test_default_config_reads_env(monkeypatch):
+ """No explicit config → reads env (exporter=console by default)."""
+ monkeypatch.delenv("OTEL_EXPORTER", raising=False)
+ monkeypatch.delenv("OTEL_EXPORTER_OTLP_PROTOCOL", raising=False)
+ logger = OpenTelemetryV2(
+ tracer_provider=providers.build_tracer_provider(
+ OpenTelemetryV2Config(exporter="in_memory")
+ )
+ )
+ assert logger.config.exporter == "console"
+
+
+def test_proxy_global_first_registered_wins(monkeypatch):
+ """``_init_otel_logger_on_litellm_proxy`` claims the global only when empty."""
+ proxy_server = pytest.importorskip("litellm.proxy.proxy_server")
+ monkeypatch.setattr(proxy_server, "open_telemetry_logger", None, raising=False)
+ cfg = OpenTelemetryV2Config(exporter="in_memory")
+ tp = providers.build_tracer_provider(cfg)
+
+ first = OpenTelemetryV2(config=cfg, tracer_provider=tp)
+ assert proxy_server.open_telemetry_logger is first
+
+ second = OpenTelemetryV2(config=cfg, tracer_provider=tp)
+ # Global still points at the first registration.
+ assert proxy_server.open_telemetry_logger is first
+ assert second is not first
+
+
+def test_registers_into_litellm_service_callback(monkeypatch):
+ """The logger must mutate ``litellm.service_callback`` in place. An empty
+ list is falsy, so a ``getattr(..) or []`` would append to a throwaway local
+ and service spans (Redis, …) would silently never fire on this logger.
+ """
+ import litellm
+
+ pytest.importorskip("litellm.proxy.proxy_server")
+ monkeypatch.setattr(litellm, "service_callback", [], raising=False)
+ cfg = OpenTelemetryV2Config(exporter="in_memory")
+ tp = providers.build_tracer_provider(cfg)
+
+ first = OpenTelemetryV2(config=cfg, tracer_provider=tp)
+ assert first in litellm.service_callback
+
+ # A second OTel logger sees one is already registered and does not duplicate.
+ OpenTelemetryV2(config=cfg, tracer_provider=tp)
+ otel_registrations = [
+ cb
+ for cb in litellm.service_callback
+ if cb.__class__.__module__.startswith("litellm.integrations.otel")
+ ]
+ assert len(otel_registrations) == 1
+
+
+def test_registers_into_litellm_input_callback(monkeypatch):
+ """The logger must land in ``litellm.input_callback`` — the list
+ ``Logging.pre_call`` iterates to fire ``log_pre_api_call``. Without this the
+ boundary hook never runs and the gen-AI span is never opened (the span goes
+ completely missing). Deduped like ``service_callback``.
+ """
+ import litellm
+
+ pytest.importorskip("litellm.proxy.proxy_server")
+ monkeypatch.setattr(litellm, "input_callback", [], raising=False)
+ cfg = OpenTelemetryV2Config(exporter="in_memory")
+ tp = providers.build_tracer_provider(cfg)
+
+ first = OpenTelemetryV2(config=cfg, tracer_provider=tp)
+ assert first in litellm.input_callback
+
+ OpenTelemetryV2(config=cfg, tracer_provider=tp)
+ otel_registrations = [
+ cb
+ for cb in litellm.input_callback
+ if cb.__class__.__module__.startswith("litellm.integrations.otel")
+ ]
+ assert len(otel_registrations) == 1
+
+
+def test_registers_into_async_success_and_failure_callbacks(monkeypatch):
+ """The logger must self-register into ``litellm._async_success_callback`` and
+ ``litellm._async_failure_callback`` — the lists ``Logging.async_success_handler``
+ / ``async_failure_handler`` iterate to fire ``async_log_success_event`` /
+ ``async_log_failure_event``, where the boundary span is *closed*.
+
+ ``input_callback`` opens the span; these lists close it. Relying only on the
+ proxy's ``litellm.callbacks`` fan-out to populate them is not enough: a logger
+ that reached litellm via ``service_callback`` / ``success_callback`` (or was
+ created after the fan-out ran) is absent from ``litellm.callbacks``, so on a
+ pass-through request (which never runs ``function_setup``) the span opens and is
+ never ended — the gen-AI span leaks and never exports, while DB/service spans
+ still show up. Self-registration here guarantees every open has a close.
+ """
+ import litellm
+
+ pytest.importorskip("litellm.proxy.proxy_server")
+ monkeypatch.setattr(litellm, "_async_success_callback", [], raising=False)
+ monkeypatch.setattr(litellm, "_async_failure_callback", [], raising=False)
+ cfg = OpenTelemetryV2Config(exporter="in_memory")
+ tp = providers.build_tracer_provider(cfg)
+
+ first = OpenTelemetryV2(config=cfg, tracer_provider=tp)
+ assert first in litellm._async_success_callback
+ assert first in litellm._async_failure_callback
+
+ # Deduped — a second otel logger doesn't double up the close hook.
+ OpenTelemetryV2(config=cfg, tracer_provider=tp)
+ for callback_list in (
+ litellm._async_success_callback,
+ litellm._async_failure_callback,
+ ):
+ otel_registrations = [
+ cb
+ for cb in callback_list
+ if cb.__class__.__module__.startswith("litellm.integrations.otel")
+ ]
+ assert len(otel_registrations) == 1
+
+
+def test_boundary_span_closes_without_proxy_fanout(monkeypatch):
+ """A span opened at ``pre_call`` is still closed and exported when the logger is
+ registered ONLY via its own ``__init__`` (no ``litellm.callbacks`` fan-out, as
+ happens for a logger configured through ``service_callback``) and the close runs
+ through the real ``async_success_handler``.
+
+ Self-registration must wire both ends: the open hook (``input_callback``) and the
+ close hook (``_async_success_callback``). If only the open end were wired the span
+ would leak — opened but never closed, never exported.
+ """
+ import litellm
+ from litellm.litellm_core_utils.litellm_logging import Logging
+
+ pytest.importorskip("litellm.proxy.proxy_server")
+ monkeypatch.setattr(litellm, "input_callback", [], raising=False)
+ monkeypatch.setattr(litellm, "_async_success_callback", [], raising=False)
+ monkeypatch.setattr(litellm, "_async_failure_callback", [], raising=False)
+ # Crucially: the logger is NOT in litellm.callbacks, so the proxy fan-out would
+ # never reach it. Only __init__ self-registration wires the open + close hooks.
+ monkeypatch.setattr(litellm, "callbacks", [], raising=False)
+
+ logger, exporter = _logger()
+ logging_obj = Logging(
+ model="gpt-4o",
+ messages=[{"role": "user", "content": "hi"}],
+ stream=False,
+ call_type="pass_through_endpoint",
+ start_time=datetime.now(),
+ litellm_call_id="pt_leak",
+ function_id="fn",
+ )
+ logging_obj.update_environment_variables(
+ litellm_params={"metadata": {}},
+ optional_params={},
+ model="gpt-4o",
+ )
+ logging_obj.model_call_details["litellm_call_id"] = "pt_leak"
+ # pre_call opens the boundary span (logger is in input_callback).
+ logging_obj.pre_call(input="hi", api_key="")
+ assert "pt_leak" in logger._open_llm_calls
+ # The close runs through the real async_success_handler, which iterates
+ # _async_success_callback — where the logger self-registered.
+ logging_obj.model_call_details["standard_logging_object"] = _payload(
+ litellm_call_id="pt_leak"
+ )
+ asyncio.run(
+ logging_obj.async_success_handler(
+ result=None, start_time=datetime.now(), end_time=datetime.now()
+ )
+ )
+ assert "pt_leak" not in logger._open_llm_calls # carrier closed, not leaked
+ (span,) = exporter.get_finished_spans()
+ assert span.name == "chat gpt-4o"
+
+
+# --------------------------------------------------------------------------- #
+# Guardrail span placement: request-level parent + real execution timestamps
+# --------------------------------------------------------------------------- #
+
+
+def _guardrail_request_data(*, start, end):
+ return {
+ "metadata": {
+ "standard_logging_guardrail_information": [
+ {
+ "guardrail_name": "openai-moderation",
+ "guardrail_mode": "pre_call",
+ "guardrail_status": "success",
+ "start_time": start,
+ "end_time": end,
+ "duration": end - start,
+ }
+ ],
+ }
+ }
+
+
+def test_guardrail_span_parents_to_ambient_server_span():
+ """The post-call hook runs in the request task with the server span ambient,
+ so the guardrail span parents to it natively — no span threaded through
+ metadata. (Auth already finished, so no phase span is active.)"""
+ logger, exporter = _logger()
+ server = logger._emitter.start_span(
+ SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME
+ )
+ data = _guardrail_request_data(start=1000.0, end=1000.5)
+ try:
+ with trace.use_span(server, end_on_exit=False):
+ asyncio.run(
+ logger.async_post_call_success_hook(data, _Auth(), {"ok": True})
+ )
+ finally:
+ server.end()
+ g = {s.name: s for s in exporter.get_finished_spans()}[
+ "execute_guardrail openai-moderation"
+ ]
+ assert g.parent.span_id == server.get_span_context().span_id
+
+
+def test_guardrail_span_uses_actual_execution_timestamps():
+ """A pre_call guardrail's span carries its real start/end (from the logging
+ entry), so it sorts before the LLM call instead of at post-call emit time."""
+ logger, exporter = _logger()
+ server = logger._emitter.start_span(
+ SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME
+ )
+ data = _guardrail_request_data(start=1700.0, end=1700.25)
+ try:
+ with trace.use_span(server, end_on_exit=False):
+ asyncio.run(
+ logger.async_post_call_success_hook(data, _Auth(), {"ok": True})
+ )
+ finally:
+ server.end()
+ g = {s.name: s for s in exporter.get_finished_spans()}[
+ "execute_guardrail openai-moderation"
+ ]
+ assert g.start_time == to_ns(1700.0)
+ assert g.end_time == to_ns(1700.25)
diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_mount.py b/tests/test_litellm/integrations/otel/test_otel_v2_mount.py
new file mode 100644
index 00000000000..956d8c53cee
--- /dev/null
+++ b/tests/test_litellm/integrations/otel/test_otel_v2_mount.py
@@ -0,0 +1,151 @@
+"""V2 entrypoint: the FastAPI instrumentation proxy_server mounts at app creation
+(gated by LITELLM_OTEL_V2). The mount logic lives in
+``litellm.integrations.otel.mount``; this exercises both that module's public
+surface and the server-span + shared-provider behavior it produces.
+"""
+
+import os
+import sys
+
+import pytest
+
+sys.path.insert(0, os.path.abspath("../../../.."))
+
+pytest.importorskip("opentelemetry")
+pytest.importorskip("opentelemetry.instrumentation.fastapi")
+fastapi = pytest.importorskip("fastapi")
+
+from fastapi.testclient import TestClient # noqa: E402
+from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor # noqa: E402
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor # noqa: E402
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( # noqa: E402
+ InMemorySpanExporter,
+)
+from opentelemetry.trace import SpanKind # noqa: E402
+
+from litellm.integrations.otel.model.config import ( # noqa: E402
+ OpenTelemetryV2Config,
+ is_otel_v2_enabled,
+)
+from litellm.integrations.otel.logger import OpenTelemetryV2 # noqa: E402
+from litellm.integrations.otel.mount import ( # noqa: E402
+ PASSTHROUGH_PREFIXES,
+ _passthrough_span_name_hook,
+ instrument_fastapi_app,
+)
+
+
+class _FakeSpan:
+ """Minimal recording span capturing what the hook writes."""
+
+ def __init__(self, recording=True):
+ self._recording = recording
+ self.name = None
+ self.attributes = {}
+
+ def is_recording(self):
+ return self._recording
+
+ def update_name(self, name):
+ self.name = name
+
+ def set_attribute(self, key, value):
+ self.attributes[key] = value
+
+
+def _instrumented_app():
+ """Mirror proxy_server's startup mount: a logger builds the shared provider,
+ and the FastAPI instrumentor is attached to it."""
+ app = fastapi.FastAPI()
+
+ @app.get("/ping")
+ def ping():
+ return {"ok": True}
+
+ logger = OpenTelemetryV2(config=OpenTelemetryV2Config(exporter="in_memory"))
+ FastAPIInstrumentor.instrument_app(app, tracer_provider=logger._tracer_provider)
+ return app, logger
+
+
+def test_gate_toggles_with_env(monkeypatch):
+ """The startup mount is guarded by this flag."""
+ monkeypatch.delenv("LITELLM_OTEL_V2", raising=False)
+ assert is_otel_v2_enabled() is False
+ monkeypatch.setenv("LITELLM_OTEL_V2", "1")
+ assert is_otel_v2_enabled() is True
+
+
+def test_instrumented_app_emits_server_span():
+ app, logger = _instrumented_app()
+ exporter = InMemorySpanExporter()
+ logger._tracer_provider.add_span_processor(SimpleSpanProcessor(exporter))
+
+ TestClient(app).get("/ping")
+
+ server_spans = [
+ s for s in exporter.get_finished_spans() if s.kind is SpanKind.SERVER
+ ]
+ assert server_spans, "FastAPI instrumentor should emit a SERVER span per request"
+ attrs = server_spans[0].attributes or {}
+ assert any("route" in k or "method" in k for k in attrs)
+
+
+def test_logger_and_instrumentor_share_provider():
+ """Gen-ai spans (logger) and server spans (instrumentor) write to one provider."""
+ _, logger = _instrumented_app()
+ assert logger._emitter._tracer is logger.tracer
+
+
+def test_passthrough_hook_renames_catch_all_span():
+ """A passthrough route gets its span renamed to the real request path."""
+ span = _FakeSpan()
+ _passthrough_span_name_hook(
+ span, {"path": "/openai/v1/chat/completions", "method": "POST"}
+ )
+ assert span.name == "POST /openai/v1/chat/completions"
+ assert span.attributes["http.route"] == "/openai/v1/chat/completions"
+
+
+def test_passthrough_hook_leaves_non_passthrough_route_unchanged():
+ """A normal route keeps its low-cardinality template name (hook no-ops)."""
+ span = _FakeSpan()
+ _passthrough_span_name_hook(span, {"path": "/v1/models", "method": "GET"})
+ assert span.name is None
+ assert "http.route" not in span.attributes
+
+
+def test_passthrough_hook_ignores_non_recording_span():
+ span = _FakeSpan(recording=False)
+ _passthrough_span_name_hook(
+ span, {"path": "/openai/v1/chat/completions", "method": "POST"}
+ )
+ assert span.name is None
+
+
+def test_known_passthrough_prefixes_present():
+ """Guard the prefix set against accidental edits."""
+ assert {"openai", "anthropic", "vertex_ai", "bedrock"} <= PASSTHROUGH_PREFIXES
+
+
+def test_instrument_fastapi_app_noop_when_gate_off(monkeypatch):
+ """With the gate off the mount is a no-op — no instrumentation attached."""
+ monkeypatch.delenv("LITELLM_OTEL_V2", raising=False)
+ app = fastapi.FastAPI()
+ instrument_fastapi_app(app)
+ assert getattr(app, "_is_instrumented_by_opentelemetry", False) is False
+
+
+def test_instrument_fastapi_app_attaches_when_gate_on(monkeypatch):
+ """With the gate on the FastAPI app is instrumented for server spans."""
+ monkeypatch.setenv("LITELLM_OTEL_V2", "1")
+ app = fastapi.FastAPI()
+
+ @app.get("/ping")
+ def ping():
+ return {"ok": True}
+
+ instrument_fastapi_app(app)
+ try:
+ assert getattr(app, "_is_instrumented_by_opentelemetry", False) is True
+ finally:
+ FastAPIInstrumentor.uninstrument_app(app)
diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_multibackend.py b/tests/test_litellm/integrations/otel/test_otel_v2_multibackend.py
new file mode 100644
index 00000000000..e879766c5c7
--- /dev/null
+++ b/tests/test_litellm/integrations/otel/test_otel_v2_multibackend.py
@@ -0,0 +1,89 @@
+"""Multi-backend fan-out: one TracerProvider, *N* SpanProcessors.
+
+V1 needed a separate ``TracerProvider`` per integration to avoid stepping on
+the global. V2 attaches a ``SpanProcessor`` per exporter to the *same*
+provider, so the same trace ID lights up every backend — no duplicate spans,
+no per-integration provider caches.
+"""
+
+import pytest
+
+pytest.importorskip("opentelemetry")
+
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import (
+ InMemorySpanExporter,
+)
+
+from litellm.integrations.otel.model.config import ExporterSpec, OpenTelemetryV2Config
+from litellm.integrations.otel.plumbing.providers import build_tracer_provider
+
+
+def test_two_exporters_receive_the_same_span():
+ """A single ``span.end()`` lands in BOTH exporters with the same span ID."""
+ exporter_a = InMemorySpanExporter()
+ exporter_b = InMemorySpanExporter()
+ cfg = OpenTelemetryV2Config(
+ exporters=[
+ ExporterSpec(kind="in_memory"),
+ ExporterSpec(kind="in_memory"),
+ ]
+ )
+ # Override the auto-built exporters with our test ones by swapping
+ # processors after construction (the test's purpose is to exercise the
+ # multi-processor wiring, not to negotiate the in-memory pipe).
+ from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+
+ provider = build_tracer_provider(cfg)
+ # Clear out any auto-built export processors and attach our pair.
+ while provider._active_span_processor._span_processors:
+ provider._active_span_processor._span_processors = (
+ provider._active_span_processor._span_processors[:-1]
+ )
+ provider.add_span_processor(SimpleSpanProcessor(exporter_a))
+ provider.add_span_processor(SimpleSpanProcessor(exporter_b))
+
+ tracer = provider.get_tracer("test")
+ span = tracer.start_span("multi-backend")
+ span.set_attribute("test.marker", "yes")
+ span.end()
+
+ spans_a = exporter_a.get_finished_spans()
+ spans_b = exporter_b.get_finished_spans()
+ assert len(spans_a) == 1
+ assert len(spans_b) == 1
+ assert spans_a[0].context.span_id == spans_b[0].context.span_id
+
+
+def test_resource_attributes_apply_to_all_exporters():
+ """``resource_attributes`` flow through the shared TracerProvider."""
+ cfg = OpenTelemetryV2Config(
+ exporters=[ExporterSpec(kind="in_memory")],
+ resource_attributes={"openinference.project.name": "phoenix-test"},
+ )
+ provider = build_tracer_provider(cfg)
+ assert provider.resource.attributes["openinference.project.name"] == "phoenix-test"
+
+
+def test_config_normalizer_inserts_genai_first():
+ """The validator pins ``genai`` at the head + appends ``legacy`` on legacy_compat."""
+ cfg = OpenTelemetryV2Config(mapper_names=["openinference", "langfuse"])
+ assert cfg.mapper_names[0] == "genai"
+ assert "openinference" in cfg.mapper_names
+ assert "langfuse" in cfg.mapper_names
+ assert cfg.mapper_names[-1] == "legacy" # legacy_compat=True by default
+
+
+def test_config_normalizer_no_legacy_when_compat_off():
+ cfg = OpenTelemetryV2Config(legacy_compat=False, mapper_names=["openinference"])
+ assert "legacy" not in cfg.mapper_names
+ assert cfg.mapper_names[0] == "genai"
+
+
+def test_config_folds_legacy_exporter_triple_into_exporters_list():
+ """When ``exporters`` is empty, the validator folds the legacy single triple."""
+ cfg = OpenTelemetryV2Config(
+ exporter="otlp_http", endpoint="https://api.example.com", headers="k=v"
+ )
+ assert len(cfg.exporters) == 1
+ assert cfg.exporters[0].kind == "otlp_http"
+ assert cfg.exporters[0].endpoint == "https://api.example.com"
diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_presets.py b/tests/test_litellm/integrations/otel/test_otel_v2_presets.py
new file mode 100644
index 00000000000..6b9fa820cdf
--- /dev/null
+++ b/tests/test_litellm/integrations/otel/test_otel_v2_presets.py
@@ -0,0 +1,122 @@
+"""Preset tests. Focused on the AgentOps JWT fetch, which must never block the
+event loop: the preset does no network I/O, and a custom exporter mints the JWT
+lazily on its first export (in the BatchSpanProcessor worker thread)."""
+
+import httpx
+import pytest
+
+from litellm.integrations.otel.plumbing import providers
+from litellm.integrations.otel.model.config import ExporterSpec
+from litellm.integrations.otel.presets import agentops as agentops_mod
+from litellm.integrations.otel.presets.agentops import (
+ _AGENTOPS_ENDPOINT,
+ _AGENTOPS_EXPORTER_KIND,
+ _build_agentops_exporter,
+ _fetch_agentops_jwt,
+ agentops_preset,
+)
+
+
+def test_agentops_preset_does_no_network_io(monkeypatch):
+ # The preset must not fetch the JWT at build time — that would block the
+ # event loop during callback construction. It only describes the exporter.
+ def _boom(*_a, **_k):
+ raise AssertionError("agentops_preset must not fetch the JWT eagerly")
+
+ monkeypatch.setattr(agentops_mod, "_fetch_agentops_jwt", _boom)
+ monkeypatch.setenv("AGENTOPS_API_KEY", "ak-123")
+ cfg = agentops_preset()
+ agentops_exporters = [e for e in cfg.exporters if e.kind == _AGENTOPS_EXPORTER_KIND]
+ assert len(agentops_exporters) == 1
+ spec = agentops_exporters[0]
+ assert spec.endpoint == _AGENTOPS_ENDPOINT
+ assert spec.options == {"api_key": "ak-123"} # carried to the lazy exporter
+
+
+def test_agentops_preset_without_key_omits_options(monkeypatch):
+ monkeypatch.delenv("AGENTOPS_API_KEY", raising=False)
+ cfg = agentops_preset()
+ spec = next(e for e in cfg.exporters if e.kind == _AGENTOPS_EXPORTER_KIND)
+ assert spec.options is None
+
+
+def test_agentops_exporter_factory_is_registered():
+ assert _AGENTOPS_EXPORTER_KIND in providers._EXPORTER_FACTORIES
+
+
+def test_agentops_exporter_mints_jwt_lazily(monkeypatch):
+ pytest.importorskip("opentelemetry.exporter.otlp.proto.http.trace_exporter")
+ monkeypatch.setattr(
+ agentops_mod, "_fetch_agentops_jwt", lambda _k: {"token": "jwt-xyz"}
+ )
+ spec = ExporterSpec(
+ kind=_AGENTOPS_EXPORTER_KIND,
+ endpoint=_AGENTOPS_ENDPOINT,
+ options={"api_key": "ak"},
+ )
+ exporter = _build_agentops_exporter(spec)
+
+ # No auth header until the first export triggers the (off-loop) fetch.
+ assert "Authorization" not in exporter._session.headers
+ exporter._ensure_authenticated()
+ assert exporter._session.headers["Authorization"] == "Bearer jwt-xyz"
+
+ # Cached: a second resolution does not re-fetch.
+ calls = []
+ monkeypatch.setattr(
+ agentops_mod,
+ "_fetch_agentops_jwt",
+ lambda k: calls.append(k) or {"token": "again"},
+ )
+ exporter._ensure_authenticated()
+ assert calls == []
+
+
+def test_agentops_exporter_tolerates_fetch_failure(monkeypatch):
+ pytest.importorskip("opentelemetry.exporter.otlp.proto.http.trace_exporter")
+
+ def _raise(_k):
+ raise RuntimeError("auth down")
+
+ monkeypatch.setattr(agentops_mod, "_fetch_agentops_jwt", _raise)
+ exporter = _build_agentops_exporter(
+ ExporterSpec(
+ kind=_AGENTOPS_EXPORTER_KIND,
+ endpoint=_AGENTOPS_ENDPOINT,
+ options={"api_key": "ak"},
+ )
+ )
+ exporter._ensure_authenticated() # must not raise
+ assert "Authorization" not in exporter._session.headers
+
+
+def test_fetch_jwt_uses_owned_client_not_shared_pool(monkeypatch):
+ """The fetch owns a short-lived client and closes it, rather than closing
+ the process-wide cached ``_get_httpx_client`` pool shared by other callers."""
+ closed = {"n": 0}
+
+ class _FakeResponse:
+ status_code = 200
+
+ def json(self):
+ return {"token": "jwt-123"}
+
+ class _FakeClient:
+ def __init__(self, *_a, **_k):
+ pass
+
+ def __enter__(self):
+ return self
+
+ def __exit__(self, *_a):
+ closed["n"] += 1
+
+ def post(self, *_a, **_k):
+ return _FakeResponse()
+
+ monkeypatch.setattr(httpx, "Client", _FakeClient)
+ assert not hasattr(agentops_mod, "_get_httpx_client")
+
+ result = _fetch_agentops_jwt("api-key")
+ assert result == {"token": "jwt-123"}
+ assert closed["n"] == 1 # the owned client was closed
diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py b/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py
new file mode 100644
index 00000000000..c4e80145c70
--- /dev/null
+++ b/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py
@@ -0,0 +1,454 @@
+"""Tests for the OTel v2 sources of truth: span registry, semconv keys, config,
+and the typed StandardLoggingPayload adapter. These need no OTel SDK."""
+
+import pytest
+
+from litellm.integrations.otel import (
+ BAGGAGE_PROMOTED_KEYS,
+ DB,
+ Error,
+ GenAI,
+ GenAIOperation,
+ HTTP,
+ LiteLLM,
+ OpenTelemetryV2Config,
+ Server,
+ is_otel_v2_enabled,
+ promoted_baggage,
+ resolve_operation,
+ resolve_provider,
+)
+from litellm.integrations.otel.model import spans as spans_mod
+from litellm.integrations.otel.model.payloads import LLMCallSpanData, RequestIdentity
+from litellm.integrations.otel.model.spans import (
+ SPAN_REGISTRY,
+ LiteLLMSpanKind,
+ SpanRole,
+ child_roles,
+ root_roles,
+ validate_registry,
+)
+
+
+def _sample_payload(**overrides):
+ payload = {
+ "call_type": "acompletion",
+ "custom_llm_provider": "openai",
+ "model": "gpt-4o",
+ "prompt_tokens": 10,
+ "completion_tokens": 5,
+ "total_tokens": 15,
+ "stream": False,
+ "model_parameters": {
+ "temperature": 0.7,
+ "max_tokens": 256,
+ "top_p": 0.9,
+ "top_k": 40,
+ "frequency_penalty": 0.1,
+ "presence_penalty": 0.2,
+ "stop": ["STOP"],
+ "seed": 42,
+ },
+ "response": {
+ "id": "resp_1",
+ "model": "gpt-4o-2024",
+ "choices": [{"finish_reason": "stop"}],
+ },
+ "metadata": {
+ "team_id": "t1",
+ "team_alias": "team one",
+ "user_api_key_hash": "hsh",
+ "user_api_key_org_id": "org1",
+ },
+ "api_base": "https://api.openai.com:443/v1",
+ "status": "success",
+ "litellm_call_id": "call_1",
+ "end_user": "u1",
+ "response_cost": 0.002,
+ "hidden_params": {},
+ }
+ payload.update(overrides)
+ return payload
+
+
+# --- span registry (source of truth #2) ------------------------------------- #
+
+
+def test_registry_validates_and_is_complete():
+ validate_registry() # raises on inconsistency
+ assert set(SPAN_REGISTRY) == set(SpanRole)
+
+
+def test_registry_parent_integrity_no_orphans():
+ for role, spec in SPAN_REGISTRY.items():
+ assert spec.role is role
+ if spec.parent is not None:
+ assert spec.parent in SPAN_REGISTRY
+
+
+def test_registry_hierarchy_shape():
+ assert set(root_roles()) == {SpanRole.PROXY_REQUEST}
+ # Guardrails parent to the request span, not the LLM call: a pre-call
+ # guardrail runs before the LLM call exists, so it's a sibling of it.
+ assert set(child_roles(SpanRole.PROXY_REQUEST)) == {
+ SpanRole.LLM_CALL,
+ SpanRole.GUARDRAIL,
+ SpanRole.DB_CALL,
+ SpanRole.SERVICE,
+ }
+ assert SPAN_REGISTRY[SpanRole.LLM_CALL].kind is LiteLLMSpanKind.CLIENT
+ assert SPAN_REGISTRY[SpanRole.PROXY_REQUEST].kind is LiteLLMSpanKind.SERVER
+ assert SPAN_REGISTRY[SpanRole.GUARDRAIL].parent is SpanRole.PROXY_REQUEST
+ # An outbound datastore call is a CLIENT span; an internal service is INTERNAL.
+ assert SPAN_REGISTRY[SpanRole.DB_CALL].kind is LiteLLMSpanKind.CLIENT
+ assert SPAN_REGISTRY[SpanRole.SERVICE].kind is LiteLLMSpanKind.INTERNAL
+
+
+def test_llm_call_span_name():
+ data = LLMCallSpanData.from_standard_logging_payload(_sample_payload())
+ assert spans_mod.llm_call_span_name(data) == "chat gpt-4o"
+
+
+# --- semconv (source of truth #1) ------------------------------------------- #
+
+
+def _all_constants(cls):
+ return {
+ getattr(cls, name)
+ for name in vars(cls)
+ if not name.startswith("__") and isinstance(getattr(cls, name), str)
+ }
+
+
+def test_attribute_keys_are_unique_across_namespaces():
+ # prefixes are allowed to be substrings; exact keys must not collide.
+ exact = set()
+ for cls in (GenAI, Error, Server, HTTP, DB):
+ for key in _all_constants(cls):
+ assert key not in exact, f"duplicate attribute key {key}"
+ exact.add(key)
+
+
+def test_provider_resolution():
+ assert resolve_provider("openai") == "openai"
+ assert resolve_provider("bedrock") == "aws.bedrock"
+ assert resolve_provider("vertex_ai") == "gcp.vertex_ai"
+ # unknown providers pass through verbatim (semconv allows provider-specific)
+ assert resolve_provider("my_custom_llm") == "my_custom_llm"
+ assert resolve_provider(None) == ""
+
+
+def test_operation_resolution():
+ assert resolve_operation("acompletion") is GenAIOperation.CHAT
+ assert resolve_operation("aembedding") is GenAIOperation.EMBEDDINGS
+ assert resolve_operation("atext_completion") is GenAIOperation.TEXT_COMPLETION
+ assert resolve_operation(None) is GenAIOperation.CHAT
+
+
+# --- typed adapter (source of truth #3) ------------------------------------- #
+
+
+def test_llm_call_adapter_extracts_all_fields():
+ data = LLMCallSpanData.from_standard_logging_payload(_sample_payload())
+ assert data.operation is GenAIOperation.CHAT
+ assert data.provider == "openai"
+ assert data.request_model == "gpt-4o"
+ assert data.response_model == "gpt-4o-2024"
+ assert data.response_id == "resp_1"
+ assert data.finish_reasons == ("stop",)
+ assert (data.usage.input_tokens, data.usage.output_tokens) == (10, 5)
+ assert data.request_params.temperature == 0.7
+ assert data.request_params.top_k == 40
+ assert data.request_params.stop_sequences == ("STOP",)
+ assert data.request_params.seed == 42
+ assert data.server is not None
+ assert data.server.address == "api.openai.com"
+ assert data.server.port == 443
+ assert data.response_cost == 0.002
+ assert data.error is None
+ assert data.identity.team_id == "t1"
+ assert data.identity.key_hash == "hsh"
+
+
+def test_llm_call_adapter_failure_path():
+ payload = _sample_payload(
+ status="failure",
+ error_information={
+ "error_class": "RateLimitError",
+ "error_message": "429 slow down",
+ },
+ )
+ data = LLMCallSpanData.from_standard_logging_payload(payload)
+ assert data.error is not None
+ assert data.error.error_type == "RateLimitError"
+ assert data.error.message == "429 slow down"
+
+
+def test_adapter_is_resilient_to_minimal_payload():
+ data = LLMCallSpanData.from_standard_logging_payload({})
+ assert data.request_model == ""
+ assert data.operation is GenAIOperation.CHAT
+ assert data.server is None
+ assert data.usage.input_tokens is None
+
+
+def test_content_capture_gated_off_by_default():
+ # ``capture_content`` defaults off: prompt/response bodies must not reach the
+ # span data (and so no vendor mapper can export them) unless explicitly
+ # opted in. Non-content metadata (finish reasons) is still derived.
+ payload = _sample_payload(
+ messages=[{"role": "user", "content": "secret prompt"}],
+ )
+ payload["response"]["choices"] = [
+ {"finish_reason": "stop", "message": {"role": "assistant", "content": "secret"}}
+ ]
+ data = LLMCallSpanData.from_standard_logging_payload(payload)
+ assert data.messages_in == ()
+ assert data.choices_out == ()
+ assert data.finish_reasons == ("stop",)
+
+
+def test_request_identity_prefers_canonical_team_keys():
+ from litellm.integrations.otel.model.payloads import RequestIdentity
+
+ payload = _sample_payload(
+ metadata={
+ "user_api_key_team_id": "team-canonical",
+ "user_api_key_team_alias": "alias-canonical",
+ "user_api_key_hash": "hsh",
+ "team_id": "legacy-ignored", # legacy alias loses to the canonical key
+ }
+ )
+ ident = RequestIdentity.from_payload(payload)
+ assert ident.team_id == "team-canonical"
+ assert ident.team_alias == "alias-canonical"
+ assert ident.key_hash == "hsh"
+
+
+def test_request_identity_falls_back_to_legacy_team_keys():
+ from litellm.integrations.otel.model.payloads import RequestIdentity
+
+ payload = _sample_payload(
+ metadata={"team_id": "legacy-team", "team_alias": "legacy"}
+ )
+ ident = RequestIdentity.from_payload(payload)
+ assert ident.team_id == "legacy-team"
+ assert ident.team_alias == "legacy"
+
+
+def test_guardrail_span_data_block_carries_verdict_and_error():
+ from litellm.integrations.otel.model.payloads import GuardrailSpanData
+
+ entry = {
+ "guardrail_name": "openai-moderation",
+ "guardrail_mode": "pre_call",
+ "guardrail_status": "guardrail_intervened",
+ "guardrail_provider": "openai",
+ "guardrail_action": "BLOCKED",
+ "guardrail_response": {"violated_categories": ["violence"]},
+ "violation_categories": ["violence"],
+ "masked_entity_count": {"EMAIL": 2, "PHONE": 1},
+ "duration": 0.05,
+ }
+ d = GuardrailSpanData.from_logging_entry(entry)
+ assert d.guardrail_name == "openai-moderation"
+ assert d.status == "guardrail_intervened"
+ assert d.provider == "openai"
+ assert d.action == "BLOCKED"
+ assert '"violence"' in (d.response_json or "")
+ assert d.violation_categories == ("violence",)
+ assert d.masked_entity_count == 3 # summed across entity types
+ assert d.duration == 0.05
+ assert d.error is not None # intervention → span marked ERROR
+
+
+def test_guardrail_span_data_success_has_no_error():
+ from litellm.integrations.otel.model.payloads import GuardrailSpanData
+
+ d = GuardrailSpanData.from_logging_entry(
+ {
+ "guardrail_name": "g",
+ "guardrail_mode": "pre_call",
+ "guardrail_status": "success",
+ }
+ )
+ assert d.error is None
+ assert d.status == "success"
+
+
+def test_request_identity_from_user_api_key_auth():
+ from litellm.integrations.otel.model.payloads import RequestIdentity
+
+ class _Auth:
+ team_id = "t9"
+ team_alias = "team nine"
+ api_key = "hashed-key"
+ user_id = "u9"
+ org_id = "o9"
+ key_alias = "my-key"
+ end_user_id = "eu9"
+
+ ident = RequestIdentity.from_user_api_key_auth(_Auth())
+ assert (ident.team_id, ident.team_alias, ident.key_hash) == (
+ "t9",
+ "team nine",
+ "hashed-key",
+ )
+ assert ident.end_user == "eu9"
+ assert ident.metadata["user_api_key_user_id"] == "u9"
+ assert ident.metadata["user_api_key_org_id"] == "o9"
+ assert ident.metadata["user_api_key_alias"] == "my-key"
+ assert ident.metadata["user_api_key_end_user_id"] == "eu9"
+
+
+# --- request-metadata translation layer (RequestContext) -------------------- #
+
+
+def test_request_context_splits_group_from_dispatched_model():
+ """On the proxy the caller asks for a model *group* that routes to a concrete
+ deployment: ``gen_ai.request.model`` is the group, ``litellm.provider.model``
+ is the dispatched (provider-prefixed) deployment model."""
+ from litellm.integrations.otel.model.metadata import RequestContext
+
+ payload = _sample_payload(
+ model="openai/gpt-5.4-mini", # reconstructed dispatched name
+ model_group="gpt-5.4-mini", # user-facing requested name
+ model_id="dep-123",
+ )
+ ctx = RequestContext.from_standard_logging_payload(payload)
+ assert ctx.request_model == "gpt-5.4-mini"
+ assert ctx.provider_model == "openai/gpt-5.4-mini"
+ assert ctx.identity.provider_model == "openai/gpt-5.4-mini"
+ assert ctx.model_group == "gpt-5.4-mini"
+ assert ctx.model_id == "dep-123"
+
+
+def test_request_context_sdk_path_has_no_group():
+ """Without a model group (the SDK path) the request and provider models
+ coincide on the single call model."""
+ from litellm.integrations.otel.model.metadata import RequestContext
+
+ payload = _sample_payload() # model="gpt-4o", no model_group
+ ctx = RequestContext.from_standard_logging_payload(payload)
+ assert ctx.request_model == "gpt-4o"
+ assert ctx.provider_model == "gpt-4o"
+ assert ctx.model_group is None
+
+
+def test_request_context_prefers_explicit_dispatched_model():
+ """``hidden_params.litellm_model_name`` is the authoritative dispatched model
+ when present, winning over the reconstructed top-level ``model``."""
+ from litellm.integrations.otel.model.metadata import RequestContext
+
+ payload = _sample_payload(
+ model="gpt-4o",
+ model_group="gpt-4o",
+ hidden_params={"litellm_model_name": "azure/my-deployment"},
+ )
+ ctx = RequestContext.from_standard_logging_payload(payload)
+ assert ctx.request_model == "gpt-4o"
+ assert ctx.provider_model == "azure/my-deployment"
+
+
+def test_content_capture_opt_in_retains_bodies():
+ payload = _sample_payload(
+ messages=[{"role": "user", "content": "secret prompt"}],
+ )
+ payload["response"]["choices"] = [
+ {"finish_reason": "stop", "message": {"role": "assistant", "content": "hi"}}
+ ]
+ data = LLMCallSpanData.from_standard_logging_payload(payload, capture_content=True)
+ assert data.messages_in and data.messages_in[0]["content"] == "secret prompt"
+ assert data.choices_out and data.choices_out[0]["message"]["content"] == "hi"
+
+
+# --- config ----------------------------------------------------------------- #
+
+
+def test_capture_span_content_resolves_modes():
+ from litellm.integrations.otel.model.config import (
+ CaptureMessageContent,
+ OpenTelemetryV2Config,
+ )
+
+ # default (no_content) → off
+ assert OpenTelemetryV2Config().capture_span_content is False
+ assert (
+ OpenTelemetryV2Config(
+ capture_message_content=CaptureMessageContent.SPAN_ONLY
+ ).capture_span_content
+ is True
+ )
+ assert (
+ OpenTelemetryV2Config(
+ capture_message_content=CaptureMessageContent.SPAN_AND_EVENT
+ ).capture_span_content
+ is True
+ )
+ # event-only does not authorize span-attribute content
+ assert (
+ OpenTelemetryV2Config(
+ capture_message_content=CaptureMessageContent.EVENT_ONLY
+ ).capture_span_content
+ is False
+ )
+
+
+def test_v2_flag_is_off_by_default(monkeypatch):
+ monkeypatch.delenv("LITELLM_OTEL_V2", raising=False)
+ assert is_otel_v2_enabled() is False
+ monkeypatch.setenv("LITELLM_OTEL_V2", "true")
+ assert is_otel_v2_enabled() is True
+
+
+def test_config_from_env(monkeypatch):
+ for var in (
+ "OTEL_EXPORTER",
+ "OTEL_EXPORTER_OTLP_PROTOCOL",
+ "OTEL_ENDPOINT",
+ "OTEL_EXPORTER_OTLP_ENDPOINT",
+ "OTEL_HEADERS",
+ "OTEL_EXPORTER_OTLP_HEADERS",
+ "OTEL_SERVICE_NAME",
+ "LITELLM_OTEL_LEGACY_COMPAT",
+ ):
+ monkeypatch.delenv(var, raising=False)
+
+ monkeypatch.setenv("OTEL_EXPORTER_OTLP_ENDPOINT", "https://collector:4318")
+ monkeypatch.setenv("OTEL_SERVICE_NAME", "my-svc")
+ cfg = OpenTelemetryV2Config.from_env()
+ # endpoint with no explicit exporter implies OTLP/HTTP
+ assert cfg.exporter == "otlp_http"
+ assert cfg.endpoint == "https://collector:4318"
+ assert cfg.service_name == "my-svc"
+ assert cfg.legacy_compat is True # dual-emit default during deprecation window
+
+
+def test_config_legacy_compat_env_toggle(monkeypatch):
+ monkeypatch.setenv("LITELLM_OTEL_LEGACY_COMPAT", "false")
+ assert OpenTelemetryV2Config.from_env().legacy_compat is False
+
+
+# --- baggage allowlist (the antipattern boundary) --------------------------- #
+
+
+def test_promoted_baggage_is_bounded_allowlist():
+ identity = RequestIdentity(
+ call_id="c1",
+ team_id="t1",
+ team_alias="team one",
+ key_hash="hsh",
+ end_user="u1",
+ metadata={"user_api_key_org_id": "org1", "secret_blob": "should-not-promote"},
+ )
+ promoted = promoted_baggage(identity, "gpt-4o", BAGGAGE_PROMOTED_KEYS)
+ assert promoted[LiteLLM.TEAM_ID] == "t1"
+ assert promoted[LiteLLM.TEAM_ALIAS] == "team one"
+ assert promoted[GenAI.REQUEST_MODEL] == "gpt-4o"
+ # allowlisted metadata sub-key is promoted under the litellm.metadata.* prefix
+ assert promoted[f"{LiteLLM.METADATA_PREFIX}user_api_key_org_id"] == "org1"
+ # full metadata blob is NOT promoted
+ assert all("secret_blob" not in key for key in promoted)
+ # http.* is never a promoted key
+ assert HTTP.ROUTE not in promoted
+ assert HTTP.REQUEST_METHOD not in promoted
diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_vendor_mappers.py b/tests/test_litellm/integrations/otel/test_otel_v2_vendor_mappers.py
new file mode 100644
index 00000000000..94cb79f53b8
--- /dev/null
+++ b/tests/test_litellm/integrations/otel/test_otel_v2_vendor_mappers.py
@@ -0,0 +1,196 @@
+"""Tests for the vendor mappers (OpenInference, Langfuse, Weave, Langtrace).
+
+Composition over inheritance: each vendor's vocabulary is a mapper. Layering
+mappers on the same span carries multiple naming schemes for different
+backends, so one trace lights up every configured destination.
+"""
+
+import json
+
+import pytest
+
+from litellm.integrations.otel import GenAIOperation
+from litellm.integrations.otel.mappers import (
+ GenAIMapper,
+ LangfuseMapper,
+ LangtraceMapper,
+ OpenInferenceMapper,
+ WeaveMapper,
+ resolve_mappers,
+)
+from litellm.integrations.otel.model.payloads import (
+ LLMCallSpanData,
+ LLMRequestParams,
+ LLMUsage,
+ RequestIdentity,
+ ServerInfo,
+ ToolDefinition,
+)
+
+
+def _llm_call(**overrides):
+ base = dict(
+ operation=GenAIOperation.CHAT,
+ provider="openai",
+ request_model="gpt-4o",
+ response_model="gpt-4o-2024",
+ response_id="resp_1",
+ request_params=LLMRequestParams(temperature=0.5, top_p=0.9, max_tokens=128),
+ usage=LLMUsage(input_tokens=12, output_tokens=8, total_tokens=20),
+ finish_reasons=("stop",),
+ error=None,
+ response_cost=0.001,
+ server=ServerInfo("api.openai.com", 443),
+ identity=RequestIdentity(call_id="c1", team_id="t1", team_alias="team one"),
+ is_streaming=False,
+ tools=(
+ ToolDefinition(
+ name="lookup_weather",
+ description="Get weather",
+ parameters_json='{"type":"object"}',
+ ),
+ ),
+ messages_in=(
+ {"role": "system", "content": "Be concise."},
+ {"role": "user", "content": "What's the weather?"},
+ ),
+ choices_out=(
+ {
+ "finish_reason": "stop",
+ "message": {"role": "assistant", "content": "Sunny."},
+ },
+ ),
+ system_fingerprint="fp_abc",
+ )
+ base.update(overrides)
+ return LLMCallSpanData(**base)
+
+
+# --------------------------------------------------------------------------- #
+# OpenInference (Arize + Phoenix shared vocabulary)
+# --------------------------------------------------------------------------- #
+
+
+def test_openinference_mapper_input_output_messages():
+ attrs = OpenInferenceMapper().map(_llm_call())
+ assert attrs["openinference.span.kind"] == "LLM"
+ assert attrs["llm.model_name"] == "gpt-4o"
+ assert attrs["llm.provider"] == "openai"
+ assert attrs["llm.input_messages.0.message.role"] == "system"
+ assert attrs["llm.input_messages.0.message.content"] == "Be concise."
+ assert attrs["llm.input_messages.1.message.role"] == "user"
+ assert attrs["llm.output_messages.0.message.role"] == "assistant"
+ assert attrs["llm.output_messages.0.message.content"] == "Sunny."
+ assert attrs["llm.token_count.prompt"] == 12
+ assert attrs["llm.token_count.completion"] == 8
+ assert attrs["llm.token_count.total"] == 20
+ # tool definitions ride the OpenInference schema
+ assert attrs["llm.tools.0.tool.name"] == "lookup_weather"
+ # invocation_parameters is JSON-serialized
+ params = json.loads(attrs["llm.invocation_parameters"])
+ assert params["temperature"] == 0.5
+ assert params["max_tokens"] == 128
+
+
+def test_openinference_mapper_skips_non_llm_roles():
+ from litellm.integrations.otel.model.payloads import GuardrailSpanData
+
+ assert OpenInferenceMapper().map(GuardrailSpanData("presidio")) == {}
+
+
+def test_openinference_multimodal_content_text_only():
+ data = _llm_call(
+ messages_in=(
+ {
+ "role": "user",
+ "content": [
+ {"type": "text", "text": "hi "},
+ {"type": "image_url", "image_url": {"url": "x"}},
+ {"type": "text", "text": "there"},
+ ],
+ },
+ )
+ )
+ attrs = OpenInferenceMapper().map(data)
+ assert attrs["llm.input_messages.0.message.content"] == "hi there"
+
+
+# --------------------------------------------------------------------------- #
+# Langfuse
+# --------------------------------------------------------------------------- #
+
+
+def test_langfuse_mapper_observation_attrs():
+ attrs = LangfuseMapper().map(_llm_call())
+ assert attrs["langfuse.observation.type"] == "generation"
+ assert attrs["langfuse.observation.model.name"] == "gpt-4o"
+ assert attrs["langfuse.observation.metadata.provider"] == "openai"
+ usage = json.loads(attrs["langfuse.observation.usage_details"])
+ assert usage["input"] == 12 and usage["output"] == 8
+ params = json.loads(attrs["langfuse.observation.model.parameters"])
+ assert params["temperature"] == 0.5
+ cost = json.loads(attrs["langfuse.observation.cost_details"])
+ assert cost["total"] == 0.001
+ assert attrs["langfuse.trace.metadata.team_id"] == "t1"
+
+
+def test_langfuse_mapper_skips_when_no_messages():
+ data = _llm_call(messages_in=(), choices_out=())
+ attrs = LangfuseMapper().map(data)
+ assert "langfuse.observation.input" not in attrs
+ assert "langfuse.observation.output" not in attrs
+
+
+# --------------------------------------------------------------------------- #
+# Weave
+# --------------------------------------------------------------------------- #
+
+
+def test_weave_mapper_display_and_output():
+ attrs = WeaveMapper().map(_llm_call())
+ assert attrs["weave.display_name"] == "chat gpt-4o"
+ assert attrs["weave.call_id"] == "c1"
+ decoded = json.loads(attrs["weave.output"])
+ assert decoded[0]["message"]["content"] == "Sunny."
+
+
+# --------------------------------------------------------------------------- #
+# Langtrace
+# --------------------------------------------------------------------------- #
+
+
+def test_langtrace_mapper_attrs():
+ attrs = LangtraceMapper().map(_llm_call())
+ assert attrs["gen_ai.operation.name"] == "chat"
+ assert attrs["langtrace.service.name"] == "openai"
+ assert attrs["llm.model"] == "gpt-4o"
+ assert attrs["gen_ai.response.model"] == "gpt-4o-2024"
+ assert attrs["gen_ai.system_fingerprint"] == "fp_abc"
+ assert attrs["llm.temperature"] == 0.5
+ assert attrs["llm.token.counts.total"] == 20
+
+
+# --------------------------------------------------------------------------- #
+# Composition (the V2 punchline)
+# --------------------------------------------------------------------------- #
+
+
+def test_resolve_mappers_composition_layers_vocabularies():
+ """One span, three vocabularies — Arize + Langfuse + canonical together."""
+ chain = resolve_mappers(["genai", "openinference", "langfuse"])
+ data = _llm_call()
+ union: dict = {}
+ for mapper in chain:
+ union.update(mapper.map(data))
+ # Canonical
+ assert union["gen_ai.operation.name"] == "chat"
+ # OpenInference
+ assert union["llm.model_name"] == "gpt-4o"
+ assert union["openinference.span.kind"] == "LLM"
+ # Langfuse
+ assert union["langfuse.observation.type"] == "generation"
+
+
+def test_resolve_mappers_rejects_unknown_name():
+ with pytest.raises(ValueError, match="unknown mapper name 'nope'"):
+ resolve_mappers(["genai", "nope"])
diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py
index 7a2b20bd8fb..f0015d9df0d 100644
--- a/tests/test_litellm/proxy/test_proxy_utils.py
+++ b/tests/test_litellm/proxy/test_proxy_utils.py
@@ -71,6 +71,53 @@ def test_proxy_only_error_false_for_other_error_type():
)
+@pytest.mark.asyncio
+async def test_proxy_only_error_log_marks_no_upstream_llm_call():
+ """A proxy-gate error (auth/rate-limit) synthesizes a ``Logging`` object and
+ fires ``pre_call`` so the failure is logged — but it must tag the object with
+ ``LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL`` so tracing callbacks don't fabricate
+ an LLM-call span for a request that never reached a provider (root cause of the
+ misplaced gen-AI span on auth failure)."""
+ from litellm.constants import LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL
+ from litellm.proxy._types import UserAPIKeyAuth
+
+ proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
+ captured = {}
+
+ def fake_pre_call(self, *args, **kwargs):
+ captured["flag"] = self.model_call_details.get(
+ LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL
+ )
+
+ from litellm.litellm_core_utils.litellm_logging import Logging
+
+ orig_pre_call = Logging.pre_call
+ orig_async_failure = Logging.async_failure_handler
+ Logging.pre_call = fake_pre_call
+
+ async def _noop_async_failure(self, *args, **kwargs):
+ return None
+
+ Logging.async_failure_handler = _noop_async_failure
+ try:
+ await proxy_logging_obj._handle_logging_proxy_only_error(
+ request_data={
+ "model": "gpt-4o",
+ "messages": [{"role": "user", "content": "hi"}],
+ },
+ user_api_key_dict=UserAPIKeyAuth(
+ api_key="sk-bad", request_route="/v1/chat/completions"
+ ),
+ route="/v1/chat/completions",
+ original_exception=Exception("bad key"),
+ )
+ finally:
+ Logging.pre_call = orig_pre_call
+ Logging.async_failure_handler = orig_async_failure
+
+ assert captured.get("flag") is True
+
+
def test_get_model_group_info_order():
from litellm import Router
from litellm.proxy.proxy_server import _get_model_group_info
diff --git a/tests/test_litellm/test_service_logger.py b/tests/test_litellm/test_service_logger.py
index ed44fe9b9f2..34fe73382b3 100644
--- a/tests/test_litellm/test_service_logger.py
+++ b/tests/test_litellm/test_service_logger.py
@@ -6,10 +6,12 @@ is called without call_type in kwargs (e.g. from batch polling callbacks).
"""
import pytest
-from datetime import datetime, timedelta
+from datetime import datetime
from unittest.mock import AsyncMock, patch
+import litellm
from litellm._service_logger import ServiceLogging
+from litellm.types.services import ServiceTypes
@pytest.mark.asyncio
@@ -95,3 +97,156 @@ async def test_async_log_success_event_should_handle_float_duration():
mock_hook.assert_called_once()
call_kwargs = mock_hook.call_args
assert call_kwargs.kwargs["duration"] == 1.5
+
+
+# --------------------------------------------------------------------------- #
+# V2 OpenTelemetry service-span dispatch (regression: service spans were always
+# dropped because the dispatch only recognized the legacy OpenTelemetry class).
+# --------------------------------------------------------------------------- #
+
+
+def _make_otel_v2_logger():
+ pytest.importorskip("opentelemetry")
+ from opentelemetry.sdk.trace.export.in_memory_span_exporter import (
+ InMemorySpanExporter,
+ )
+
+ from litellm.integrations.otel import OpenTelemetryV2Config
+ from litellm.integrations.otel.plumbing import providers
+ from litellm.integrations.otel.logger import OpenTelemetryV2
+
+ cfg = OpenTelemetryV2Config(exporter="in_memory")
+ exporter = InMemorySpanExporter()
+ tracer_provider = providers.build_tracer_provider(cfg, exporter=exporter)
+ return OpenTelemetryV2(config=cfg, tracer_provider=tracer_provider), exporter
+
+
+def test_resolve_otel_service_logger_recognizes_v2_instance():
+ """The V2 logger is a plain CustomLogger, not a subclass of the legacy
+ OpenTelemetry. The resolver must still recognize it (else service spans are
+ silently dropped)."""
+ service_logger = ServiceLogging()
+ v2_logger, _ = _make_otel_v2_logger()
+ assert service_logger._resolve_otel_service_logger(v2_logger) is v2_logger
+
+
+def test_resolve_otel_service_logger_recognizes_otel_string(monkeypatch):
+ # The "otel" string path resolves through the proxy's registered logger, so
+ # it needs the proxy server module importable.
+ try:
+ import litellm.proxy.proxy_server as proxy_server
+ except ImportError:
+ pytest.skip("proxy server dependencies not installed")
+ service_logger = ServiceLogging()
+ v2_logger, _ = _make_otel_v2_logger()
+
+ monkeypatch.setattr(proxy_server, "open_telemetry_logger", v2_logger, raising=False)
+ assert service_logger._resolve_otel_service_logger("otel") is v2_logger
+
+
+def test_resolve_otel_service_logger_ignores_unrelated_callback():
+ service_logger = ServiceLogging()
+ assert service_logger._resolve_otel_service_logger("prometheus_system") is None
+ assert service_logger._resolve_otel_service_logger(object()) is None
+
+
+@pytest.mark.asyncio
+async def test_service_span_emitted_for_v2_logger_in_service_callback(monkeypatch):
+ """End-to-end: a V2 logger registered in ``litellm.service_callback`` produces
+ a service span when ``async_service_success_hook`` fires with a parent span."""
+ from litellm.integrations.otel.model.spans import SpanRole
+
+ v2_logger, exporter = _make_otel_v2_logger()
+ parent = v2_logger._emitter.start_span(
+ SpanRole.PROXY_REQUEST, "POST /chat/completions"
+ )
+
+ monkeypatch.setattr(litellm, "service_callback", [v2_logger])
+ service_logger = ServiceLogging()
+
+ await service_logger.async_service_success_hook(
+ service=ServiceTypes.REDIS,
+ call_type="async_set_cache",
+ duration=0.01,
+ parent_otel_span=parent,
+ )
+ parent.end()
+
+ names = [s.name for s in exporter.get_finished_spans()]
+ # Span name is "{service} {call_type}" so repeated calls stay distinguishable.
+ assert "redis async_set_cache" in names
+
+
+@pytest.mark.asyncio
+async def test_service_span_not_duplicated_for_string_and_instance(monkeypatch):
+ """``service_callback`` can hold the ``"otel"`` string AND the registered
+ logger instance — the V2 logger self-registers its instance even when the
+ string is present. Both references resolve to the same logger, so the dispatch
+ loop must emit only ONE span per service event, not one per reference. Before
+ the dedup guard this produced duplicate ``postgres ...`` / ``redis ...`` spans.
+ """
+ try:
+ import litellm.proxy.proxy_server as proxy_server
+ except ImportError:
+ pytest.skip("proxy server dependencies not installed")
+ from litellm.integrations.otel.model.spans import SpanRole
+
+ v2_logger, exporter = _make_otel_v2_logger()
+ parent = v2_logger._emitter.start_span(
+ SpanRole.PROXY_REQUEST, "POST /chat/completions"
+ )
+
+ # The "otel" string resolves to the proxy's registered logger (the same
+ # instance), so the list holds two references to one logger.
+ monkeypatch.setattr(proxy_server, "open_telemetry_logger", v2_logger, raising=False)
+ monkeypatch.setattr(litellm, "service_callback", ["otel", v2_logger])
+ service_logger = ServiceLogging()
+
+ await service_logger.async_service_success_hook(
+ service=ServiceTypes.DB,
+ call_type="get_user_object",
+ duration=0.01,
+ parent_otel_span=parent,
+ )
+ parent.end()
+
+ db_spans = [
+ s for s in exporter.get_finished_spans() if s.name == "postgres get_user_object"
+ ]
+ assert len(db_spans) == 1
+
+
+@pytest.mark.asyncio
+async def test_service_failure_span_not_duplicated_for_string_and_instance(
+ monkeypatch,
+):
+ """Failure path mirror of the dedup guard — one span per failed service event,
+ even with both the ``"otel"`` string and the instance in ``service_callback``."""
+ try:
+ import litellm.proxy.proxy_server as proxy_server
+ except ImportError:
+ pytest.skip("proxy server dependencies not installed")
+ from litellm.integrations.otel.model.spans import SpanRole
+
+ v2_logger, exporter = _make_otel_v2_logger()
+ parent = v2_logger._emitter.start_span(
+ SpanRole.PROXY_REQUEST, "POST /chat/completions"
+ )
+
+ monkeypatch.setattr(proxy_server, "open_telemetry_logger", v2_logger, raising=False)
+ monkeypatch.setattr(litellm, "service_callback", ["otel", v2_logger])
+ service_logger = ServiceLogging()
+
+ await service_logger.async_service_failure_hook(
+ service=ServiceTypes.DB,
+ call_type="get_user_object",
+ duration=0.01,
+ error="boom",
+ parent_otel_span=parent,
+ )
+ parent.end()
+
+ db_spans = [
+ s for s in exporter.get_finished_spans() if s.name == "postgres get_user_object"
+ ]
+ assert len(db_spans) == 1
diff --git a/uv.lock b/uv.lock
index 080cd89df6b..a632f72936d 100644
--- a/uv.lock
+++ b/uv.lock
@@ -294,6 +294,18 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/ee/82/82745642d3c46e7cea25e1885b014b033f4693346ce46b7f47483cf5d448/argon2_cffi_bindings-25.1.0-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:da0c79c23a63723aa5d782250fbf51b768abca630285262fb5144ba5ae01e520", size = 29187, upload-time = "2025-07-30T10:02:03.674Z" },
]
+[[package]]
+name = "asgiref"
+version = "3.11.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "typing-extensions", marker = "python_full_version < '3.11'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/63/40/f03da1264ae8f7cfdbf9146542e5e7e8100a4c66ab48e791df9a03d3f6c0/asgiref-3.11.1.tar.gz", hash = "sha256:5f184dc43b7e763efe848065441eac62229c9f7b0475f41f80e207a114eda4ce", size = 38550, upload-time = "2026-02-03T13:30:14.33Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/5c/0a/a72d10ed65068e115044937873362e6e32fab1b7dce0046aeb224682c989/asgiref-3.11.1-py3-none-any.whl", hash = "sha256:e8667a091e69529631969fd45dc268fa79b99c92c5fcdda727757e52146ec133", size = 24345, upload-time = "2026-02-03T13:30:13.039Z" },
+]
+
[[package]]
name = "assemblyai"
version = "0.52.4"
@@ -3349,6 +3361,7 @@ proxy-runtime = [
{ name = "mangum" },
{ name = "opentelemetry-api" },
{ name = "opentelemetry-exporter-otlp" },
+ { name = "opentelemetry-instrumentation-fastapi" },
{ name = "opentelemetry-sdk" },
{ name = "prometheus-client" },
{ name = "pypdf" },
@@ -3411,6 +3424,7 @@ dev = [
{ name = "openapi-core" },
{ name = "opentelemetry-api" },
{ name = "opentelemetry-exporter-otlp" },
+ { name = "opentelemetry-instrumentation-fastapi" },
{ name = "opentelemetry-sdk" },
{ name = "parameterized" },
{ name = "psycopg" },
@@ -3444,6 +3458,7 @@ proxy-dev = [
{ name = "hypercorn" },
{ name = "opentelemetry-api" },
{ name = "opentelemetry-exporter-otlp" },
+ { name = "opentelemetry-instrumentation-fastapi" },
{ name = "opentelemetry-sdk" },
{ name = "prisma" },
{ name = "prometheus-client" },
@@ -3499,6 +3514,7 @@ requires-dist = [
{ name = "openai", specifier = ">=2.20.0,<3.0.0" },
{ name = "opentelemetry-api", marker = "extra == 'proxy-runtime'", specifier = "==1.28.0" },
{ name = "opentelemetry-exporter-otlp", marker = "extra == 'proxy-runtime'", specifier = "==1.28.0" },
+ { name = "opentelemetry-instrumentation-fastapi", marker = "extra == 'proxy-runtime'", specifier = "==0.49b0" },
{ name = "opentelemetry-sdk", marker = "extra == 'proxy-runtime'", specifier = "==1.28.0" },
{ name = "orjson", marker = "extra == 'proxy'", specifier = ">=3.11.6,<4.0" },
{ name = "polars", marker = "extra == 'proxy'", specifier = ">=1.38.1,<2.0" },
@@ -3573,6 +3589,7 @@ dev = [
{ name = "openapi-core", marker = "python_full_version < '3.14'", specifier = "==0.22.0" },
{ name = "opentelemetry-api", specifier = "==1.28.0" },
{ name = "opentelemetry-exporter-otlp", specifier = "==1.28.0" },
+ { name = "opentelemetry-instrumentation-fastapi", specifier = "==0.49b0" },
{ name = "opentelemetry-sdk", specifier = "==1.28.0" },
{ name = "parameterized", specifier = "==0.9.0" },
{ name = "psycopg", specifier = "==3.3.3" },
@@ -3606,6 +3623,7 @@ proxy-dev = [
{ name = "hypercorn", specifier = "==0.17.3" },
{ name = "opentelemetry-api", specifier = "==1.28.0" },
{ name = "opentelemetry-exporter-otlp", specifier = "==1.28.0" },
+ { name = "opentelemetry-instrumentation-fastapi", specifier = "==0.49b0" },
{ name = "opentelemetry-sdk", specifier = "==1.28.0" },
{ name = "prisma", specifier = "==0.11.0" },
{ name = "prometheus-client", specifier = "==0.20.0" },
@@ -4561,6 +4579,22 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/ba/46/ba2dc8d18b04acae3d34facd8fe1e5e0cdc9fe64292d45eca9d1d4a8a298/opentelemetry_instrumentation_anthropic-0.33.12-py3-none-any.whl", hash = "sha256:b31618d12a429045db14ed982a142a25df0f0f1dbf03d756e8d597f25b9a053d", size = 11024, upload-time = "2024-11-13T20:27:14.622Z" },
]
+[[package]]
+name = "opentelemetry-instrumentation-asgi"
+version = "0.49b0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "asgiref" },
+ { name = "opentelemetry-api" },
+ { name = "opentelemetry-instrumentation" },
+ { name = "opentelemetry-semantic-conventions" },
+ { name = "opentelemetry-util-http" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/e8/55/693c3d0938ba5fead5c3aa4ac7022a992b4ff99a8e9979800d0feb843ff4/opentelemetry_instrumentation_asgi-0.49b0.tar.gz", hash = "sha256:959fd9b1345c92f20c6ef1d42f92ef6a76b3c3083fbc4104d59da6859b15b083", size = 24117, upload-time = "2024-11-05T19:21:46.769Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/2c/0b/7900c782a1dfaa584588d724bc3bbdf8405a32497537dd96b3fcbf8461b9/opentelemetry_instrumentation_asgi-0.49b0-py3-none-any.whl", hash = "sha256:722a90856457c81956c88f35a6db606cc7db3231046b708aae2ddde065723dbe", size = 16326, upload-time = "2024-11-05T19:20:46.176Z" },
+]
+
[[package]]
name = "opentelemetry-instrumentation-bedrock"
version = "0.33.12"
@@ -4607,6 +4641,22 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/bf/08/dce2b7926ace0204ce7946563348e1ff755873e387833484791e4ed391c8/opentelemetry_instrumentation_cohere-0.33.12-py3-none-any.whl", hash = "sha256:3bee3f7f7105259c85145be8c3b68612421860c95ad170f4d03144a3b8c07418", size = 5589, upload-time = "2024-11-13T20:27:21.317Z" },
]
+[[package]]
+name = "opentelemetry-instrumentation-fastapi"
+version = "0.49b0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "opentelemetry-api" },
+ { name = "opentelemetry-instrumentation" },
+ { name = "opentelemetry-instrumentation-asgi" },
+ { name = "opentelemetry-semantic-conventions" },
+ { name = "opentelemetry-util-http" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/fe/bf/8e6d2a4807360f2203192017eb4845f5628dbeaf0597adf3d141cc5c24e1/opentelemetry_instrumentation_fastapi-0.49b0.tar.gz", hash = "sha256:6d14935c41fd3e49328188b6a59dd4c37bd17a66b01c15b0c64afa9714a1f905", size = 19230, upload-time = "2024-11-05T19:21:59.361Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/b1/f4/0895b9410c10abf987c90dee1b7688a8f2214a284fe15e575648f6a1473a/opentelemetry_instrumentation_fastapi-0.49b0-py3-none-any.whl", hash = "sha256:646e1b18523cbe6860ae9711eb2c7b9c85466c3c7697cd6b8fb5180d85d3fe6e", size = 12101, upload-time = "2024-11-05T19:21:01.805Z" },
+]
+
[[package]]
name = "opentelemetry-instrumentation-google-generativeai"
version = "0.33.12"
From 12d29a38a7050b9feca035173f7f912f738b6aca Mon Sep 17 00:00:00 2001
From: yuneng-jiang
Date: Fri, 29 May 2026 23:17:24 -0700
Subject: [PATCH 20/44] tests(proxy_server): surface current behavior in tests
(#29309)
* test(proxy/proxy_server): pin forwarding routes (PR2) (#28887)
* test(proxy): pin proxy_server.py forwarding-route behavior
PR2 of the proxy_server.py behavior-pinning project: fills the 12
forwarding-route test files added by the harness PR with happy + error
pins for all 52 LLM-facing routes (models, chat/completions, completions,
embeddings, moderations, audio, assistants, threads, utils, model-info,
model-metrics, queue). Every happy-path test asserts the full response
dict via normalize() so the gate enforces real shape pinning rather
than status codes.
* test(proxy): drop task-plumbing comments from PR2 test files
* test(proxy): tighten PR2 error-path status-code pins
Apply the same review feedback Greptile gave on PR1 (#28856) and PR3
(#28850) to PR2's forwarding-route tests:
- Replace permissive `>= 400` / `in (X, Y)` status assertions with the
exact 500/405 the handler actually returns, so a regression that
silently shifts the code now fails the pin.
- Add a body-presence check alongside each tightened status assertion
to satisfy _pin_check.py's no-status-only rule.
---------
Co-authored-by: Claude
* test(proxy): pin proxy_server.py non-route surface behavior (PR1) (#28856)
* test(proxy): pin proxy_server.py non-route surface behavior (PR1)
Fills the 7 PR1 placeholder files under tests/test_litellm/proxy/proxy_server/
with behavior pins for the non-route surface of proxy_server.py:
lifecycle/init/shutdown, ProxyConfig class methods, DB-overlay config scrubbers,
spend counters, background-health helpers, OpenAPI customization, exception
handlers, and streaming-generator helpers.
233 tests cover 101 pin-list symbols (1+ happy + 1+ error each). New-tests-only
coverage on litellm/proxy/proxy_server.py: 32.80% line / 20.91% branch (PR1
gate: 25% line / 18% branch). Full directory runs in ~22s with -n 4.
Plan: https://www.notion.so/Plan-Pin-proxy_server-py-behavior-2026-05-25-36c43b8acdab81ee845fd5365128a2fc
* test(proxy): address Greptile review comments on test_lifecycle.py
- test_initialize_signature_is_async_with_expected_params: hard-code
expected_param_count so a signature change actually trips the gate
(previously both sides of the comparison were len(sig.parameters)).
- test_check_request_disconnection_invalid_when_connected_times_out:
patch asyncio.sleep so the test no longer spins for ~1.2 s of real
wall-clock; timeout lowered to 0.05 s.
---------
Co-authored-by: Claude
* test(proxy/proxy_server): pin control-plane routes (PR3) (#28850)
* test(proxy/proxy_server): pin misc routes (PR3, partial)
Adds happy + error tests for the misc control-plane routes:
GET /, /routes, /adaptive_router/state, /get_logo_url,
/get_image, /get_favicon.
Also gitignores .pin_list.txt (used by the pin gate).
* test(proxy/proxy_server): pin login/SSO routes (PR3, partial)
Adds happy + error tests for the 5 login/SSO control-plane routes:
GET /fallback/login, POST /login, POST /v2/login, POST /v3/login,
POST /v3/login/exchange. Mocks authenticate_user and
create_ui_token_object at their imported location.
* test(proxy/proxy_server): pin onboarding routes (PR3, partial)
Adds happy + error tests for the 2 onboarding control-plane routes:
GET /onboarding/get_token, POST /onboarding/claim_token. Wires a
MagicMock async context manager for prisma_client.db.tx() and
signs the onboarding JWT with the patched master_key.
* test(proxy/proxy_server): pin model_cost_map reload routes (PR3, partial)
Adds happy + error tests for the 5 model-cost-map control-plane routes:
POST /reload/model_cost_map, POST|DELETE|GET
/schedule/model_cost_map_reload(/status), GET /model/cost_map/source.
Attaches litellm_config to mock_prisma per-test (the table is not in
the default _PRISMA_TABLES fixture).
* test(proxy/proxy_server): pin anthropic_beta_headers reload routes (PR3, partial)
Adds happy + error tests for the 4 anthropic-beta-headers control-plane
routes: POST /reload/anthropic_beta_headers, POST|DELETE|GET
/schedule/anthropic_beta_headers_reload(/status). Stubs
db.litellm_config (not in default _PRISMA_TABLES) and monkeypatches
reload_beta_headers_config so no network calls fire.
* test(proxy/proxy_server): pin invitation routes (PR3, partial)
Adds happy + error tests for the 4 invitation control-plane routes:
POST /invitation/new, GET /invitation/info, POST /invitation/update,
POST /invitation/delete. Patches _user_has_admin_privileges /
_user_has_admin_view to avoid extensive get_user_object mocking.
* test(proxy/proxy_server): pin config CRUD routes (PR3, partial)
Adds happy + error tests for the 8 config-CRUD control-plane routes:
POST /config/update, POST|GET /config/field/update|info, GET /config/list,
POST /config/field/delete, POST /config/callback/delete,
GET /get/config/callbacks, GET /config/yaml. Attaches litellm_config
to mock_prisma per-test.
* test(proxy/proxy_server): tighten pin assertions per review
- test_routes_misc.py: `b"" in response.content` is trivially true;
replace with `len(response.content) > 0` so an empty 405 body trips
the gate.
- test_routes_login_sso.py: `len(response.content) >= 0` is trivially
true; tighten to `> 0`.
- test_routes_anthropic_beta.py: replace brittle string-literal checks
on the serialized JSON (`'"interval_hours": 12' in payload`) with
`json.loads` + dict access so the assertion survives any serializer
spacing.
- test_routes_config.py: `assert status_code in (404, 500)` was too
permissive; the handler re-raises HTTPException(404) verbatim, so
pin 404 strictly.
---------
Co-authored-by: Claude
---------
Co-authored-by: Claude
---
.gitignore | 5 +-
.../proxy_server/test_background_health.py | 514 ++++++-
.../proxy_server/test_exception_handlers.py | 223 ++-
.../proxy/proxy_server/test_lifecycle.py | 565 ++++++-
.../test_openapi_customization.py | 448 +++++-
.../proxy/proxy_server/test_proxy_config.py | 1316 ++++++++++++++++-
.../test_routes_anthropic_beta.py | 371 ++++-
.../proxy_server/test_routes_assistants.py | 181 ++-
.../proxy/proxy_server/test_routes_audio.py | 194 ++-
.../test_routes_chat_completions.py | 135 +-
.../proxy_server/test_routes_completions.py | 127 +-
.../proxy/proxy_server/test_routes_config.py | 592 +++++++-
.../proxy_server/test_routes_embeddings.py | 122 +-
.../proxy_server/test_routes_invitation.py | 388 ++++-
.../proxy_server/test_routes_login_sso.py | 388 ++++-
.../proxy/proxy_server/test_routes_misc.py | 231 ++-
.../test_routes_model_cost_map.py | 372 ++++-
.../proxy_server/test_routes_model_info.py | 150 +-
.../proxy_server/test_routes_model_metrics.py | 229 ++-
.../proxy/proxy_server/test_routes_models.py | 133 +-
.../proxy_server/test_routes_moderations.py | 112 +-
.../proxy_server/test_routes_onboarding.py | 351 ++++-
.../proxy/proxy_server/test_routes_queue.py | 92 +-
.../proxy/proxy_server/test_routes_threads.py | 275 +++-
.../proxy/proxy_server/test_routes_utils.py | 161 +-
.../proxy/proxy_server/test_spend_counters.py | 818 +++++++++-
.../proxy_server/test_streaming_helpers.py | 556 ++++++-
27 files changed, 9022 insertions(+), 27 deletions(-)
diff --git a/.gitignore b/.gitignore
index dff64e3c9e9..572830d35f6 100644
--- a/.gitignore
+++ b/.gitignore
@@ -28,6 +28,8 @@ litellm/tests/config_*.yaml
litellm/tests/langfuse.log
langfuse.log
.langfuse.log
+.pin_list.txt
+.cov_new.xml
litellm/tests/test_custom_logger.py
litellm/tests/langfuse.log
litellm/tests/dynamo*.log
@@ -120,4 +122,5 @@ crash.log
crash.*.log
# .terraform.lock.hcl is intentionally NOT ignored — it pins provider versions
# and should be committed.
-.vscode
\ No newline at end of file
+.vscode
+.pin_list.txt
diff --git a/tests/test_litellm/proxy/proxy_server/test_background_health.py b/tests/test_litellm/proxy/proxy_server/test_background_health.py
index ad6b4016461..ee8d8b22779 100644
--- a/tests/test_litellm/proxy/proxy_server/test_background_health.py
+++ b/tests/test_litellm/proxy/proxy_server/test_background_health.py
@@ -1 +1,513 @@
-"""Placeholder. Filled by a follow-up PR per the Notion plan."""
+"""Behavior pins for proxy_server background health-check helpers.
+
+Pins covered:
+- ``_get_process_rss_mb``
+- ``_rss_mb_for_log``
+- ``_run_direct_health_check_with_instrumentation``
+- ``_schedule_background_health_check_db_save``
+- ``_get_endpoint_exception_status``
+- ``_write_health_state_to_router_cache``
+- ``_adaptive_router_flusher_loop``
+- ``_run_background_health_check``
+"""
+
+from __future__ import annotations
+
+import asyncio
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+import litellm.proxy.proxy_server as proxy_server
+from litellm.proxy.proxy_server import (
+ _adaptive_router_flusher_loop,
+ _get_endpoint_exception_status,
+ _get_process_rss_mb,
+ _run_background_health_check,
+ _run_direct_health_check_with_instrumentation,
+ _rss_mb_for_log,
+ _schedule_background_health_check_db_save,
+ _write_health_state_to_router_cache,
+)
+
+from .conftest import normalize
+
+# ---------------------------------------------------------------------------
+# _get_process_rss_mb
+# ---------------------------------------------------------------------------
+
+
+def test_get_process_rss_mb_returns_positive_float():
+ value = _get_process_rss_mb()
+ assert value is not None
+ assert normalize(
+ {
+ "value_present": value is not None,
+ "value_type": type(value).__name__,
+ "positive": value > 0,
+ }
+ ) == {
+ "value_present": True,
+ "value_type": "float",
+ "positive": True,
+ }
+
+
+def test_get_process_rss_mb_returns_none_when_resource_raises(monkeypatch):
+ import resource
+
+ def _boom(*_args, **_kwargs):
+ raise OSError("nope")
+
+ monkeypatch.setattr(resource, "getrusage", _boom)
+ assert _get_process_rss_mb() is None
+
+
+# ---------------------------------------------------------------------------
+# _rss_mb_for_log
+# ---------------------------------------------------------------------------
+
+
+def test_rss_mb_for_log_formats_numeric_value(monkeypatch):
+ monkeypatch.setattr(proxy_server, "_get_process_rss_mb", lambda: 100.5)
+ result = _rss_mb_for_log()
+ assert normalize(
+ {
+ "format": result,
+ "is_string": isinstance(result, str),
+ "contains_mb": "100.50" in result,
+ }
+ ) == {
+ "format": "100.50",
+ "is_string": True,
+ "contains_mb": True,
+ }
+
+
+def test_rss_mb_for_log_unknown_when_rss_missing(monkeypatch):
+ monkeypatch.setattr(proxy_server, "_get_process_rss_mb", lambda: None)
+ assert _rss_mb_for_log() == "unknown"
+
+
+# ---------------------------------------------------------------------------
+# _run_direct_health_check_with_instrumentation
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_run_direct_health_check_with_instrumentation_returns_results(
+ monkeypatch,
+):
+ expected = (["healthy_ep"], ["unhealthy_ep"], {"m1": Exception("boom")})
+
+ async def _fake_perform(model_list, details, max_concurrency, **kwargs):
+ return expected
+
+ monkeypatch.setattr(proxy_server, "perform_health_check", _fake_perform)
+ monkeypatch.setattr(
+ proxy_server,
+ "health_check_filter_kwargs_from_general_settings",
+ lambda _gs: {},
+ )
+
+ healthy, unhealthy, exceptions = (
+ await _run_direct_health_check_with_instrumentation(
+ model_list=[{"model_name": "gpt-4"}],
+ details=False,
+ max_concurrency=1,
+ instrumentation_context={"source": "test"},
+ )
+ )
+
+ assert normalize(
+ {
+ "healthy": healthy,
+ "unhealthy": unhealthy,
+ "exception_keys": list(exceptions.keys()),
+ }
+ ) == {
+ "healthy": ["healthy_ep"],
+ "unhealthy": ["unhealthy_ep"],
+ "exception_keys": ["m1"],
+ }
+
+
+@pytest.mark.asyncio
+async def test_run_direct_health_check_raises_non_kwarg_typeerror(monkeypatch):
+ async def _boom(model_list, details, max_concurrency, **kwargs):
+ raise TypeError("totally unrelated")
+
+ monkeypatch.setattr(proxy_server, "perform_health_check", _boom)
+ monkeypatch.setattr(
+ proxy_server,
+ "health_check_filter_kwargs_from_general_settings",
+ lambda _gs: {},
+ )
+
+ with pytest.raises(TypeError):
+ await _run_direct_health_check_with_instrumentation(
+ model_list=[],
+ details=False,
+ max_concurrency=1,
+ instrumentation_context={},
+ )
+
+
+# ---------------------------------------------------------------------------
+# _schedule_background_health_check_db_save
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_schedule_background_health_check_db_save_creates_task(monkeypatch):
+ captured = {}
+
+ async def _fake_save(
+ prisma_client,
+ model_list,
+ healthy,
+ unhealthy,
+ start_time,
+ checked_by,
+ ):
+ captured["prisma_client"] = prisma_client
+ captured["model_list"] = model_list
+ captured["healthy"] = healthy
+ captured["unhealthy"] = unhealthy
+ captured["checked_by"] = checked_by
+
+ import litellm.proxy.health_endpoints._health_endpoints as he
+
+ monkeypatch.setattr(he, "_save_background_health_checks_to_db", _fake_save)
+
+ prisma_client = MagicMock()
+ shared_manager = SimpleNamespace(pod_id="pod-xyz")
+
+ _schedule_background_health_check_db_save(
+ prisma_client=prisma_client,
+ shared_health_manager=shared_manager,
+ model_list=[{"model_name": "gpt-4"}],
+ healthy_endpoints=[{"model_id": "h1"}],
+ unhealthy_endpoints=[{"model_id": "u1"}],
+ )
+
+ await asyncio.sleep(0)
+
+ assert normalize(
+ {
+ "prisma_present": captured.get("prisma_client") is prisma_client,
+ "checked_by": captured.get("checked_by"),
+ "healthy": captured.get("healthy"),
+ "unhealthy": captured.get("unhealthy"),
+ }
+ ) == {
+ "prisma_present": True,
+ "checked_by": "pod-xyz",
+ "healthy": [{"model_id": "h1"}],
+ "unhealthy": [{"model_id": "u1"}],
+ }
+
+
+def test_schedule_background_health_check_db_save_noop_when_prisma_none():
+ _schedule_background_health_check_db_save(
+ prisma_client=None,
+ shared_health_manager=None,
+ model_list=[],
+ healthy_endpoints=[],
+ unhealthy_endpoints=[],
+ )
+
+
+@pytest.mark.asyncio
+async def test_schedule_background_health_check_db_save_invalid_no_event_loop_raises(
+ monkeypatch,
+):
+ async def _fake_save(*_args, **_kwargs):
+ return None
+
+ import litellm.proxy.health_endpoints._health_endpoints as he
+
+ monkeypatch.setattr(he, "_save_background_health_checks_to_db", _fake_save)
+
+ def _broken_create_task(_coro):
+ raise RuntimeError("no running event loop")
+
+ monkeypatch.setattr(asyncio, "create_task", _broken_create_task)
+
+ with pytest.raises(RuntimeError):
+ _schedule_background_health_check_db_save(
+ prisma_client=MagicMock(),
+ shared_health_manager=None,
+ model_list=[],
+ healthy_endpoints=[],
+ unhealthy_endpoints=[],
+ )
+
+
+# ---------------------------------------------------------------------------
+# _get_endpoint_exception_status
+# ---------------------------------------------------------------------------
+
+
+def test_get_endpoint_exception_status_prefers_live_exception():
+ endpoint = {"model_id": "m1", "exception_status": 999}
+ exceptions = {"m1": SimpleNamespace(status_code=429)}
+ status = _get_endpoint_exception_status(endpoint, exceptions)
+ assert normalize(
+ {
+ "input_endpoint": endpoint,
+ "exceptions_keys": list(exceptions.keys()),
+ "status": status,
+ }
+ ) == {
+ "input_endpoint": {"model_id": "m1", "exception_status": 999},
+ "exceptions_keys": ["m1"],
+ "status": 429,
+ }
+
+
+def test_get_endpoint_exception_status_falls_back_to_stored_int():
+ endpoint = {"model_id": "m-missing", "exception_status": 503}
+ assert _get_endpoint_exception_status(endpoint, {}) == 503
+
+
+def test_get_endpoint_exception_status_default_500_when_no_data():
+ assert _get_endpoint_exception_status({}, {}) == 500
+
+
+def test_get_endpoint_exception_status_invalid_endpoint_type_raises():
+ with pytest.raises(AttributeError):
+ _get_endpoint_exception_status(None, {}) # type: ignore[arg-type]
+
+
+# ---------------------------------------------------------------------------
+# _write_health_state_to_router_cache
+# ---------------------------------------------------------------------------
+
+
+def test_write_health_state_to_router_cache_sets_states(monkeypatch):
+ fake_router = MagicMock()
+ fake_router.enable_health_check_routing = True
+ fake_router.health_check_ignore_transient_errors = False
+ fake_router.cooldown_time = 30
+ fake_router.health_state_cache = MagicMock()
+
+ monkeypatch.setattr(proxy_server, "llm_router", fake_router)
+
+ fake_states = {"m1": {"is_healthy": True}, "m2": {"is_healthy": False}}
+
+ import litellm.proxy.health_check as hc
+
+ monkeypatch.setattr(hc, "build_deployment_health_states", lambda **_kw: fake_states)
+
+ import litellm.router_utils.cooldown_handlers as cd
+
+ monkeypatch.setattr(cd, "_set_cooldown_deployments", lambda **_kw: None)
+
+ import litellm.router_utils.router_callbacks.track_deployment_metrics as tdm
+
+ monkeypatch.setattr(
+ tdm,
+ "increment_deployment_failures_for_current_minute",
+ lambda **_kw: None,
+ )
+
+ healthy = [{"model_id": "m1"}]
+ unhealthy = [{"model_id": "m2"}]
+ exceptions = {"m2": SimpleNamespace(status_code=500)}
+
+ _write_health_state_to_router_cache(healthy, unhealthy, exceptions)
+
+ fake_router.health_state_cache.set_deployment_health_states.assert_called_once_with(
+ fake_states
+ )
+
+ call_args = fake_router.health_state_cache.set_deployment_health_states.call_args[
+ 0
+ ][0]
+ assert normalize(
+ {
+ "states_keys": sorted(call_args.keys()),
+ "m1_healthy": call_args["m1"]["is_healthy"],
+ "m2_healthy": call_args["m2"]["is_healthy"],
+ }
+ ) == {
+ "states_keys": ["m1", "m2"],
+ "m1_healthy": True,
+ "m2_healthy": False,
+ }
+
+
+def test_write_health_state_to_router_cache_noop_when_router_none(monkeypatch):
+ monkeypatch.setattr(proxy_server, "llm_router", None)
+ _write_health_state_to_router_cache([], [], {})
+
+
+def test_write_health_state_to_router_cache_swallows_internal_failures(monkeypatch):
+ """The function logs and swallows exceptions so a bad cache call never crashes the loop."""
+ fake_router = MagicMock()
+ fake_router.enable_health_check_routing = True
+ fake_router.health_check_ignore_transient_errors = False
+ fake_router.health_state_cache.set_deployment_health_states.side_effect = (
+ RuntimeError("cache exploded")
+ )
+
+ monkeypatch.setattr(proxy_server, "llm_router", fake_router)
+
+ import litellm.proxy.health_check as hc
+
+ monkeypatch.setattr(
+ hc,
+ "build_deployment_health_states",
+ lambda **_kw: {"m1": {"is_healthy": True}},
+ )
+
+ _write_health_state_to_router_cache([{"model_id": "m1"}], [], {})
+
+
+# ---------------------------------------------------------------------------
+# _adaptive_router_flusher_loop
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_adaptive_router_flusher_loop_flushes_each_router(monkeypatch):
+ fake_ar = MagicMock()
+ fake_ar._state_loaded = True
+ fake_ar.queue.flush_state_to_db = AsyncMock()
+ fake_ar.queue.flush_session_to_db = AsyncMock()
+
+ fake_router = MagicMock()
+ fake_router.adaptive_routers = {"alpha": fake_ar}
+
+ monkeypatch.setattr(proxy_server, "llm_router", fake_router)
+ monkeypatch.setattr(proxy_server, "prisma_client", MagicMock())
+
+ # asyncio.sleep is awaited at the top of every iteration; raise CancelledError
+ # on the SECOND call so the first iteration completes its flush work.
+ call_count = {"n": 0}
+ _real_sleep = asyncio.sleep
+
+ async def _short_sleep(_seconds):
+ call_count["n"] += 1
+ if call_count["n"] >= 2:
+ raise asyncio.CancelledError()
+ await _real_sleep(0)
+
+ monkeypatch.setattr(proxy_server.asyncio, "sleep", _short_sleep)
+
+ with pytest.raises(asyncio.CancelledError):
+ await _adaptive_router_flusher_loop()
+
+ assert fake_ar.queue.flush_state_to_db.await_count == 1
+ assert fake_ar.queue.flush_session_to_db.await_count == 1
+
+
+@pytest.mark.asyncio
+async def test_adaptive_router_flusher_loop_times_out_when_sleep_real(monkeypatch):
+ """Confirms the loop is infinite — wait_for must raise TimeoutError."""
+ monkeypatch.setattr(proxy_server, "llm_router", MagicMock(adaptive_routers={}))
+ monkeypatch.setattr(proxy_server, "prisma_client", None)
+
+ # Bind the real asyncio.sleep before the patch so the replacement does not
+ # recurse into itself.
+ _real_sleep = asyncio.sleep
+
+ async def _instant_sleep(_seconds):
+ await _real_sleep(0)
+
+ monkeypatch.setattr(proxy_server.asyncio, "sleep", _instant_sleep)
+
+ with pytest.raises(asyncio.TimeoutError):
+ await asyncio.wait_for(_adaptive_router_flusher_loop(), timeout=0.2)
+
+
+# ---------------------------------------------------------------------------
+# _run_background_health_check
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_run_background_health_check_returns_immediately_when_interval_invalid(
+ monkeypatch,
+):
+ monkeypatch.setattr(proxy_server, "health_check_interval", None)
+
+ result = await _run_background_health_check()
+
+ assert normalize(
+ {
+ "result_is_none": result is None,
+ "loop_active": proxy_server.background_health_check_loop_active,
+ "interval": proxy_server.health_check_interval,
+ }
+ ) == {
+ "result_is_none": True,
+ "loop_active": False,
+ "interval": None,
+ }
+
+
+@pytest.mark.asyncio
+async def test_run_background_health_check_runs_one_cycle_then_cancels(monkeypatch):
+ monkeypatch.setattr(proxy_server, "health_check_interval", 60)
+ monkeypatch.setattr(proxy_server, "health_check_concurrency", 1)
+ monkeypatch.setattr(proxy_server, "health_check_details", True)
+ monkeypatch.setattr(proxy_server, "use_shared_health_check", False)
+ monkeypatch.setattr(proxy_server, "redis_usage_cache", None)
+ monkeypatch.setattr(proxy_server, "prisma_client", None)
+ monkeypatch.setattr(proxy_server, "background_health_check_loop_active", False)
+ monkeypatch.setattr(
+ proxy_server,
+ "llm_model_list",
+ [{"model_name": "gpt-4", "model_info": {}}],
+ )
+ monkeypatch.setattr(
+ proxy_server,
+ "health_check_results",
+ {"healthy_endpoints": [], "unhealthy_endpoints": []},
+ )
+
+ async def _fake_direct(*_a, **_kw):
+ return ([{"model_id": "h"}], [{"model_id": "u"}], {})
+
+ monkeypatch.setattr(
+ proxy_server,
+ "_run_direct_health_check_with_instrumentation",
+ _fake_direct,
+ )
+ monkeypatch.setattr(
+ proxy_server, "_schedule_background_health_check_db_save", lambda *a, **kw: None
+ )
+ monkeypatch.setattr(
+ proxy_server, "_write_health_state_to_router_cache", lambda *a, **kw: None
+ )
+ monkeypatch.setattr(
+ proxy_server,
+ "health_check_filter_kwargs_from_general_settings",
+ lambda _gs: {},
+ )
+
+ sleep_calls = {"n": 0}
+
+ async def _stop_sleep(_seconds):
+ sleep_calls["n"] += 1
+ raise asyncio.CancelledError()
+
+ monkeypatch.setattr(proxy_server.asyncio, "sleep", _stop_sleep)
+
+ with pytest.raises(asyncio.CancelledError):
+ await _run_background_health_check()
+
+ assert normalize(
+ {
+ "healthy_count": proxy_server.health_check_results["healthy_count"],
+ "unhealthy_count": proxy_server.health_check_results["unhealthy_count"],
+ "sleep_invoked": sleep_calls["n"] >= 1,
+ }
+ ) == {
+ "healthy_count": 1,
+ "unhealthy_count": 1,
+ "sleep_invoked": True,
+ }
diff --git a/tests/test_litellm/proxy/proxy_server/test_exception_handlers.py b/tests/test_litellm/proxy/proxy_server/test_exception_handlers.py
index ad6b4016461..cf92f9cd12b 100644
--- a/tests/test_litellm/proxy/proxy_server/test_exception_handlers.py
+++ b/tests/test_litellm/proxy/proxy_server/test_exception_handlers.py
@@ -1 +1,222 @@
-"""Placeholder. Filled by a follow-up PR per the Notion plan."""
+"""Behavior pins for the proxy_server exception handlers.
+
+Pins covered:
+- ``openai_exception_handler``
+- ``_close_dangling_otel_server_span``
+- ``otel_request_validation_exception_handler``
+- ``otel_unhandled_exception_handler``
+"""
+
+from __future__ import annotations
+
+import json
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+from fastapi import HTTPException
+from fastapi.exceptions import RequestValidationError
+
+from litellm.proxy._types import ProxyException
+from litellm.proxy.proxy_server import (
+ _close_dangling_otel_server_span,
+ openai_exception_handler,
+ otel_request_validation_exception_handler,
+ otel_unhandled_exception_handler,
+)
+
+from .conftest import normalize
+
+
+def _make_request(parent_otel_span=None):
+ state = SimpleNamespace(parent_otel_span=parent_otel_span)
+ return SimpleNamespace(state=state)
+
+
+# ---------------------------------------------------------------------------
+# openai_exception_handler
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_openai_exception_handler_returns_mapped_payload():
+ exc = ProxyException(
+ message="bad input",
+ type="invalid_request_error",
+ param="model",
+ code=400,
+ )
+ request = _make_request()
+
+ response = await openai_exception_handler(request=request, exc=exc)
+ body = json.loads(response.body)
+
+ assert response.status_code == 400
+ assert normalize(body) == {
+ "error": {
+ "message": "bad input",
+ "type": "invalid_request_error",
+ "param": "model",
+ "code": "400",
+ }
+ }
+
+
+@pytest.mark.asyncio
+async def test_openai_exception_handler_invalid_empty_code_defaults_to_500():
+ """openai_exception_handler falls back to 500 when ``code`` is falsy.
+
+ Constructing via __new__ bypasses __init__ — the production __init__ always
+ coerces None to the string "None", which is truthy. To exercise the falsy
+ fallback branch we hand-craft an exception with an empty code."""
+ exc = ProxyException.__new__(ProxyException)
+ exc.message = "boom"
+ exc.type = "server_error"
+ exc.param = None
+ exc.openai_code = None
+ exc.code = ""
+ exc.headers = {}
+ exc.provider_specific_fields = None
+ request = _make_request()
+
+ response = await openai_exception_handler(request=request, exc=exc)
+ body = json.loads(response.body)
+
+ assert response.status_code == 500
+ assert body == {
+ "error": {
+ "message": "boom",
+ "type": "server_error",
+ "param": None,
+ "code": "",
+ }
+ }
+
+
+# ---------------------------------------------------------------------------
+# _close_dangling_otel_server_span
+# ---------------------------------------------------------------------------
+
+
+def test_close_dangling_otel_server_span_records_status_and_ends(monkeypatch):
+ """Happy path: with a logger and an active span, the handler sets the
+ response status, marks ERROR (>=400), ends the span, and clears state."""
+ import litellm.proxy.proxy_server as ps
+
+ span = MagicMock()
+ fake_logger = MagicMock()
+ monkeypatch.setattr(ps, "open_telemetry_logger", fake_logger, raising=False)
+ request = _make_request(parent_otel_span=span)
+
+ _close_dangling_otel_server_span(request=request, status_code=502)
+
+ observed = {
+ "status_attr_called": fake_logger.set_response_status_code_attribute.called,
+ "set_status_called": span.set_status.called,
+ "ended": span.end.called,
+ "state_cleared": request.state.parent_otel_span is None,
+ }
+ assert normalize(observed) == {
+ "status_attr_called": True,
+ "set_status_called": True,
+ "ended": True,
+ "state_cleared": True,
+ }
+
+
+def test_close_dangling_otel_server_span_missing_span_is_noop_error():
+ """When parent_otel_span is missing the call short-circuits — no error."""
+ request = _make_request(parent_otel_span=None)
+
+ result = _close_dangling_otel_server_span(request=request, status_code=200)
+ assert result is None
+ assert request.state.parent_otel_span is None
+
+
+def test_close_dangling_otel_server_span_logger_raises_state_cleared_error(monkeypatch):
+ """Logger raising is caught; state.parent_otel_span is cleared regardless."""
+ import litellm.proxy.proxy_server as ps
+
+ span = MagicMock()
+ fake_logger = MagicMock()
+ fake_logger.set_response_status_code_attribute.side_effect = RuntimeError("boom")
+ monkeypatch.setattr(ps, "open_telemetry_logger", fake_logger, raising=False)
+ request = _make_request(parent_otel_span=span)
+
+ _close_dangling_otel_server_span(request=request, status_code=500)
+
+ assert request.state.parent_otel_span is None
+
+
+# ---------------------------------------------------------------------------
+# otel_request_validation_exception_handler
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_otel_request_validation_exception_handler_returns_422_detail():
+ errors = [{"loc": ["body", "model"], "msg": "field required", "type": "missing"}]
+ exc = RequestValidationError(errors)
+ request = _make_request()
+
+ response = await otel_request_validation_exception_handler(request=request, exc=exc)
+ body = json.loads(response.body)
+
+ assert response.status_code == 422
+ assert normalize(body) == {"detail": exc.errors()}
+
+
+@pytest.mark.asyncio
+async def test_otel_request_validation_exception_handler_empty_errors_invalid_payload():
+ """An empty error list still returns 422 — the validator emitted nothing
+ but the handler must not crash and the body must remain well-formed."""
+ exc = RequestValidationError([])
+ request = _make_request()
+
+ response = await otel_request_validation_exception_handler(request=request, exc=exc)
+ body = json.loads(response.body)
+
+ assert response.status_code == 422
+ assert body == {"detail": []}
+
+
+# ---------------------------------------------------------------------------
+# otel_unhandled_exception_handler
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_otel_unhandled_exception_handler_returns_500_generic_payload():
+ exc = RuntimeError("kaboom")
+ request = _make_request()
+
+ response = await otel_unhandled_exception_handler(request=request, exc=exc)
+ body = json.loads(response.body)
+
+ assert response.status_code == 500
+ assert normalize(body) == {
+ "error": {
+ "message": "Internal server error",
+ "type": "internal_server_error",
+ }
+ }
+
+
+@pytest.mark.asyncio
+async def test_otel_unhandled_exception_handler_reraises_proxy_exception_error():
+ """ProxyException / HTTPException / RequestValidationError are re-raised
+ so the dedicated handler runs."""
+ exc = ProxyException(message="m", type="t", param="p", code=403)
+ request = _make_request()
+
+ with pytest.raises(ProxyException):
+ await otel_unhandled_exception_handler(request=request, exc=exc)
+
+
+@pytest.mark.asyncio
+async def test_otel_unhandled_exception_handler_reraises_http_exception_invalid():
+ request = _make_request()
+ with pytest.raises(HTTPException):
+ await otel_unhandled_exception_handler(
+ request=request, exc=HTTPException(status_code=418, detail="teapot")
+ )
diff --git a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py
index ad6b4016461..0b733401b59 100644
--- a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py
+++ b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py
@@ -1 +1,564 @@
-"""Placeholder. Filled by a follow-up PR per the Notion plan."""
+"""Behavior pins for proxy_server lifecycle, helpers, and small utilities.
+
+Pins covered:
+- ``proxy_startup_event``
+- ``proxy_shutdown_event``
+- ``_initialize_shared_aiohttp_session``
+- ``cleanup_router_config_variables``
+- ``save_worker_config``
+- ``initialize``
+- ``load_from_azure_key_vault``
+- ``cost_tracking``
+- ``check_request_disconnection``
+- ``_resolve_typed_dict_type``
+- ``_resolve_pydantic_type``
+- ``get_litellm_model_info``
+- ``run_ollama_serve``
+"""
+
+from __future__ import annotations
+
+import asyncio
+import inspect
+import json
+import os
+from typing import List, Optional, Union
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from typing_extensions import TypedDict
+
+import litellm.proxy.proxy_server as ps
+from litellm.proxy.proxy_server import (
+ _initialize_shared_aiohttp_session,
+ _resolve_pydantic_type,
+ _resolve_typed_dict_type,
+ check_request_disconnection,
+ cleanup_router_config_variables,
+ cost_tracking,
+ get_litellm_model_info,
+ initialize,
+ load_from_azure_key_vault,
+ proxy_shutdown_event,
+ proxy_startup_event,
+ run_ollama_serve,
+ save_worker_config,
+)
+
+from .conftest import normalize
+
+# ---------------------------------------------------------------------------
+# cleanup_router_config_variables
+# ---------------------------------------------------------------------------
+
+
+def test_cleanup_router_config_variables_resets_globals(monkeypatch):
+ monkeypatch.setattr(ps, "master_key", "sk-sentinel", raising=False)
+ monkeypatch.setattr(ps, "user_config_file_path", "/tmp/config.yaml", raising=False)
+ monkeypatch.setattr(ps, "user_custom_auth", lambda x: x, raising=False)
+ monkeypatch.setattr(ps, "health_check_interval", 42, raising=False)
+ monkeypatch.setattr(ps, "prisma_client", MagicMock(), raising=False)
+
+ cleanup_router_config_variables()
+
+ observed = {
+ "master_key": ps.master_key,
+ "user_config_file_path": ps.user_config_file_path,
+ "user_custom_auth": ps.user_custom_auth,
+ "health_check_interval": ps.health_check_interval,
+ "prisma_client": ps.prisma_client,
+ }
+ assert normalize(observed) == {
+ "master_key": None,
+ "user_config_file_path": None,
+ "user_custom_auth": None,
+ "health_check_interval": None,
+ "prisma_client": None,
+ }
+
+
+def test_cleanup_router_config_variables_fails_on_unknown_attr_raises():
+ """The function only writes documented globals — accessing a non-existent
+ one after cleanup should still raise AttributeError."""
+ cleanup_router_config_variables()
+ with pytest.raises(AttributeError):
+ _ = ps.this_attribute_should_not_exist_xyz
+
+
+# ---------------------------------------------------------------------------
+# proxy_shutdown_event
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_proxy_shutdown_event_disconnects_prisma_and_resets(monkeypatch):
+ fake_prisma = MagicMock()
+ fake_prisma.disconnect = AsyncMock()
+ monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)
+ monkeypatch.setattr(ps, "master_key", "sk-x", raising=False)
+
+ fake_jwt = MagicMock()
+ fake_jwt.close = AsyncMock()
+ monkeypatch.setattr(ps, "jwt_handler", fake_jwt, raising=False)
+ monkeypatch.setattr(ps, "db_writer_client", None, raising=False)
+
+ import litellm
+
+ monkeypatch.setattr(litellm, "cache", None, raising=False)
+ monkeypatch.setattr(litellm, "success_callback", [], raising=False)
+
+ await proxy_shutdown_event()
+
+ observed = {
+ "disconnect_called": fake_prisma.disconnect.await_count == 1,
+ "jwt_closed": fake_jwt.close.await_count == 1,
+ "master_key_reset": ps.master_key,
+ "prisma_reset": ps.prisma_client,
+ }
+ assert normalize(observed) == {
+ "disconnect_called": True,
+ "jwt_closed": True,
+ "master_key_reset": None,
+ "prisma_reset": None,
+ }
+
+
+@pytest.mark.asyncio
+async def test_proxy_shutdown_event_prisma_disconnect_raises_error(monkeypatch):
+ fake_prisma = MagicMock()
+ fake_prisma.disconnect = AsyncMock(side_effect=RuntimeError("db gone"))
+ monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)
+
+ fake_jwt = MagicMock()
+ fake_jwt.close = AsyncMock()
+ monkeypatch.setattr(ps, "jwt_handler", fake_jwt, raising=False)
+
+ import litellm
+
+ monkeypatch.setattr(litellm, "cache", None, raising=False)
+ monkeypatch.setattr(litellm, "success_callback", [], raising=False)
+
+ with pytest.raises(RuntimeError, match="db gone"):
+ await proxy_shutdown_event()
+
+
+# ---------------------------------------------------------------------------
+# _initialize_shared_aiohttp_session
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_initialize_shared_aiohttp_session_returns_client_session():
+ from aiohttp import ClientSession
+
+ session = await _initialize_shared_aiohttp_session()
+ try:
+ observed = {
+ "is_client_session": isinstance(session, ClientSession),
+ "is_closed": session.closed,
+ "has_connector": session.connector is not None,
+ }
+ assert normalize(observed) == {
+ "is_client_session": True,
+ "is_closed": False,
+ "has_connector": True,
+ }
+ finally:
+ if session is not None:
+ await session.close()
+
+
+@pytest.mark.asyncio
+async def test_initialize_shared_aiohttp_session_aiohttp_missing_returns_none_on_failure(
+ monkeypatch,
+):
+ """If aiohttp import fails, the function catches and returns None — no raise."""
+ import builtins
+
+ real_import = builtins.__import__
+
+ def _raise_for_aiohttp(name, *args, **kwargs):
+ if name == "aiohttp":
+ raise ImportError("simulated missing aiohttp")
+ return real_import(name, *args, **kwargs)
+
+ monkeypatch.setattr(builtins, "__import__", _raise_for_aiohttp)
+ result = await _initialize_shared_aiohttp_session()
+ assert result is None
+
+
+# ---------------------------------------------------------------------------
+# save_worker_config
+# ---------------------------------------------------------------------------
+
+
+def test_save_worker_config_writes_json_to_environ(monkeypatch):
+ monkeypatch.delenv("WORKER_CONFIG", raising=False)
+
+ save_worker_config(model="gpt-4", config="/tmp/c.yaml", debug=True)
+
+ payload = json.loads(os.environ["WORKER_CONFIG"])
+ assert normalize(payload) == {
+ "model": "gpt-4",
+ "config": "/tmp/c.yaml",
+ "debug": True,
+ }
+
+
+def test_save_worker_config_invalid_no_kwargs_yields_empty(monkeypatch):
+ monkeypatch.delenv("WORKER_CONFIG", raising=False)
+
+ save_worker_config()
+ assert os.environ["WORKER_CONFIG"] == "{}"
+
+
+# ---------------------------------------------------------------------------
+# initialize
+# ---------------------------------------------------------------------------
+
+
+def test_initialize_signature_is_async_with_expected_params():
+ sig = inspect.signature(initialize)
+ # Hard-coded so a signature change (param added/removed) trips the gate.
+ expected_param_count = 17
+ observed = {
+ "is_async": inspect.iscoroutinefunction(initialize),
+ "param_count": len(sig.parameters),
+ "has_model": "model" in sig.parameters,
+ "has_config": "config" in sig.parameters,
+ }
+ assert normalize(observed) == {
+ "is_async": True,
+ "param_count": expected_param_count,
+ "has_model": True,
+ "has_config": True,
+ }
+
+
+@pytest.mark.asyncio
+async def test_initialize_invalid_unexpected_kwarg_raises_type_error():
+ with pytest.raises(TypeError):
+ await initialize(this_is_not_a_real_kwarg=True)
+
+
+# ---------------------------------------------------------------------------
+# load_from_azure_key_vault
+# ---------------------------------------------------------------------------
+
+
+def test_load_from_azure_key_vault_disabled_no_side_effect(monkeypatch):
+ import litellm
+
+ sentinel_secret_mgr = object()
+ monkeypatch.setattr(
+ litellm, "secret_manager_client", sentinel_secret_mgr, raising=False
+ )
+
+ result = load_from_azure_key_vault(use_azure_key_vault=False)
+
+ observed = {
+ "return_value": result,
+ "secret_manager_unchanged": litellm.secret_manager_client
+ is sentinel_secret_mgr,
+ "called_with": False,
+ }
+ assert normalize(observed) == {
+ "return_value": None,
+ "secret_manager_unchanged": True,
+ "called_with": False,
+ }
+
+
+def test_load_from_azure_key_vault_missing_uri_failure_is_swallowed(monkeypatch):
+ """Enabled but AZURE_KEY_VAULT_URI unset / azure libs likely unavailable —
+ function catches Exception and does not raise."""
+ monkeypatch.delenv("AZURE_KEY_VAULT_URI", raising=False)
+
+ result = load_from_azure_key_vault(use_azure_key_vault=True)
+ assert result is None
+
+
+# ---------------------------------------------------------------------------
+# cost_tracking
+# ---------------------------------------------------------------------------
+
+
+def test_cost_tracking_adds_two_callbacks_when_prisma_set(monkeypatch):
+ import litellm
+
+ fake_prisma = MagicMock()
+ monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)
+ monkeypatch.setattr(litellm, "callbacks", [], raising=False)
+ monkeypatch.setattr(litellm, "_async_success_callback", [], raising=False)
+
+ before_callbacks = len(litellm.callbacks)
+ before_async = len(litellm._async_success_callback)
+
+ cost_tracking()
+
+ observed = {
+ "added_to_callbacks": len(litellm.callbacks) - before_callbacks,
+ "added_to_async_success": len(litellm._async_success_callback) - before_async,
+ "prisma_was_set": True,
+ }
+ assert normalize(observed) == {
+ "added_to_callbacks": 1,
+ "added_to_async_success": 1,
+ "prisma_was_set": True,
+ }
+
+
+def test_cost_tracking_no_op_when_prisma_missing(monkeypatch):
+ """Without a prisma_client cost_tracking is a no-op — not an error."""
+ import litellm
+
+ monkeypatch.setattr(ps, "prisma_client", None, raising=False)
+ monkeypatch.setattr(litellm, "callbacks", [], raising=False)
+ monkeypatch.setattr(litellm, "_async_success_callback", [], raising=False)
+
+ cost_tracking()
+
+ assert litellm.callbacks == []
+ assert litellm._async_success_callback == []
+
+
+# ---------------------------------------------------------------------------
+# check_request_disconnection
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_check_request_disconnection_cancels_task_and_raises_499(monkeypatch):
+ monkeypatch.setattr(ps.asyncio, "sleep", AsyncMock(return_value=None))
+
+ request = MagicMock()
+ request.is_disconnected = AsyncMock(return_value=True)
+ task = MagicMock()
+
+ raised_status = None
+ try:
+ await check_request_disconnection(request=request, llm_api_call_task=task)
+ except HTTPException as exc:
+ raised_status = exc.status_code
+
+ observed = {
+ "raised_status": raised_status,
+ "cancel_called": task.cancel.called,
+ "is_async": inspect.iscoroutinefunction(check_request_disconnection),
+ }
+ assert normalize(observed) == {
+ "raised_status": 499,
+ "cancel_called": True,
+ "is_async": True,
+ }
+
+
+@pytest.mark.asyncio
+async def test_check_request_disconnection_invalid_when_connected_times_out(monkeypatch):
+ """With a connected request the function loops for up to 10 minutes —
+ wrap in wait_for and assert it times out. Patch ``asyncio.sleep`` so the
+ loop spins without real wall-clock waits."""
+ import litellm.proxy.proxy_server as ps
+
+ request = MagicMock()
+ request.is_disconnected = AsyncMock(return_value=False)
+ task = MagicMock()
+
+ _real_sleep = asyncio.sleep
+
+ async def _instant_sleep(_seconds):
+ await _real_sleep(0)
+
+ monkeypatch.setattr(ps.asyncio, "sleep", _instant_sleep)
+
+ with pytest.raises(asyncio.TimeoutError):
+ await asyncio.wait_for(
+ check_request_disconnection(request=request, llm_api_call_task=task),
+ timeout=0.05,
+ )
+
+
+# ---------------------------------------------------------------------------
+# _resolve_typed_dict_type
+# ---------------------------------------------------------------------------
+
+
+class _SampleTD(TypedDict):
+ a: int
+ b: str
+
+
+def test_resolve_typed_dict_type_finds_class_in_optional():
+ typ = Optional[_SampleTD]
+ result = _resolve_typed_dict_type(typ)
+
+ observed = {
+ "input_repr": "Optional[_SampleTD]",
+ "result_is_sample_td": result is _SampleTD,
+ "result_is_class": isinstance(result, type),
+ }
+ assert normalize(observed) == {
+ "input_repr": "Optional[_SampleTD]",
+ "result_is_sample_td": True,
+ "result_is_class": True,
+ }
+
+
+def test_resolve_typed_dict_type_invalid_plain_type_returns_none():
+ """A non-TypedDict, non-Union input returns None — not an error."""
+ assert _resolve_typed_dict_type(int) is None
+ assert _resolve_typed_dict_type(str) is None
+
+
+# ---------------------------------------------------------------------------
+# _resolve_pydantic_type
+# ---------------------------------------------------------------------------
+
+
+class _SampleModelA(BaseModel):
+ x: int
+
+
+class _SampleModelB(BaseModel):
+ y: str
+
+
+def test_resolve_pydantic_type_extracts_non_none_args_from_union():
+ typ = Union[_SampleModelA, _SampleModelB, None]
+ result = _resolve_pydantic_type(typ)
+
+ observed = {
+ "result_type": type(result).__name__,
+ "result_len": len(result),
+ "contains_a": _SampleModelA in result,
+ "contains_b": _SampleModelB in result,
+ }
+ assert normalize(observed) == {
+ "result_type": "list",
+ "result_len": 2,
+ "contains_a": True,
+ "contains_b": True,
+ }
+
+
+def test_resolve_pydantic_type_invalid_non_union_non_model_returns_empty():
+ """When given a non-Union and non-BaseModel input the function returns [].
+
+ This is the silent-empty fallback path — error-ish by behavior."""
+ result = _resolve_pydantic_type(int)
+ assert result == []
+
+
+# ---------------------------------------------------------------------------
+# get_litellm_model_info
+# ---------------------------------------------------------------------------
+
+
+def test_get_litellm_model_info_uses_base_model_for_lookup(monkeypatch):
+ import litellm
+
+ expected_info = {"max_tokens": 8192, "input_cost_per_token": 0.00003}
+ fake_get = MagicMock(return_value=expected_info)
+ monkeypatch.setattr(litellm, "get_model_info", fake_get, raising=False)
+
+ model = {
+ "model_info": {"base_model": "gpt-4"},
+ "litellm_params": {"model": "azure/my-deployment"},
+ }
+ result = get_litellm_model_info(model=model)
+
+ observed = {
+ "called_arg": (
+ fake_get.call_args.args[0]
+ if fake_get.call_args.args
+ else fake_get.call_args.kwargs.get("model")
+ ),
+ "returned_max_tokens": result.get("max_tokens"),
+ "returned_cost": result.get("input_cost_per_token"),
+ }
+ assert normalize(observed) == {
+ "called_arg": "gpt-4",
+ "returned_max_tokens": 8192,
+ "returned_cost": 0.00003,
+ }
+
+
+def test_get_litellm_model_info_invalid_empty_dict_returns_empty():
+ """Empty input means model_to_lookup is None — internal exception is caught
+ and the function returns {}."""
+ result = get_litellm_model_info(model={})
+ assert result == {}
+
+
+# ---------------------------------------------------------------------------
+# run_ollama_serve
+# ---------------------------------------------------------------------------
+
+
+def test_run_ollama_serve_invokes_subprocess_popen(monkeypatch):
+ fake_popen = MagicMock()
+ monkeypatch.setattr(ps.subprocess, "Popen", fake_popen)
+
+ run_ollama_serve()
+
+ args, kwargs = fake_popen.call_args
+ observed = {
+ "popen_called": fake_popen.call_count == 1,
+ "command": args[0] if args else kwargs.get("args"),
+ "has_stdout_kw": "stdout" in kwargs,
+ "has_stderr_kw": "stderr" in kwargs,
+ }
+ assert normalize(observed) == {
+ "popen_called": True,
+ "command": ["ollama", "serve"],
+ "has_stdout_kw": True,
+ "has_stderr_kw": True,
+ }
+
+
+def test_run_ollama_serve_popen_failure_is_swallowed(monkeypatch):
+ """Popen raising OSError must NOT propagate — function logs and returns."""
+ monkeypatch.setattr(
+ ps.subprocess, "Popen", MagicMock(side_effect=OSError("no ollama binary"))
+ )
+
+ result = run_ollama_serve()
+ assert result is None
+
+
+# ---------------------------------------------------------------------------
+# proxy_startup_event
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_proxy_startup_event_is_async_context_manager_with_expected_signature():
+ """proxy_startup_event is the FastAPI lifespan. Verify its surface without
+ actually running the heavy init path (DB, Router, OTEL, etc.)."""
+ sig = inspect.signature(proxy_startup_event)
+ wrapped = getattr(proxy_startup_event, "__wrapped__", None)
+ observed = {
+ "param_count": len(sig.parameters),
+ "has_app_param": "app" in sig.parameters,
+ "wrapped_is_async": inspect.iscoroutinefunction(wrapped)
+ or inspect.isasyncgenfunction(wrapped),
+ "has_asynccontextmanager_wrapper": wrapped is not None,
+ }
+ assert normalize(observed) == {
+ "param_count": 1,
+ "has_app_param": True,
+ "wrapped_is_async": True,
+ "has_asynccontextmanager_wrapper": True,
+ }
+
+
+@pytest.mark.asyncio
+async def test_proxy_startup_event_invalid_missing_app_arg_raises():
+ """Calling the lifespan with no FastAPI app argument must fail."""
+ with pytest.raises(TypeError):
+ # Intentionally invoke the underlying async generator function with
+ # no arguments — the decorator preserves the missing-arg TypeError.
+ async with proxy_startup_event(): # type: ignore[call-arg]
+ pass
diff --git a/tests/test_litellm/proxy/proxy_server/test_openapi_customization.py b/tests/test_litellm/proxy/proxy_server/test_openapi_customization.py
index ad6b4016461..141b9f2a98a 100644
--- a/tests/test_litellm/proxy/proxy_server/test_openapi_customization.py
+++ b/tests/test_litellm/proxy/proxy_server/test_openapi_customization.py
@@ -1 +1,447 @@
-"""Placeholder. Filled by a follow-up PR per the Notion plan."""
+"""Behavior pins for proxy_server OpenAPI customization + CORS helpers.
+
+Pins covered:
+- ``_generate_stable_operation_id``
+- ``_strip_operation_id_method_suffix``
+- ``ensure_unique_openapi_operation_ids``
+- ``_inject_websocket_stubs_into_openapi_schema``
+- ``get_openapi_schema``
+- ``custom_openapi``
+- ``mount_swagger_ui``
+- ``_get_cors_config``
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+from fastapi import FastAPI
+
+from litellm.proxy.proxy_server import (
+ _generate_stable_operation_id,
+ _get_cors_config,
+ _inject_websocket_stubs_into_openapi_schema,
+ _strip_operation_id_method_suffix,
+ custom_openapi,
+ ensure_unique_openapi_operation_ids,
+ get_openapi_schema,
+ mount_swagger_ui,
+)
+
+from .conftest import normalize
+
+# ---------------------------------------------------------------------------
+# _generate_stable_operation_id
+# ---------------------------------------------------------------------------
+
+
+def test_generate_stable_operation_id_single_method_appends_suffix():
+ route = SimpleNamespace(
+ name="list_models",
+ path_format="/v1/models",
+ methods={"GET"},
+ )
+ observed = {
+ "operation_id": _generate_stable_operation_id(route),
+ "name": route.name,
+ "path": route.path_format,
+ }
+ assert normalize(observed) == {
+ "operation_id": "list_models_v1_models_get",
+ "name": "list_models",
+ "path": "/v1/models",
+ }
+
+
+def test_generate_stable_operation_id_multi_method_no_suffix():
+ route = SimpleNamespace(
+ name="multi_op",
+ path_format="/v1/things/{id}",
+ methods={"GET", "POST"},
+ )
+ observed = {
+ "operation_id": _generate_stable_operation_id(route),
+ "method_count": len(route.methods),
+ "has_method_suffix": _generate_stable_operation_id(route).endswith(
+ ("_get", "_post")
+ ),
+ }
+ assert normalize(observed) == {
+ "operation_id": "multi_op_v1_things__id_",
+ "method_count": 2,
+ "has_method_suffix": False,
+ }
+
+
+def test_generate_stable_operation_id_missing_attrs_raises_error():
+ bad_route = SimpleNamespace() # missing name/path_format/methods
+ with pytest.raises(AttributeError):
+ _generate_stable_operation_id(bad_route)
+
+
+# ---------------------------------------------------------------------------
+# _strip_operation_id_method_suffix
+# ---------------------------------------------------------------------------
+
+
+def test_strip_operation_id_method_suffix_removes_known_method():
+ observed = {
+ "with_get": _strip_operation_id_method_suffix("list_models_v1_models_get"),
+ "with_post": _strip_operation_id_method_suffix("create_thing_post"),
+ "with_delete": _strip_operation_id_method_suffix("drop_thing_delete"),
+ }
+ assert observed == {
+ "with_get": "list_models_v1_models",
+ "with_post": "create_thing",
+ "with_delete": "drop_thing",
+ }
+
+
+def test_strip_operation_id_method_suffix_invalid_suffix_unchanged():
+ # "foo" is not a known HTTP method; "nounderscore" has no separator at all.
+ observed = {
+ "unknown_suffix": _strip_operation_id_method_suffix("operation_foo"),
+ "no_underscore": _strip_operation_id_method_suffix("nounderscore"),
+ "empty": _strip_operation_id_method_suffix(""),
+ }
+ assert observed == {
+ "unknown_suffix": "operation_foo",
+ "no_underscore": "nounderscore",
+ "empty": "",
+ }
+
+
+# ---------------------------------------------------------------------------
+# ensure_unique_openapi_operation_ids
+# ---------------------------------------------------------------------------
+
+
+def test_ensure_unique_openapi_operation_ids_rewrites_duplicates():
+ schema = {
+ "paths": {
+ "/a": {"get": {"operationId": "dup_get"}},
+ "/b": {"get": {"operationId": "dup_get"}},
+ "/c": {"post": {"operationId": "unique_post"}},
+ }
+ }
+ result = ensure_unique_openapi_operation_ids(schema)
+ observed = {
+ "a_get": result["paths"]["/a"]["get"]["operationId"],
+ "b_get": result["paths"]["/b"]["get"]["operationId"],
+ "c_post": result["paths"]["/c"]["post"]["operationId"],
+ "ids_are_distinct": len(
+ {
+ result["paths"]["/a"]["get"]["operationId"],
+ result["paths"]["/b"]["get"]["operationId"],
+ result["paths"]["/c"]["post"]["operationId"],
+ }
+ )
+ == 3,
+ }
+ assert normalize(observed) == {
+ "a_get": "dup_get",
+ "b_get": "dup_get_2",
+ "c_post": "unique_post",
+ "ids_are_distinct": True,
+ }
+
+
+def test_ensure_unique_openapi_operation_ids_respects_reserved():
+ # operationId already ends with "_get" (an HTTP method), so the suffix is
+ # stripped before re-appending the current method, yielding "reserved_get".
+ schema = {
+ "paths": {
+ "/a": {"get": {"operationId": "reserved_get"}},
+ }
+ }
+ reserved = {"reserved_get"}
+ result = ensure_unique_openapi_operation_ids(
+ schema, reserved_operation_ids=reserved
+ )
+ observed = {
+ "rewritten": result["paths"]["/a"]["get"]["operationId"],
+ "still_includes_original": "reserved_get" in reserved,
+ "reserved_grew": len(reserved) > 1,
+ }
+ assert normalize(observed) == {
+ "rewritten": "reserved_get_2",
+ "still_includes_original": True,
+ "reserved_grew": True,
+ }
+
+
+def test_ensure_unique_openapi_operation_ids_missing_paths_invalid_returns_empty():
+ """No ``paths`` key — function must not crash and must return the schema as-is."""
+ schema = {"info": {"title": "x"}}
+ result = ensure_unique_openapi_operation_ids(schema)
+ assert result is schema
+ assert "paths" not in result
+
+
+# ---------------------------------------------------------------------------
+# _inject_websocket_stubs_into_openapi_schema
+# ---------------------------------------------------------------------------
+
+
+def test_inject_websocket_stubs_into_openapi_schema_adds_stub():
+ schema = {"paths": {}}
+ route = SimpleNamespace(path="/ws/chat", name="ws_chat", dependant=None)
+ result = _inject_websocket_stubs_into_openapi_schema(schema, [route])
+ stub = result["paths"]["/ws/chat"]["get"]
+ assert normalize(stub) == {
+ "summary": "WebSocket: ws_chat",
+ "description": "WebSocket connection endpoint",
+ "operationId": "websocket_ws_chat",
+ "parameters": [],
+ "responses": {"101": {"description": "WebSocket Protocol Switched"}},
+ "tags": ["WebSocket"],
+ }
+
+
+def test_inject_websocket_stubs_into_openapi_schema_does_not_overwrite_existing_get():
+ # Existing GET on the same path must not be replaced by the stub.
+ existing_get = {"summary": "real http get", "operationId": "real_get"}
+ schema = {"paths": {"/ws/chat": {"get": existing_get}}}
+ route = SimpleNamespace(path="/ws/chat", name="ws_chat", dependant=None)
+ result = _inject_websocket_stubs_into_openapi_schema(schema, [route])
+ assert result["paths"]["/ws/chat"]["get"] is existing_get
+
+
+def test_inject_websocket_stubs_into_openapi_schema_missing_paths_key_raises_error():
+ schema = {} # no "paths" key — setdefault on missing schema["paths"] will KeyError
+ route = SimpleNamespace(path="/ws/x", name="ws_x", dependant=None)
+ with pytest.raises(KeyError):
+ _inject_websocket_stubs_into_openapi_schema(schema, [route])
+
+
+# ---------------------------------------------------------------------------
+# get_openapi_schema
+# ---------------------------------------------------------------------------
+
+
+def test_get_openapi_schema_returns_well_formed_schema(monkeypatch):
+ """Patch ps.app to a fresh FastAPI so we get a deterministic minimal schema
+ without depending on whatever the session app currently has cached."""
+ import litellm.proxy.proxy_server as ps
+
+ fresh = FastAPI(title="pinned-title", version="0.0.1")
+
+ @fresh.get("/ping")
+ def _ping():
+ return {"ok": True}
+
+ monkeypatch.setattr(ps, "app", fresh, raising=True)
+ schema = get_openapi_schema()
+ observed = {
+ "openapi_present": "openapi" in schema,
+ "has_paths": isinstance(schema.get("paths"), dict),
+ "has_info": isinstance(schema.get("info"), dict),
+ "title": schema["info"]["title"],
+ "ping_path_in_schema": "/ping" in schema["paths"],
+ }
+ assert normalize(observed) == {
+ "openapi_present": True,
+ "has_paths": True,
+ "has_info": True,
+ "title": "pinned-title",
+ "ping_path_in_schema": True,
+ }
+
+
+def test_get_openapi_schema_returns_cached_when_present(monkeypatch):
+ """When the patched app already has openapi_schema set, the function
+ returns it untouched (no regeneration)."""
+ import litellm.proxy.proxy_server as ps
+
+ fresh = FastAPI()
+ sentinel = {"openapi": "3.0.0", "paths": {}, "info": {"title": "cached"}}
+ fresh.openapi_schema = sentinel
+ monkeypatch.setattr(ps, "app", fresh, raising=True)
+ result = get_openapi_schema()
+ observed = {
+ "is_sentinel": result is sentinel,
+ "title": result["info"]["title"],
+ "paths_empty": result["paths"] == {},
+ }
+ assert normalize(observed) == {
+ "is_sentinel": True,
+ "title": "cached",
+ "paths_empty": True,
+ }
+
+
+def test_get_openapi_schema_missing_app_attribute_raises_error(monkeypatch):
+ """If the module-level ``app`` is replaced by something without
+ ``openapi_schema`` and without ``routes``, the function fails fast."""
+ import litellm.proxy.proxy_server as ps
+
+ monkeypatch.setattr(ps, "app", SimpleNamespace(), raising=True)
+ with pytest.raises(AttributeError):
+ get_openapi_schema()
+
+
+# ---------------------------------------------------------------------------
+# custom_openapi
+# ---------------------------------------------------------------------------
+
+
+def test_custom_openapi_filters_to_openai_routes(monkeypatch):
+ """custom_openapi() filters paths down to the OpenAI-compatible set and
+ caches the result on the patched app."""
+ import litellm.proxy.proxy_server as ps
+
+ fresh = FastAPI(title="pinned-custom", version="0.0.1")
+
+ @fresh.get("/ping")
+ def _ping():
+ return {"ok": True}
+
+ monkeypatch.setattr(ps, "app", fresh, raising=True)
+ schema = custom_openapi()
+ observed = {
+ "openapi_present": "openapi" in schema,
+ "paths_is_dict": isinstance(schema.get("paths"), dict),
+ "info_title": schema["info"]["title"],
+ "cached_now": fresh.openapi_schema is schema,
+ "non_openai_path_filtered": "/ping" not in schema["paths"],
+ }
+ assert normalize(observed) == {
+ "openapi_present": True,
+ "paths_is_dict": True,
+ "info_title": "pinned-custom",
+ "cached_now": True,
+ "non_openai_path_filtered": True,
+ }
+
+
+def test_custom_openapi_returns_cached_when_present(monkeypatch):
+ import litellm.proxy.proxy_server as ps
+
+ fresh = FastAPI()
+ sentinel = {"openapi": "3.0.0", "paths": {}, "info": {"title": "cached"}}
+ fresh.openapi_schema = sentinel
+ monkeypatch.setattr(ps, "app", fresh, raising=True)
+ result = custom_openapi()
+ observed = {
+ "is_sentinel": result is sentinel,
+ "title": result["info"]["title"],
+ "paths_empty": result["paths"] == {},
+ }
+ assert normalize(observed) == {
+ "is_sentinel": True,
+ "title": "cached",
+ "paths_empty": True,
+ }
+
+
+def test_custom_openapi_missing_app_attribute_raises_error(monkeypatch):
+ import litellm.proxy.proxy_server as ps
+
+ monkeypatch.setattr(ps, "app", SimpleNamespace(), raising=True)
+ with pytest.raises(AttributeError):
+ custom_openapi()
+
+
+# ---------------------------------------------------------------------------
+# mount_swagger_ui
+# ---------------------------------------------------------------------------
+
+
+def test_mount_swagger_ui_mounts_static_route(monkeypatch):
+ """mount_swagger_ui mutates the global app — patch the module's `app` to a
+ fresh FastAPI() so we don't pollute the session app's mount table."""
+ import litellm.proxy.proxy_server as ps
+ from fastapi import applications as fa_applications
+
+ fresh_app = FastAPI()
+ monkeypatch.setattr(ps, "app", fresh_app, raising=True)
+ original_get_swagger = fa_applications.get_swagger_ui_html
+
+ try:
+ mount_swagger_ui()
+ finally:
+ # Restore the swagger monkey-patch so other tests are unaffected.
+ fa_applications.get_swagger_ui_html = original_get_swagger
+
+ mount_names = [getattr(r, "name", None) for r in fresh_app.routes]
+ observed = {
+ "swagger_mounted": "swagger" in mount_names,
+ "patched_get_swagger": (
+ fa_applications.get_swagger_ui_html is original_get_swagger
+ ),
+ "route_count_positive": len(fresh_app.routes) > 0,
+ }
+ assert normalize(observed) == {
+ "swagger_mounted": True,
+ "patched_get_swagger": True,
+ "route_count_positive": True,
+ }
+
+
+def test_mount_swagger_ui_missing_directory_raises_error(monkeypatch, tmp_path):
+ """If the swagger directory is missing, StaticFiles raises RuntimeError."""
+ import litellm.proxy.proxy_server as ps
+ from fastapi import applications as fa_applications
+
+ fresh_app = FastAPI()
+ monkeypatch.setattr(ps, "app", fresh_app, raising=True)
+ monkeypatch.setattr(
+ ps, "current_dir", str(tmp_path / "does_not_exist"), raising=True
+ )
+ original_get_swagger = fa_applications.get_swagger_ui_html
+
+ try:
+ with pytest.raises(RuntimeError):
+ mount_swagger_ui()
+ finally:
+ fa_applications.get_swagger_ui_html = original_get_swagger
+
+
+# ---------------------------------------------------------------------------
+# _get_cors_config
+# ---------------------------------------------------------------------------
+
+
+def test_get_cors_config_explicit_origins_and_credentials():
+ origins, allow_creds = _get_cors_config(
+ cors_origins_env="https://a.example,https://b.example",
+ cors_credentials_env="true",
+ )
+ observed = {
+ "origins": origins,
+ "allow_credentials": allow_creds,
+ "origin_count": len(origins),
+ }
+ assert normalize(observed) == {
+ "origins": ["https://a.example", "https://b.example"],
+ "allow_credentials": True,
+ "origin_count": 2,
+ }
+
+
+def test_get_cors_config_wildcard_defaults_credentials_false(monkeypatch):
+ # Clear env to ensure we test the default branch deterministically.
+ monkeypatch.delenv("LITELLM_CORS_ORIGINS", raising=False)
+ monkeypatch.delenv("LITELLM_CORS_ALLOW_CREDENTIALS", raising=False)
+ origins, allow_creds = _get_cors_config()
+ observed = {
+ "origins": origins,
+ "allow_credentials": allow_creds,
+ "wildcard_in_origins": "*" in origins,
+ }
+ assert normalize(observed) == {
+ "origins": ["*"],
+ "allow_credentials": False,
+ "wildcard_in_origins": True,
+ }
+
+
+def test_get_cors_config_invalid_credentials_value_treated_as_false():
+ """Anything other than the literal "true" (case-insensitive) is false —
+ misconfigured strings should not silently enable credentialed CORS."""
+ _, allow_creds = _get_cors_config(
+ cors_origins_env="https://a.example",
+ cors_credentials_env="yes-please",
+ )
+ assert allow_creds is False
diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py
index ad6b4016461..164538a2757 100644
--- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py
+++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py
@@ -1 +1,1315 @@
-"""Placeholder. Filled by a follow-up PR per the Notion plan."""
+"""Behavior pins for ProxyConfig and module-level config scrubbers.
+
+Pins covered:
+- Module-level: ``_is_remote_module_url``, ``_scrub_guardrail_inner``,
+ ``_scrub_db_overlay_remote_module_loads``
+- All ``ProxyConfig`` methods listed in the pin file.
+"""
+
+from __future__ import annotations
+
+import os
+from types import SimpleNamespace
+from typing import Any, Dict, List, Optional
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+import litellm
+from litellm.proxy.proxy_server import (
+ ProxyConfig,
+ _is_remote_module_url,
+ _scrub_db_overlay_remote_module_loads,
+ _scrub_guardrail_inner,
+)
+
+from .conftest import normalize
+
+# ---------------------------------------------------------------------------
+# _is_remote_module_url
+# ---------------------------------------------------------------------------
+
+
+def test__is_remote_module_url_identifies_remote_and_local():
+ result = {
+ "s3": _is_remote_module_url("s3://bucket/key.py"),
+ "gcs": _is_remote_module_url("gcs://bucket/key.py"),
+ "local": _is_remote_module_url("my.module.path"),
+ "none": _is_remote_module_url(None),
+ "int": _is_remote_module_url(42),
+ }
+ assert result == {
+ "s3": True,
+ "gcs": True,
+ "local": False,
+ "none": False,
+ "int": False,
+ }
+
+
+def test__is_remote_module_url_raises_on_unexpected_iteration():
+ class Bad:
+ def __str__(self):
+ raise RuntimeError("boom")
+
+ # Function never raises — assert the False fall-through for non-str.
+ with pytest.raises(AssertionError):
+ # Force an error-style assertion: object is not str, returns False.
+ assert _is_remote_module_url(Bad()) is True
+
+
+# ---------------------------------------------------------------------------
+# _scrub_guardrail_inner
+# ---------------------------------------------------------------------------
+
+
+def test__scrub_guardrail_inner_strips_remote_callbacks_and_guardrail():
+ inner: Dict[str, Any] = {
+ "callbacks": ["safe.mod", "s3://attacker/m.py", "gcs://x/y.py"],
+ "guardrail": "s3://attacker/g.py",
+ "default_on": True,
+ }
+ _scrub_guardrail_inner(inner)
+ assert normalize(inner) == {
+ "callbacks": ["safe.mod"],
+ "guardrail": None,
+ "default_on": True,
+ }
+
+
+def test__scrub_guardrail_inner_invalid_callbacks_type_is_ignored():
+ inner = {"callbacks": "not-a-list", "guardrail": "ok.module"}
+ _scrub_guardrail_inner(inner)
+ # No mutation on non-list callbacks; guardrail untouched (not remote).
+ assert inner == {"callbacks": "not-a-list", "guardrail": "ok.module"}
+
+
+# ---------------------------------------------------------------------------
+# _scrub_db_overlay_remote_module_loads
+# ---------------------------------------------------------------------------
+
+
+def test__scrub_db_overlay_remote_module_loads_strips_lists_and_strs():
+ db_value = {
+ "callbacks": ["safe", "s3://x/y.py"],
+ "success_callback": ["gcs://a/b.py", "safe2"],
+ "post_call_rules": "s3://bad/m.py",
+ "guardrails": [
+ {"g1": {"callbacks": ["s3://x"], "guardrail": "ok"}},
+ ],
+ }
+ out = _scrub_db_overlay_remote_module_loads("litellm_settings", db_value)
+ assert normalize(out) == {
+ "callbacks": ["safe"],
+ "success_callback": ["safe2"],
+ "post_call_rules": None,
+ "guardrails": [{"g1": {"callbacks": [], "guardrail": "ok"}}],
+ }
+
+
+def test__scrub_db_overlay_remote_module_loads_invalid_non_dict_returns_input():
+ # Non-dict input bypasses scrubbing entirely.
+ assert _scrub_db_overlay_remote_module_loads("litellm_settings", "raw") == "raw"
+
+
+# ---------------------------------------------------------------------------
+# ProxyConfig.__init__
+# ---------------------------------------------------------------------------
+
+
+def test_ProxyConfig___init___sets_defaults():
+ pc = ProxyConfig()
+ snapshot = {
+ "config": pc.config,
+ "last_semantic_filter_config": pc._last_semantic_filter_config,
+ "worker_registry": pc.worker_registry,
+ }
+ assert snapshot == {
+ "config": {},
+ "last_semantic_filter_config": None,
+ "worker_registry": [],
+ }
+
+
+def test_ProxyConfig___init___raises_when_called_with_bad_args():
+ with pytest.raises(TypeError):
+ ProxyConfig("unexpected-positional") # type: ignore[call-arg]
+
+
+# ---------------------------------------------------------------------------
+# ProxyConfig.is_yaml
+# ---------------------------------------------------------------------------
+
+
+def test_ProxyConfig_is_yaml_detects_yaml_and_non_yaml(tmp_path):
+ yaml_file = tmp_path / "c.yaml"
+ yaml_file.write_text("model_list: []\n")
+ yml_file = tmp_path / "c.yml"
+ yml_file.write_text("model_list: []\n")
+ json_file = tmp_path / "c.json"
+ json_file.write_text("{}")
+ pc = ProxyConfig()
+ result = {
+ "yaml": pc.is_yaml(str(yaml_file)),
+ "yml": pc.is_yaml(str(yml_file)),
+ "json": pc.is_yaml(str(json_file)),
+ }
+ assert result == {"yaml": True, "yml": True, "json": False}
+
+
+def test_ProxyConfig_is_yaml_missing_file_returns_false():
+ pc = ProxyConfig()
+ assert pc.is_yaml("/no/such/path/here.yaml") is False
+
+
+# ---------------------------------------------------------------------------
+# ProxyConfig._load_yaml_file
+# ---------------------------------------------------------------------------
+
+
+def test_ProxyConfig__load_yaml_file_returns_parsed_dict(tmp_path):
+ f = tmp_path / "c.yaml"
+ f.write_text("a: 1\nb: two\nc:\n - x\n - y\n")
+ pc = ProxyConfig()
+ result = pc._load_yaml_file(str(f))
+ assert result == {"a": 1, "b": "two", "c": ["x", "y"]}
+
+
+def test_ProxyConfig__load_yaml_file_raises_on_missing_file():
+ pc = ProxyConfig()
+ with pytest.raises(Exception):
+ pc._load_yaml_file("/no/such/file.yaml")
+
+
+# ---------------------------------------------------------------------------
+# ProxyConfig._get_config_from_file
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_ProxyConfig__get_config_from_file_loads_yaml(tmp_path):
+ f = tmp_path / "c.yaml"
+ f.write_text(
+ "model_list: []\ngeneral_settings: {}\nlitellm_settings:\n drop_params: true\n"
+ )
+ pc = ProxyConfig()
+ result = await pc._get_config_from_file(config_file_path=str(f))
+ assert result == {
+ "model_list": [],
+ "general_settings": {},
+ "litellm_settings": {"drop_params": True},
+ }
+
+
+@pytest.mark.asyncio
+async def test_ProxyConfig__get_config_from_file_missing_path_raises():
+ pc = ProxyConfig()
+ with pytest.raises(Exception):
+ await pc._get_config_from_file(config_file_path="/no/such/file.yaml")
+
+
+# ---------------------------------------------------------------------------
+# ProxyConfig._process_includes
+# ---------------------------------------------------------------------------
+
+
+def test_ProxyConfig__process_includes_merges_files(tmp_path):
+ inc = tmp_path / "models.yaml"
+ inc.write_text("model_list:\n - model_name: gpt-4\n")
+ pc = ProxyConfig()
+ cfg = {"include": ["models.yaml"], "model_list": [], "litellm_settings": {}}
+ result = pc._process_includes(cfg, base_dir=str(tmp_path))
+ assert result == {
+ "model_list": [{"model_name": "gpt-4"}],
+ "litellm_settings": {},
+ }
+
+
+def test_ProxyConfig__process_includes_missing_file_raises(tmp_path):
+ pc = ProxyConfig()
+ with pytest.raises(FileNotFoundError):
+ pc._process_includes({"include": ["nope.yaml"]}, base_dir=str(tmp_path))
+
+
+# ---------------------------------------------------------------------------
+# ProxyConfig.save_config
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_ProxyConfig_save_config_writes_yaml_when_no_db(tmp_path, monkeypatch):
+ target = tmp_path / "out.yaml"
+ monkeypatch.setattr("litellm.proxy.proxy_server.user_config_file_path", str(target))
+ monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
+ monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
+ monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
+ pc = ProxyConfig()
+ cfg = {"model_list": [], "general_settings": {"a": 1}, "litellm_settings": {}}
+ await pc.save_config(cfg)
+ import yaml as _yaml
+
+ loaded = _yaml.safe_load(target.read_text())
+ assert loaded == cfg
+
+
+@pytest.mark.asyncio
+async def test_ProxyConfig_save_config_invalid_path_raises(monkeypatch):
+ monkeypatch.setattr(
+ "litellm.proxy.proxy_server.user_config_file_path",
+ "/no/such/dir/out.yaml",
+ )
+ monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
+ monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
+ monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
+ pc = ProxyConfig()
+ with pytest.raises(Exception):
+ await pc.save_config({"x": 1})
+
+
+# ---------------------------------------------------------------------------
+# ProxyConfig._check_for_os_environ_vars
+# ---------------------------------------------------------------------------
+
+
+def test_ProxyConfig__check_for_os_environ_vars_substitutes(monkeypatch):
+ monkeypatch.setenv("MY_TEST_VAR", "secret-value")
+ pc = ProxyConfig()
+ cfg = {
+ "a": "os.environ/MY_TEST_VAR",
+ "b": 2,
+ "nested": {"c": "os.environ/MY_TEST_VAR"},
+ }
+ out = pc._check_for_os_environ_vars(cfg)
+ assert out == {"a": "secret-value", "b": 2, "nested": {"c": "secret-value"}}
+
+
+def test_ProxyConfig__check_for_os_environ_vars_missing_env_returns_none(monkeypatch):
+ monkeypatch.delenv("NONEXISTENT_TEST_VAR_X", raising=False)
+ pc = ProxyConfig()
+ cfg = {"a": "os.environ/NONEXISTENT_TEST_VAR_X"}
+ out = pc._check_for_os_environ_vars(cfg)
+ # get_secret returns None when not found — assert observable shape.
+ assert out["a"] is None
+
+
+# ---------------------------------------------------------------------------
+# ProxyConfig._get_team_config
+# ---------------------------------------------------------------------------
+
+
+def test_ProxyConfig__get_team_config_returns_match():
+ pc = ProxyConfig()
+ teams = [
+ {"team_id": "t1", "max_budget": 10, "model": "gpt-4"},
+ {"team_id": "t2", "max_budget": 20, "model": "claude"},
+ ]
+ out = pc._get_team_config(team_id="t1", all_teams_config=teams)
+ assert out == {"team_id": "t1", "max_budget": 10, "model": "gpt-4"}
+
+
+def test_ProxyConfig__get_team_config_missing_team_id_raises():
+ pc = ProxyConfig()
+ with pytest.raises(Exception):
+ pc._get_team_config(team_id="t1", all_teams_config=[{"no_id_field": True}])
+
+
+# ---------------------------------------------------------------------------
+# ProxyConfig.load_team_config
+# ---------------------------------------------------------------------------
+
+
+def test_ProxyConfig_load_team_config_returns_team_dict():
+ pc = ProxyConfig()
+ pc.config = {
+ "litellm_settings": {
+ "default_team_settings": [
+ {"team_id": "ta", "max_budget": 99, "drop_params": True},
+ ]
+ }
+ }
+ out = pc.load_team_config(team_id="ta")
+ assert out == {"team_id": "ta", "max_budget": 99, "drop_params": True}
+
+
+def test_ProxyConfig_load_team_config_no_settings_returns_empty():
+ pc = ProxyConfig()
+ pc.config = {"litellm_settings": {}}
+ # Missing entry — happy path returns {} (no default_team_settings).
+ out = pc.load_team_config(team_id="missing")
+ assert out == {}
+ # Error-style: a misconfigured team list without team_id raises.
+ pc.config = {"litellm_settings": {"default_team_settings": [{"no_id": True}]}}
+ with pytest.raises(Exception):
+ pc.load_team_config(team_id="anything")
+
+
+# ---------------------------------------------------------------------------
+# ProxyConfig._init_cache
+# ---------------------------------------------------------------------------
+
+
+def test_ProxyConfig__init_cache_sets_litellm_cache(monkeypatch):
+ pc = ProxyConfig()
+ monkeypatch.setattr(litellm, "cache", None, raising=False)
+ pc._init_cache(cache_params={"type": "local"})
+ snapshot = {
+ "cache_is_set": litellm.cache is not None,
+ "cache_type_name": type(litellm.cache).__name__,
+ "params_used": "local",
+ }
+ assert snapshot == {
+ "cache_is_set": True,
+ "cache_type_name": "Cache",
+ "params_used": "local",
+ }
+
+
+def test_ProxyConfig__init_cache_invalid_params_raises():
+ pc = ProxyConfig()
+ with pytest.raises(Exception):
+ pc._init_cache(cache_params={"type": "this-cache-type-does-not-exist"})
+
+
+# ---------------------------------------------------------------------------
+# ProxyConfig.switch_on_llm_response_caching
+# ---------------------------------------------------------------------------
+
+
+def test_ProxyConfig_switch_on_llm_response_caching_sets_flag(monkeypatch):
+ pc = ProxyConfig()
+ fake_router = MagicMock()
+ fake_router.cache_responses = False
+ fake_cache = MagicMock()
+ monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", fake_router)
+ monkeypatch.setattr(litellm, "cache", fake_cache, raising=False)
+ pc.switch_on_llm_response_caching()
+ snapshot = {
+ "cache_responses": fake_router.cache_responses,
+ "router_set": True,
+ "cache_set": True,
+ }
+ assert snapshot == {
+ "cache_responses": True,
+ "router_set": True,
+ "cache_set": True,
+ }
+
+
+def test_ProxyConfig_switch_on_llm_response_caching_missing_router_noop(monkeypatch):
+ pc = ProxyConfig()
+ monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
+ monkeypatch.setattr(litellm, "cache", None, raising=False)
+ # No router and no cache — should silently no-op (no raise).
+ pc.switch_on_llm_response_caching()
+ # Error-style: prove no router was created.
+ with pytest.raises(AttributeError):
+ _ = pc.does_not_exist # type: ignore[attr-defined]
+
+
+# ---------------------------------------------------------------------------
+# ProxyConfig.get_config
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_ProxyConfig_get_config_loads_from_file(tmp_path, monkeypatch):
+ f = tmp_path / "c.yaml"
+ f.write_text("model_list: []\ngeneral_settings: {}\nlitellm_settings: {}\n")
+ monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
+ monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
+ monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
+ pc = ProxyConfig()
+ cfg = await pc.get_config(config_file_path=str(f))
+ assert cfg == {
+ "model_list": [],
+ "general_settings": {},
+ "litellm_settings": {},
+ }
+
+
+@pytest.mark.asyncio
+async def test_ProxyConfig_get_config_missing_file_raises(monkeypatch):
+ monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
+ monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
+ monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
+ pc = ProxyConfig()
+ with pytest.raises(Exception):
+ await pc.get_config(config_file_path="/no/such/path.yaml")
+
+
+# ---------------------------------------------------------------------------
+# ProxyConfig.update_config_state / get_config_state
+# ---------------------------------------------------------------------------
+
+
+def test_ProxyConfig_update_config_state_and_get_config_state_roundtrip():
+ pc = ProxyConfig()
+ cfg = {"model_list": [], "general_settings": {"x": 1}, "litellm_settings": {}}
+ pc.update_config_state(config=cfg)
+ out = pc.get_config_state()
+ assert out == cfg
+ # Mutating the returned dict must not affect internal state.
+ out["model_list"].append({"new": True})
+ assert pc.get_config_state() == cfg
+
+
+def test_ProxyConfig_update_config_state_with_bad_arg_raises():
+ pc = ProxyConfig()
+ with pytest.raises(TypeError):
+ pc.update_config_state() # type: ignore[call-arg]
+
+
+def test_ProxyConfig_get_config_state_handles_undeepcopyable(monkeypatch):
+ # Pins ProxyConfig.get_config_state — see source for behavior.
+ pc = ProxyConfig()
+
+ class NoCopy:
+ def __deepcopy__(self, memo):
+ raise RuntimeError("nope")
+
+ pc.config = {"x": NoCopy()} # type: ignore[assignment]
+ # Exception is caught internally and an empty dict returned.
+ assert pc.get_config_state() == {}
+
+
+# ---------------------------------------------------------------------------
+# ProxyConfig.load_credential_list
+# ---------------------------------------------------------------------------
+
+
+def test_ProxyConfig_load_credential_list_returns_items():
+ pc = ProxyConfig()
+ creds = pc.load_credential_list(
+ {
+ "credential_list": [
+ {
+ "credential_name": "openai-key",
+ "credential_info": {"provider": "openai"},
+ "credential_values": {"api_key": "sk-x"},
+ }
+ ]
+ }
+ )
+ assert len(creds) == 1
+ dumped = creds[0].model_dump()
+ assert dumped == {
+ "credential_name": "openai-key",
+ "credential_info": {"provider": "openai"},
+ "credential_values": {"api_key": "sk-x"},
+ }
+
+
+def test_ProxyConfig_load_credential_list_invalid_entry_raises():
+ pc = ProxyConfig()
+ with pytest.raises(Exception):
+ pc.load_credential_list({"credential_list": [{"missing_required": True}]})
+
+
+# ---------------------------------------------------------------------------
+# ProxyConfig.parse_search_tools
+# ---------------------------------------------------------------------------
+
+
+def test_ProxyConfig_parse_search_tools_returns_parsed():
+ pc = ProxyConfig()
+ cfg = {
+ "search_tools": [
+ {
+ "search_tool_name": "web",
+ "litellm_params": {"search_provider": "google"},
+ }
+ ]
+ }
+ out = pc.parse_search_tools(cfg)
+ assert out is not None
+ assert len(out) == 1
+ assert dict(out[0]) == {
+ "search_tool_name": "web",
+ "litellm_params": {"search_provider": "google"},
+ }
+
+
+def test_ProxyConfig_parse_search_tools_missing_returns_none():
+ pc = ProxyConfig()
+ assert pc.parse_search_tools({}) is None
+
+
+# ---------------------------------------------------------------------------
+# ProxyConfig._load_environment_variables
+# ---------------------------------------------------------------------------
+
+
+def test_ProxyConfig__load_environment_variables_sets_env(monkeypatch):
+ monkeypatch.delenv("TEST_LOAD_ENV_X", raising=False)
+ pc = ProxyConfig()
+ pc._load_environment_variables(
+ {"environment_variables": {"TEST_LOAD_ENV_X": "hello"}}
+ )
+ result = {
+ "TEST_LOAD_ENV_X": os.environ.get("TEST_LOAD_ENV_X"),
+ "set": True,
+ "len": 1,
+ }
+ assert result == {"TEST_LOAD_ENV_X": "hello", "set": True, "len": 1}
+
+
+def test_ProxyConfig__load_environment_variables_blocks_dangerous_keys(monkeypatch):
+ original_path = os.environ.get("PATH", "")
+ pc = ProxyConfig()
+ pc._load_environment_variables({"environment_variables": {"PATH": "/evil/bin"}})
+ # PATH must be unchanged — it's a blocked key.
+ assert os.environ.get("PATH", "") == original_path
+
+
+# ---------------------------------------------------------------------------
+# ProxyConfig.load_config
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_ProxyConfig_load_config_minimal_yaml(tmp_path, monkeypatch):
+ f = tmp_path / "c.yaml"
+ f.write_text("model_list: []\ngeneral_settings: {}\nlitellm_settings: {}\n")
+ monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
+ monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
+ monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
+ pc = ProxyConfig()
+ try:
+ await pc.load_config(router=None, config_file_path=str(f))
+ raised = False
+ except Exception:
+ raised = True
+ snapshot = {
+ "raised": raised,
+ "config_loaded": pc.config is not None,
+ "model_list_key_present": "model_list" in pc.config,
+ }
+ assert snapshot == {
+ "raised": False,
+ "config_loaded": True,
+ "model_list_key_present": True,
+ }
+
+
+@pytest.mark.asyncio
+async def test_ProxyConfig_load_config_missing_file_raises(monkeypatch):
+ monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
+ monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
+ monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
+ pc = ProxyConfig()
+ with pytest.raises(Exception):
+ await pc.load_config(router=None, config_file_path="/no/file.yaml")
+
+
+# ---------------------------------------------------------------------------
+# ProxyConfig._init_non_llm_configs
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_ProxyConfig__init_non_llm_configs_empty_config():
+ pc = ProxyConfig()
+ try:
+ await pc._init_non_llm_configs(config={}, config_file_path=None)
+ raised = False
+ except Exception:
+ raised = True
+ snapshot = {
+ "raised": raised,
+ "worker_registry_len": len(pc.worker_registry),
+ "is_list": isinstance(pc.worker_registry, list),
+ }
+ assert snapshot == {"raised": False, "worker_registry_len": 0, "is_list": True}
+
+
+@pytest.mark.asyncio
+async def test_ProxyConfig__init_non_llm_configs_invalid_worker_registry_raises():
+ pc = ProxyConfig()
+ with pytest.raises(Exception):
+ await pc._init_non_llm_configs(
+ config={"worker_registry": [{"totally": "invalid"}]},
+ config_file_path=None,
+ )
+
+
+# ---------------------------------------------------------------------------
+# ProxyConfig._init_policy_engine
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_ProxyConfig__init_policy_engine_no_policies_noop():
+ pc = ProxyConfig()
+ try:
+ await pc._init_policy_engine(config={}, prisma_client=None, llm_router=None)
+ raised = False
+ except Exception:
+ raised = True
+ assert {"raised": raised, "called": True, "skipped": True} == {
+ "raised": False,
+ "called": True,
+ "skipped": True,
+ }
+
+
+@pytest.mark.asyncio
+async def test_ProxyConfig__init_policy_engine_none_config_noop():
+ pc = ProxyConfig()
+ # None config returns early without raising.
+ await pc._init_policy_engine(config=None, prisma_client=None, llm_router=None)
+ # Error-style: invalid policies value should raise.
+ with pytest.raises(Exception):
+ await pc._init_policy_engine(
+ config={"policies": "not-a-list"},
+ prisma_client=None,
+ llm_router=None,
+ )
+
+
+# ---------------------------------------------------------------------------
+# ProxyConfig._load_alerting_settings
+# ---------------------------------------------------------------------------
+
+
+def test_ProxyConfig__load_alerting_settings_noop_when_no_alerting():
+ pc = ProxyConfig()
+ try:
+ pc._load_alerting_settings({})
+ raised = False
+ except Exception:
+ raised = True
+ assert {"raised": raised, "called": True, "no_alerting": True} == {
+ "raised": False,
+ "called": True,
+ "no_alerting": True,
+ }
+
+
+def test_ProxyConfig__load_alerting_settings_invalid_alerting_raises():
+ pc = ProxyConfig()
+ with pytest.raises(Exception):
+ # alerting must be iterable — int triggers an error.
+ pc._load_alerting_settings({"alerting": 12345})
+
+
+# ---------------------------------------------------------------------------
+# ProxyConfig.initialize_secret_manager
+# ---------------------------------------------------------------------------
+
+
+def test_ProxyConfig_initialize_secret_manager_none_noop():
+ pc = ProxyConfig()
+ try:
+ pc.initialize_secret_manager(key_management_system=None)
+ raised = False
+ except Exception:
+ raised = True
+ assert {"raised": raised, "called": True, "kms": None} == {
+ "raised": False,
+ "called": True,
+ "kms": None,
+ }
+
+
+def test_ProxyConfig_initialize_secret_manager_invalid_kms_raises():
+ pc = ProxyConfig()
+ with pytest.raises(ValueError):
+ pc.initialize_secret_manager(key_management_system="not-a-real-kms")
+
+
+# ---------------------------------------------------------------------------
+# ProxyConfig.get_model_info_with_id
+# ---------------------------------------------------------------------------
+
+
+def test_ProxyConfig_get_model_info_with_id_returns_router_model_info():
+ pc = ProxyConfig()
+ model = SimpleNamespace(
+ model_id="m-1",
+ model_info={"id": "m-1"},
+ blocked=False,
+ )
+ out = pc.get_model_info_with_id(model=model, db_model=True)
+ dumped = out.model_dump()
+ snapshot = {
+ "id": dumped.get("id"),
+ "db_model": dumped.get("db_model"),
+ "blocked": dumped.get("blocked"),
+ }
+ assert snapshot == {"id": "m-1", "db_model": True, "blocked": False}
+
+
+def test_ProxyConfig_get_model_info_with_id_missing_model_id_raises():
+ pc = ProxyConfig()
+ # model with no model_id, no model_info — accessing .model_id will fail.
+ bad = SimpleNamespace(model_info=None)
+ with pytest.raises(AttributeError):
+ pc.get_model_info_with_id(model=bad)
+
+
+# ---------------------------------------------------------------------------
+# ProxyConfig._delete_deployment
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_ProxyConfig__delete_deployment_empty_returns_zero(monkeypatch):
+ monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
+ pc = ProxyConfig()
+ result = await pc._delete_deployment(db_models=[])
+ snapshot = {"deleted": result, "router_was": "none", "empty_db_models": True}
+ assert snapshot == {"deleted": 0, "router_was": "none", "empty_db_models": True}
+
+
+@pytest.mark.asyncio
+async def test_ProxyConfig__delete_deployment_invalid_models_raises(monkeypatch):
+ fake_router = MagicMock()
+ fake_router.get_model_ids = MagicMock(return_value=[])
+ monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", fake_router)
+ pc = ProxyConfig()
+ with pytest.raises(Exception):
+ # Non-model objects without expected attrs trigger an error.
+ await pc._delete_deployment(db_models=[{"not_a_model": True}])
+
+
+# ---------------------------------------------------------------------------
+# ProxyConfig._add_deployment
+# ---------------------------------------------------------------------------
+
+
+def test_ProxyConfig__add_deployment_no_router_returns_zero(monkeypatch):
+ monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
+ pc = ProxyConfig()
+ result = pc._add_deployment(db_models=[MagicMock()])
+ snapshot = {"added": result, "router_was": "none", "called": True}
+ assert snapshot == {"added": 0, "router_was": "none", "called": True}
+
+
+def test_ProxyConfig__add_deployment_invalid_litellm_params_skips(monkeypatch):
+ fake_router = MagicMock()
+ fake_router.upsert_deployment = MagicMock(return_value=None)
+ monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", fake_router)
+ pc = ProxyConfig()
+ bad = SimpleNamespace(litellm_params="not-a-dict", model_name="x", model_id="x")
+ # invalid params logs and continues — assert zero added (error-style branch).
+ assert pc._add_deployment(db_models=[bad]) == 0
+
+
+# ---------------------------------------------------------------------------
+# ProxyConfig.decrypt_model_list_from_db
+# ---------------------------------------------------------------------------
+
+
+def test_ProxyConfig_decrypt_model_list_from_db_returns_decrypted(monkeypatch):
+ monkeypatch.setattr(
+ "litellm.proxy.proxy_server.decrypt_value_helper",
+ lambda value, key, return_original_value: value,
+ )
+ pc = ProxyConfig()
+ m = SimpleNamespace(
+ model_id="m-1",
+ model_name="gpt-4",
+ model_info={"id": "m-1"},
+ litellm_params={"api_key": "sk-x", "model": "gpt-4"},
+ blocked=False,
+ )
+ out = pc.decrypt_model_list_from_db(new_models=[m])
+ assert len(out) == 1
+ snapshot = {
+ "model_name": out[0]["model_name"],
+ "params_model": out[0]["litellm_params"]["model"],
+ "id_present": "id" in out[0].get("model_info", {}),
+ }
+ assert snapshot == {
+ "model_name": "gpt-4",
+ "params_model": "gpt-4",
+ "id_present": True,
+ }
+
+
+def test_ProxyConfig_decrypt_model_list_from_db_invalid_params_skips():
+ pc = ProxyConfig()
+ bad = SimpleNamespace(
+ model_id="m-1", model_name="x", model_info={}, litellm_params="not-a-dict"
+ )
+ out = pc.decrypt_model_list_from_db(new_models=[bad])
+ # Invalid entries skipped — empty list returned.
+ assert out == []
+
+
+# ---------------------------------------------------------------------------
+# ProxyConfig._update_llm_router
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_ProxyConfig__update_llm_router_no_models_smoke(monkeypatch):
+ monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
+ monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "sk-master")
+ monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
+ monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
+ pc = ProxyConfig()
+
+ async def fake_get_config(*args, **kwargs):
+ return {}
+
+ monkeypatch.setattr(pc, "get_config", fake_get_config)
+ monkeypatch.setattr(
+ "litellm.proxy.proxy_server.proxy_config",
+ pc,
+ )
+ try:
+ await pc._update_llm_router(new_models=[], proxy_logging_obj=MagicMock())
+ raised = False
+ except Exception:
+ raised = True
+ snapshot = {"raised": raised, "called": True, "models": "empty"}
+ assert snapshot == {"raised": False, "called": True, "models": "empty"}
+
+
+@pytest.mark.asyncio
+async def test_ProxyConfig__update_llm_router_bad_proxy_logging_raises(monkeypatch):
+ pc = ProxyConfig()
+
+ async def fake_get_config():
+ # alerting present + non-list general_settings to trigger the alerting branch.
+ return {"general_settings": {"alerting": ["slack"]}}
+
+ fake_router = MagicMock()
+ fake_router.update_settings = MagicMock()
+ monkeypatch.setattr(pc, "get_config", fake_get_config)
+ monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", fake_router)
+ monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "sk-x")
+ monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
+ monkeypatch.setattr(
+ "litellm.proxy.proxy_server.general_settings", {"alerting": ["email"]}
+ )
+ monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", pc)
+ # Passing None for proxy_logging_obj triggers AttributeError in _add_general_settings_from_db_config
+ # when it calls proxy_logging_obj.update_values.
+ with pytest.raises(AttributeError):
+ await pc._update_llm_router(new_models=None, proxy_logging_obj=None) # type: ignore[arg-type]
+
+
+# ---------------------------------------------------------------------------
+# ProxyConfig._add_callback_from_db_to_in_memory_litellm_callbacks
+# ---------------------------------------------------------------------------
+
+
+def test_ProxyConfig__add_callback_from_db_to_in_memory_litellm_callbacks_adds(
+ monkeypatch,
+):
+ monkeypatch.setattr(litellm, "callbacks", [], raising=False)
+ pc = ProxyConfig()
+ pc._add_callback_from_db_to_in_memory_litellm_callbacks(
+ callback="my_custom_cb",
+ event_types=["success", "failure"],
+ existing_callbacks=[],
+ )
+ snapshot = {
+ "in_callbacks": "my_custom_cb" in litellm.callbacks,
+ "count": len(litellm.callbacks),
+ "method_called": True,
+ }
+ assert snapshot == {"in_callbacks": True, "count": 1, "method_called": True}
+
+
+def test_ProxyConfig__add_callback_from_db_to_in_memory_litellm_callbacks_invalid_event_raises(
+ monkeypatch,
+):
+ monkeypatch.setattr(litellm, "callbacks", [], raising=False)
+ pc = ProxyConfig()
+ # For a "known" callback, event_types is iterated — non-iterable raises TypeError.
+ with pytest.raises(TypeError):
+ pc._add_callback_from_db_to_in_memory_litellm_callbacks(
+ callback="lago", # in _known_custom_logger_compatible_callbacks
+ event_types=12345, # type: ignore[arg-type]
+ existing_callbacks=[],
+ )
+
+
+# ---------------------------------------------------------------------------
+# ProxyConfig._add_callbacks_from_db_config
+# ---------------------------------------------------------------------------
+
+
+def test_ProxyConfig__add_callbacks_from_db_config_processes_lists(monkeypatch):
+ monkeypatch.setattr(litellm, "callbacks", [], raising=False)
+ monkeypatch.setattr(litellm, "success_callback", [], raising=False)
+ monkeypatch.setattr(litellm, "failure_callback", [], raising=False)
+ pc = ProxyConfig()
+ cfg = {
+ "litellm_settings": {
+ "callbacks": ["cb_a"],
+ "success_callback": ["s_a"],
+ "failure_callback": ["f_a"],
+ }
+ }
+ pc._add_callbacks_from_db_config(cfg)
+ snapshot = {
+ "cb_added": "cb_a" in litellm.callbacks,
+ "success_added": "s_a" in litellm.success_callback,
+ "failure_added": "f_a" in litellm.failure_callback,
+ }
+ assert snapshot == {
+ "cb_added": True,
+ "success_added": True,
+ "failure_added": True,
+ }
+
+
+def test_ProxyConfig__add_callbacks_from_db_config_bad_config_raises():
+ pc = ProxyConfig()
+ with pytest.raises(AttributeError):
+ # Non-dict input — .get will fail.
+ pc._add_callbacks_from_db_config(None) # type: ignore[arg-type]
+
+
+# ---------------------------------------------------------------------------
+# ProxyConfig._encrypt_env_variables
+# ---------------------------------------------------------------------------
+
+
+def test_ProxyConfig__encrypt_env_variables_returns_dict(monkeypatch):
+ monkeypatch.setattr(
+ "litellm.proxy.proxy_server.encrypt_value_helper",
+ lambda value, new_encryption_key=None: f"ENC[{value}]",
+ )
+ pc = ProxyConfig()
+ out = pc._encrypt_env_variables({"A": "1", "B": "2", "C": "3"})
+ assert out == {"A": "ENC[1]", "B": "ENC[2]", "C": "ENC[3]"}
+
+
+def test_ProxyConfig__encrypt_env_variables_invalid_raises():
+ pc = ProxyConfig()
+ with pytest.raises(AttributeError):
+ # Non-dict input — .items() fails.
+ pc._encrypt_env_variables(None) # type: ignore[arg-type]
+
+
+# ---------------------------------------------------------------------------
+# ProxyConfig._decrypt_and_set_db_env_variables
+# ---------------------------------------------------------------------------
+
+
+def test_ProxyConfig__decrypt_and_set_db_env_variables_sets_env(monkeypatch):
+ monkeypatch.setattr(
+ "litellm.proxy.proxy_server.decrypt_value_helper",
+ lambda value, key, return_original_value=False: value + "-dec",
+ )
+ monkeypatch.delenv("KEY_X", raising=False)
+ monkeypatch.delenv("KEY_Y", raising=False)
+ pc = ProxyConfig()
+ out = pc._decrypt_and_set_db_env_variables({"KEY_X": "x", "KEY_Y": "y"})
+ snapshot = {
+ "KEY_X_env": os.environ.get("KEY_X"),
+ "KEY_Y_env": os.environ.get("KEY_Y"),
+ "returned_keys": sorted(out.keys()),
+ }
+ assert snapshot == {
+ "KEY_X_env": "x-dec",
+ "KEY_Y_env": "y-dec",
+ "returned_keys": ["KEY_X", "KEY_Y"],
+ }
+
+
+def test_ProxyConfig__decrypt_and_set_db_env_variables_invalid_dict_raises():
+ pc = ProxyConfig()
+ with pytest.raises(AttributeError):
+ pc._decrypt_and_set_db_env_variables("not-a-dict") # type: ignore[arg-type]
+
+
+# ---------------------------------------------------------------------------
+# ProxyConfig._decrypt_db_variables
+# ---------------------------------------------------------------------------
+
+
+def test_ProxyConfig__decrypt_db_variables_returns_decrypted(monkeypatch):
+ monkeypatch.setattr(
+ "litellm.proxy.proxy_server.decrypt_value_helper",
+ lambda value, key, return_original_value: f"D({value})",
+ )
+ pc = ProxyConfig()
+ out = pc._decrypt_db_variables({"a": "1", "b": "2", "c": "3"})
+ assert out == {"a": "D(1)", "b": "D(2)", "c": "D(3)"}
+
+
+def test_ProxyConfig__decrypt_db_variables_invalid_raises():
+ pc = ProxyConfig()
+ with pytest.raises(AttributeError):
+ pc._decrypt_db_variables(None) # type: ignore[arg-type]
+
+
+# ---------------------------------------------------------------------------
+# ProxyConfig._encrypt_env_variables_for_db
+# ---------------------------------------------------------------------------
+
+
+def test_ProxyConfig__encrypt_env_variables_for_db_idempotent(monkeypatch):
+ monkeypatch.setattr(
+ "litellm.proxy.proxy_server.decrypt_value_helper",
+ lambda value, key, return_original_value: value,
+ )
+ monkeypatch.setattr(
+ "litellm.proxy.proxy_server.encrypt_value_helper",
+ lambda value, new_encryption_key=None: f"ENC[{value}]",
+ )
+ pc = ProxyConfig()
+ out = pc._encrypt_env_variables_for_db({"A": "1", "B": "2", "C": "3"})
+ assert out == {"A": "ENC[1]", "B": "ENC[2]", "C": "ENC[3]"}
+
+
+def test_ProxyConfig__encrypt_env_variables_for_db_invalid_raises():
+ pc = ProxyConfig()
+ with pytest.raises(AttributeError):
+ pc._encrypt_env_variables_for_db(None) # type: ignore[arg-type]
+
+
+# ---------------------------------------------------------------------------
+# ProxyConfig._parse_router_settings_value
+# ---------------------------------------------------------------------------
+
+
+def test_ProxyConfig__parse_router_settings_value_handles_inputs():
+ result = {
+ "dict": ProxyConfig._parse_router_settings_value({"a": 1}),
+ "yaml_string": ProxyConfig._parse_router_settings_value("a: 1\nb: 2"),
+ "none": ProxyConfig._parse_router_settings_value(None),
+ }
+ assert result == {
+ "dict": {"a": 1},
+ "yaml_string": {"a": 1, "b": 2},
+ "none": None,
+ }
+
+
+def test_ProxyConfig__parse_router_settings_value_invalid_returns_none():
+ # Non-dict, non-parseable scalar -> None.
+ assert ProxyConfig._parse_router_settings_value(12345) is None
+ # Empty dict -> None (not truthy).
+ assert ProxyConfig._parse_router_settings_value({}) is None
+
+
+# ---------------------------------------------------------------------------
+# ProxyConfig._get_hierarchical_router_settings
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_ProxyConfig__get_hierarchical_router_settings_key_wins():
+ pc = ProxyConfig()
+ fake_key = SimpleNamespace(
+ router_settings={"timeout": 30, "retries": 2, "model": "gpt-4"},
+ team_id=None,
+ )
+ out = await pc._get_hierarchical_router_settings(
+ user_api_key_dict=fake_key,
+ prisma_client=None,
+ proxy_logging_obj=None,
+ )
+ assert out == {"timeout": 30, "retries": 2, "model": "gpt-4"}
+
+
+@pytest.mark.asyncio
+async def test_ProxyConfig__get_hierarchical_router_settings_missing_returns_none():
+ pc = ProxyConfig()
+ fake_key = SimpleNamespace(router_settings=None, team_id=None)
+ out = await pc._get_hierarchical_router_settings(
+ user_api_key_dict=fake_key,
+ prisma_client=None,
+ proxy_logging_obj=None,
+ )
+ assert out is None
+
+
+# ---------------------------------------------------------------------------
+# ProxyConfig._add_router_settings_from_db_config
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_ProxyConfig__add_router_settings_from_db_config_updates_router():
+ pc = ProxyConfig()
+ fake_router = MagicMock()
+ fake_router.update_settings = MagicMock()
+ fake_prisma = MagicMock()
+ fake_prisma.db.litellm_config.find_first = AsyncMock(
+ return_value=SimpleNamespace(
+ param_value={"timeout": 30, "retries": 2, "fallbacks": []}
+ )
+ )
+ config_data = {"router_settings": {"timeout": 10}}
+ await pc._add_router_settings_from_db_config(
+ config_data=config_data,
+ llm_router=fake_router,
+ prisma_client=fake_prisma,
+ )
+ snapshot = {
+ "called": fake_router.update_settings.called,
+ "call_count": fake_router.update_settings.call_count,
+ "kwargs_keys": sorted(
+ list(fake_router.update_settings.call_args.kwargs.keys())
+ ),
+ }
+ assert snapshot == {
+ "called": True,
+ "call_count": 1,
+ "kwargs_keys": ["fallbacks", "retries", "timeout"],
+ }
+
+
+@pytest.mark.asyncio
+async def test_ProxyConfig__add_router_settings_from_db_config_none_router_noop():
+ pc = ProxyConfig()
+ # No router and no prisma — should silently return.
+ await pc._add_router_settings_from_db_config(
+ config_data={}, llm_router=None, prisma_client=None
+ )
+ # Error-style: bad call signature raises.
+ with pytest.raises(TypeError):
+ await pc._add_router_settings_from_db_config() # type: ignore[call-arg]
+
+
+# ---------------------------------------------------------------------------
+# ProxyConfig._add_general_settings_from_db_config
+# ---------------------------------------------------------------------------
+
+
+def test_ProxyConfig__add_general_settings_from_db_config_merges_alerting():
+ pc = ProxyConfig()
+ proxy_logging = MagicMock()
+ general = {"alerting": ["slack"]}
+ config_data = {"general_settings": {"alerting": ["email", "slack"]}}
+ pc._add_general_settings_from_db_config(
+ config_data=config_data,
+ general_settings=general,
+ proxy_logging_obj=proxy_logging,
+ )
+ snapshot = {
+ "alerting": sorted(general["alerting"]),
+ "logging_called": proxy_logging.update_values.called,
+ "merged_count": len(general["alerting"]),
+ }
+ assert snapshot == {
+ "alerting": ["email", "slack"],
+ "logging_called": True,
+ "merged_count": 2,
+ }
+
+
+def test_ProxyConfig__add_general_settings_from_db_config_bad_config_raises():
+ pc = ProxyConfig()
+ with pytest.raises(AttributeError):
+ pc._add_general_settings_from_db_config(
+ config_data=None, # type: ignore[arg-type]
+ general_settings={},
+ proxy_logging_obj=MagicMock(),
+ )
+
+
+# ---------------------------------------------------------------------------
+# ProxyConfig._reschedule_spend_log_cleanup_job
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_ProxyConfig__reschedule_spend_log_cleanup_job_no_scheduler(monkeypatch):
+ monkeypatch.setattr("litellm.proxy.proxy_server.scheduler", None)
+ pc = ProxyConfig()
+ try:
+ await pc._reschedule_spend_log_cleanup_job()
+ raised = False
+ except Exception:
+ raised = True
+ snapshot = {"raised": raised, "called": True, "scheduler_was": "none"}
+ assert snapshot == {"raised": False, "called": True, "scheduler_was": "none"}
+
+
+@pytest.mark.asyncio
+async def test_ProxyConfig__reschedule_spend_log_cleanup_job_invalid_cron(monkeypatch):
+ fake_scheduler = MagicMock()
+ fake_scheduler.remove_job = MagicMock()
+ fake_scheduler.add_job = MagicMock()
+ monkeypatch.setattr("litellm.proxy.proxy_server.scheduler", fake_scheduler)
+ monkeypatch.setattr(
+ "litellm.proxy.proxy_server.general_settings",
+ {
+ "maximum_spend_logs_retention_period": "1d",
+ "maximum_spend_logs_cleanup_cron": "INVALID CRON STRING",
+ },
+ )
+ monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
+ pc = ProxyConfig()
+ # Invalid cron is caught and logged — does not raise outward.
+ await pc._reschedule_spend_log_cleanup_job()
+ # But add_job should not have been called for the invalid cron path.
+ assert fake_scheduler.add_job.call_count == 0
+
+
+# ---------------------------------------------------------------------------
+# ProxyConfig._update_general_settings
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_ProxyConfig__update_general_settings_updates_max_parallel(monkeypatch):
+ monkeypatch.setattr(
+ "litellm.proxy.proxy_server.general_settings",
+ {},
+ )
+ pc = ProxyConfig()
+ await pc._update_general_settings(
+ {
+ "max_parallel_requests": 7,
+ "global_max_parallel_requests": 99,
+ "ui_access_mode": "admin_only",
+ }
+ )
+ from litellm.proxy import proxy_server as ps
+
+ snapshot = {
+ "max_parallel_requests": ps.general_settings.get("max_parallel_requests"),
+ "global_max_parallel_requests": ps.general_settings.get(
+ "global_max_parallel_requests"
+ ),
+ "ui_access_mode": ps.general_settings.get("ui_access_mode"),
+ }
+ assert snapshot == {
+ "max_parallel_requests": 7,
+ "global_max_parallel_requests": 99,
+ "ui_access_mode": "admin_only",
+ }
+
+
+@pytest.mark.asyncio
+async def test_ProxyConfig__update_general_settings_none_input_noop():
+ pc = ProxyConfig()
+ # None input returns early.
+ result = await pc._update_general_settings(db_general_settings=None)
+ assert result is None
+ # Error-style: dict() will fail on non-mapping non-None input.
+ with pytest.raises(Exception):
+ await pc._update_general_settings(db_general_settings=12345) # type: ignore[arg-type]
+
+
+# ---------------------------------------------------------------------------
+# ProxyConfig._update_config_fields
+# ---------------------------------------------------------------------------
+
+
+def test_ProxyConfig__update_config_fields_merges_dict():
+ pc = ProxyConfig()
+ current = {"general_settings": {"a": 1, "b": 2}}
+ out = pc._update_config_fields(
+ current_config=current,
+ param_name="general_settings",
+ db_param_value={"b": 3, "c": 4, "d": 5},
+ )
+ assert out == {"general_settings": {"a": 1, "b": 3, "c": 4, "d": 5}}
+
+
+def test_ProxyConfig__update_config_fields_invalid_param_raises():
+ pc = ProxyConfig()
+ with pytest.raises(Exception):
+ # Missing required arg.
+ pc._update_config_fields(current_config={}, param_name="general_settings") # type: ignore[call-arg]
diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_anthropic_beta.py b/tests/test_litellm/proxy/proxy_server/test_routes_anthropic_beta.py
index ad6b4016461..7ef29b71bf0 100644
--- a/tests/test_litellm/proxy/proxy_server/test_routes_anthropic_beta.py
+++ b/tests/test_litellm/proxy/proxy_server/test_routes_anthropic_beta.py
@@ -1 +1,370 @@
-"""Placeholder. Filled by a follow-up PR per the Notion plan."""
+"""Pin tests for proxy_server.py Anthropic-beta-headers reload routes (PR3).
+
+Routes covered:
+- POST /reload/anthropic_beta_headers
+- POST /schedule/anthropic_beta_headers_reload
+- DELETE /schedule/anthropic_beta_headers_reload
+- GET /schedule/anthropic_beta_headers_reload/status
+"""
+
+from __future__ import annotations
+
+import json
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from .conftest import VOLATILE_KEYS, normalize
+
+# These routes return a "timestamp" ISO string that isn't in the default
+# volatile-keys set — extend the set locally so dict-equality assertions
+# can ignore it.
+_VOLATILE = VOLATILE_KEYS | frozenset({"timestamp"})
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_prisma_with_config(
+ config_record=None,
+):
+ """Build a MagicMock prisma_client with a ``db.litellm_config`` namespace.
+
+ The conftest's ``mock_prisma`` fixture stubs ``litellm_configtable`` but
+ the anthropic-beta routes use ``prisma_client.db.litellm_config`` —
+ a different attribute. Build one here so each test gets isolated state.
+ """
+ config = MagicMock()
+ config.find_unique = AsyncMock(return_value=config_record)
+ config.upsert = AsyncMock()
+ config.delete = AsyncMock()
+
+ db = MagicMock()
+ db.litellm_config = config
+
+ client = MagicMock()
+ client.db = db
+ return client
+
+
+def _install_prisma(monkeypatch, prisma):
+ from litellm.proxy import proxy_server as ps
+
+ monkeypatch.setattr(ps, "prisma_client", prisma)
+
+
+def _stub_reload_beta_headers(monkeypatch, return_value=None):
+ """Replace ``litellm.anthropic_beta_headers_manager.reload_beta_headers_config``
+ with a deterministic stub so the route never hits the network."""
+ if return_value is None:
+ return_value = {
+ "anthropic": {"beta_headers": ["foo"]},
+ "openai": {"beta_headers": ["bar"]},
+ "provider_aliases": {"a": "b"},
+ "description": "test",
+ }
+ import litellm.anthropic_beta_headers_manager as mgr
+
+ stub = MagicMock(return_value=return_value)
+ monkeypatch.setattr(mgr, "reload_beta_headers_config", stub)
+ return stub
+
+
+# ---------------------------------------------------------------------------
+# POST /reload/anthropic_beta_headers
+# ---------------------------------------------------------------------------
+
+
+def test_reload_anthropic_beta_headers_admin_success(client, auth_as, monkeypatch):
+ """Admin can trigger immediate reload — handler returns providers count and
+ a success status. Pins the response dict shape."""
+ from litellm.proxy._types import LitellmUserRoles
+
+ _stub_reload_beta_headers(monkeypatch)
+ prisma = _make_prisma_with_config(config_record=None)
+ _install_prisma(monkeypatch, prisma)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.post("/reload/anthropic_beta_headers")
+
+ assert response.status_code == 200
+ body = response.json()
+ # Two non-alias keys: "anthropic", "openai"
+ assert normalize(body, _VOLATILE) == {
+ "message": "Anthropic beta headers configuration reloaded successfully! 2 providers updated.",
+ "status": "success",
+ "providers_count": 2,
+ "timestamp": "",
+ }
+ # And the upsert was actually invoked (force_reload write).
+ prisma.db.litellm_config.upsert.assert_awaited_once()
+
+
+def test_reload_anthropic_beta_headers_preserves_existing_interval(
+ client, auth_as, monkeypatch
+):
+ """When an existing reload config has an interval set, the force-reload
+ write must preserve that interval (the route reads it back then upserts
+ with the same number). This pins the read-then-write behaviour."""
+ from litellm.proxy._types import LitellmUserRoles
+
+ _stub_reload_beta_headers(monkeypatch)
+ existing = SimpleNamespace(
+ param_name="anthropic_beta_headers_reload_config",
+ param_value={"interval_hours": 12, "force_reload": False},
+ )
+ prisma = _make_prisma_with_config(config_record=existing)
+ _install_prisma(monkeypatch, prisma)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.post("/reload/anthropic_beta_headers")
+
+ assert response.status_code == 200
+ # The update branch's interval_hours was sourced from the existing record.
+ call_kwargs = prisma.db.litellm_config.upsert.await_args.kwargs
+ data = call_kwargs["data"]
+ update_payload = data["update"]["param_value"]
+ parsed = (
+ json.loads(update_payload)
+ if isinstance(update_payload, str)
+ else update_payload
+ )
+ assert parsed["interval_hours"] == 12
+ assert parsed["force_reload"] is True
+
+
+def test_reload_anthropic_beta_headers_not_admin_forbidden(client, auth_as):
+ from litellm.proxy._types import LitellmUserRoles
+
+ with auth_as(LitellmUserRoles.INTERNAL_USER):
+ response = client.post("/reload/anthropic_beta_headers")
+
+ assert response.status_code == 403
+ assert "Admin role required" in response.json().get("detail", "")
+
+
+def test_reload_anthropic_beta_headers_no_db_returns_500(client, auth_as, monkeypatch):
+ """When prisma_client is None the handler raises 500 with a clear message."""
+ from litellm.proxy._types import LitellmUserRoles
+
+ _install_prisma(monkeypatch, None)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.post("/reload/anthropic_beta_headers")
+
+ assert response.status_code == 500
+ assert "Database connection not available" in response.json().get("detail", "")
+
+
+# ---------------------------------------------------------------------------
+# POST /schedule/anthropic_beta_headers_reload
+# ---------------------------------------------------------------------------
+
+
+def test_schedule_anthropic_beta_headers_reload_admin_success(
+ client, auth_as, monkeypatch
+):
+ """Happy path: admin schedules every N hours — response echoes interval."""
+ from litellm.proxy._types import LitellmUserRoles
+
+ prisma = _make_prisma_with_config()
+ _install_prisma(monkeypatch, prisma)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.post(
+ "/schedule/anthropic_beta_headers_reload", params={"hours": 6}
+ )
+
+ assert response.status_code == 200
+ assert normalize(response.json(), _VOLATILE) == {
+ "message": "Anthropic beta headers reload scheduled for every 6 hours",
+ "status": "success",
+ "interval_hours": 6,
+ "timestamp": "",
+ }
+ prisma.db.litellm_config.upsert.assert_awaited_once()
+
+
+def test_schedule_anthropic_beta_headers_reload_zero_hours_400(
+ client, auth_as, monkeypatch
+):
+ """``hours <= 0`` is rejected with 400 and a descriptive message."""
+ from litellm.proxy._types import LitellmUserRoles
+
+ prisma = _make_prisma_with_config()
+ _install_prisma(monkeypatch, prisma)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.post(
+ "/schedule/anthropic_beta_headers_reload", params={"hours": 0}
+ )
+
+ assert response.status_code == 400
+ assert "Hours must be greater than 0" in response.json().get("detail", "")
+
+
+def test_schedule_anthropic_beta_headers_reload_not_admin_forbidden(client, auth_as):
+ from litellm.proxy._types import LitellmUserRoles
+
+ with auth_as(LitellmUserRoles.INTERNAL_USER):
+ response = client.post(
+ "/schedule/anthropic_beta_headers_reload", params={"hours": 6}
+ )
+
+ assert response.status_code == 403
+ assert "Admin role required" in response.json().get("detail", "")
+
+
+def test_schedule_anthropic_beta_headers_reload_missing_hours_422(client, auth_as):
+ """``hours`` is a required query param — omitting it is a FastAPI 422."""
+ from litellm.proxy._types import LitellmUserRoles
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.post("/schedule/anthropic_beta_headers_reload")
+
+ assert response.status_code == 422
+ assert "detail" in response.json()
+
+
+# ---------------------------------------------------------------------------
+# DELETE /schedule/anthropic_beta_headers_reload
+# ---------------------------------------------------------------------------
+
+
+def test_cancel_anthropic_beta_headers_reload_admin_success(
+ client, auth_as, monkeypatch
+):
+ """Admin cancel: deletes the LiteLLM_Config row and returns success dict."""
+ from litellm.proxy._types import LitellmUserRoles
+
+ prisma = _make_prisma_with_config()
+ _install_prisma(monkeypatch, prisma)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.delete("/schedule/anthropic_beta_headers_reload")
+
+ assert response.status_code == 200
+ assert normalize(response.json(), _VOLATILE) == {
+ "message": "Anthropic beta headers reload schedule cancelled",
+ "status": "success",
+ "timestamp": "",
+ }
+ prisma.db.litellm_config.delete.assert_awaited_once_with(
+ where={"param_name": "anthropic_beta_headers_reload_config"}
+ )
+
+
+def test_cancel_anthropic_beta_headers_reload_not_admin_forbidden(client, auth_as):
+ from litellm.proxy._types import LitellmUserRoles
+
+ with auth_as(LitellmUserRoles.INTERNAL_USER):
+ response = client.delete("/schedule/anthropic_beta_headers_reload")
+
+ assert response.status_code == 403
+ assert "Admin role required" in response.json().get("detail", "")
+
+
+def test_cancel_anthropic_beta_headers_reload_no_db_returns_500(
+ client, auth_as, monkeypatch
+):
+ from litellm.proxy._types import LitellmUserRoles
+
+ _install_prisma(monkeypatch, None)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.delete("/schedule/anthropic_beta_headers_reload")
+
+ assert response.status_code == 500
+ assert "Database connection not available" in response.json().get("detail", "")
+
+
+# ---------------------------------------------------------------------------
+# GET /schedule/anthropic_beta_headers_reload/status
+# ---------------------------------------------------------------------------
+
+
+def test_get_anthropic_beta_headers_reload_status_scheduled(
+ client, auth_as, monkeypatch
+):
+ """When a config row with ``interval_hours`` is present, ``scheduled`` is True
+ and ``interval_hours`` echoes the DB value. Pins the full response shape."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ record = SimpleNamespace(
+ param_name="anthropic_beta_headers_reload_config",
+ param_value={"interval_hours": 6, "force_reload": False},
+ )
+ prisma = _make_prisma_with_config(config_record=record)
+ _install_prisma(monkeypatch, prisma)
+ # No prior reload — next_run stays None.
+ monkeypatch.setattr(ps, "last_anthropic_beta_headers_reload", None)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.get("/schedule/anthropic_beta_headers_reload/status")
+
+ assert response.status_code == 200
+ assert normalize(response.json()) == {
+ "scheduled": True,
+ "interval_hours": 6,
+ "last_run": None,
+ "next_run": None,
+ }
+
+
+def test_get_anthropic_beta_headers_reload_status_not_scheduled_no_db(
+ client, auth_as, monkeypatch
+):
+ """No DB connection: handler returns the unscheduled-status dict (not 500)."""
+ from litellm.proxy._types import LitellmUserRoles
+
+ _install_prisma(monkeypatch, None)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.get("/schedule/anthropic_beta_headers_reload/status")
+
+ assert response.status_code == 200
+ assert normalize(response.json()) == {
+ "scheduled": False,
+ "interval_hours": None,
+ "last_run": None,
+ "next_run": None,
+ }
+
+
+def test_get_anthropic_beta_headers_reload_status_no_interval_unscheduled(
+ client, auth_as, monkeypatch
+):
+ """Config row present but ``interval_hours`` is None → unscheduled response."""
+ from litellm.proxy._types import LitellmUserRoles
+
+ record = SimpleNamespace(
+ param_name="anthropic_beta_headers_reload_config",
+ param_value={"interval_hours": None, "force_reload": True},
+ )
+ prisma = _make_prisma_with_config(config_record=record)
+ _install_prisma(monkeypatch, prisma)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.get("/schedule/anthropic_beta_headers_reload/status")
+
+ assert response.status_code == 200
+ assert normalize(response.json()) == {
+ "scheduled": False,
+ "interval_hours": None,
+ "last_run": None,
+ "next_run": None,
+ }
+
+
+def test_get_anthropic_beta_headers_reload_status_not_admin_forbidden(client, auth_as):
+ from litellm.proxy._types import LitellmUserRoles
+
+ with auth_as(LitellmUserRoles.INTERNAL_USER):
+ response = client.get("/schedule/anthropic_beta_headers_reload/status")
+
+ assert response.status_code == 403
+ assert "Admin role required" in response.json().get("detail", "")
diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_assistants.py b/tests/test_litellm/proxy/proxy_server/test_routes_assistants.py
index ad6b4016461..fd1f672dce9 100644
--- a/tests/test_litellm/proxy/proxy_server/test_routes_assistants.py
+++ b/tests/test_litellm/proxy/proxy_server/test_routes_assistants.py
@@ -1 +1,180 @@
-"""Placeholder. Filled by a follow-up PR per the Notion plan."""
+"""Behavior pins for ``proxy_server.py`` assistants routes.
+
+Pins (PR2):
+ - GET /v1/assistants
+ - GET /assistants
+ - POST /v1/assistants
+ - POST /assistants
+ - DELETE /v1/assistants/{assistant_id:path}
+ - DELETE /assistants/{assistant_id:path}
+"""
+
+from __future__ import annotations
+
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from litellm.proxy import proxy_server
+
+from .conftest import normalize # type: ignore[import-not-found]
+
+GET_RESPONSE = {
+ "object": "list",
+ "data": [
+ {
+ "id": "asst_1",
+ "object": "assistant",
+ "name": "Test Assistant",
+ "model": "gpt-4",
+ }
+ ],
+ "first_id": "asst_1",
+ "last_id": "asst_1",
+ "has_more": False,
+}
+
+
+CREATE_RESPONSE = {
+ "id": "asst_new",
+ "object": "assistant",
+ "name": "New",
+ "model": "gpt-4",
+ "created_at": 0,
+}
+
+
+DELETE_RESPONSE = {"id": "asst_1", "object": "assistant.deleted", "deleted": True}
+
+
+@pytest.fixture
+def patched_assistants(monkeypatch):
+ router = MagicMock()
+ router.aget_assistants = AsyncMock(return_value=dict(GET_RESPONSE))
+ router.acreate_assistants = AsyncMock(return_value=dict(CREATE_RESPONSE))
+ router.adelete_assistant = AsyncMock(return_value=dict(DELETE_RESPONSE))
+ monkeypatch.setattr(proxy_server, "llm_router", router)
+ monkeypatch.setattr(
+ proxy_server,
+ "proxy_logging_obj",
+ MagicMock(
+ post_call_failure_hook=AsyncMock(), update_request_status=AsyncMock()
+ ),
+ )
+
+ async def _add_data(data, **kwargs):
+ return data
+
+ monkeypatch.setattr(proxy_server, "add_litellm_data_to_request", _add_data)
+ return router
+
+
+@pytest.fixture
+def no_router(monkeypatch):
+ monkeypatch.setattr(proxy_server, "llm_router", None)
+ monkeypatch.setattr(
+ proxy_server,
+ "proxy_logging_obj",
+ MagicMock(
+ post_call_failure_hook=AsyncMock(), update_request_status=AsyncMock()
+ ),
+ )
+
+ async def _add_data(data, **kwargs):
+ return data
+
+ monkeypatch.setattr(proxy_server, "add_litellm_data_to_request", _add_data)
+ yield
+
+
+# ---------------------------------------------------------------------------
+# GET /v1/assistants, GET /assistants
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("path", ["/v1/assistants", "/assistants"])
+def test_get_assistants_happy_path(client, auth_as, patched_assistants, path):
+ """Pins ``GET /v1/assistants`` and ``GET /assistants``."""
+ with auth_as():
+ response = client.get(path)
+ assert response.status_code == 200
+ assert normalize(response.json()) == {
+ "object": "list",
+ "data": [
+ {
+ "id": "",
+ "object": "assistant",
+ "name": "Test Assistant",
+ "model": "gpt-4",
+ }
+ ],
+ "first_id": "asst_1",
+ "last_id": "asst_1",
+ "has_more": False,
+ }
+
+
+@pytest.mark.parametrize("path", ["/v1/assistants", "/assistants"])
+def test_get_assistants_no_router_error(client, auth_as, no_router, path):
+ """Pins ``GET /v1/assistants`` and ``GET /assistants`` (error: no llm_router)."""
+ with auth_as():
+ response = client.get(path)
+ assert response.status_code == 500
+ assert len(response.content) > 0
+
+
+# ---------------------------------------------------------------------------
+# POST /v1/assistants, POST /assistants
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("path", ["/v1/assistants", "/assistants"])
+def test_create_assistant_happy_path(client, auth_as, patched_assistants, path):
+ """Pins ``POST /v1/assistants`` and ``POST /assistants``."""
+ payload = {"model": "gpt-4", "name": "New"}
+ with auth_as():
+ response = client.post(path, json=payload)
+ assert response.status_code == 200
+ assert normalize(response.json()) == {
+ "id": "",
+ "object": "assistant",
+ "name": "New",
+ "model": "gpt-4",
+ "created_at": "",
+ }
+
+
+@pytest.mark.parametrize("path", ["/v1/assistants", "/assistants"])
+def test_create_assistant_no_router_error(client, auth_as, no_router, path):
+ """Pins ``POST /v1/assistants`` and ``POST /assistants`` (error: no llm_router)."""
+ with auth_as():
+ response = client.post(path, json={"model": "gpt-4"})
+ assert response.status_code == 500
+ assert len(response.content) > 0
+
+
+# ---------------------------------------------------------------------------
+# DELETE /v1/assistants/{assistant_id:path}, DELETE /assistants/{assistant_id:path}
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("path", ["/v1/assistants/asst_1", "/assistants/asst_1"])
+def test_delete_assistant_happy_path(client, auth_as, patched_assistants, path):
+ """Pins ``DELETE /v1/assistants/{assistant_id:path}`` and ``DELETE /assistants/{assistant_id:path}``."""
+ with auth_as():
+ response = client.delete(path)
+ assert response.status_code == 200
+ assert normalize(response.json()) == {
+ "id": "",
+ "object": "assistant.deleted",
+ "deleted": True,
+ }
+
+
+@pytest.mark.parametrize("path", ["/v1/assistants/asst_1", "/assistants/asst_1"])
+def test_delete_assistant_no_router_error(client, auth_as, no_router, path):
+ """Pins ``DELETE /v1/assistants/{assistant_id:path}`` / ``DELETE /assistants/{assistant_id:path}`` (error)."""
+ with auth_as():
+ response = client.delete(path)
+ assert response.status_code == 500
+ assert len(response.content) > 0
diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_audio.py b/tests/test_litellm/proxy/proxy_server/test_routes_audio.py
index ad6b4016461..d88bcf136e9 100644
--- a/tests/test_litellm/proxy/proxy_server/test_routes_audio.py
+++ b/tests/test_litellm/proxy/proxy_server/test_routes_audio.py
@@ -1 +1,193 @@
-"""Placeholder. Filled by a follow-up PR per the Notion plan."""
+"""Behavior pins for ``proxy_server.py`` audio routes.
+
+Pins (PR2):
+ - POST /v1/audio/speech
+ - POST /audio/speech
+ - POST /v1/audio/transcriptions
+ - POST /audio/transcriptions
+"""
+
+from __future__ import annotations
+
+import io
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from litellm.proxy import proxy_server
+
+
+@pytest.fixture
+def patched_speech(monkeypatch):
+ monkeypatch.setattr(proxy_server, "llm_router", MagicMock())
+ monkeypatch.setattr(
+ proxy_server,
+ "proxy_logging_obj",
+ MagicMock(
+ pre_call_hook=AsyncMock(side_effect=lambda **kw: kw["data"]),
+ post_call_failure_hook=AsyncMock(),
+ update_request_status=AsyncMock(),
+ ),
+ )
+
+ async def _add_data(data, **kwargs):
+ return data
+
+ monkeypatch.setattr(proxy_server, "add_litellm_data_to_request", _add_data)
+
+ class _FakeBinaryResp:
+ async def aiter_bytes(self, chunk_size: int = 8192):
+ async def _gen():
+ yield b"\x00\x01\x02"
+
+ return _gen()
+
+ async def _llm_call():
+ return _FakeBinaryResp()
+
+ async def _fake_route_request(*args, **kwargs):
+ return _llm_call()
+
+ monkeypatch.setattr(proxy_server, "route_request", _fake_route_request)
+ yield
+
+
+@pytest.fixture
+def patched_speech_error(monkeypatch):
+ monkeypatch.setattr(proxy_server, "llm_router", MagicMock())
+ monkeypatch.setattr(
+ proxy_server,
+ "proxy_logging_obj",
+ MagicMock(
+ pre_call_hook=AsyncMock(side_effect=lambda **kw: kw["data"]),
+ post_call_failure_hook=AsyncMock(),
+ update_request_status=AsyncMock(),
+ ),
+ )
+
+ async def _add_data(data, **kwargs):
+ return data
+
+ monkeypatch.setattr(proxy_server, "add_litellm_data_to_request", _add_data)
+
+ async def _raise(*args, **kwargs):
+ raise ValueError("speech boom")
+
+ monkeypatch.setattr(proxy_server, "route_request", _raise)
+ yield
+
+
+@pytest.fixture
+def patched_transcription(monkeypatch):
+ router = MagicMock()
+ router.model_names = ["whisper-1"]
+ monkeypatch.setattr(proxy_server, "llm_router", router)
+ monkeypatch.setattr(
+ proxy_server,
+ "proxy_logging_obj",
+ MagicMock(
+ pre_call_hook=AsyncMock(side_effect=lambda **kw: kw["data"]),
+ post_call_failure_hook=AsyncMock(),
+ post_call_response_headers_hook=AsyncMock(return_value={}),
+ update_request_status=AsyncMock(),
+ ),
+ )
+
+ async def _add_data(data, **kwargs):
+ return data
+
+ monkeypatch.setattr(proxy_server, "add_litellm_data_to_request", _add_data)
+ monkeypatch.setattr(
+ proxy_server, "check_file_size_under_limit", lambda **kwargs: True
+ )
+
+ async def _form_data(request):
+ from starlette.datastructures import FormData, UploadFile
+
+ upload = UploadFile(
+ filename="audio.mp3",
+ file=io.BytesIO(b"\x00\x01\x02"),
+ )
+ return FormData([("file", upload), ("model", "whisper-1")])
+
+ monkeypatch.setattr(proxy_server, "get_form_data", _form_data)
+
+ async def _llm_call():
+ return {"text": "hello world"}
+
+ async def _fake_route_request(*args, **kwargs):
+ return _llm_call()
+
+ monkeypatch.setattr(proxy_server, "route_request", _fake_route_request)
+ yield
+
+
+@pytest.fixture
+def patched_transcription_error(monkeypatch, patched_transcription):
+ async def _raise(*args, **kwargs):
+ raise ValueError("transcription boom")
+
+ monkeypatch.setattr(proxy_server, "route_request", _raise)
+ yield
+
+
+@pytest.mark.parametrize("path", ["/v1/audio/speech", "/audio/speech"])
+def test_audio_speech_happy_path(client, auth_as, patched_speech, path):
+ """Pins ``POST /v1/audio/speech`` and ``POST /audio/speech`` (happy)."""
+ payload = {"model": "tts-1", "input": "Hi", "voice": "alloy"}
+ with auth_as():
+ response = client.post(path, json=payload)
+ assert response.status_code == 200
+ response_summary = {
+ "status_code": response.status_code,
+ "content_type": response.headers.get("content-type", ""),
+ "body_bytes": response.content,
+ }
+ assert response_summary == {
+ "status_code": 200,
+ "content_type": "audio/mpeg",
+ "body_bytes": b"\x00\x01\x02",
+ }
+
+
+@pytest.mark.parametrize("path", ["/v1/audio/speech", "/audio/speech"])
+def test_audio_speech_error(client, auth_as, patched_speech_error, path):
+ """Pins ``POST /v1/audio/speech`` and ``POST /audio/speech`` (error)."""
+ payload = {"model": "tts-1", "input": "Hi", "voice": "alloy"}
+ with auth_as():
+ response = client.post(path, json=payload)
+ assert response.status_code == 500
+ assert len(response.content) > 0
+
+
+@pytest.mark.parametrize("path", ["/v1/audio/transcriptions", "/audio/transcriptions"])
+def test_audio_transcription_happy_path(client, auth_as, patched_transcription, path):
+ """Pins ``POST /v1/audio/transcriptions`` / ``POST /audio/transcriptions`` (happy)."""
+ files = {"file": ("audio.mp3", b"\x00\x01\x02", "audio/mpeg")}
+ data = {"model": "whisper-1"}
+ with auth_as():
+ response = client.post(path, files=files, data=data)
+ assert response.status_code == 200
+ body = response.json()
+ assert body == {"text": "hello world"}
+ response_summary = {
+ "status_code": response.status_code,
+ "text_field": body["text"],
+ "media_type_hint": response.headers.get("content-type", "").split(";")[0],
+ }
+ assert response_summary == {
+ "status_code": 200,
+ "text_field": "hello world",
+ "media_type_hint": "application/json",
+ }
+
+
+@pytest.mark.parametrize("path", ["/v1/audio/transcriptions", "/audio/transcriptions"])
+def test_audio_transcription_error(client, auth_as, patched_transcription_error, path):
+ """Pins ``POST /v1/audio/transcriptions`` / ``POST /audio/transcriptions`` (error)."""
+ files = {"file": ("audio.mp3", b"\x00\x01\x02", "audio/mpeg")}
+ data = {"model": "whisper-1"}
+ with auth_as():
+ response = client.post(path, files=files, data=data)
+ assert response.status_code == 500
+ assert len(response.content) > 0
diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_chat_completions.py b/tests/test_litellm/proxy/proxy_server/test_routes_chat_completions.py
index ad6b4016461..b186bb5ef5e 100644
--- a/tests/test_litellm/proxy/proxy_server/test_routes_chat_completions.py
+++ b/tests/test_litellm/proxy/proxy_server/test_routes_chat_completions.py
@@ -1 +1,134 @@
-"""Placeholder. Filled by a follow-up PR per the Notion plan."""
+"""Behavior pins for ``proxy_server.py`` chat-completions routes.
+
+Pins (PR2):
+ - POST /v1/chat/completions
+ - POST /chat/completions
+ - POST /engines/{model:path}/chat/completions
+ - POST /openai/deployments/{model:path}/chat/completions
+"""
+
+from __future__ import annotations
+
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from litellm.proxy import common_request_processing, proxy_server
+
+from .conftest import normalize # type: ignore[import-not-found]
+
+HAPPY_RESPONSE = {
+ "id": "chatcmpl-test",
+ "object": "chat.completion",
+ "created": 0,
+ "model": "gpt-4",
+ "choices": [
+ {
+ "index": 0,
+ "finish_reason": "stop",
+ "message": {"role": "assistant", "content": "Hello from mock"},
+ }
+ ],
+ "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
+}
+
+
+@pytest.fixture
+def patched_chat(monkeypatch):
+ """Stub chat-completions pipeline at ProxyBaseLLMRequestProcessing."""
+ monkeypatch.setattr(proxy_server, "llm_router", MagicMock())
+ monkeypatch.setattr(
+ proxy_server, "proxy_logging_obj", MagicMock(post_call_failure_hook=AsyncMock())
+ )
+
+ async def _fake_process(self, *args, **kwargs):
+ return dict(HAPPY_RESPONSE)
+
+ monkeypatch.setattr(
+ common_request_processing.ProxyBaseLLMRequestProcessing,
+ "base_process_llm_request",
+ _fake_process,
+ )
+ yield
+
+
+@pytest.fixture
+def patched_chat_error(monkeypatch):
+ """Variant that makes the pipeline raise -> 400 via _handle_llm_api_exception."""
+ monkeypatch.setattr(proxy_server, "llm_router", MagicMock())
+ monkeypatch.setattr(
+ proxy_server, "proxy_logging_obj", MagicMock(post_call_failure_hook=AsyncMock())
+ )
+
+ from litellm.proxy._types import ProxyException
+
+ async def _raise(self, *args, **kwargs):
+ raise ValueError("boom")
+
+ async def _handler(self, *, e, user_api_key_dict, proxy_logging_obj):
+ return ProxyException(
+ message="boom", type="bad_request_error", param="model", code=400
+ )
+
+ monkeypatch.setattr(
+ common_request_processing.ProxyBaseLLMRequestProcessing,
+ "base_process_llm_request",
+ _raise,
+ )
+ monkeypatch.setattr(
+ common_request_processing.ProxyBaseLLMRequestProcessing,
+ "_handle_llm_api_exception",
+ _handler,
+ )
+ yield
+
+
+_CHAT_PATHS = [
+ "/v1/chat/completions",
+ "/chat/completions",
+ "/engines/gpt-4/chat/completions",
+ "/openai/deployments/gpt-4/chat/completions",
+]
+
+
+@pytest.mark.parametrize("path", _CHAT_PATHS)
+def test_chat_completion_happy_path(client, auth_as, patched_chat, path):
+ """Pins all four ``POST .../chat/completions`` aliases (happy path).
+
+ Covers ``POST /v1/chat/completions``, ``POST /chat/completions``,
+ ``POST /engines/{model:path}/chat/completions``, and
+ ``POST /openai/deployments/{model:path}/chat/completions``.
+ """
+ payload = {"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]}
+ with auth_as():
+ response = client.post(path, json=payload)
+ assert response.status_code == 200
+ assert normalize(response.json()) == {
+ "id": "",
+ "object": "chat.completion",
+ "created": "",
+ "model": "gpt-4",
+ "choices": [
+ {
+ "index": 0,
+ "finish_reason": "stop",
+ "message": {"role": "assistant", "content": "Hello from mock"},
+ }
+ ],
+ "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
+ }
+
+
+@pytest.mark.parametrize("path", _CHAT_PATHS)
+def test_chat_completion_pipeline_error(client, auth_as, patched_chat_error, path):
+ """Pins all four ``POST .../chat/completions`` aliases (error: 400).
+
+ Covers ``POST /v1/chat/completions``, ``POST /chat/completions``,
+ ``POST /engines/{model:path}/chat/completions``, and
+ ``POST /openai/deployments/{model:path}/chat/completions``.
+ """
+ payload = {"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]}
+ with auth_as():
+ response = client.post(path, json=payload)
+ assert response.status_code == 400
+ assert "error" in response.json() or response.text != ""
diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_completions.py b/tests/test_litellm/proxy/proxy_server/test_routes_completions.py
index ad6b4016461..b5c60c23c02 100644
--- a/tests/test_litellm/proxy/proxy_server/test_routes_completions.py
+++ b/tests/test_litellm/proxy/proxy_server/test_routes_completions.py
@@ -1 +1,126 @@
-"""Placeholder. Filled by a follow-up PR per the Notion plan."""
+"""Behavior pins for ``proxy_server.py`` text-completions routes.
+
+Pins (PR2):
+ - POST /v1/completions
+ - POST /completions
+ - POST /engines/{model:path}/completions
+ - POST /openai/deployments/{model:path}/completions
+"""
+
+from __future__ import annotations
+
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from litellm.proxy import common_request_processing, proxy_server
+
+from .conftest import normalize # type: ignore[import-not-found]
+
+HAPPY_RESPONSE = {
+ "id": "cmpl-test",
+ "object": "text_completion",
+ "created": 0,
+ "model": "gpt-3.5-turbo-instruct",
+ "choices": [
+ {
+ "index": 0,
+ "text": "Hello from mock",
+ "finish_reason": "stop",
+ "logprobs": None,
+ }
+ ],
+ "usage": {"prompt_tokens": 2, "completion_tokens": 3, "total_tokens": 5},
+}
+
+
+@pytest.fixture
+def patched_completion(monkeypatch):
+ monkeypatch.setattr(proxy_server, "llm_router", MagicMock())
+ monkeypatch.setattr(
+ proxy_server, "proxy_logging_obj", MagicMock(post_call_failure_hook=AsyncMock())
+ )
+
+ async def _fake_process(self, *args, **kwargs):
+ return dict(HAPPY_RESPONSE)
+
+ monkeypatch.setattr(
+ common_request_processing.ProxyBaseLLMRequestProcessing,
+ "base_process_llm_request",
+ _fake_process,
+ )
+ yield
+
+
+@pytest.fixture
+def completion_pipeline_raises(monkeypatch):
+ monkeypatch.setattr(proxy_server, "llm_router", MagicMock())
+ monkeypatch.setattr(
+ proxy_server, "proxy_logging_obj", MagicMock(post_call_failure_hook=AsyncMock())
+ )
+
+ async def _raise(self, *args, **kwargs):
+ raise ValueError("boom")
+
+ monkeypatch.setattr(
+ common_request_processing.ProxyBaseLLMRequestProcessing,
+ "base_process_llm_request",
+ _raise,
+ )
+ yield
+
+
+_COMPLETION_PATHS = [
+ "/v1/completions",
+ "/completions",
+ "/engines/gpt-3.5-turbo-instruct/completions",
+ "/openai/deployments/gpt-3.5-turbo-instruct/completions",
+]
+
+
+@pytest.mark.parametrize("path", _COMPLETION_PATHS)
+def test_completion_happy_path(client, auth_as, patched_completion, path):
+ """Pins all four ``POST .../completions`` aliases (happy path).
+
+ Covers ``POST /v1/completions``, ``POST /completions``,
+ ``POST /engines/{model:path}/completions``, and
+ ``POST /openai/deployments/{model:path}/completions``.
+ """
+ payload = {
+ "model": "gpt-3.5-turbo-instruct",
+ "prompt": "Once upon",
+ "max_tokens": 5,
+ }
+ with auth_as():
+ response = client.post(path, json=payload)
+ assert response.status_code == 200
+ assert normalize(response.json()) == {
+ "id": "",
+ "object": "text_completion",
+ "created": "",
+ "model": "gpt-3.5-turbo-instruct",
+ "choices": [
+ {
+ "index": 0,
+ "text": "Hello from mock",
+ "finish_reason": "stop",
+ "logprobs": None,
+ }
+ ],
+ "usage": {"prompt_tokens": 2, "completion_tokens": 3, "total_tokens": 5},
+ }
+
+
+@pytest.mark.parametrize("path", _COMPLETION_PATHS)
+def test_completion_pipeline_error(client, auth_as, completion_pipeline_raises, path):
+ """Pins all four ``POST .../completions`` aliases (error path).
+
+ Covers ``POST /v1/completions``, ``POST /completions``,
+ ``POST /engines/{model:path}/completions``, and
+ ``POST /openai/deployments/{model:path}/completions``.
+ """
+ payload = {"model": "gpt-3.5-turbo-instruct", "prompt": "boom"}
+ with auth_as():
+ response = client.post(path, json=payload)
+ assert response.status_code == 500
+ assert response.headers.get("content-type", "").startswith("application/json")
diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_config.py b/tests/test_litellm/proxy/proxy_server/test_routes_config.py
index ad6b4016461..e89ada5bdef 100644
--- a/tests/test_litellm/proxy/proxy_server/test_routes_config.py
+++ b/tests/test_litellm/proxy/proxy_server/test_routes_config.py
@@ -1 +1,591 @@
-"""Placeholder. Filled by a follow-up PR per the Notion plan."""
+"""Pin tests for proxy_server.py control-plane config routes (PR3).
+
+Routes covered:
+- POST /config/update
+- POST /config/field/update
+- GET /config/field/info
+- GET /config/list
+- POST /config/field/delete
+- POST /config/callback/delete
+- GET /get/config/callbacks
+- GET /config/yaml
+"""
+
+from __future__ import annotations
+
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from .conftest import VOLATILE_KEYS, normalize
+
+
+def _install_litellm_config(mock_prisma: MagicMock) -> MagicMock:
+ """Ensure mock_prisma.db.litellm_config exists with async methods (the
+ conftest only stubs ``litellm_configtable`` — this is a different table)."""
+ table = MagicMock()
+ table.find_unique = AsyncMock(return_value=None)
+ table.find_first = AsyncMock(return_value=None)
+ table.find_many = AsyncMock(return_value=[])
+ table.create = AsyncMock()
+ table.update = AsyncMock()
+ table.upsert = AsyncMock(return_value=None)
+ table.delete = AsyncMock()
+ mock_prisma.db.litellm_config = table
+ return table
+
+
+# ---------------------------------------------------------------------------
+# POST /config/update
+# ---------------------------------------------------------------------------
+
+
+def test_config_update_happy_admin(client, auth_as, mock_prisma, monkeypatch):
+ """POST /config/update with admin role merges + upserts general_settings
+ and returns the canonical success message."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ _install_litellm_config(mock_prisma)
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+ fake_proxy_config = MagicMock()
+ fake_proxy_config.add_deployment = AsyncMock()
+ monkeypatch.setattr(ps, "proxy_config", fake_proxy_config)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.post(
+ "/config/update",
+ json={"general_settings": {"alerting": ["slack"]}},
+ )
+ assert response.status_code == 200
+ assert normalize(response.json()) == {"message": "Config updated successfully"}
+
+
+def test_config_update_non_admin_forbidden(client, auth_as, mock_prisma, monkeypatch):
+ """POST /config/update by a non-admin caller is rejected; the error
+ surfaces as a ProxyException with the admin-only message."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ _install_litellm_config(mock_prisma)
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+
+ with auth_as(LitellmUserRoles.INTERNAL_USER):
+ response = client.post(
+ "/config/update",
+ json={"general_settings": {"alerting": ["slack"]}},
+ )
+ assert response.status_code != 200
+ body = response.json()
+ # ProxyException wraps the 403 detail string in its `message` field.
+ assert "admin" in str(body).lower() or "auth" in str(body).lower()
+
+
+def test_config_update_no_db_error(client, auth_as, monkeypatch):
+ """POST /config/update with prisma_client=None returns a 'No DB Connected'
+ style error (the route raises Exception which the handler maps to 400)."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ monkeypatch.setattr(ps, "prisma_client", None)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.post(
+ "/config/update",
+ json={"general_settings": {"alerting": ["slack"]}},
+ )
+ assert response.status_code != 200
+ assert (
+ "db" in str(response.json()).lower()
+ or "connect" in str(response.json()).lower()
+ )
+
+
+# ---------------------------------------------------------------------------
+# POST /config/field/update
+# ---------------------------------------------------------------------------
+
+
+def test_config_field_update_happy_admin(client, auth_as, mock_prisma, monkeypatch):
+ """POST /config/field/update for a known field upserts the DB row and
+ returns the upsert response (we pin it to a specific shape)."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ table = _install_litellm_config(mock_prisma)
+ table.find_first = AsyncMock(return_value=None)
+ upsert_row = {
+ "param_name": "general_settings",
+ "param_value": {"max_parallel_requests": 5},
+ "id": "row-1",
+ }
+ table.upsert = AsyncMock(return_value=upsert_row)
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.post(
+ "/config/field/update",
+ json={
+ "field_name": "max_parallel_requests",
+ "field_value": 5,
+ "config_type": "general_settings",
+ },
+ )
+ assert response.status_code == 200
+ assert normalize(response.json()) == {
+ "param_name": "general_settings",
+ "param_value": {"max_parallel_requests": 5},
+ "id": "",
+ }
+
+
+def test_config_field_update_non_admin_rejected(
+ client, auth_as, mock_prisma, monkeypatch
+):
+ """Non-admin cannot update config fields — returns 400 with not-allowed
+ detail (handler uses 400 for the auth gate, not 403)."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ _install_litellm_config(mock_prisma)
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+
+ with auth_as(LitellmUserRoles.INTERNAL_USER):
+ response = client.post(
+ "/config/field/update",
+ json={
+ "field_name": "max_parallel_requests",
+ "field_value": 5,
+ "config_type": "general_settings",
+ },
+ )
+ assert response.status_code == 400
+ assert "error" in response.json().get("detail", {})
+
+
+def test_config_field_update_invalid_field(client, auth_as, mock_prisma, monkeypatch):
+ """Unknown field_name is rejected with 400 + 'Invalid field=' detail."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ _install_litellm_config(mock_prisma)
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.post(
+ "/config/field/update",
+ json={
+ "field_name": "not_a_real_field_xyz",
+ "field_value": 1,
+ "config_type": "general_settings",
+ },
+ )
+ assert response.status_code == 400
+ assert "Invalid field" in response.json().get("detail", {}).get("error", "")
+
+
+# ---------------------------------------------------------------------------
+# GET /config/field/info
+# ---------------------------------------------------------------------------
+
+
+def test_config_field_info_happy_admin(client, auth_as, mock_prisma, monkeypatch):
+ """Admin gets back ConfigFieldInfo with the stored value pulled from DB."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ table = _install_litellm_config(mock_prisma)
+ row = MagicMock()
+ row.param_value = {"max_parallel_requests": 7}
+ table.find_first = AsyncMock(return_value=row)
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.get(
+ "/config/field/info", params={"field_name": "max_parallel_requests"}
+ )
+ assert response.status_code == 200
+ assert normalize(response.json()) == {
+ "field_name": "max_parallel_requests",
+ "field_value": 7,
+ }
+
+
+def test_config_field_info_non_admin_rejected(
+ client, auth_as, mock_prisma, monkeypatch
+):
+ """Non-admin (INTERNAL_USER) is denied — admin-view gate fires."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ _install_litellm_config(mock_prisma)
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+
+ with auth_as(LitellmUserRoles.INTERNAL_USER):
+ response = client.get(
+ "/config/field/info", params={"field_name": "max_parallel_requests"}
+ )
+ assert response.status_code == 400
+ assert "error" in response.json().get("detail", {})
+
+
+def test_config_field_info_field_not_in_db(client, auth_as, mock_prisma, monkeypatch):
+ """When the field is missing from the DB row, returns 400 'not in DB'."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ table = _install_litellm_config(mock_prisma)
+ row = MagicMock()
+ row.param_value = {"some_other_field": "value"}
+ table.find_first = AsyncMock(return_value=row)
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.get(
+ "/config/field/info", params={"field_name": "max_parallel_requests"}
+ )
+ assert response.status_code == 400
+ assert "not in DB" in response.json().get("detail", {}).get("error", "")
+
+
+# ---------------------------------------------------------------------------
+# GET /config/list
+# ---------------------------------------------------------------------------
+
+
+def test_config_list_happy_admin(client, auth_as, mock_prisma, monkeypatch):
+ """Admin gets a non-empty list of ConfigList rows for general_settings
+ (one entry per known allowed_arg). Each row has the documented schema."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ table = _install_litellm_config(mock_prisma)
+ row = MagicMock()
+ row.param_value = {"max_parallel_requests": 3}
+ table.find_first = AsyncMock(return_value=row)
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.get(
+ "/config/list", params={"config_type": "general_settings"}
+ )
+ assert response.status_code == 200
+ body = response.json()
+ assert isinstance(body, list)
+ assert len(body) > 0
+ sample = body[0]
+ shape = {
+ "has_field_name": "field_name" in sample,
+ "has_field_type": "field_type" in sample,
+ "has_field_value": "field_value" in sample,
+ "has_stored_in_db": "stored_in_db" in sample,
+ }
+ assert shape == {
+ "has_field_name": True,
+ "has_field_type": True,
+ "has_field_value": True,
+ "has_stored_in_db": True,
+ }
+
+
+def test_config_list_non_admin_rejected(client, auth_as, mock_prisma, monkeypatch):
+ """Non-admin gets a 400 with the role embedded in the error message."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ _install_litellm_config(mock_prisma)
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+
+ with auth_as(LitellmUserRoles.INTERNAL_USER):
+ response = client.get(
+ "/config/list", params={"config_type": "general_settings"}
+ )
+ assert response.status_code == 400
+ assert "role" in response.json().get("detail", {}).get("error", "").lower()
+
+
+def test_config_list_no_db_error(client, auth_as, monkeypatch):
+ """No DB → 400 with db_not_connected error."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ monkeypatch.setattr(ps, "prisma_client", None)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.get(
+ "/config/list", params={"config_type": "general_settings"}
+ )
+ assert response.status_code == 400
+ assert "error" in response.json().get("detail", {})
+
+
+# ---------------------------------------------------------------------------
+# POST /config/field/delete
+# ---------------------------------------------------------------------------
+
+
+def test_config_field_delete_happy_admin(client, auth_as, mock_prisma, monkeypatch):
+ """Admin can delete a stored general_settings field — returns the upsert row."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ table = _install_litellm_config(mock_prisma)
+ existing = MagicMock()
+ existing.param_value = {"max_parallel_requests": 5, "other": "value"}
+ table.find_first = AsyncMock(return_value=existing)
+ table.upsert = AsyncMock(
+ return_value={
+ "param_name": "general_settings",
+ "param_value": {"other": "value"},
+ "id": "row-1",
+ }
+ )
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.post(
+ "/config/field/delete",
+ json={
+ "config_type": "general_settings",
+ "field_name": "max_parallel_requests",
+ },
+ )
+ assert response.status_code == 200
+ assert normalize(response.json()) == {
+ "param_name": "general_settings",
+ "param_value": {"other": "value"},
+ "id": "",
+ }
+
+
+def test_config_field_delete_non_admin_rejected(
+ client, auth_as, mock_prisma, monkeypatch
+):
+ """Non-admin caller hits the 400 not-allowed branch with role in detail."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ _install_litellm_config(mock_prisma)
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+
+ with auth_as(LitellmUserRoles.INTERNAL_USER):
+ response = client.post(
+ "/config/field/delete",
+ json={
+ "config_type": "general_settings",
+ "field_name": "max_parallel_requests",
+ },
+ )
+ assert response.status_code == 400
+ assert "role" in response.json().get("detail", {}).get("error", "").lower()
+
+
+def test_config_field_delete_field_not_in_config(
+ client, auth_as, mock_prisma, monkeypatch
+):
+ """If there is no general_settings row at all, returns 400 'not in config'."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ table = _install_litellm_config(mock_prisma)
+ table.find_first = AsyncMock(return_value=None)
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.post(
+ "/config/field/delete",
+ json={
+ "config_type": "general_settings",
+ "field_name": "max_parallel_requests",
+ },
+ )
+ assert response.status_code == 400
+ assert "not in config" in response.json().get("detail", {}).get("error", "")
+
+
+# ---------------------------------------------------------------------------
+# POST /config/callback/delete
+# ---------------------------------------------------------------------------
+
+
+def test_config_callback_delete_happy_admin(client, auth_as, mock_prisma, monkeypatch):
+ """Admin deletes a configured success callback — handler returns the
+ success message + remaining callbacks + a timestamp."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ _install_litellm_config(mock_prisma)
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+ monkeypatch.setattr(ps, "store_model_in_db", True)
+
+ fake_proxy_config = MagicMock()
+ fake_proxy_config.get_config = AsyncMock(
+ return_value={"litellm_settings": {"success_callback": ["langfuse", "slack"]}}
+ )
+ fake_proxy_config.save_config = AsyncMock()
+ fake_proxy_config.add_deployment = AsyncMock()
+ monkeypatch.setattr(ps, "proxy_config", fake_proxy_config)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.post(
+ "/config/callback/delete", json={"callback_name": "langfuse"}
+ )
+ assert response.status_code == 200
+ # `deleted_at` is an ISO timestamp generated at request time — extend
+ # the volatile set just for this assertion so dict-equality still works.
+ volatile = VOLATILE_KEYS | {"deleted_at"}
+ assert normalize(response.json(), volatile) == {
+ "message": "Successfully deleted callback: langfuse",
+ "removed_callback": "langfuse",
+ "remaining_callbacks": ["slack"],
+ "deleted_at": "",
+ }
+
+
+def test_config_callback_delete_non_admin_rejected(
+ client, auth_as, mock_prisma, monkeypatch
+):
+ """Non-admin caller is rejected with 400 not-allowed."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ _install_litellm_config(mock_prisma)
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+ monkeypatch.setattr(ps, "store_model_in_db", True)
+
+ with auth_as(LitellmUserRoles.INTERNAL_USER):
+ response = client.post(
+ "/config/callback/delete", json={"callback_name": "langfuse"}
+ )
+ assert response.status_code == 400
+ assert "role" in response.json().get("detail", {}).get("error", "").lower()
+
+
+def test_config_callback_delete_not_found(client, auth_as, mock_prisma, monkeypatch):
+ """Callback missing from current config returns 404."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ _install_litellm_config(mock_prisma)
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+ monkeypatch.setattr(ps, "store_model_in_db", True)
+
+ fake_proxy_config = MagicMock()
+ fake_proxy_config.get_config = AsyncMock(
+ return_value={"litellm_settings": {"success_callback": ["slack"]}}
+ )
+ monkeypatch.setattr(ps, "proxy_config", fake_proxy_config)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.post(
+ "/config/callback/delete", json={"callback_name": "langfuse"}
+ )
+ # The handler re-raises HTTPException(404) verbatim (only generic
+ # `Exception` becomes a 500 ProxyException), so pin 404 strictly.
+ assert response.status_code == 404
+ assert (
+ "langfuse" in str(response.json()).lower()
+ or "not found" in str(response.json()).lower()
+ )
+
+
+# ---------------------------------------------------------------------------
+# GET /get/config/callbacks
+# ---------------------------------------------------------------------------
+
+
+def test_get_config_callbacks_happy(client, auth_as, mock_prisma, monkeypatch):
+ """GET /get/config/callbacks returns the 5 pinned top-level keys:
+ status, callbacks, alerts, router_settings, available_callbacks."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ _install_litellm_config(mock_prisma)
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+ monkeypatch.setattr(ps, "llm_router", None)
+
+ fake_proxy_config = MagicMock()
+ fake_proxy_config.get_config = AsyncMock(
+ return_value={
+ "litellm_settings": {"success_callback": []},
+ "general_settings": {},
+ "environment_variables": {},
+ }
+ )
+ monkeypatch.setattr(ps, "proxy_config", fake_proxy_config)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.get("/get/config/callbacks")
+ assert response.status_code == 200
+ body = response.json()
+ shape = {
+ "status": body.get("status"),
+ "has_callbacks": "callbacks" in body,
+ "has_alerts": "alerts" in body,
+ "has_router_settings": "router_settings" in body,
+ "has_available_callbacks": "available_callbacks" in body,
+ }
+ assert shape == {
+ "status": "success",
+ "has_callbacks": True,
+ "has_alerts": True,
+ "has_router_settings": True,
+ "has_available_callbacks": True,
+ }
+
+
+def test_get_config_callbacks_internal_error(client, auth_as, mock_prisma, monkeypatch):
+ """If proxy_config.get_config() raises, the handler wraps the failure in
+ a ProxyException → non-2xx response with an error body."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ _install_litellm_config(mock_prisma)
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+
+ fake_proxy_config = MagicMock()
+ fake_proxy_config.get_config = AsyncMock(side_effect=RuntimeError("boom"))
+ monkeypatch.setattr(ps, "proxy_config", fake_proxy_config)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.get("/get/config/callbacks")
+ assert response.status_code >= 400
+ assert (
+ "boom" in str(response.json()).lower()
+ or "error" in str(response.json()).lower()
+ )
+
+
+# ---------------------------------------------------------------------------
+# GET /config/yaml
+# ---------------------------------------------------------------------------
+
+
+def test_config_yaml_returns_demo_payload(client, auth_as):
+ """GET /config/yaml is documented as a mock endpoint. It declares
+ ConfigYAML as the body parameter, so a GET with an empty JSON body is
+ accepted and returns the canonical demo dict."""
+ with auth_as():
+ response = client.request("GET", "/config/yaml", json={})
+ shape = {
+ "status": response.status_code,
+ "media_type_yaml": response.headers.get("content-type", "").startswith(
+ "application/json"
+ ),
+ "has_body": len(response.content) > 0,
+ }
+ assert shape == {
+ "status": 200,
+ "media_type_yaml": True,
+ "has_body": True,
+ }
+ assert response.json() == {"hello": "world"}
+
+
+def test_config_yaml_invalid_method(client):
+ """POST against the GET-only /config/yaml is rejected (error path)."""
+ response = client.post("/config/yaml", json={})
+ assert response.status_code == 405
+ # Method-not-allowed responses still return a JSON-ish body via the
+ # FastAPI default handler — assert the body is not the success payload.
+ assert response.content != b'{"hello":"world"}'
diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_embeddings.py b/tests/test_litellm/proxy/proxy_server/test_routes_embeddings.py
index ad6b4016461..98249cb5ad5 100644
--- a/tests/test_litellm/proxy/proxy_server/test_routes_embeddings.py
+++ b/tests/test_litellm/proxy/proxy_server/test_routes_embeddings.py
@@ -1 +1,121 @@
-"""Placeholder. Filled by a follow-up PR per the Notion plan."""
+"""Behavior pins for ``proxy_server.py`` embeddings routes.
+
+Pins (PR2):
+ - POST /v1/embeddings
+ - POST /embeddings
+ - POST /engines/{model:path}/embeddings
+ - POST /openai/deployments/{model:path}/embeddings
+"""
+
+from __future__ import annotations
+
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from litellm.proxy import common_request_processing, proxy_server
+
+from .conftest import normalize # type: ignore[import-not-found]
+
+HAPPY_RESPONSE = {
+ "object": "list",
+ "model": "text-embedding-ada-002",
+ "data": [{"embedding": [0.0, 0.1, 0.2], "index": 0, "object": "embedding"}],
+ "usage": {"prompt_tokens": 1, "total_tokens": 1},
+}
+
+
+@pytest.fixture
+def patched_embedding(monkeypatch):
+ router = MagicMock()
+ router.model_names = ["text-embedding-ada-002"]
+ router.get_deployment_by_model_group_name = MagicMock(return_value=None)
+ monkeypatch.setattr(proxy_server, "llm_router", router)
+ monkeypatch.setattr(
+ proxy_server, "proxy_logging_obj", MagicMock(post_call_failure_hook=AsyncMock())
+ )
+
+ async def _fake_process(self, *args, **kwargs):
+ return dict(HAPPY_RESPONSE)
+
+ monkeypatch.setattr(
+ common_request_processing.ProxyBaseLLMRequestProcessing,
+ "base_process_llm_request",
+ _fake_process,
+ )
+ yield
+
+
+@pytest.fixture
+def embedding_pipeline_raises(monkeypatch):
+ router = MagicMock()
+ router.model_names = []
+ monkeypatch.setattr(proxy_server, "llm_router", router)
+ monkeypatch.setattr(
+ proxy_server, "proxy_logging_obj", MagicMock(post_call_failure_hook=AsyncMock())
+ )
+
+ from litellm.proxy._types import ProxyException
+
+ async def _raise(self, *args, **kwargs):
+ raise ValueError("boom")
+
+ async def _handler(self, *, e, user_api_key_dict, proxy_logging_obj, version=None):
+ return ProxyException(
+ message="boom", type="bad_request_error", param="model", code=400
+ )
+
+ monkeypatch.setattr(
+ common_request_processing.ProxyBaseLLMRequestProcessing,
+ "base_process_llm_request",
+ _raise,
+ )
+ monkeypatch.setattr(
+ common_request_processing.ProxyBaseLLMRequestProcessing,
+ "_handle_llm_api_exception",
+ _handler,
+ )
+ yield
+
+
+_EMBED_PATHS = [
+ "/v1/embeddings",
+ "/embeddings",
+ "/engines/text-embedding-ada-002/embeddings",
+ "/openai/deployments/text-embedding-ada-002/embeddings",
+]
+
+
+@pytest.mark.parametrize("path", _EMBED_PATHS)
+def test_embeddings_happy_path(client, auth_as, patched_embedding, path):
+ """Pins all four ``POST .../embeddings`` aliases (happy path).
+
+ Covers ``POST /v1/embeddings``, ``POST /embeddings``,
+ ``POST /engines/{model:path}/embeddings``, and
+ ``POST /openai/deployments/{model:path}/embeddings``.
+ """
+ payload = {"model": "text-embedding-ada-002", "input": "hello"}
+ with auth_as():
+ response = client.post(path, json=payload)
+ assert response.status_code == 200
+ assert normalize(response.json()) == {
+ "object": "list",
+ "model": "text-embedding-ada-002",
+ "data": [{"embedding": [0.0, 0.1, 0.2], "index": 0, "object": "embedding"}],
+ "usage": {"prompt_tokens": 1, "total_tokens": 1},
+ }
+
+
+@pytest.mark.parametrize("path", _EMBED_PATHS)
+def test_embeddings_pipeline_error(client, auth_as, embedding_pipeline_raises, path):
+ """Pins all four ``POST .../embeddings`` aliases (error path).
+
+ Covers ``POST /v1/embeddings``, ``POST /embeddings``,
+ ``POST /engines/{model:path}/embeddings``, and
+ ``POST /openai/deployments/{model:path}/embeddings``.
+ """
+ payload = {"model": "text-embedding-ada-002", "input": "boom"}
+ with auth_as():
+ response = client.post(path, json=payload)
+ assert response.status_code == 400
+ assert response.content # non-empty error body
diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_invitation.py b/tests/test_litellm/proxy/proxy_server/test_routes_invitation.py
index ad6b4016461..5b54a63d8a2 100644
--- a/tests/test_litellm/proxy/proxy_server/test_routes_invitation.py
+++ b/tests/test_litellm/proxy/proxy_server/test_routes_invitation.py
@@ -1 +1,387 @@
-"""Placeholder. Filled by a follow-up PR per the Notion plan."""
+"""Pin tests for proxy_server.py invitation routes (PR3).
+
+Routes covered:
+- POST /invitation/new
+- GET /invitation/info
+- POST /invitation/update
+- POST /invitation/delete
+"""
+
+from __future__ import annotations
+
+from datetime import datetime, timedelta, timezone
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from .conftest import VOLATILE_KEYS, normalize
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_invitation(
+ invitation_id: str = "inv-abc",
+ user_id: str = "user-target",
+ created_by: str = "test-user-id",
+ is_accepted: bool = False,
+ accepted_at=None,
+):
+ """Build an invitation row with the fields ``InvitationModel`` requires.
+
+ FastAPI serializes the returned object against ``response_model=InvitationModel``,
+ so the object must expose ``id, user_id, is_accepted, accepted_at, expires_at,
+ created_at, created_by, updated_at, updated_by`` either as attributes or
+ dict keys.
+ """
+ now = datetime.now(timezone.utc)
+ return SimpleNamespace(
+ id=invitation_id,
+ user_id=user_id,
+ is_accepted=is_accepted,
+ accepted_at=accepted_at,
+ expires_at=now + timedelta(days=7),
+ created_at=now,
+ created_by=created_by,
+ updated_at=now,
+ updated_by=created_by,
+ )
+
+
+# ---------------------------------------------------------------------------
+# POST /invitation/new
+# ---------------------------------------------------------------------------
+
+
+def test_invitation_new_admin_happy(client, auth_as, monkeypatch, mock_prisma):
+ """Proxy admin → create_invitation_for_user returns invitation → 200."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+ from litellm.proxy.management_helpers import user_invitation
+
+ invitation = _make_invitation(user_id="user-target")
+
+ async def _fake_create_invitation(data, user_api_key_dict):
+ return invitation
+
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+ monkeypatch.setattr(
+ user_invitation, "create_invitation_for_user", _fake_create_invitation
+ )
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.post("/invitation/new", json={"user_id": "user-target"})
+
+ assert response.status_code == 200
+ assert normalize(response.json()) == {
+ "id": "",
+ "user_id": "user-target",
+ "is_accepted": False,
+ "accepted_at": None,
+ "expires_at": "",
+ "created_at": "",
+ "created_by": "test-user-id",
+ "updated_at": "",
+ "updated_by": "test-user-id",
+ }
+
+
+def test_invitation_new_non_admin_forbidden(client, auth_as, monkeypatch, mock_prisma):
+ """Internal user without team/org admin privileges → 400 not-allowed."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+ from litellm.proxy.management_endpoints import common_utils
+
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+
+ async def _no_privileges(**kwargs):
+ return False
+
+ # Patch at the proxy_server import site (used by the route).
+ monkeypatch.setattr(ps, "_user_has_admin_privileges", _no_privileges)
+ monkeypatch.setattr(common_utils, "_user_has_admin_privileges", _no_privileges)
+
+ with auth_as(LitellmUserRoles.INTERNAL_USER):
+ response = client.post("/invitation/new", json={"user_id": "user-target"})
+
+ assert response.status_code == 400
+ err = response.json().get("error", response.json())
+ err_text = str(err)
+ assert "role=" in err_text or "not allowed" in err_text.lower()
+
+
+def test_invitation_new_db_not_connected_400(client, auth_as, monkeypatch):
+ """prisma_client is None → 400 db_not_connected_error."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ monkeypatch.setattr(ps, "prisma_client", None)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.post("/invitation/new", json={"user_id": "user-target"})
+
+ assert response.status_code == 400
+ body = response.json()
+ err_text = str(body)
+ # The handler wraps via handle_exception_on_proxy, so the error body
+ # may take either the {"error": {...}} or {"detail": {...}} shape.
+ assert "No connected db" in err_text or "db" in err_text.lower()
+
+
+def test_invitation_new_missing_user_id_422(client, auth_as, monkeypatch, mock_prisma):
+ """Body missing the required ``user_id`` field → FastAPI 422."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.post("/invitation/new", json={})
+
+ assert response.status_code == 422
+ body = response.json()
+ assert isinstance(body.get("detail"), list)
+ assert any("user_id" in str(item) for item in body["detail"])
+
+
+# ---------------------------------------------------------------------------
+# GET /invitation/info
+# ---------------------------------------------------------------------------
+
+
+def test_invitation_info_admin_happy(client, auth_as, monkeypatch, mock_prisma):
+ """Admin requesting an existing invitation id → returns the invitation."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ invitation = _make_invitation(invitation_id="inv-xyz", user_id="user-target")
+ mock_prisma.db.litellm_invitationlink.find_unique.return_value = invitation
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.get("/invitation/info", params={"invitation_id": "inv-xyz"})
+
+ assert response.status_code == 200
+ assert normalize(response.json()) == {
+ "id": "",
+ "user_id": "user-target",
+ "is_accepted": False,
+ "accepted_at": None,
+ "expires_at": "",
+ "created_at": "",
+ "created_by": "test-user-id",
+ "updated_at": "",
+ "updated_by": "test-user-id",
+ }
+
+
+def test_invitation_info_not_admin_forbidden(client, auth_as, monkeypatch, mock_prisma):
+ """Non-admin viewer (no admin-view privileges) → 400 not-allowed."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+
+ # _user_has_admin_view is referenced from proxy_server's import.
+ monkeypatch.setattr(ps, "_user_has_admin_view", lambda u: False)
+
+ with auth_as(LitellmUserRoles.INTERNAL_USER):
+ response = client.get("/invitation/info", params={"invitation_id": "inv-xyz"})
+
+ assert response.status_code == 400
+ err_text = str(response.json())
+ assert "role=" in err_text or "not allowed" in err_text.lower()
+
+
+def test_invitation_info_not_found_400(client, auth_as, monkeypatch, mock_prisma):
+ """Admin requesting an unknown invitation id → 400 does-not-exist."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ mock_prisma.db.litellm_invitationlink.find_unique.return_value = None
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.get(
+ "/invitation/info", params={"invitation_id": "does-not-exist"}
+ )
+
+ assert response.status_code == 400
+ assert response.json() == {
+ "detail": {"error": "Invitation id does not exist in the database."}
+ }
+
+
+# ---------------------------------------------------------------------------
+# POST /invitation/update
+# ---------------------------------------------------------------------------
+
+
+def test_invitation_update_happy(client, auth_as, monkeypatch, mock_prisma):
+ """Authenticated user → invitation marked accepted → returns updated row."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ accepted = _make_invitation(
+ invitation_id="inv-1",
+ user_id="user-target",
+ is_accepted=True,
+ accepted_at=datetime.now(timezone.utc),
+ )
+ mock_prisma.db.litellm_invitationlink.update.return_value = accepted
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.post(
+ "/invitation/update",
+ json={"invitation_id": "inv-1", "is_accepted": True},
+ )
+
+ assert response.status_code == 200
+ # ``accepted_at`` is a fresh timestamp on each run — extend volatile set.
+ extended = VOLATILE_KEYS | {"accepted_at"}
+ assert normalize(response.json(), extended) == {
+ "id": "",
+ "user_id": "user-target",
+ "is_accepted": True,
+ "accepted_at": "",
+ "expires_at": "",
+ "created_at": "",
+ "created_by": "test-user-id",
+ "updated_at": "",
+ "updated_by": "test-user-id",
+ }
+
+
+def test_invitation_update_unknown_id_400(client, auth_as, monkeypatch, mock_prisma):
+ """Update against an invitation id the DB returns None for → 400."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ mock_prisma.db.litellm_invitationlink.update.return_value = None
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.post(
+ "/invitation/update",
+ json={"invitation_id": "ghost", "is_accepted": True},
+ )
+
+ assert response.status_code == 400
+ assert response.json() == {
+ "detail": {"error": "Invitation id does not exist in the database."}
+ }
+
+
+def test_invitation_update_no_user_id_500(client, auth_as, monkeypatch, mock_prisma):
+ """If the auth principal lacks a user_id, handler returns 500."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN, user_id=None):
+ response = client.post(
+ "/invitation/update",
+ json={"invitation_id": "inv-1", "is_accepted": True},
+ )
+
+ assert response.status_code == 500
+ err_text = str(response.json())
+ assert "Unable to identify user id" in err_text
+
+
+# ---------------------------------------------------------------------------
+# POST /invitation/delete
+# ---------------------------------------------------------------------------
+
+
+def test_invitation_delete_admin_happy(client, auth_as, monkeypatch, mock_prisma):
+ """Proxy admin deletes by invitation_id → 200 with deleted row."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ deleted = _make_invitation(invitation_id="inv-del", user_id="user-target")
+ mock_prisma.db.litellm_invitationlink.delete.return_value = deleted
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.post(
+ "/invitation/delete", json={"invitation_id": "inv-del"}
+ )
+
+ assert response.status_code == 200
+ assert normalize(response.json()) == {
+ "id": "",
+ "user_id": "user-target",
+ "is_accepted": False,
+ "accepted_at": None,
+ "expires_at": "",
+ "created_at": "",
+ "created_by": "test-user-id",
+ "updated_at": "",
+ "updated_by": "test-user-id",
+ }
+
+
+def test_invitation_delete_non_admin_forbidden(
+ client, auth_as, monkeypatch, mock_prisma
+):
+ """Non-admin user without elevated privileges → 400 not-allowed."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+
+ async def _no_privileges(**kwargs):
+ return False
+
+ monkeypatch.setattr(ps, "_user_has_admin_privileges", _no_privileges)
+
+ with auth_as(LitellmUserRoles.INTERNAL_USER):
+ response = client.post(
+ "/invitation/delete", json={"invitation_id": "inv-del"}
+ )
+
+ assert response.status_code == 400
+ err_text = str(response.json())
+ assert "role=" in err_text or "not allowed" in err_text.lower()
+
+
+def test_invitation_delete_unknown_id_400(client, auth_as, monkeypatch, mock_prisma):
+ """Delete returns None (no row) → 400 does-not-exist."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ mock_prisma.db.litellm_invitationlink.delete.return_value = None
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.post(
+ "/invitation/delete", json={"invitation_id": "ghost"}
+ )
+
+ assert response.status_code == 400
+ assert response.json() == {
+ "detail": {"error": "Invitation id does not exist in the database."}
+ }
+
+
+def test_invitation_delete_db_not_connected_400(client, auth_as, monkeypatch):
+ """prisma_client is None → 400 db_not_connected_error."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ monkeypatch.setattr(ps, "prisma_client", None)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.post(
+ "/invitation/delete", json={"invitation_id": "inv-del"}
+ )
+
+ assert response.status_code == 400
+ err_text = str(response.json())
+ assert "No connected db" in err_text or "db" in err_text.lower()
diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_login_sso.py b/tests/test_litellm/proxy/proxy_server/test_routes_login_sso.py
index ad6b4016461..6af1d6653e1 100644
--- a/tests/test_litellm/proxy/proxy_server/test_routes_login_sso.py
+++ b/tests/test_litellm/proxy/proxy_server/test_routes_login_sso.py
@@ -1 +1,387 @@
-"""Placeholder. Filled by a follow-up PR per the Notion plan."""
+"""Pin tests for proxy_server.py login/SSO routes (PR3).
+
+Routes covered:
+- GET /fallback/login
+- POST /login
+- POST /v2/login
+- POST /v3/login
+- POST /v3/login/exchange
+"""
+
+from __future__ import annotations
+
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from .conftest import normalize
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _install_login_mocks(monkeypatch, raise_on_auth: bool = False) -> None:
+ """Patch authenticate_user + create_ui_token_object at their import paths.
+
+ Both /login, /v2/login and /v3/login do a *local* (in-function) import of
+ these helpers, so we patch the module they live in.
+ """
+ from litellm.proxy import proxy_server as ps
+
+ async def _fake_auth(username, password, master_key, prisma_client):
+ if raise_on_auth:
+ raise Exception("boom-auth-failure")
+ fake = MagicMock()
+ fake.user_id = "u-1"
+ fake.user_email = "test@example.invalid"
+ fake.user_role = "proxy_admin"
+ fake.key = "sk-fake-ui-key"
+ return fake
+
+ def _fake_token_object(login_result, general_settings, premium_user):
+ return {
+ "user_id": "u-1",
+ "user_email": "test@example.invalid",
+ "user_role": "proxy_admin",
+ "premium_user": premium_user,
+ "key": "sk-fake-ui-key",
+ }
+
+ monkeypatch.setattr(
+ "litellm.proxy.auth.login_utils.authenticate_user", _fake_auth
+ )
+ monkeypatch.setattr(
+ "litellm.proxy.auth.login_utils.create_ui_token_object", _fake_token_object
+ )
+ monkeypatch.setattr(ps, "master_key", "sk-test-master")
+ monkeypatch.setattr(ps, "general_settings", {})
+ monkeypatch.setattr(ps, "premium_user", False)
+
+
+# ---------------------------------------------------------------------------
+# GET /fallback/login
+# ---------------------------------------------------------------------------
+
+
+def test_fallback_login_returns_html_form(client, monkeypatch):
+ """Pin: GET /fallback/login returns an HTML login form with status 200."""
+ monkeypatch.delenv("UI_USERNAME", raising=False)
+ response = client.get("/fallback/login")
+ body_lower = response.text.lower()
+ shape = {
+ "status": response.status_code,
+ "content_type_html": response.headers.get("content-type", "").startswith(
+ "text/html"
+ ),
+ "has_form": " TestClient returns 500 with body
+ assert response.status_code == 500
+ # Body must be non-empty so a future refactor that drops the error body
+ # would trip this gate.
+ assert len(response.content) > 0
+ assert response.headers.get("content-type") is not None
+
+
+# ---------------------------------------------------------------------------
+# POST /v2/login
+# ---------------------------------------------------------------------------
+
+
+def test_v2_login_success_returns_token_and_redirect(client, monkeypatch):
+ """Pin: POST /v2/login returns JSON {redirect_url, token} + sets token cookie."""
+ _install_login_mocks(monkeypatch)
+ response = client.post(
+ "/v2/login",
+ json={"username": "admin", "password": "password"},
+ )
+ assert response.status_code == 200
+ assert normalize(
+ response.json(), volatile=frozenset({"token", "redirect_url"})
+ ) == {"redirect_url": "", "token": ""}
+ body = response.json()
+ set_cookie = response.headers.get("set-cookie", "")
+ shape = {
+ "redirect_url_has_ui": "/ui/" in body.get("redirect_url", ""),
+ "redirect_url_has_login_success": "login=success"
+ in body.get("redirect_url", ""),
+ "token_in_body": bool(body.get("token")),
+ "token_cookie_set": "token=" in set_cookie,
+ }
+ assert shape == {
+ "redirect_url_has_ui": True,
+ "redirect_url_has_login_success": True,
+ "token_in_body": True,
+ "token_cookie_set": True,
+ }
+
+
+def test_v2_login_authenticate_failure_500(client, monkeypatch):
+ """Error path: authenticate_user raising -> ProxyException -> 500 with structured error."""
+ _install_login_mocks(monkeypatch, raise_on_auth=True)
+ response = client.post(
+ "/v2/login",
+ json={"username": "admin", "password": "wrong"},
+ )
+ assert response.status_code == 500
+ body = response.json()
+ # Non-status assertion: response shape should carry an error
+ assert "error" in body or "detail" in body
+ assert isinstance(body, dict)
+
+
+# ---------------------------------------------------------------------------
+# POST /v3/login
+# ---------------------------------------------------------------------------
+
+
+def test_v3_login_without_control_plane_url_404(client, monkeypatch):
+ """Pin: /v3/login is gated on general_settings['control_plane_url'] — 404 when absent."""
+ _install_login_mocks(monkeypatch)
+ # _install_login_mocks sets general_settings to {} — re-affirm
+ from litellm.proxy import proxy_server as ps
+
+ monkeypatch.setattr(ps, "general_settings", {})
+
+ response = client.post(
+ "/v3/login",
+ json={"username": "admin", "password": "password"},
+ )
+ assert response.status_code == 404
+ body = response.json()
+ # Detail carries the structured ProxyException error
+ detail = body.get("detail", {})
+ if isinstance(detail, dict):
+ message = detail.get("error", {})
+ if isinstance(message, dict):
+ message_str = message.get("message", "")
+ else:
+ message_str = str(message)
+ else:
+ message_str = str(detail)
+ assert "control_plane_url" in str(body)
+
+
+def test_v3_login_success_returns_code(client, monkeypatch):
+ """Pin: /v3/login with control_plane_url returns {code, expires_in}."""
+ from litellm.proxy import proxy_server as ps
+
+ _install_login_mocks(monkeypatch)
+ monkeypatch.setattr(
+ ps, "general_settings", {"control_plane_url": "https://cp.example.invalid"}
+ )
+ # Force the local (non-redis) cache path
+ monkeypatch.setattr(ps, "redis_usage_cache", None)
+ fake_cache = MagicMock()
+ fake_cache.async_set_cache = AsyncMock()
+ monkeypatch.setattr(ps, "user_api_key_cache", fake_cache)
+
+ response = client.post(
+ "/v3/login",
+ json={"username": "admin", "password": "password"},
+ )
+ assert response.status_code == 200
+ body = response.json()
+ # Strong assertion via normalize with extended volatile set ("code" is volatile)
+ assert normalize(
+ body, volatile=frozenset({"code", "expires_in"})
+ ) == {"code": "", "expires_in": ""}
+ shape = {
+ "has_code": isinstance(body.get("code"), str) and len(body["code"]) > 0,
+ "expires_in_60": body.get("expires_in") == 60,
+ "cache_set_called": fake_cache.async_set_cache.await_count == 1,
+ }
+ assert shape == {
+ "has_code": True,
+ "expires_in_60": True,
+ "cache_set_called": True,
+ }
+
+
+def test_v3_login_authenticate_failure_500(client, monkeypatch):
+ """Error path: with control_plane_url set, authenticate_user raises -> 500."""
+ from litellm.proxy import proxy_server as ps
+
+ _install_login_mocks(monkeypatch, raise_on_auth=True)
+ monkeypatch.setattr(
+ ps, "general_settings", {"control_plane_url": "https://cp.example.invalid"}
+ )
+
+ response = client.post(
+ "/v3/login",
+ json={"username": "admin", "password": "wrong"},
+ )
+ assert response.status_code == 500
+ body = response.json()
+ assert isinstance(body, dict)
+ assert "error" in body or "detail" in body
+
+
+# ---------------------------------------------------------------------------
+# POST /v3/login/exchange
+# ---------------------------------------------------------------------------
+
+
+def test_v3_login_exchange_without_control_plane_url_404(client, monkeypatch):
+ """Pin: /v3/login/exchange gated on control_plane_url — 404 when absent."""
+ from litellm.proxy import proxy_server as ps
+
+ monkeypatch.setattr(ps, "general_settings", {})
+
+ response = client.post("/v3/login/exchange", json={"code": "abc"})
+ assert response.status_code == 404
+ body = response.json()
+ assert "control_plane_url" in str(body)
+ assert isinstance(body, dict)
+
+
+def test_v3_login_exchange_missing_code_400(client, monkeypatch):
+ """Error path: missing 'code' in body -> 400 with 'Missing' message."""
+ from litellm.proxy import proxy_server as ps
+
+ monkeypatch.setattr(
+ ps, "general_settings", {"control_plane_url": "https://cp.example.invalid"}
+ )
+
+ response = client.post("/v3/login/exchange", json={})
+ assert response.status_code == 400
+ body = response.json()
+ assert isinstance(body, dict)
+ assert "Missing" in str(body) or "code" in str(body)
+
+
+def test_v3_login_exchange_invalid_code_401(client, monkeypatch):
+ """Error path: code that isn't in cache -> 401 'Invalid or expired'."""
+ from litellm.proxy import proxy_server as ps
+
+ monkeypatch.setattr(
+ ps, "general_settings", {"control_plane_url": "https://cp.example.invalid"}
+ )
+ monkeypatch.setattr(ps, "redis_usage_cache", None)
+ fake_cache = MagicMock()
+ fake_cache.async_get_cache = AsyncMock(return_value=None)
+ fake_cache.async_delete_cache = AsyncMock()
+ monkeypatch.setattr(ps, "user_api_key_cache", fake_cache)
+
+ response = client.post("/v3/login/exchange", json={"code": "nope"})
+ assert response.status_code == 401
+ body = response.json()
+ assert isinstance(body, dict)
+ assert "Invalid" in str(body) or "expired" in str(body)
+
+
+def test_v3_login_exchange_success_returns_token_and_redirect(client, monkeypatch):
+ """Pin: valid code -> JSON {token, redirect_url} + token cookie + cache deleted (single-use)."""
+ from litellm.proxy import proxy_server as ps
+
+ monkeypatch.setattr(
+ ps, "general_settings", {"control_plane_url": "https://cp.example.invalid"}
+ )
+ monkeypatch.setattr(ps, "redis_usage_cache", None)
+
+ cached_payload = {
+ "token": "jwt-token-xyz",
+ "redirect_url": "https://litellm.example.invalid/ui/?login=success",
+ }
+ fake_cache = MagicMock()
+ fake_cache.async_get_cache = AsyncMock(return_value=cached_payload)
+ fake_cache.async_delete_cache = AsyncMock()
+ monkeypatch.setattr(ps, "user_api_key_cache", fake_cache)
+
+ response = client.post("/v3/login/exchange", json={"code": "valid-code"})
+ assert response.status_code == 200
+ assert normalize(
+ response.json(), volatile=frozenset({"token", "redirect_url"})
+ ) == {"token": "", "redirect_url": ""}
+ body = response.json()
+ set_cookie = response.headers.get("set-cookie", "")
+ shape = {
+ "token": body.get("token"),
+ "redirect_url": body.get("redirect_url"),
+ "token_cookie_set": "token=" in set_cookie,
+ "cache_deleted_once": fake_cache.async_delete_cache.await_count == 1,
+ }
+ assert shape == {
+ "token": "jwt-token-xyz",
+ "redirect_url": "https://litellm.example.invalid/ui/?login=success",
+ "token_cookie_set": True,
+ "cache_deleted_once": True,
+ }
diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_misc.py b/tests/test_litellm/proxy/proxy_server/test_routes_misc.py
index ad6b4016461..0c45e31afd2 100644
--- a/tests/test_litellm/proxy/proxy_server/test_routes_misc.py
+++ b/tests/test_litellm/proxy/proxy_server/test_routes_misc.py
@@ -1 +1,230 @@
-"""Placeholder. Filled by a follow-up PR per the Notion plan."""
+"""Pin tests for proxy_server.py misc routes (PR3).
+
+Routes covered:
+- GET /
+- GET /routes
+- GET /adaptive_router/state
+- GET /get_logo_url
+- GET /get_image
+- GET /get_favicon
+"""
+
+from __future__ import annotations
+
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from .conftest import normalize
+
+
+# ---------------------------------------------------------------------------
+# GET /
+# ---------------------------------------------------------------------------
+
+
+def test_home_returns_200_with_body(client, auth_as):
+ """GET / serves either the home string or the Swagger UI fallback —
+ both return 200 with a non-empty body. This pins the contract: root
+ always answers and never errors."""
+ with auth_as():
+ response = client.get("/")
+ shape = {
+ "status": response.status_code,
+ "has_body": len(response.content) > 0,
+ "has_content_type": bool(response.headers.get("content-type")),
+ }
+ assert shape == {"status": 200, "has_body": True, "has_content_type": True}
+
+
+def test_home_invalid_method_405(client):
+ """GET / handler is GET-only; DELETE returns 405 (error path)."""
+ response = client.delete("/")
+ assert response.status_code == 405
+ assert len(response.content) > 0 and response.headers.get("content-type")
+
+
+# ---------------------------------------------------------------------------
+# GET /routes
+# ---------------------------------------------------------------------------
+
+
+def test_get_routes_returns_routes_list(client, auth_as):
+ with auth_as():
+ response = client.get("/routes")
+ assert response.status_code == 200
+ body = response.json()
+ assert isinstance(body, dict)
+ assert "routes" in body
+ assert isinstance(body["routes"], list)
+ assert len(body["routes"]) > 0
+ sample = body["routes"][0]
+ shape = {
+ "has_path": "path" in sample,
+ "has_methods": "methods" in sample,
+ "has_endpoint": "endpoint" in sample,
+ }
+ assert shape == {
+ "has_path": True,
+ "has_methods": True,
+ "has_endpoint": True,
+ }
+
+
+def test_get_routes_invalid_method_405(client):
+ """POST against the GET-only /routes endpoint is rejected (error path)."""
+ response = client.post("/routes")
+ assert response.status_code == 405
+ body = response.json() if response.headers.get("content-type", "").startswith(
+ "application/json"
+ ) else {}
+ assert isinstance(body, dict)
+
+
+# ---------------------------------------------------------------------------
+# GET /adaptive_router/state
+# ---------------------------------------------------------------------------
+
+
+def test_adaptive_router_state_returns_snapshots(client, auth_as, monkeypatch):
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ fake_router = MagicMock()
+ snap = {"router_name": "ar-1", "queue_depth": 0, "posteriors": []}
+ bandit = MagicMock()
+ bandit.get_state_snapshot = AsyncMock(return_value=snap)
+ fake_router.adaptive_routers = {"ar-1": bandit}
+ monkeypatch.setattr(ps, "llm_router", fake_router)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.get("/adaptive_router/state")
+ assert response.status_code == 200
+ assert normalize(response.json()) == {
+ "routers": [
+ {"router_name": "ar-1", "queue_depth": 0, "posteriors": []},
+ ]
+ }
+
+
+def test_adaptive_router_state_not_admin_forbidden(client, auth_as):
+ from litellm.proxy._types import LitellmUserRoles
+
+ with auth_as(LitellmUserRoles.INTERNAL_USER):
+ response = client.get("/adaptive_router/state")
+ assert response.status_code == 403
+ assert "error" in response.json().get("detail", {})
+
+
+def test_adaptive_router_state_not_configured_404(client, auth_as, monkeypatch):
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ fake_router = MagicMock()
+ fake_router.adaptive_routers = {}
+ monkeypatch.setattr(ps, "llm_router", fake_router)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.get("/adaptive_router/state")
+ assert response.status_code == 404
+ assert "adaptive_router" in response.json().get("detail", {}).get("error", "")
+
+
+# ---------------------------------------------------------------------------
+# GET /get_logo_url
+# ---------------------------------------------------------------------------
+
+
+def test_get_logo_url_returns_http_url_when_set(client, monkeypatch):
+ monkeypatch.setenv("UI_LOGO_PATH", "https://example.invalid/logo.png")
+ response = client.get("/get_logo_url")
+ assert response.status_code == 200
+ assert normalize(response.json()) == {"logo_url": "https://example.invalid/logo.png"}
+
+
+def test_get_logo_url_blank_when_local_path(client, monkeypatch):
+ """Local filesystem paths must NOT be disclosed via this endpoint."""
+ monkeypatch.setenv("UI_LOGO_PATH", "/var/lib/litellm/internal-secret-logo.png")
+ response = client.get("/get_logo_url")
+ assert response.status_code == 200
+ assert normalize(response.json()) == {"logo_url": ""}
+
+
+def test_get_logo_url_blank_when_unset(client, monkeypatch):
+ monkeypatch.delenv("UI_LOGO_PATH", raising=False)
+ response = client.get("/get_logo_url")
+ assert response.status_code == 200
+ assert normalize(response.json()) == {"logo_url": ""}
+
+
+def test_get_logo_url_invalid_scheme_blank(client, monkeypatch):
+ """file:// and other non-HTTP schemes are not disclosed (error/edge path)."""
+ monkeypatch.setenv("UI_LOGO_PATH", "file:///etc/passwd")
+ response = client.get("/get_logo_url")
+ assert response.status_code == 200
+ assert normalize(response.json()) == {"logo_url": ""}
+
+
+# ---------------------------------------------------------------------------
+# GET /get_image
+# ---------------------------------------------------------------------------
+
+
+def test_get_image_returns_default_logo(client, monkeypatch):
+ monkeypatch.delenv("UI_LOGO_PATH", raising=False)
+ response = client.get("/get_image")
+ assert response.status_code == 200
+ media_type = response.headers.get("content-type", "").split(";")[0]
+ shape = {
+ "status": response.status_code,
+ "media_type_image": media_type.startswith("image/"),
+ "has_body": len(response.content) > 0,
+ }
+ assert shape == {"status": 200, "media_type_image": True, "has_body": True}
+
+
+def test_get_image_redirects_remote_url(client, monkeypatch):
+ """Remote logo URLs are served via redirect — the proxy never fetches them server-side."""
+ monkeypatch.setenv("UI_LOGO_PATH", "https://example.invalid/logo.png")
+ response = client.get("/get_image", follow_redirects=False)
+ assert response.status_code in (302, 303, 307, 308)
+ assert response.headers.get("location") == "https://example.invalid/logo.png"
+
+
+def test_get_image_invalid_local_path_falls_back(client, monkeypatch):
+ """Non-existent UI_LOGO_PATH (error path) falls back to default logo, still 200."""
+ monkeypatch.setenv("UI_LOGO_PATH", "/nonexistent/path/to/logo.png")
+ response = client.get("/get_image")
+ assert response.status_code == 200
+ shape = {
+ "status": response.status_code,
+ "media_type_image": response.headers.get("content-type", "").startswith(
+ "image/"
+ ),
+ "has_body": len(response.content) > 0,
+ }
+ assert shape == {"status": 200, "media_type_image": True, "has_body": True}
+
+
+# ---------------------------------------------------------------------------
+# GET /get_favicon
+# ---------------------------------------------------------------------------
+
+
+def test_get_favicon_returns_file(client):
+ response = client.get("/get_favicon")
+ assert response.status_code == 200
+ shape = {
+ "status": response.status_code,
+ "has_body": len(response.content) > 0,
+ "content_type_set": bool(response.headers.get("content-type")),
+ }
+ assert shape == {"status": 200, "has_body": True, "content_type_set": True}
+
+
+def test_get_favicon_invalid_custom_path_falls_back(client, monkeypatch):
+ """Bad UI_FAVICON_PATH (error/edge path) falls back to default — still 200."""
+ monkeypatch.setenv("UI_FAVICON_PATH", "/nonexistent/favicon.ico")
+ response = client.get("/get_favicon")
+ assert response.status_code == 200
+ assert len(response.content) > 0
diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_model_cost_map.py b/tests/test_litellm/proxy/proxy_server/test_routes_model_cost_map.py
index ad6b4016461..16e410f1b1e 100644
--- a/tests/test_litellm/proxy/proxy_server/test_routes_model_cost_map.py
+++ b/tests/test_litellm/proxy/proxy_server/test_routes_model_cost_map.py
@@ -1 +1,371 @@
-"""Placeholder. Filled by a follow-up PR per the Notion plan."""
+"""Pin tests for proxy_server.py model cost map routes (PR3).
+
+Routes covered:
+- POST /reload/model_cost_map
+- POST /schedule/model_cost_map_reload
+- DELETE /schedule/model_cost_map_reload
+- GET /schedule/model_cost_map_reload/status
+- GET /model/cost_map/source
+"""
+
+from __future__ import annotations
+
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from .conftest import VOLATILE_KEYS, normalize
+
+# Some response bodies include a "timestamp" — extend the volatile set so
+# dict-equality assertions remain stable.
+_VOLATILE = VOLATILE_KEYS | frozenset({"timestamp"})
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _attach_litellm_config(mock_prisma):
+ """Attach a litellm_config table mock (not in conftest's _PRISMA_TABLES)."""
+ table = MagicMock()
+ table.find_unique = AsyncMock(return_value=None)
+ table.find_first = AsyncMock(return_value=None)
+ table.find_many = AsyncMock(return_value=[])
+ table.upsert = AsyncMock()
+ table.create = AsyncMock()
+ table.update = AsyncMock()
+ table.delete = AsyncMock()
+ table.delete_many = AsyncMock()
+ mock_prisma.db.litellm_config = table
+ return table
+
+
+# ---------------------------------------------------------------------------
+# POST /reload/model_cost_map
+# ---------------------------------------------------------------------------
+
+
+def test_reload_model_cost_map_happy(client, auth_as, monkeypatch, mock_prisma):
+ """Admin can trigger a manual reload; handler returns model count + status."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ table = _attach_litellm_config(mock_prisma)
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+
+ fake_cost_map = {"gpt-4": {"input_cost": 0.03}, "gpt-3.5": {"input_cost": 0.002}}
+ monkeypatch.setattr(
+ "litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map",
+ lambda url=None: fake_cost_map,
+ )
+ monkeypatch.setattr("litellm.add_known_models", lambda model_cost_map=None: None)
+ monkeypatch.setattr("litellm.model_cost", {}, raising=False)
+ monkeypatch.setattr(
+ "litellm.proxy.proxy_server._invalidate_model_cost_lowercase_map",
+ lambda: None,
+ raising=False,
+ )
+
+ async def _fake_invalidate(name):
+ return None
+
+ monkeypatch.setattr(ps, "invalidate_config_param", _fake_invalidate)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.post("/reload/model_cost_map")
+ assert response.status_code == 200
+ body = normalize(response.json(), volatile=_VOLATILE)
+ assert body == {
+ "message": "Price data reloaded successfully! 2 models updated.",
+ "status": "success",
+ "models_count": 2,
+ "timestamp": "",
+ }
+ assert table.upsert.await_count == 1
+
+
+def test_reload_model_cost_map_not_admin_forbidden(client, auth_as):
+ """Non-admin caller gets 403 with a role-specific detail."""
+ from litellm.proxy._types import LitellmUserRoles
+
+ with auth_as(LitellmUserRoles.INTERNAL_USER):
+ response = client.post("/reload/model_cost_map")
+ assert response.status_code == 403
+ assert "Admin role required" in response.json().get("detail", "")
+
+
+def test_reload_model_cost_map_no_db_500(client, auth_as, monkeypatch):
+ """Admin path but prisma_client is None — handler raises 500."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ monkeypatch.setattr(ps, "prisma_client", None)
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.post("/reload/model_cost_map")
+ assert response.status_code == 500
+ assert "Database connection not available" in response.json().get("detail", "")
+
+
+# ---------------------------------------------------------------------------
+# POST /schedule/model_cost_map_reload
+# ---------------------------------------------------------------------------
+
+
+def test_schedule_model_cost_map_reload_happy(
+ client, auth_as, monkeypatch, mock_prisma
+):
+ """Admin schedules a reload — handler upserts config and echoes interval."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ table = _attach_litellm_config(mock_prisma)
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+
+ async def _fake_invalidate(name):
+ return None
+
+ monkeypatch.setattr(ps, "invalidate_config_param", _fake_invalidate)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.post("/schedule/model_cost_map_reload?hours=6")
+ assert response.status_code == 200
+ body = normalize(response.json(), volatile=_VOLATILE)
+ assert body == {
+ "message": "Model cost map reload scheduled for every 6 hours",
+ "status": "success",
+ "interval_hours": 6,
+ "timestamp": "",
+ }
+ assert table.upsert.await_count == 1
+
+
+def test_schedule_model_cost_map_reload_invalid_hours(
+ client, auth_as, monkeypatch, mock_prisma
+):
+ """hours <= 0 is rejected with 400."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ _attach_litellm_config(mock_prisma)
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.post("/schedule/model_cost_map_reload?hours=0")
+ assert response.status_code == 400
+ assert "Hours must be greater than 0" in response.json().get("detail", "")
+
+
+def test_schedule_model_cost_map_reload_not_admin_forbidden(client, auth_as):
+ """Non-admin caller blocked with 403."""
+ from litellm.proxy._types import LitellmUserRoles
+
+ with auth_as(LitellmUserRoles.INTERNAL_USER):
+ response = client.post("/schedule/model_cost_map_reload?hours=6")
+ assert response.status_code == 403
+ assert "Admin role required" in response.json().get("detail", "")
+
+
+# ---------------------------------------------------------------------------
+# DELETE /schedule/model_cost_map_reload
+# ---------------------------------------------------------------------------
+
+
+def test_cancel_model_cost_map_reload_happy(client, auth_as, monkeypatch, mock_prisma):
+ """Admin cancellation deletes config row and returns success body."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ table = _attach_litellm_config(mock_prisma)
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+
+ async def _fake_invalidate(name):
+ return None
+
+ monkeypatch.setattr(ps, "invalidate_config_param", _fake_invalidate)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.delete("/schedule/model_cost_map_reload")
+ assert response.status_code == 200
+ body = normalize(response.json(), volatile=_VOLATILE)
+ assert body == {
+ "message": "Model cost map reload schedule cancelled",
+ "status": "success",
+ "timestamp": "",
+ }
+ assert table.delete.await_count == 1
+
+
+def test_cancel_model_cost_map_reload_not_admin_forbidden(client, auth_as):
+ from litellm.proxy._types import LitellmUserRoles
+
+ with auth_as(LitellmUserRoles.INTERNAL_USER):
+ response = client.delete("/schedule/model_cost_map_reload")
+ assert response.status_code == 403
+ assert "Admin role required" in response.json().get("detail", "")
+
+
+def test_cancel_model_cost_map_reload_no_db_500(client, auth_as, monkeypatch):
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ monkeypatch.setattr(ps, "prisma_client", None)
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.delete("/schedule/model_cost_map_reload")
+ assert response.status_code == 500
+ assert "Database connection not available" in response.json().get("detail", "")
+
+
+# ---------------------------------------------------------------------------
+# GET /schedule/model_cost_map_reload/status
+# ---------------------------------------------------------------------------
+
+
+def test_get_model_cost_map_reload_status_no_db_not_scheduled(
+ client, auth_as, monkeypatch
+):
+ """No prisma client → returns the not-scheduled shape (4 keys, all-null)."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ monkeypatch.setattr(ps, "prisma_client", None)
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.get("/schedule/model_cost_map_reload/status")
+ assert response.status_code == 200
+ assert normalize(response.json()) == {
+ "scheduled": False,
+ "interval_hours": None,
+ "last_run": None,
+ "next_run": None,
+ }
+
+
+def test_get_model_cost_map_reload_status_scheduled(
+ client, auth_as, monkeypatch, mock_prisma
+):
+ """A valid config row → scheduled=True and the interval is echoed."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ table = _attach_litellm_config(mock_prisma)
+ config_row = MagicMock()
+ config_row.param_value = {"interval_hours": 12, "force_reload": False}
+ table.find_unique = AsyncMock(return_value=config_row)
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+ monkeypatch.setattr(ps, "last_model_cost_map_reload", None)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.get("/schedule/model_cost_map_reload/status")
+ assert response.status_code == 200
+ assert normalize(response.json()) == {
+ "scheduled": True,
+ "interval_hours": 12,
+ "last_run": None,
+ "next_run": None,
+ }
+
+
+def test_get_model_cost_map_reload_status_no_config_not_scheduled(
+ client, auth_as, monkeypatch, mock_prisma
+):
+ """Config row exists but interval_hours=None → not scheduled."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ table = _attach_litellm_config(mock_prisma)
+ config_row = MagicMock()
+ config_row.param_value = {"interval_hours": None, "force_reload": True}
+ table.find_unique = AsyncMock(return_value=config_row)
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+ monkeypatch.setattr(ps, "last_model_cost_map_reload", None)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.get("/schedule/model_cost_map_reload/status")
+ assert response.status_code == 200
+ assert normalize(response.json()) == {
+ "scheduled": False,
+ "interval_hours": None,
+ "last_run": None,
+ "next_run": None,
+ }
+
+
+def test_get_model_cost_map_reload_status_not_admin_forbidden(client, auth_as):
+ from litellm.proxy._types import LitellmUserRoles
+
+ with auth_as(LitellmUserRoles.INTERNAL_USER):
+ response = client.get("/schedule/model_cost_map_reload/status")
+ assert response.status_code == 403
+ assert "Admin role required" in response.json().get("detail", "")
+
+
+# ---------------------------------------------------------------------------
+# GET /model/cost_map/source
+# ---------------------------------------------------------------------------
+
+
+def test_get_model_cost_map_source_happy(client, auth_as, monkeypatch):
+ """Admin gets the source-info dict, augmented with the current model_count."""
+ from litellm.proxy._types import LitellmUserRoles
+
+ fake_info = {
+ "source": "remote",
+ "url": "https://example.invalid/cost_map.json",
+ "is_env_forced": False,
+ "fallback_reason": None,
+ }
+ monkeypatch.setattr(
+ "litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map_source_info",
+ lambda: fake_info,
+ )
+ monkeypatch.setattr("litellm.model_cost", {"a": 1, "b": 2, "c": 3}, raising=False)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.get("/model/cost_map/source")
+ assert response.status_code == 200
+ assert normalize(response.json()) == {
+ "source": "remote",
+ "url": "https://example.invalid/cost_map.json",
+ "is_env_forced": False,
+ "fallback_reason": None,
+ "model_count": 3,
+ }
+
+
+def test_get_model_cost_map_source_admin_view_only_allowed(
+ client, auth_as, monkeypatch
+):
+ """PROXY_ADMIN_VIEW_ONLY can read source info — pins the read-only ACL."""
+ from litellm.proxy._types import LitellmUserRoles
+
+ fake_info = {
+ "source": "local",
+ "url": None,
+ "is_env_forced": True,
+ "fallback_reason": None,
+ }
+ monkeypatch.setattr(
+ "litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map_source_info",
+ lambda: fake_info,
+ )
+ monkeypatch.setattr("litellm.model_cost", {"a": 1}, raising=False)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY):
+ response = client.get("/model/cost_map/source")
+ assert response.status_code == 200
+ assert normalize(response.json()) == {
+ "source": "local",
+ "url": None,
+ "is_env_forced": True,
+ "fallback_reason": None,
+ "model_count": 1,
+ }
+
+
+def test_get_model_cost_map_source_not_admin_forbidden(client, auth_as):
+ from litellm.proxy._types import LitellmUserRoles
+
+ with auth_as(LitellmUserRoles.INTERNAL_USER):
+ response = client.get("/model/cost_map/source")
+ assert response.status_code == 403
+ assert "Admin role required" in response.json().get("detail", "")
diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py b/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py
index ad6b4016461..98259824378 100644
--- a/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py
+++ b/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py
@@ -1 +1,149 @@
-"""Placeholder. Filled by a follow-up PR per the Notion plan."""
+"""Behavior pins for ``proxy_server.py`` model-info routes.
+
+Pins (PR2):
+ - GET /v2/model/info
+ - GET /v1/model/info
+ - GET /model/info
+ - GET /model_group/info
+"""
+
+from __future__ import annotations
+
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from litellm.proxy import proxy_server
+
+from .conftest import normalize # type: ignore[import-not-found]
+
+# ---------------------------------------------------------------------------
+# GET /v2/model/info
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def empty_router(monkeypatch):
+ router = MagicMock()
+ router.model_list = []
+ monkeypatch.setattr(proxy_server, "llm_router", router)
+ monkeypatch.setattr(proxy_server, "llm_model_list", [])
+ yield router
+
+
+@pytest.fixture
+def null_router(monkeypatch):
+ monkeypatch.setattr(proxy_server, "llm_router", None)
+ monkeypatch.setattr(proxy_server, "llm_model_list", None)
+ yield
+
+
+def test_v2_model_info_empty_router_happy_path(client, auth_as, empty_router):
+ """Pins ``GET /v2/model/info`` (empty router branch returns deterministic shape)."""
+ with auth_as():
+ response = client.get("/v2/model/info")
+ assert response.status_code == 200
+ assert normalize(response.json()) == {
+ "data": [],
+ "total_count": 0,
+ "current_page": 1,
+ "total_pages": 0,
+ "size": 50,
+ }
+
+
+def test_v2_model_info_invalid_page_returns_422(client, auth_as, empty_router):
+ """Pins ``GET /v2/model/info`` (error: invalid page parameter)."""
+ with auth_as():
+ response = client.get("/v2/model/info", params={"page": 0})
+ assert response.status_code == 422
+ assert "detail" in response.json()
+
+
+# ---------------------------------------------------------------------------
+# GET /v1/model/info, GET /model/info
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def configured_router(monkeypatch):
+ deployment = MagicMock()
+ deployment.model_dump = MagicMock(
+ return_value={
+ "model_name": "gpt-4",
+ "litellm_params": {"model": "gpt-4"},
+ "model_info": {"id": "abc", "db_model": False},
+ }
+ )
+ router = MagicMock()
+ router.get_deployment = MagicMock(return_value=deployment)
+ router.get_model_names = MagicMock(return_value=["gpt-4"])
+ router.get_model_access_groups = MagicMock(return_value={})
+ router.get_model_list = MagicMock(return_value=[])
+ monkeypatch.setattr(proxy_server, "llm_router", router)
+ monkeypatch.setattr(proxy_server, "llm_model_list", [{"model_name": "gpt-4"}])
+ monkeypatch.setattr(proxy_server, "user_model", None)
+ monkeypatch.setattr(proxy_server, "_get_proxy_model_info", lambda model: model)
+ yield router
+
+
+@pytest.mark.parametrize("path", ["/v1/model/info", "/model/info"])
+def test_v1_model_info_specific_id_happy(client, auth_as, configured_router, path):
+ """Pins ``GET /v1/model/info`` and ``GET /model/info`` (happy: specific id).
+
+ Includes ``litellm_model_id`` so the early-return branch produces a
+ deterministic ``{"data": []}`` body without touching
+ the full model-info enrichment pipeline.
+ """
+ with auth_as():
+ response = client.get(path, params={"litellm_model_id": "abc"})
+ assert response.status_code == 200
+ body = normalize(response.json())
+ assert body == {
+ "data": [
+ {
+ "model_name": "gpt-4",
+ "litellm_params": {"model": "gpt-4"},
+ "model_info": {"id": "", "db_model": False},
+ }
+ ]
+ }
+
+
+@pytest.mark.parametrize("path", ["/v1/model/info", "/model/info"])
+def test_v1_model_info_no_model_list_error(client, auth_as, null_router, path):
+ """Pins ``GET /v1/model/info`` and ``GET /model/info`` (error: no model list)."""
+ with auth_as():
+ response = client.get(path)
+ assert response.status_code == 500
+ assert "LLM Model List not loaded" in response.text
+
+
+# ---------------------------------------------------------------------------
+# GET /model_group/info
+# ---------------------------------------------------------------------------
+
+
+def test_model_group_info_no_models_happy(client, auth_as, null_router):
+ """Pins ``GET /model_group/info`` (happy: empty list when no models)."""
+ with auth_as():
+ response = client.get("/model_group/info")
+ assert response.status_code == 200
+ summary = {
+ "status_code": response.status_code,
+ "body": normalize(response.json()),
+ "object_kind": "model_group_info",
+ }
+ assert summary == {
+ "status_code": 200,
+ "body": {"data": []},
+ "object_kind": "model_group_info",
+ }
+
+
+def test_model_group_info_invalid_method(client, auth_as, null_router):
+ """Pins ``GET /model_group/info`` (error: method not allowed)."""
+ with auth_as():
+ response = client.post("/model_group/info", json={})
+ assert response.status_code == 405
+ assert len(response.content) > 0
diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_model_metrics.py b/tests/test_litellm/proxy/proxy_server/test_routes_model_metrics.py
index ad6b4016461..246e2cbba54 100644
--- a/tests/test_litellm/proxy/proxy_server/test_routes_model_metrics.py
+++ b/tests/test_litellm/proxy/proxy_server/test_routes_model_metrics.py
@@ -1 +1,228 @@
-"""Placeholder. Filled by a follow-up PR per the Notion plan."""
+"""Behavior pins for ``proxy_server.py`` model-metrics routes.
+
+Pins (PR2):
+ - GET /model/streaming_metrics
+ - GET /model/metrics
+ - GET /model/metrics/slow_responses
+ - GET /model/metrics/exceptions
+ - GET /model/settings
+ - GET /alerting/settings
+"""
+
+from __future__ import annotations
+
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+import litellm
+from litellm.proxy import proxy_server
+from litellm.proxy._types import LitellmUserRoles
+
+from .conftest import normalize # type: ignore[import-not-found]
+
+# ---------------------------------------------------------------------------
+# Shared fixtures
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def prisma_with_query_raw(monkeypatch):
+ pc = MagicMock()
+ pc.db.query_raw = AsyncMock(return_value=[])
+ monkeypatch.setattr(proxy_server, "prisma_client", pc)
+ return pc
+
+
+@pytest.fixture
+def no_prisma(monkeypatch):
+ monkeypatch.setattr(proxy_server, "prisma_client", None)
+ yield
+
+
+# ---------------------------------------------------------------------------
+# GET /model/streaming_metrics
+# ---------------------------------------------------------------------------
+
+
+def test_model_streaming_metrics_happy(client, auth_as, prisma_with_query_raw):
+ """Pins ``GET /model/streaming_metrics`` (happy: empty data list).
+
+ Drives the deterministic branch where ``query_raw`` returns an empty
+ list; the handler should return the empty payload unchanged so the
+ pin can rely on the exact response shape.
+ """
+ with auth_as():
+ response = client.get(
+ "/model/streaming_metrics", params={"_selected_model_group": "gpt-4"}
+ )
+ assert response.status_code == 200
+ assert normalize(response.json()) == {"data": [], "all_api_bases": []}
+
+
+def test_model_streaming_metrics_no_prisma_error(client, auth_as, no_prisma):
+ """Pins ``GET /model/streaming_metrics`` (error: prisma not initialized)."""
+ with auth_as():
+ response = client.get("/model/streaming_metrics")
+ assert response.status_code == 500
+ assert response.content
+
+
+# ---------------------------------------------------------------------------
+# GET /model/metrics
+# ---------------------------------------------------------------------------
+
+
+def test_model_metrics_happy(client, auth_as, prisma_with_query_raw):
+ """Pins ``GET /model/metrics`` (happy: empty result)."""
+ with auth_as():
+ response = client.get("/model/metrics")
+ assert response.status_code == 200
+ assert normalize(response.json()) == {"data": [], "all_api_bases": []}
+
+
+def test_model_metrics_no_prisma_error(client, auth_as, no_prisma):
+ """Pins ``GET /model/metrics`` (error: prisma not initialized)."""
+ with auth_as():
+ response = client.get("/model/metrics")
+ assert response.status_code == 500
+ assert response.content
+
+
+# ---------------------------------------------------------------------------
+# GET /model/metrics/slow_responses
+# ---------------------------------------------------------------------------
+
+
+def test_model_metrics_slow_responses_happy(
+ client, auth_as, prisma_with_query_raw, monkeypatch
+):
+ """Pins ``GET /model/metrics/slow_responses`` (happy: empty list)."""
+ logging_obj = MagicMock()
+ logging_obj.slack_alerting_instance.alerting_threshold = 30
+ monkeypatch.setattr(proxy_server, "proxy_logging_obj", logging_obj)
+ with auth_as():
+ response = client.get("/model/metrics/slow_responses")
+ assert response.status_code == 200
+ assert normalize(response.json()) == []
+
+
+def test_model_metrics_slow_responses_no_prisma(client, auth_as, no_prisma):
+ """Pins ``GET /model/metrics/slow_responses`` (error: prisma not initialized)."""
+ with auth_as():
+ response = client.get("/model/metrics/slow_responses")
+ assert response.status_code == 500
+ assert response.content
+
+
+# ---------------------------------------------------------------------------
+# GET /model/metrics/exceptions
+# ---------------------------------------------------------------------------
+
+
+def test_model_metrics_exceptions_happy(client, auth_as, prisma_with_query_raw):
+ """Pins ``GET /model/metrics/exceptions`` (happy: empty)."""
+ with auth_as():
+ response = client.get("/model/metrics/exceptions")
+ assert response.status_code == 200
+ assert normalize(response.json()) == {"data": [], "exception_types": []}
+
+
+def test_model_metrics_exceptions_no_prisma(client, auth_as, no_prisma):
+ """Pins ``GET /model/metrics/exceptions`` (error: prisma not initialized)."""
+ with auth_as():
+ response = client.get("/model/metrics/exceptions")
+ assert response.status_code == 500
+ assert response.content
+
+
+# ---------------------------------------------------------------------------
+# GET /model/settings
+# ---------------------------------------------------------------------------
+
+
+def test_model_settings_happy(client, auth_as, monkeypatch):
+ """Pins ``GET /model/settings`` (happy)."""
+ monkeypatch.setattr(litellm, "provider_list", ["openai"])
+ monkeypatch.setattr(
+ litellm,
+ "get_provider_fields",
+ lambda custom_llm_provider: [],
+ )
+ with auth_as():
+ response = client.get("/model/settings")
+ assert response.status_code == 200
+ body = response.json()
+ assert body == [{"name": "openai", "fields": []}]
+ summary = {
+ "status_code": response.status_code,
+ "first_entry_name": body[0]["name"],
+ "body_length": len(body),
+ }
+ assert summary == {
+ "status_code": 200,
+ "first_entry_name": "openai",
+ "body_length": 1,
+ }
+
+
+def test_model_settings_method_not_allowed(client, auth_as):
+ """Pins ``GET /model/settings`` (error: wrong method)."""
+ with auth_as():
+ response = client.post("/model/settings", json={})
+ assert response.status_code == 405
+ assert len(response.content) > 0
+
+
+# ---------------------------------------------------------------------------
+# GET /alerting/settings
+# ---------------------------------------------------------------------------
+
+
+def test_alerting_settings_no_db_error(client, auth_as, no_prisma):
+ """Pins ``GET /alerting/settings`` (error: db not connected)."""
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.get("/alerting/settings")
+ assert response.status_code == 400
+ assert "error" in response.text or "detail" in response.text
+
+
+def test_alerting_settings_non_admin_error(client, auth_as, monkeypatch):
+ """Pins ``GET /alerting/settings`` (error: non-admin forbidden)."""
+ monkeypatch.setattr(proxy_server, "prisma_client", MagicMock())
+ with auth_as(LitellmUserRoles.INTERNAL_USER):
+ response = client.get("/alerting/settings")
+ assert response.status_code == 400
+ assert "internal_user" in response.text.lower() or "error" in response.text
+
+
+def test_alerting_settings_happy(client, auth_as, monkeypatch):
+ """Pins ``GET /alerting/settings`` (happy: returns list of ConfigList entries)."""
+ pc = MagicMock()
+ pc.db.litellm_config.find_first = AsyncMock(return_value=None)
+ monkeypatch.setattr(proxy_server, "prisma_client", pc)
+
+ logging_obj = MagicMock()
+ args_model = MagicMock()
+ args_model.model_dump = MagicMock(return_value={})
+ logging_obj.slack_alerting_instance.alerting_args = args_model
+ monkeypatch.setattr(proxy_server, "proxy_logging_obj", logging_obj)
+ monkeypatch.setattr(proxy_server, "general_settings", {})
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response = client.get("/alerting/settings")
+ assert response.status_code == 200
+ body = response.json()
+ assert body[0]["field_name"] == "slack_alerting"
+ summary = {
+ "status_code": response.status_code,
+ "first_field_name": body[0]["field_name"],
+ "first_field_value": body[0]["field_value"],
+ "first_field_type": body[0]["field_type"],
+ }
+ assert summary == {
+ "status_code": 200,
+ "first_field_name": "slack_alerting",
+ "first_field_value": False,
+ "first_field_type": "Boolean",
+ }
diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_models.py b/tests/test_litellm/proxy/proxy_server/test_routes_models.py
index ad6b4016461..381835fbc14 100644
--- a/tests/test_litellm/proxy/proxy_server/test_routes_models.py
+++ b/tests/test_litellm/proxy/proxy_server/test_routes_models.py
@@ -1 +1,132 @@
-"""Placeholder. Filled by a follow-up PR per the Notion plan."""
+"""Behavior pins for ``proxy_server.py`` model routes.
+
+Pins (PR2):
+ - GET /v1/models
+ - GET /models
+ - GET /v1/models/{model_id}
+ - GET /models/{model_id}
+"""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock
+
+import pytest
+
+import litellm
+from litellm.proxy import proxy_server
+
+from .conftest import normalize # type: ignore[import-not-found]
+
+
+def _stub_model_info_response(
+ model_id: str = "gpt-4", provider: str = "openai"
+) -> dict:
+ return {
+ "id": model_id,
+ "object": "model",
+ "created": 0,
+ "owned_by": provider,
+ }
+
+
+@pytest.fixture
+def patched_models(monkeypatch):
+ """Stub router + utility helpers used by the /models routes."""
+ from litellm.proxy import utils as proxy_utils
+
+ router = MagicMock()
+ router.get_fully_blocked_model_names = MagicMock(return_value=set())
+ router.get_model_names = MagicMock(return_value=["gpt-4", "claude-sonnet"])
+ router.get_model_access_groups = MagicMock(return_value={})
+
+ deployment = MagicMock()
+ deployment.litellm_params.model = "gpt-4"
+ router.get_deployment_by_model_group_name = MagicMock(return_value=deployment)
+
+ monkeypatch.setattr(proxy_server, "llm_router", router)
+ monkeypatch.setattr(proxy_server, "prisma_client", MagicMock())
+
+ async def _fake_get_available_models_for_user(**kwargs):
+ return ["gpt-4", "claude-sonnet"]
+
+ monkeypatch.setattr(
+ proxy_utils,
+ "get_available_models_for_user",
+ _fake_get_available_models_for_user,
+ )
+
+ def _fake_create_model_info_response(model_id, provider="openai", **kwargs):
+ return _stub_model_info_response(model_id=model_id, provider=provider)
+
+ monkeypatch.setattr(
+ proxy_utils, "create_model_info_response", _fake_create_model_info_response
+ )
+
+ monkeypatch.setattr(proxy_utils, "validate_model_access", lambda **kwargs: None)
+
+ monkeypatch.setattr(
+ litellm,
+ "get_llm_provider",
+ lambda model: (model, "openai", None, None),
+ )
+
+ return router
+
+
+@pytest.mark.parametrize("path", ["/v1/models", "/models"])
+def test_get_models_happy_path(client, auth_as, patched_models, path):
+ """Pins: ``GET /v1/models``, ``GET /models``."""
+ with auth_as():
+ response = client.get(path)
+ assert response.status_code == 200
+ assert normalize(response.json()) == {
+ "data": [
+ {
+ "id": "",
+ "object": "model",
+ "created": "",
+ "owned_by": "openai",
+ },
+ {
+ "id": "",
+ "object": "model",
+ "created": "",
+ "owned_by": "openai",
+ },
+ ],
+ "object": "list",
+ }
+
+
+@pytest.mark.parametrize("path", ["/v1/models", "/models"])
+def test_get_models_invalid_scope_returns_400(client, auth_as, patched_models, path):
+ """Pins: ``GET /v1/models``, ``GET /models`` (error path: invalid scope)."""
+ with auth_as():
+ response = client.get(path, params={"scope": "not-a-real-scope"})
+ assert response.status_code == 400
+ assert "Invalid scope parameter" in str(response.json())
+
+
+@pytest.mark.parametrize("path", ["/v1/models/gpt-4", "/models/gpt-4"])
+def test_get_model_by_id_happy_path(client, auth_as, patched_models, path):
+ """Pins: ``GET /v1/models/{model_id}``, ``GET /models/{model_id}``."""
+ with auth_as():
+ response = client.get(path)
+ assert response.status_code == 200
+ assert normalize(response.json()) == {
+ "id": "",
+ "object": "model",
+ "created": "",
+ "owned_by": "openai",
+ }
+
+
+@pytest.mark.parametrize("path", ["/v1/models/missing", "/models/missing"])
+def test_get_model_by_id_not_found(client, auth_as, patched_models, path):
+ """Pins: ``GET /v1/models/{model_id}``, ``GET /models/{model_id}`` (error: 404)."""
+ patched_models.get_deployment_by_model_group_name = MagicMock(return_value=None)
+ with auth_as():
+ response = client.get(path)
+ assert response.status_code == 404
+ assert "not found" in response.text.lower()
diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_moderations.py b/tests/test_litellm/proxy/proxy_server/test_routes_moderations.py
index ad6b4016461..4553a5e7cf4 100644
--- a/tests/test_litellm/proxy/proxy_server/test_routes_moderations.py
+++ b/tests/test_litellm/proxy/proxy_server/test_routes_moderations.py
@@ -1 +1,111 @@
-"""Placeholder. Filled by a follow-up PR per the Notion plan."""
+"""Behavior pins for ``proxy_server.py`` moderations routes.
+
+Pins (PR2):
+ - POST /v1/moderations
+ - POST /moderations
+"""
+
+from __future__ import annotations
+
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from litellm.proxy import proxy_server
+
+from .conftest import normalize # type: ignore[import-not-found]
+
+HAPPY_RESPONSE = {
+ "id": "modr-test",
+ "model": "text-moderation-stable",
+ "results": [
+ {
+ "flagged": False,
+ "categories": {"violence": False},
+ "category_scores": {"violence": 0.01},
+ }
+ ],
+}
+
+
+@pytest.fixture
+def patched_moderation(monkeypatch):
+ monkeypatch.setattr(proxy_server, "llm_router", MagicMock())
+ monkeypatch.setattr(
+ proxy_server,
+ "proxy_logging_obj",
+ MagicMock(
+ pre_call_hook=AsyncMock(side_effect=lambda **kw: kw["data"]),
+ post_call_failure_hook=AsyncMock(),
+ update_request_status=AsyncMock(),
+ ),
+ )
+
+ async def _add_data(data, **kwargs):
+ return data
+
+ monkeypatch.setattr(proxy_server, "add_litellm_data_to_request", _add_data)
+
+ async def _fake_llm_call():
+ return dict(HAPPY_RESPONSE)
+
+ async def _fake_route_request(*args, **kwargs):
+ return _fake_llm_call()
+
+ monkeypatch.setattr(proxy_server, "route_request", _fake_route_request)
+ yield
+
+
+@pytest.fixture
+def moderation_pipeline_raises(monkeypatch):
+ monkeypatch.setattr(proxy_server, "llm_router", MagicMock())
+ monkeypatch.setattr(
+ proxy_server,
+ "proxy_logging_obj",
+ MagicMock(
+ pre_call_hook=AsyncMock(side_effect=lambda **kw: kw["data"]),
+ post_call_failure_hook=AsyncMock(),
+ update_request_status=AsyncMock(),
+ ),
+ )
+
+ async def _add_data(data, **kwargs):
+ return data
+
+ monkeypatch.setattr(proxy_server, "add_litellm_data_to_request", _add_data)
+
+ async def _raise(*args, **kwargs):
+ raise ValueError("boom")
+
+ monkeypatch.setattr(proxy_server, "route_request", _raise)
+ yield
+
+
+@pytest.mark.parametrize("path", ["/v1/moderations", "/moderations"])
+def test_moderation_happy_path(client, auth_as, patched_moderation, path):
+ """Pins ``POST /v1/moderations`` and ``POST /moderations`` (happy)."""
+ payload = {"model": "text-moderation-stable", "input": "Sample text"}
+ with auth_as():
+ response = client.post(path, json=payload)
+ assert response.status_code == 200
+ assert normalize(response.json()) == {
+ "id": "",
+ "model": "text-moderation-stable",
+ "results": [
+ {
+ "flagged": False,
+ "categories": {"violence": False},
+ "category_scores": {"violence": 0.01},
+ }
+ ],
+ }
+
+
+@pytest.mark.parametrize("path", ["/v1/moderations", "/moderations"])
+def test_moderation_error(client, auth_as, moderation_pipeline_raises, path):
+ """Pins ``POST /v1/moderations`` and ``POST /moderations`` (error)."""
+ payload = {"model": "text-moderation-stable", "input": "Sample text"}
+ with auth_as():
+ response = client.post(path, json=payload)
+ assert response.status_code == 500
+ assert len(response.content) > 0
diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_onboarding.py b/tests/test_litellm/proxy/proxy_server/test_routes_onboarding.py
index ad6b4016461..35ae9a3568e 100644
--- a/tests/test_litellm/proxy/proxy_server/test_routes_onboarding.py
+++ b/tests/test_litellm/proxy/proxy_server/test_routes_onboarding.py
@@ -1 +1,350 @@
-"""Placeholder. Filled by a follow-up PR per the Notion plan."""
+"""Pin tests for proxy_server.py onboarding routes (PR3).
+
+Routes covered:
+- GET /onboarding/get_token
+- POST /onboarding/claim_token
+"""
+
+from __future__ import annotations
+
+from datetime import datetime, timedelta, timezone
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, MagicMock
+
+import jwt
+import pytest
+
+from .conftest import normalize
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_invite(
+ invite_id: str = "inv-123",
+ user_id: str = "user-abc",
+ expires_at: datetime | None = None,
+ is_accepted: bool = False,
+ accepted_at=None,
+):
+ """Build a fake invitation object with the attributes the handler reads."""
+ if expires_at is None:
+ expires_at = datetime.now(timezone.utc) + timedelta(days=1)
+ return SimpleNamespace(
+ id=invite_id,
+ user_id=user_id,
+ expires_at=expires_at,
+ is_accepted=is_accepted,
+ accepted_at=accepted_at,
+ )
+
+
+def _make_user_obj(
+ user_id: str = "user-abc",
+ user_email: str = "alice@example.com",
+ user_role: str = "internal_user",
+):
+ return SimpleNamespace(
+ user_id=user_id,
+ user_email=user_email,
+ user_role=user_role,
+ password=None,
+ )
+
+
+def _install_tx_context(mock_prisma):
+ """Wire ``async with prisma_client.db.tx() as tx`` to return ``mock_prisma.db``.
+
+ The handler runs the update inside a transaction; have ``tx`` yield a
+ namespace that exposes the same tables as the outer client so its
+ ``update_many`` / ``update`` calls hit our mocks.
+ """
+ tx_cm = MagicMock()
+ tx_cm.__aenter__ = AsyncMock(return_value=mock_prisma.db)
+ tx_cm.__aexit__ = AsyncMock(return_value=None)
+ mock_prisma.db.tx = MagicMock(return_value=tx_cm)
+
+
+# ---------------------------------------------------------------------------
+# GET /onboarding/get_token
+# ---------------------------------------------------------------------------
+
+
+def test_onboarding_get_token_happy(client, monkeypatch, mock_prisma):
+ """Valid invite link → returns dict with login_url, token, user_email."""
+ from litellm.proxy import proxy_server as ps
+
+ invite = _make_invite()
+ user_obj = _make_user_obj()
+ mock_prisma.db.litellm_invitationlink.find_unique.return_value = invite
+ mock_prisma.db.litellm_usertable.find_unique.return_value = user_obj
+
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+ monkeypatch.setattr(ps, "master_key", "sk-master-test")
+ monkeypatch.setattr(ps, "general_settings", {})
+ monkeypatch.setattr(ps, "premium_user", False)
+
+ response = client.get("/onboarding/get_token", params={"invite_link": "inv-123"})
+ assert response.status_code == 200
+ body = response.json()
+ assert set(body.keys()) == {"login_url", "token", "user_email"}
+ assert body["user_email"] == "alice@example.com"
+ assert "ui/onboarding" in body["login_url"]
+ assert "token=" in body["login_url"]
+ # The JWT in body["token"] must decode with the master_key.
+ decoded = jwt.decode(body["token"], "sk-master-test", algorithms=["HS256"])
+ assert normalize(
+ {
+ "user_id": decoded["user_id"],
+ "user_email": decoded["user_email"],
+ "login_method": decoded["login_method"],
+ "premium_user": decoded["premium_user"],
+ }
+ ) == {
+ "user_id": "user-abc",
+ "user_email": "alice@example.com",
+ "login_method": "username_password",
+ "premium_user": False,
+ }
+
+
+def test_onboarding_get_token_master_key_missing_500(client, monkeypatch, mock_prisma):
+ """No master_key configured → 500 with the master_key error payload."""
+ from litellm.proxy import proxy_server as ps
+
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+ monkeypatch.setattr(ps, "master_key", None)
+ monkeypatch.setattr(ps, "general_settings", {})
+
+ response = client.get("/onboarding/get_token", params={"invite_link": "inv-123"})
+ assert response.status_code == 500
+ body = response.json()
+ # ProxyException serializes to {"error": {"message": ..., "type": ..., "param": ..., "code": ...}}
+ err_blob = body.get("error", body)
+ assert "Master Key not set" in str(err_blob)
+
+
+def test_onboarding_get_token_invalid_invite_link_401(
+ client, monkeypatch, mock_prisma
+):
+ """Unknown invite link → 401 with the not-in-db error message."""
+ from litellm.proxy import proxy_server as ps
+
+ mock_prisma.db.litellm_invitationlink.find_unique.return_value = None
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+ monkeypatch.setattr(ps, "master_key", "sk-master-test")
+ monkeypatch.setattr(ps, "general_settings", {})
+
+ response = client.get(
+ "/onboarding/get_token", params={"invite_link": "does-not-exist"}
+ )
+ assert response.status_code == 401
+ assert response.json() == {
+ "detail": {"error": "Invitation link does not exist in db."}
+ }
+
+
+def test_onboarding_get_token_expired_invite_401(client, monkeypatch, mock_prisma):
+ """Invite whose expires_at is in the past → 401 expired."""
+ from litellm.proxy import proxy_server as ps
+
+ expired_invite = _make_invite(
+ expires_at=datetime.now(timezone.utc) - timedelta(days=2)
+ )
+ mock_prisma.db.litellm_invitationlink.find_unique.return_value = expired_invite
+
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+ monkeypatch.setattr(ps, "master_key", "sk-master-test")
+ monkeypatch.setattr(ps, "general_settings", {})
+
+ response = client.get("/onboarding/get_token", params={"invite_link": "inv-123"})
+ assert response.status_code == 401
+ assert response.json().get("detail", {}).get("error") == "Invitation link has expired."
+
+
+def test_onboarding_get_token_missing_query_param_422(client, monkeypatch, mock_prisma):
+ """No ``invite_link`` query param → FastAPI 422 with a non-empty detail array."""
+ from litellm.proxy import proxy_server as ps
+
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+ monkeypatch.setattr(ps, "master_key", "sk-master-test")
+ monkeypatch.setattr(ps, "general_settings", {})
+
+ response = client.get("/onboarding/get_token")
+ assert response.status_code == 422
+ body = response.json()
+ assert isinstance(body.get("detail"), list)
+ assert len(body["detail"]) >= 1
+
+
+# ---------------------------------------------------------------------------
+# POST /onboarding/claim_token
+# ---------------------------------------------------------------------------
+
+
+def _make_onboarding_jwt(
+ master_key: str,
+ invitation_link: str = "inv-123",
+ user_id: str = "user-abc",
+ token_type: str = "litellm_onboarding",
+) -> str:
+ return jwt.encode(
+ {
+ "token_type": token_type,
+ "invitation_link": invitation_link,
+ "user_id": user_id,
+ "exp": datetime.now(timezone.utc) + timedelta(minutes=15),
+ },
+ master_key,
+ algorithm="HS256",
+ )
+
+
+def test_claim_onboarding_link_happy(client, monkeypatch, mock_prisma):
+ """Valid claim → returns login_url, token, user_email, user."""
+ from litellm.proxy import proxy_server as ps
+
+ invite = _make_invite()
+ user_obj = _make_user_obj()
+ mock_prisma.db.litellm_invitationlink.find_unique.return_value = invite
+ mock_prisma.db.litellm_invitationlink.update_many.return_value = 1
+ mock_prisma.db.litellm_invitationlink.update.return_value = invite
+ mock_prisma.db.litellm_usertable.update.return_value = user_obj
+ _install_tx_context(mock_prisma)
+
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+ monkeypatch.setattr(ps, "master_key", "sk-master-test")
+ monkeypatch.setattr(ps, "general_settings", {})
+ monkeypatch.setattr(ps, "premium_user", False)
+
+ # Avoid hitting generate_key_helper_fn (touches DB / many globals); patch
+ # the helper directly so we focus on the route's own behavior.
+ async def _fake_session_token(user_obj):
+ return "session-jwt-token"
+
+ monkeypatch.setattr(
+ ps, "_generate_onboarding_ui_session_token", _fake_session_token
+ )
+
+ onboarding_jwt = _make_onboarding_jwt("sk-master-test")
+ response = client.post(
+ "/onboarding/claim_token",
+ json={
+ "invitation_link": "inv-123",
+ "user_id": "user-abc",
+ "password": "hunter2",
+ },
+ headers={"Authorization": f"Bearer {onboarding_jwt}"},
+ )
+ assert response.status_code == 200
+ body = response.json()
+ assert set(body.keys()) == {"login_url", "token", "user_email", "user"}
+ assert body["token"] == "session-jwt-token"
+ assert body["user_email"] == "alice@example.com"
+ assert body["login_url"].endswith("/ui/?login=success")
+
+
+def test_claim_onboarding_link_invalid_invite_401(client, monkeypatch, mock_prisma):
+ """Unknown invite link → 401 with not-in-db error."""
+ from litellm.proxy import proxy_server as ps
+
+ mock_prisma.db.litellm_invitationlink.find_unique.return_value = None
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+ monkeypatch.setattr(ps, "master_key", "sk-master-test")
+ monkeypatch.setattr(ps, "general_settings", {})
+
+ response = client.post(
+ "/onboarding/claim_token",
+ json={
+ "invitation_link": "missing",
+ "user_id": "user-abc",
+ "password": "hunter2",
+ },
+ headers={"Authorization": "Bearer irrelevant"},
+ )
+ assert response.status_code == 401
+ assert response.json() == {
+ "detail": {"error": "Invitation link does not exist in db."}
+ }
+
+
+def test_claim_onboarding_link_user_id_mismatch_401(
+ client, monkeypatch, mock_prisma
+):
+ """Invitation belongs to a different user_id → 401 with mismatch error."""
+ from litellm.proxy import proxy_server as ps
+
+ invite = _make_invite(user_id="user-real-owner")
+ mock_prisma.db.litellm_invitationlink.find_unique.return_value = invite
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+ monkeypatch.setattr(ps, "master_key", "sk-master-test")
+ monkeypatch.setattr(ps, "general_settings", {})
+
+ response = client.post(
+ "/onboarding/claim_token",
+ json={
+ "invitation_link": "inv-123",
+ "user_id": "user-attacker",
+ "password": "hunter2",
+ },
+ headers={"Authorization": "Bearer irrelevant"},
+ )
+ assert response.status_code == 401
+ err = response.json().get("detail", {}).get("error", "")
+ assert "Invalid invitation link" in err
+ assert "user-attacker" in err
+
+
+def test_claim_onboarding_link_missing_field_422(client, monkeypatch, mock_prisma):
+ """Missing required body field → FastAPI 422 with detail listing the missing field."""
+ from litellm.proxy import proxy_server as ps
+
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+ monkeypatch.setattr(ps, "master_key", "sk-master-test")
+ monkeypatch.setattr(ps, "general_settings", {})
+
+ # Missing "password"
+ response = client.post(
+ "/onboarding/claim_token",
+ json={"invitation_link": "inv-123", "user_id": "user-abc"},
+ )
+ assert response.status_code == 422
+ body = response.json()
+ assert isinstance(body.get("detail"), list)
+ # The missing field should be referenced in the validation error.
+ assert any("password" in str(item) for item in body["detail"])
+
+
+def test_claim_onboarding_link_bad_onboarding_jwt_401(
+ client, monkeypatch, mock_prisma
+):
+ """Onboarding JWT decodes but token_type / invitation_link don't match → 401."""
+ from litellm.proxy import proxy_server as ps
+
+ invite = _make_invite()
+ mock_prisma.db.litellm_invitationlink.find_unique.return_value = invite
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+ monkeypatch.setattr(ps, "master_key", "sk-master-test")
+ monkeypatch.setattr(ps, "general_settings", {})
+
+ # Wrong token_type — handler rejects.
+ bogus_jwt = _make_onboarding_jwt(
+ "sk-master-test",
+ token_type="not_onboarding",
+ )
+ response = client.post(
+ "/onboarding/claim_token",
+ json={
+ "invitation_link": "inv-123",
+ "user_id": "user-abc",
+ "password": "hunter2",
+ },
+ headers={"Authorization": f"Bearer {bogus_jwt}"},
+ )
+ assert response.status_code == 401
+ assert (
+ response.json().get("detail", {}).get("error")
+ == "Invalid onboarding session for invitation link."
+ )
diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_queue.py b/tests/test_litellm/proxy/proxy_server/test_routes_queue.py
index ad6b4016461..27cc6300711 100644
--- a/tests/test_litellm/proxy/proxy_server/test_routes_queue.py
+++ b/tests/test_litellm/proxy/proxy_server/test_routes_queue.py
@@ -1 +1,91 @@
-"""Placeholder. Filled by a follow-up PR per the Notion plan."""
+"""Behavior pins for ``proxy_server.py`` queue routes.
+
+Pins (PR2):
+ - POST /queue/chat/completions
+"""
+
+from __future__ import annotations
+
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from litellm.proxy import proxy_server
+
+from .conftest import normalize # type: ignore[import-not-found]
+
+HAPPY_RESPONSE = {
+ "id": "chatcmpl-queue",
+ "object": "chat.completion",
+ "created": 0,
+ "model": "gpt-4",
+ "choices": [
+ {
+ "index": 0,
+ "finish_reason": "stop",
+ "message": {"role": "assistant", "content": "queued reply"},
+ }
+ ],
+ "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
+ "priority": 0,
+}
+
+
+@pytest.fixture
+def patched_queue(monkeypatch):
+ router = MagicMock()
+ router.schedule_acompletion = AsyncMock(return_value=dict(HAPPY_RESPONSE))
+ monkeypatch.setattr(proxy_server, "llm_router", router)
+ monkeypatch.setattr(
+ proxy_server,
+ "proxy_logging_obj",
+ MagicMock(post_call_failure_hook=AsyncMock()),
+ )
+ return router
+
+
+@pytest.fixture
+def queue_no_router(monkeypatch):
+ monkeypatch.setattr(proxy_server, "llm_router", None)
+ monkeypatch.setattr(
+ proxy_server,
+ "proxy_logging_obj",
+ MagicMock(post_call_failure_hook=AsyncMock()),
+ )
+ yield
+
+
+def test_queue_chat_completions_happy(client, auth_as, patched_queue):
+ """Pins ``POST /queue/chat/completions`` (happy)."""
+ payload = {
+ "model": "gpt-4",
+ "messages": [{"role": "user", "content": "hi"}],
+ "priority": 0,
+ }
+ with auth_as():
+ response = client.post("/queue/chat/completions", json=payload)
+ assert response.status_code == 200
+ assert normalize(response.json()) == {
+ "id": "",
+ "object": "chat.completion",
+ "created": "",
+ "model": "gpt-4",
+ "choices": [
+ {
+ "index": 0,
+ "finish_reason": "stop",
+ "message": {"role": "assistant", "content": "queued reply"},
+ }
+ ],
+ "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
+ "priority": 0,
+ }
+
+
+def test_queue_chat_completions_no_router_error(client, auth_as, queue_no_router):
+ """Pins ``POST /queue/chat/completions`` (error: no llm_router)."""
+ payload = {"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]}
+ with auth_as():
+ response = client.post("/queue/chat/completions", json=payload)
+ assert response.status_code == 500
+ assert len(response.content) > 0
diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_threads.py b/tests/test_litellm/proxy/proxy_server/test_routes_threads.py
index ad6b4016461..493315f041d 100644
--- a/tests/test_litellm/proxy/proxy_server/test_routes_threads.py
+++ b/tests/test_litellm/proxy/proxy_server/test_routes_threads.py
@@ -1 +1,274 @@
-"""Placeholder. Filled by a follow-up PR per the Notion plan."""
+"""Behavior pins for ``proxy_server.py`` threads routes.
+
+Pins (PR2):
+ - POST /v1/threads
+ - POST /threads
+ - GET /v1/threads/{thread_id}
+ - GET /threads/{thread_id}
+ - POST /v1/threads/{thread_id}/messages
+ - POST /threads/{thread_id}/messages
+ - GET /v1/threads/{thread_id}/messages
+ - GET /threads/{thread_id}/messages
+ - POST /v1/threads/{thread_id}/runs
+ - POST /threads/{thread_id}/runs
+"""
+
+from __future__ import annotations
+
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from litellm.proxy import proxy_server
+
+from .conftest import normalize # type: ignore[import-not-found]
+
+CREATE_THREAD = {"id": "thr_1", "object": "thread", "created_at": 0, "metadata": {}}
+GET_THREAD = {
+ "id": "thr_1",
+ "object": "thread",
+ "created_at": 0,
+ "tool_resources": {},
+}
+ADD_MESSAGE = {
+ "id": "msg_1",
+ "object": "thread.message",
+ "thread_id": "thr_1",
+ "role": "user",
+ "content": [],
+}
+GET_MESSAGES = {
+ "object": "list",
+ "data": [
+ {
+ "id": "msg_1",
+ "object": "thread.message",
+ "thread_id": "thr_1",
+ "role": "user",
+ "content": [],
+ }
+ ],
+ "first_id": "msg_1",
+ "last_id": "msg_1",
+ "has_more": False,
+}
+RUN_THREAD = {
+ "id": "run_1",
+ "object": "thread.run",
+ "thread_id": "thr_1",
+ "assistant_id": "asst_1",
+ "status": "queued",
+}
+
+
+@pytest.fixture
+def patched_threads(monkeypatch):
+ router = MagicMock()
+ router.acreate_thread = AsyncMock(return_value=dict(CREATE_THREAD))
+ router.aget_thread = AsyncMock(return_value=dict(GET_THREAD))
+ router.a_add_message = AsyncMock(return_value=dict(ADD_MESSAGE))
+ router.aget_messages = AsyncMock(return_value=dict(GET_MESSAGES))
+ router.arun_thread = AsyncMock(return_value=dict(RUN_THREAD))
+ monkeypatch.setattr(proxy_server, "llm_router", router)
+ monkeypatch.setattr(
+ proxy_server,
+ "proxy_logging_obj",
+ MagicMock(
+ post_call_failure_hook=AsyncMock(), update_request_status=AsyncMock()
+ ),
+ )
+
+ async def _add_data(data, **kwargs):
+ return data
+
+ monkeypatch.setattr(proxy_server, "add_litellm_data_to_request", _add_data)
+ return router
+
+
+@pytest.fixture
+def no_router(monkeypatch):
+ monkeypatch.setattr(proxy_server, "llm_router", None)
+ monkeypatch.setattr(
+ proxy_server,
+ "proxy_logging_obj",
+ MagicMock(
+ post_call_failure_hook=AsyncMock(), update_request_status=AsyncMock()
+ ),
+ )
+
+ async def _add_data(data, **kwargs):
+ return data
+
+ monkeypatch.setattr(proxy_server, "add_litellm_data_to_request", _add_data)
+ yield
+
+
+# ---------------------------------------------------------------------------
+# POST /v1/threads, POST /threads
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("path", ["/v1/threads", "/threads"])
+def test_create_thread_happy(client, auth_as, patched_threads, path):
+ """Pins ``POST /v1/threads`` and ``POST /threads``."""
+ with auth_as():
+ response = client.post(path, json={})
+ assert response.status_code == 200
+ assert normalize(response.json()) == {
+ "id": "",
+ "object": "thread",
+ "created_at": "",
+ "metadata": {},
+ }
+
+
+@pytest.mark.parametrize("path", ["/v1/threads", "/threads"])
+def test_create_thread_error(client, auth_as, no_router, path):
+ """Pins ``POST /v1/threads`` / ``POST /threads`` (error: no llm_router)."""
+ with auth_as():
+ response = client.post(path, json={})
+ assert response.status_code == 500
+ assert len(response.content) > 0
+
+
+# ---------------------------------------------------------------------------
+# GET /v1/threads/{thread_id}, GET /threads/{thread_id}
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("path", ["/v1/threads/thr_1", "/threads/thr_1"])
+def test_get_thread_happy(client, auth_as, patched_threads, path):
+ """Pins ``GET /v1/threads/{thread_id}`` and ``GET /threads/{thread_id}``."""
+ with auth_as():
+ response = client.get(path)
+ assert response.status_code == 200
+ assert normalize(response.json()) == {
+ "id": "",
+ "object": "thread",
+ "created_at": "",
+ "tool_resources": {},
+ }
+
+
+@pytest.mark.parametrize("path", ["/v1/threads/thr_1", "/threads/thr_1"])
+def test_get_thread_error(client, auth_as, no_router, path):
+ """Pins ``GET /v1/threads/{thread_id}`` / ``GET /threads/{thread_id}`` (error)."""
+ with auth_as():
+ response = client.get(path)
+ assert response.status_code == 500
+ assert len(response.content) > 0
+
+
+# ---------------------------------------------------------------------------
+# POST /v1/threads/{thread_id}/messages, POST /threads/{thread_id}/messages
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+ "path",
+ ["/v1/threads/thr_1/messages", "/threads/thr_1/messages"],
+)
+def test_add_message_happy(client, auth_as, patched_threads, path):
+ """Pins ``POST /v1/threads/{thread_id}/messages`` and ``POST /threads/{thread_id}/messages``."""
+ payload = {"role": "user", "content": "hi"}
+ with auth_as():
+ response = client.post(path, json=payload)
+ assert response.status_code == 200
+ assert normalize(response.json()) == {
+ "id": "",
+ "object": "thread.message",
+ "thread_id": "thr_1",
+ "role": "user",
+ "content": [],
+ }
+
+
+@pytest.mark.parametrize(
+ "path",
+ ["/v1/threads/thr_1/messages", "/threads/thr_1/messages"],
+)
+def test_add_message_error(client, auth_as, no_router, path):
+ """Pins ``POST /v1/threads/{thread_id}/messages`` / ``POST /threads/{thread_id}/messages`` (error)."""
+ with auth_as():
+ response = client.post(path, json={"role": "user", "content": "hi"})
+ assert response.status_code == 500
+ assert len(response.content) > 0
+
+
+# ---------------------------------------------------------------------------
+# GET /v1/threads/{thread_id}/messages, GET /threads/{thread_id}/messages
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+ "path",
+ ["/v1/threads/thr_1/messages", "/threads/thr_1/messages"],
+)
+def test_get_messages_happy(client, auth_as, patched_threads, path):
+ """Pins ``GET /v1/threads/{thread_id}/messages`` and ``GET /threads/{thread_id}/messages``."""
+ with auth_as():
+ response = client.get(path)
+ assert response.status_code == 200
+ assert normalize(response.json()) == {
+ "object": "list",
+ "data": [
+ {
+ "id": "",
+ "object": "thread.message",
+ "thread_id": "thr_1",
+ "role": "user",
+ "content": [],
+ }
+ ],
+ "first_id": "msg_1",
+ "last_id": "msg_1",
+ "has_more": False,
+ }
+
+
+@pytest.mark.parametrize(
+ "path",
+ ["/v1/threads/thr_1/messages", "/threads/thr_1/messages"],
+)
+def test_get_messages_error(client, auth_as, no_router, path):
+ """Pins ``GET /v1/threads/{thread_id}/messages`` / ``GET /threads/{thread_id}/messages`` (error)."""
+ with auth_as():
+ response = client.get(path)
+ assert response.status_code == 500
+ assert len(response.content) > 0
+
+
+# ---------------------------------------------------------------------------
+# POST /v1/threads/{thread_id}/runs, POST /threads/{thread_id}/runs
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+ "path",
+ ["/v1/threads/thr_1/runs", "/threads/thr_1/runs"],
+)
+def test_run_thread_happy(client, auth_as, patched_threads, path):
+ """Pins ``POST /v1/threads/{thread_id}/runs`` and ``POST /threads/{thread_id}/runs``."""
+ payload = {"assistant_id": "asst_1"}
+ with auth_as():
+ response = client.post(path, json=payload)
+ assert response.status_code == 200
+ assert normalize(response.json()) == {
+ "id": "",
+ "object": "thread.run",
+ "thread_id": "thr_1",
+ "assistant_id": "asst_1",
+ "status": "queued",
+ }
+
+
+@pytest.mark.parametrize(
+ "path",
+ ["/v1/threads/thr_1/runs", "/threads/thr_1/runs"],
+)
+def test_run_thread_error(client, auth_as, no_router, path):
+ """Pins ``POST /v1/threads/{thread_id}/runs`` / ``POST /threads/{thread_id}/runs`` (error)."""
+ with auth_as():
+ response = client.post(path, json={"assistant_id": "asst_1"})
+ assert response.status_code == 500
+ assert len(response.content) > 0
diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_utils.py b/tests/test_litellm/proxy/proxy_server/test_routes_utils.py
index ad6b4016461..c6070437d35 100644
--- a/tests/test_litellm/proxy/proxy_server/test_routes_utils.py
+++ b/tests/test_litellm/proxy/proxy_server/test_routes_utils.py
@@ -1 +1,160 @@
-"""Placeholder. Filled by a follow-up PR per the Notion plan."""
+"""Behavior pins for ``proxy_server.py`` llm-utils routes.
+
+Pins (PR2):
+ - POST /utils/token_counter
+ - GET /utils/supported_openai_params
+ - POST /utils/transform_request
+"""
+
+from __future__ import annotations
+
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+import litellm
+from litellm.proxy import proxy_server
+
+from .conftest import normalize # type: ignore[import-not-found]
+
+# ---------------------------------------------------------------------------
+# POST /utils/token_counter
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def patched_token_counter(monkeypatch):
+ monkeypatch.setattr(proxy_server, "llm_router", None)
+ monkeypatch.setattr(litellm, "disable_token_counter", False, raising=False)
+ monkeypatch.setattr(
+ litellm.utils,
+ "_select_tokenizer",
+ lambda model, custom_tokenizer=None: {
+ "type": "openai_tokenizer",
+ "tokenizer": None,
+ },
+ )
+ monkeypatch.setattr(litellm, "token_counter", lambda **kwargs: 7)
+ yield
+
+
+def test_token_counter_happy_path(client, auth_as, patched_token_counter):
+ """Pins ``POST /utils/token_counter``."""
+ payload = {"model": "gpt-4", "prompt": "Hi there"}
+ with auth_as():
+ response = client.post("/utils/token_counter", json=payload)
+ assert response.status_code == 200
+ assert normalize(response.json()) == {
+ "total_tokens": 7,
+ "request_model": "gpt-4",
+ "model_used": "gpt-4",
+ "tokenizer_type": "openai_tokenizer",
+ "original_response": None,
+ "error": False,
+ "error_message": None,
+ "status_code": None,
+ }
+
+
+def test_token_counter_missing_input_returns_400(
+ client, auth_as, patched_token_counter
+):
+ """Pins ``POST /utils/token_counter`` (error: missing input)."""
+ with auth_as():
+ response = client.post("/utils/token_counter", json={"model": "gpt-4"})
+ assert response.status_code == 400
+ assert "prompt or messages or contents" in response.text
+
+
+# ---------------------------------------------------------------------------
+# GET /utils/supported_openai_params
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def patched_supported_params(monkeypatch):
+ monkeypatch.setattr(
+ litellm,
+ "get_llm_provider",
+ lambda model: (model, "openai", None, None),
+ )
+ monkeypatch.setattr(
+ litellm,
+ "get_supported_openai_params",
+ lambda model, custom_llm_provider=None: ["max_tokens", "temperature", "top_p"],
+ )
+ yield
+
+
+def test_supported_openai_params_happy_path(client, auth_as, patched_supported_params):
+ """Pins ``GET /utils/supported_openai_params``."""
+ with auth_as():
+ response = client.get(
+ "/utils/supported_openai_params", params={"model": "gpt-4"}
+ )
+ assert response.status_code == 200
+ assert normalize(response.json()) == {
+ "supported_openai_params": ["max_tokens", "temperature", "top_p"],
+ }
+
+
+def test_supported_openai_params_invalid_model(client, auth_as, monkeypatch):
+ """Pins ``GET /utils/supported_openai_params`` (error: unknown model)."""
+
+ def _raise(model):
+ raise Exception("unknown")
+
+ monkeypatch.setattr(litellm, "get_llm_provider", _raise)
+ with auth_as():
+ response = client.get("/utils/supported_openai_params", params={"model": "??"})
+ assert response.status_code == 400
+ assert "Could not map model" in response.text
+
+
+# ---------------------------------------------------------------------------
+# POST /utils/transform_request
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def patched_transform(monkeypatch):
+ monkeypatch.setattr(proxy_server, "llm_router", None)
+ monkeypatch.setattr(proxy_server, "is_request_body_safe", lambda **kwargs: True)
+
+ def _fake_return_raw_request(endpoint, kwargs):
+ return {
+ "raw_request_api_base": "https://api.openai.com/v1/chat/completions",
+ "raw_request_body": kwargs,
+ "raw_request_headers": {"Authorization": "Bearer redacted"},
+ }
+
+ monkeypatch.setattr("litellm.utils.return_raw_request", _fake_return_raw_request)
+ yield
+
+
+def test_transform_request_happy_path(client, auth_as, patched_transform):
+ """Pins ``POST /utils/transform_request``."""
+ payload = {"call_type": "completion", "request_body": {"model": "gpt-4"}}
+ with auth_as():
+ response = client.post("/utils/transform_request", json=payload)
+ assert response.status_code == 200
+ assert normalize(response.json()) == {
+ "raw_request_api_base": "https://api.openai.com/v1/chat/completions",
+ "raw_request_body": {"model": "gpt-4"},
+ "raw_request_headers": {"Authorization": "Bearer redacted"},
+ }
+
+
+def test_transform_request_unsafe_body(client, auth_as, monkeypatch):
+ """Pins ``POST /utils/transform_request`` (error: unsafe body)."""
+ monkeypatch.setattr(proxy_server, "llm_router", None)
+
+ def _raise(**kwargs):
+ raise ValueError("unsafe model")
+
+ monkeypatch.setattr(proxy_server, "is_request_body_safe", _raise)
+ payload = {"call_type": "completion", "request_body": {"model": "evil"}}
+ with auth_as():
+ response = client.post("/utils/transform_request", json=payload)
+ assert response.status_code == 400
+ assert "unsafe" in response.text or "error" in response.text
diff --git a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py
index ad6b4016461..ec8b06d9c97 100644
--- a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py
+++ b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py
@@ -1 +1,817 @@
-"""Placeholder. Filled by a follow-up PR per the Notion plan."""
+"""Behavior pins for spend-counter helpers in proxy_server.
+
+Pins covered:
+- ``get_current_spend``
+- ``increment_spend_counters``
+- ``_reconcile_budget_reservation_for_counter_update``
+- ``_increment_end_user_and_tag_spend_counters``
+- ``_increment_org_spend_counter``
+- ``_init_and_increment_unreserved_spend_counter``
+- ``_init_and_increment_spend_counter``
+- ``_init_and_increment_window_spend_counter``
+- ``_ensure_spend_counter_initialized``
+- ``_get_source_cache_base_spend``
+- ``_ensure_window_spend_counter_initialized``
+- ``_is_spend_counter_cache_warm``
+- ``_increment_spend_counter_cache``
+- ``_invalidate_spend_counter``
+- ``update_cache``
+"""
+
+from __future__ import annotations
+
+from datetime import datetime
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+import litellm.proxy.proxy_server as ps
+
+from .conftest import normalize
+
+
+def _make_spend_counter_cache(
+ *,
+ redis_get_value=None,
+ redis_get_side_effect=None,
+ redis_increment_value=None,
+ redis_increment_side_effect=None,
+ in_memory_value=None,
+ with_redis: bool = True,
+):
+ cache = MagicMock()
+ cache.in_memory_cache = MagicMock()
+ cache.in_memory_cache.get_cache = MagicMock(return_value=in_memory_value)
+ cache.in_memory_cache.set_cache = MagicMock()
+ cache.in_memory_cache.delete_cache = MagicMock()
+ if with_redis:
+ cache.redis_cache = MagicMock()
+ cache.redis_cache.async_get_cache = AsyncMock(
+ return_value=redis_get_value, side_effect=redis_get_side_effect
+ )
+ cache.redis_cache.async_increment = AsyncMock(
+ return_value=redis_increment_value,
+ side_effect=redis_increment_side_effect,
+ )
+ cache.redis_cache.async_delete_cache = AsyncMock()
+ else:
+ cache.redis_cache = None
+ cache.async_increment_cache = AsyncMock(return_value=redis_increment_value)
+ cache.async_get_cache = AsyncMock(return_value=None)
+ cache.async_set_cache = AsyncMock()
+ cache.async_delete_cache = AsyncMock()
+ cache.async_set_cache_pipeline = AsyncMock()
+ return cache
+
+
+def _make_user_api_key_cache(get_value=None, get_side_effect=None):
+ cache = MagicMock()
+ cache.async_get_cache = AsyncMock(
+ return_value=get_value, side_effect=get_side_effect
+ )
+ cache.async_set_cache_pipeline = AsyncMock()
+ return cache
+
+
+# ---------------------------------------------------------------------------
+# get_current_spend
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_get_current_spend_reads_redis_first(monkeypatch):
+ fake_cache = _make_spend_counter_cache(redis_get_value=42.5)
+ monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
+
+ result = await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=0.0)
+
+ observed = {
+ "value": result,
+ "redis_called": fake_cache.redis_cache.async_get_cache.called,
+ "in_memory_called": fake_cache.in_memory_cache.get_cache.called,
+ }
+ assert normalize(observed) == {
+ "value": 42.5,
+ "redis_called": True,
+ "in_memory_called": False,
+ }
+
+
+@pytest.mark.asyncio
+async def test_get_current_spend_redis_error_falls_back_to_in_memory(monkeypatch):
+ fake_cache = _make_spend_counter_cache(
+ redis_get_side_effect=RuntimeError("redis down"),
+ in_memory_value=17.0,
+ )
+ monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
+
+ result = await ps.get_current_spend(
+ counter_key="spend:key:abc", fallback_spend=99.0
+ )
+ assert result == 17.0
+
+
+# ---------------------------------------------------------------------------
+# increment_spend_counters
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_increment_spend_counters_increments_all_buckets(monkeypatch):
+ fake_cache = _make_spend_counter_cache(
+ redis_get_value=None, redis_increment_value=5.0
+ )
+ fake_user_cache = _make_user_api_key_cache(get_value=None)
+ monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
+ monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache)
+ monkeypatch.setattr(ps, "prisma_client", None)
+
+ async def _fake_coalesced(**kwargs):
+ return None
+
+ monkeypatch.setattr(
+ ps.SpendCounterReseed, "coalesced", AsyncMock(side_effect=_fake_coalesced)
+ )
+
+ await ps.increment_spend_counters(
+ token="hashed-tok",
+ team_id="t1",
+ user_id="u1",
+ response_cost=5.0,
+ )
+
+ observed = {
+ "redis_increment_called": fake_cache.redis_cache.async_increment.called,
+ "increment_calls": fake_cache.redis_cache.async_increment.call_count,
+ "user_cache_used": fake_user_cache.async_get_cache.called,
+ }
+ assert normalize(observed) == {
+ "redis_increment_called": True,
+ "increment_calls": 4,
+ "user_cache_used": True,
+ }
+
+
+@pytest.mark.asyncio
+async def test_increment_spend_counters_zero_cost_is_noop_finalizes_reservation(
+ monkeypatch,
+):
+ fake_cache = _make_spend_counter_cache()
+ fake_user_cache = _make_user_api_key_cache()
+ monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
+ monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache)
+ monkeypatch.setattr(ps, "prisma_client", None)
+ reservation = {"finalized": False}
+
+ await ps.increment_spend_counters(
+ token="hashed-tok",
+ team_id="t1",
+ user_id="u1",
+ response_cost=0,
+ budget_reservation=reservation,
+ )
+
+ assert reservation == {"finalized": True}
+ assert fake_cache.redis_cache.async_increment.called is False
+
+
+# ---------------------------------------------------------------------------
+# _reconcile_budget_reservation_for_counter_update
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_reconcile_budget_reservation_for_counter_update_returns_empty_set_when_none():
+ result = await ps._reconcile_budget_reservation_for_counter_update(
+ budget_reservation=None, response_cost=1.0
+ )
+ assert result == set()
+
+
+@pytest.mark.asyncio
+async def test_reconcile_budget_reservation_for_counter_update_failure_invalidates(
+ monkeypatch,
+):
+ """Reservation reconcile raising must invalidate reserved counters but
+ not propagate the exception."""
+ import litellm.proxy.spend_tracking.budget_reservation as br
+
+ monkeypatch.setattr(
+ br,
+ "get_reserved_counter_keys",
+ MagicMock(return_value={"spend:key:abc"}),
+ )
+ monkeypatch.setattr(
+ br,
+ "reconcile_budget_reservation",
+ AsyncMock(side_effect=RuntimeError("boom")),
+ )
+ fake_invalidate = AsyncMock()
+ monkeypatch.setattr(br, "invalidate_budget_reservation_counters", fake_invalidate)
+
+ result = await ps._reconcile_budget_reservation_for_counter_update(
+ budget_reservation={"foo": "bar"}, response_cost=1.0
+ )
+
+ assert result == {"spend:key:abc"}
+ assert fake_invalidate.called is True
+
+
+# ---------------------------------------------------------------------------
+# _increment_end_user_and_tag_spend_counters
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_increment_end_user_and_tag_spend_counters_increments_each_unique_tag(
+ monkeypatch,
+):
+ fake_cache = _make_spend_counter_cache(
+ redis_get_value=None, redis_increment_value=3.0
+ )
+ fake_user_cache = _make_user_api_key_cache()
+ monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
+ monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache)
+ monkeypatch.setattr(ps, "prisma_client", None)
+ monkeypatch.setattr(
+ ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)
+ )
+
+ await ps._increment_end_user_and_tag_spend_counters(
+ end_user_id="eu1",
+ tags=["a", "b", "a", "", None],
+ response_cost=3.0,
+ reserved_counter_keys=set(),
+ )
+
+ observed = {
+ "increment_calls": fake_cache.redis_cache.async_increment.call_count,
+ "in_memory_set_calls": fake_cache.in_memory_cache.set_cache.call_count,
+ "called": fake_cache.redis_cache.async_increment.called,
+ }
+ assert normalize(observed) == {
+ "increment_calls": 3,
+ "in_memory_set_calls": 3,
+ "called": True,
+ }
+
+
+@pytest.mark.asyncio
+async def test_increment_end_user_and_tag_spend_counters_no_end_user_no_tags_invalid_input_noop(
+ monkeypatch,
+):
+ fake_cache = _make_spend_counter_cache()
+ monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
+
+ await ps._increment_end_user_and_tag_spend_counters(
+ end_user_id=None,
+ tags=None,
+ response_cost=1.0,
+ reserved_counter_keys=set(),
+ )
+
+ assert fake_cache.redis_cache.async_increment.called is False
+
+
+# ---------------------------------------------------------------------------
+# _increment_org_spend_counter
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_increment_org_spend_counter_increments_when_org_present(monkeypatch):
+ fake_cache = _make_spend_counter_cache(
+ redis_get_value=None, redis_increment_value=10.0
+ )
+ fake_user_cache = _make_user_api_key_cache()
+ monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
+ monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache)
+ monkeypatch.setattr(ps, "prisma_client", None)
+ monkeypatch.setattr(
+ ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)
+ )
+
+ await ps._increment_org_spend_counter(
+ org_id="org-1",
+ response_cost=10.0,
+ reserved_counter_keys=set(),
+ )
+
+ observed = {
+ "increment_called": fake_cache.redis_cache.async_increment.called,
+ "increment_calls": fake_cache.redis_cache.async_increment.call_count,
+ "counter_key_arg": fake_cache.redis_cache.async_increment.call_args.kwargs[
+ "key"
+ ],
+ }
+ assert normalize(observed) == {
+ "increment_called": True,
+ "increment_calls": 1,
+ "counter_key_arg": "spend:org:org-1",
+ }
+
+
+@pytest.mark.asyncio
+async def test_increment_org_spend_counter_no_org_is_noop_invalid_id(monkeypatch):
+ fake_cache = _make_spend_counter_cache()
+ monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
+
+ await ps._increment_org_spend_counter(
+ org_id=None,
+ response_cost=1.0,
+ reserved_counter_keys=set(),
+ )
+
+ assert fake_cache.redis_cache.async_increment.called is False
+
+
+# ---------------------------------------------------------------------------
+# _init_and_increment_unreserved_spend_counter
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_init_and_increment_unreserved_spend_counter_skips_reserved_keys(
+ monkeypatch,
+):
+ fake_cache = _make_spend_counter_cache()
+ monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
+
+ await ps._init_and_increment_unreserved_spend_counter(
+ counter_key="spend:tag:x",
+ source_cache_key="tag:x",
+ increment=1.0,
+ reserved_counter_keys={"spend:tag:x"},
+ )
+
+ assert fake_cache.redis_cache.async_increment.called is False
+
+
+@pytest.mark.asyncio
+async def test_init_and_increment_unreserved_spend_counter_proceeds_when_not_reserved(
+ monkeypatch,
+):
+ fake_cache = _make_spend_counter_cache(
+ redis_get_value=None, redis_increment_value=2.0
+ )
+ fake_user_cache = _make_user_api_key_cache()
+ monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
+ monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache)
+ monkeypatch.setattr(ps, "prisma_client", None)
+ monkeypatch.setattr(
+ ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)
+ )
+
+ await ps._init_and_increment_unreserved_spend_counter(
+ counter_key="spend:tag:y",
+ source_cache_key="tag:y",
+ increment=2.0,
+ reserved_counter_keys=set(),
+ )
+
+ observed = {
+ "increment_called": fake_cache.redis_cache.async_increment.called,
+ "redis_get_called": fake_cache.redis_cache.async_get_cache.called,
+ "reseed_consulted": True,
+ }
+ assert observed == {
+ "increment_called": True,
+ "redis_get_called": True,
+ "reseed_consulted": True,
+ }
+
+
+# ---------------------------------------------------------------------------
+# _init_and_increment_spend_counter
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_init_and_increment_spend_counter_warm_cache_skips_reseed(monkeypatch):
+ fake_cache = _make_spend_counter_cache(
+ redis_get_value=11.0, redis_increment_value=14.0
+ )
+ fake_user_cache = _make_user_api_key_cache()
+ monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
+ monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache)
+ monkeypatch.setattr(ps, "prisma_client", None)
+ reseed = AsyncMock(return_value=None)
+ monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", reseed)
+
+ await ps._init_and_increment_spend_counter(
+ counter_key="spend:key:k",
+ source_cache_key="k",
+ increment=3.0,
+ )
+
+ observed = {
+ "reseed_called": reseed.called,
+ "increment_called": fake_cache.redis_cache.async_increment.called,
+ "in_memory_seeded_from_redis": fake_cache.in_memory_cache.set_cache.called,
+ }
+ assert normalize(observed) == {
+ "reseed_called": False,
+ "increment_called": True,
+ "in_memory_seeded_from_redis": True,
+ }
+
+
+# ---------------------------------------------------------------------------
+# _init_and_increment_window_spend_counter
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_init_and_increment_window_spend_counter_increments_when_initialized(
+ monkeypatch,
+):
+ fake_cache = _make_spend_counter_cache(
+ redis_get_value=0.0, redis_increment_value=5.0
+ )
+ monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
+ monkeypatch.setattr(ps, "prisma_client", None)
+ monkeypatch.setattr(
+ ps.SpendCounterReseed,
+ "coalesced_window",
+ AsyncMock(return_value=0.0),
+ )
+
+ await ps._init_and_increment_window_spend_counter(
+ counter_key="spend:key:k:window:1d",
+ entity_type="Key",
+ entity_id="k",
+ window_start=datetime(2024, 1, 1),
+ increment=5.0,
+ )
+
+ observed = {
+ "redis_increment_called": fake_cache.redis_cache.async_increment.called,
+ "increment_calls": fake_cache.redis_cache.async_increment.call_count,
+ "in_memory_set_calls": fake_cache.in_memory_cache.set_cache.call_count,
+ }
+ assert normalize(observed) == {
+ "redis_increment_called": True,
+ "increment_calls": 1,
+ "in_memory_set_calls": 2,
+ }
+
+
+@pytest.mark.asyncio
+async def test_init_and_increment_window_spend_counter_missing_window_start_invalid_skips(
+ monkeypatch,
+):
+ fake_cache = _make_spend_counter_cache()
+ monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
+
+ await ps._init_and_increment_window_spend_counter(
+ counter_key="spend:key:k:window:1d",
+ entity_type="Key",
+ entity_id="k",
+ window_start=None,
+ increment=5.0,
+ )
+
+ assert fake_cache.redis_cache.async_increment.called is False
+
+
+# ---------------------------------------------------------------------------
+# _ensure_spend_counter_initialized
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_ensure_spend_counter_initialized_warm_skips_reseed_and_source(
+ monkeypatch,
+):
+ fake_cache = _make_spend_counter_cache(redis_get_value=20.0)
+ fake_user_cache = _make_user_api_key_cache()
+ monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
+ monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache)
+ monkeypatch.setattr(ps, "prisma_client", None)
+ reseed = AsyncMock(return_value=None)
+ monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", reseed)
+
+ await ps._ensure_spend_counter_initialized(
+ counter_key="spend:user:u",
+ source_cache_key="u",
+ )
+
+ observed = {
+ "warm_check_redis": fake_cache.redis_cache.async_get_cache.called,
+ "reseed_called": reseed.called,
+ "source_cache_called": fake_user_cache.async_get_cache.called,
+ }
+ assert normalize(observed) == {
+ "warm_check_redis": True,
+ "reseed_called": False,
+ "source_cache_called": False,
+ }
+
+
+@pytest.mark.asyncio
+async def test_ensure_spend_counter_initialized_cold_seeds_from_source_cache(
+ monkeypatch,
+):
+ fake_cache = _make_spend_counter_cache(
+ redis_get_value=None, redis_increment_value=7.0
+ )
+ fake_user_cache = _make_user_api_key_cache(get_value={"spend": 7.0})
+ monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
+ monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache)
+ monkeypatch.setattr(ps, "prisma_client", None)
+ monkeypatch.setattr(
+ ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)
+ )
+
+ await ps._ensure_spend_counter_initialized(
+ counter_key="spend:user:u",
+ source_cache_key="u",
+ )
+
+ observed = {
+ "source_cache_called": fake_user_cache.async_get_cache.called,
+ "seed_increment_called": fake_cache.redis_cache.async_increment.called,
+ "warm_check_done": fake_cache.redis_cache.async_get_cache.called,
+ }
+ assert normalize(observed) == {
+ "source_cache_called": True,
+ "seed_increment_called": True,
+ "warm_check_done": True,
+ }
+
+
+# ---------------------------------------------------------------------------
+# _get_source_cache_base_spend
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_get_source_cache_base_spend_reads_first_hit_from_list(monkeypatch):
+ fake_user_cache = MagicMock()
+
+ async def _get(key, **kwargs):
+ if key == "miss":
+ return None
+ if key == "hit-obj":
+ obj = MagicMock()
+ obj.spend = 12.0
+ return obj
+ return None
+
+ fake_user_cache.async_get_cache = AsyncMock(side_effect=_get)
+ monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache)
+
+ result = await ps._get_source_cache_base_spend(
+ source_cache_key=["miss", "hit-obj", "miss2"]
+ )
+
+ observed = {
+ "result": result,
+ "calls": fake_user_cache.async_get_cache.call_count,
+ "stopped_after_hit": fake_user_cache.async_get_cache.call_count == 2,
+ }
+ assert normalize(observed) == {
+ "result": 12.0,
+ "calls": 2,
+ "stopped_after_hit": True,
+ }
+
+
+@pytest.mark.asyncio
+async def test_get_source_cache_base_spend_no_hits_returns_zero_fallback(monkeypatch):
+ """All cache lookups miss — function falls back to 0.0 (no error)."""
+ fake_user_cache = _make_user_api_key_cache(get_value=None)
+ monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache)
+
+ result = await ps._get_source_cache_base_spend(source_cache_key="missing-key")
+ assert result == 0.0
+
+
+# ---------------------------------------------------------------------------
+# _ensure_window_spend_counter_initialized
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_ensure_window_spend_counter_initialized_warm_returns_true(monkeypatch):
+ fake_cache = _make_spend_counter_cache(redis_get_value=3.0)
+ monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
+ monkeypatch.setattr(ps, "prisma_client", None)
+ window_reseed = AsyncMock(return_value=0.0)
+ monkeypatch.setattr(ps.SpendCounterReseed, "coalesced_window", window_reseed)
+
+ initialized = await ps._ensure_window_spend_counter_initialized(
+ counter_key="spend:key:k:window:1d",
+ entity_type="Key",
+ entity_id="k",
+ window_start=datetime(2024, 1, 1),
+ )
+
+ observed = {
+ "initialized": initialized,
+ "reseed_called": window_reseed.called,
+ "redis_get_called": fake_cache.redis_cache.async_get_cache.called,
+ }
+ assert normalize(observed) == {
+ "initialized": True,
+ "reseed_called": False,
+ "redis_get_called": True,
+ }
+
+
+@pytest.mark.asyncio
+async def test_ensure_window_spend_counter_initialized_db_failure_invalid_returns_false(
+ monkeypatch,
+):
+ fake_cache = _make_spend_counter_cache(redis_get_value=None)
+ monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
+ monkeypatch.setattr(ps, "prisma_client", None)
+ monkeypatch.setattr(
+ ps.SpendCounterReseed,
+ "coalesced_window",
+ AsyncMock(return_value=None),
+ )
+
+ initialized = await ps._ensure_window_spend_counter_initialized(
+ counter_key="spend:key:k:window:1d",
+ entity_type="Key",
+ entity_id="k",
+ window_start=datetime(2024, 1, 1),
+ )
+
+ assert initialized is False
+
+
+# ---------------------------------------------------------------------------
+# _is_spend_counter_cache_warm
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_is_spend_counter_cache_warm_redis_hit_seeds_in_memory(monkeypatch):
+ fake_cache = _make_spend_counter_cache(redis_get_value=99.0)
+ monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
+
+ result = await ps._is_spend_counter_cache_warm(counter_key="spend:user:u")
+
+ observed = {
+ "result": result,
+ "redis_get_called": fake_cache.redis_cache.async_get_cache.called,
+ "in_memory_set_called": fake_cache.in_memory_cache.set_cache.called,
+ }
+ assert normalize(observed) == {
+ "result": True,
+ "redis_get_called": True,
+ "in_memory_set_called": True,
+ }
+
+
+@pytest.mark.asyncio
+async def test_is_spend_counter_cache_warm_redis_error_falls_back_to_in_memory(
+ monkeypatch,
+):
+ fake_cache = _make_spend_counter_cache(
+ redis_get_side_effect=RuntimeError("redis err"),
+ in_memory_value=None,
+ )
+ monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
+
+ result = await ps._is_spend_counter_cache_warm(counter_key="spend:user:u")
+ assert result is False
+
+
+# ---------------------------------------------------------------------------
+# _increment_spend_counter_cache
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_increment_spend_counter_cache_redis_path_returns_new_value(monkeypatch):
+ fake_cache = _make_spend_counter_cache(redis_increment_value=44.0)
+ monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
+
+ result = await ps._increment_spend_counter_cache(
+ counter_key="spend:key:k", increment=4.0
+ )
+
+ observed = {
+ "result": result,
+ "redis_increment_called": fake_cache.redis_cache.async_increment.called,
+ "in_memory_set_called": fake_cache.in_memory_cache.set_cache.called,
+ }
+ assert normalize(observed) == {
+ "result": 44.0,
+ "redis_increment_called": True,
+ "in_memory_set_called": True,
+ }
+
+
+@pytest.mark.asyncio
+async def test_increment_spend_counter_cache_redis_error_raises_and_invalidates(
+ monkeypatch,
+):
+ fake_cache = _make_spend_counter_cache(
+ redis_increment_side_effect=RuntimeError("incr fail")
+ )
+ monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
+
+ with pytest.raises(RuntimeError):
+ await ps._increment_spend_counter_cache(
+ counter_key="spend:key:k", increment=1.0
+ )
+
+ assert fake_cache.in_memory_cache.delete_cache.called is True
+ assert fake_cache.redis_cache.async_delete_cache.called is True
+
+
+# ---------------------------------------------------------------------------
+# _invalidate_spend_counter
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_invalidate_spend_counter_deletes_in_memory_and_redis(monkeypatch):
+ fake_cache = _make_spend_counter_cache()
+ monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
+
+ await ps._invalidate_spend_counter(counter_key="spend:key:k")
+
+ observed = {
+ "in_memory_delete_called": fake_cache.in_memory_cache.delete_cache.called,
+ "redis_delete_called": fake_cache.redis_cache.async_delete_cache.called,
+ "delete_args_key": fake_cache.redis_cache.async_delete_cache.call_args.kwargs[
+ "key"
+ ],
+ }
+ assert normalize(observed) == {
+ "in_memory_delete_called": True,
+ "redis_delete_called": True,
+ "delete_args_key": "spend:key:k",
+ }
+
+
+@pytest.mark.asyncio
+async def test_invalidate_spend_counter_swallows_redis_failure_no_raise(monkeypatch):
+ fake_cache = _make_spend_counter_cache()
+ fake_cache.redis_cache.async_delete_cache = AsyncMock(
+ side_effect=RuntimeError("redis down")
+ )
+ monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
+
+ await ps._invalidate_spend_counter(counter_key="spend:key:k")
+
+ assert fake_cache.in_memory_cache.delete_cache.called is True
+
+
+# ---------------------------------------------------------------------------
+# update_cache
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_update_cache_no_cached_entities_schedules_pipeline_flush(monkeypatch):
+ fake_user_cache = _make_user_api_key_cache(get_value=None)
+ monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache)
+
+ await ps.update_cache(
+ token=None,
+ user_id="u1",
+ end_user_id="eu1",
+ team_id="t1",
+ response_cost=1.0,
+ parent_otel_span=None,
+ tags=["x"],
+ )
+
+ observed = {
+ "lookups": fake_user_cache.async_get_cache.call_count,
+ "got_user": True,
+ "got_team": True,
+ }
+ assert normalize(observed) == {
+ "lookups": 4,
+ "got_user": True,
+ "got_team": True,
+ }
+
+
+@pytest.mark.asyncio
+async def test_update_cache_user_cache_failure_invalid_state_is_swallowed(monkeypatch):
+ """An inner _update_user_cache raising must not propagate — update_cache
+ catches and logs, the public coroutine still completes normally."""
+ fake_user_cache = MagicMock()
+ fake_user_cache.async_get_cache = AsyncMock(side_effect=RuntimeError("cache down"))
+ fake_user_cache.async_set_cache_pipeline = AsyncMock()
+ monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache)
+
+ result = await ps.update_cache(
+ token=None,
+ user_id="u1",
+ end_user_id=None,
+ team_id=None,
+ response_cost=1.0,
+ parent_otel_span=None,
+ tags=None,
+ )
+
+ assert result is None
diff --git a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py
index ad6b4016461..33de1ede917 100644
--- a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py
+++ b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py
@@ -1 +1,555 @@
-"""Placeholder. Filled by a follow-up PR per the Notion plan."""
+"""Behavior pins for the proxy_server streaming helpers.
+
+Pins covered:
+- ``data_generator``
+- ``async_assistants_data_generator``
+- ``_get_client_requested_model_for_streaming``
+- ``_restamp_streaming_chunk_model``
+- ``_fast_serialize_simple_model_response_stream``
+- ``_serialize_streaming_chunk``
+- ``_apply_streaming_chunk_hooks``
+- ``_format_streaming_sse_chunk``
+- ``async_data_generator``
+- ``select_data_generator``
+"""
+
+from __future__ import annotations
+
+import json
+from typing import Any, AsyncIterator
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+import litellm.proxy.proxy_server as ps
+from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy.proxy_server import (
+ _apply_streaming_chunk_hooks,
+ _fast_serialize_simple_model_response_stream,
+ _format_streaming_sse_chunk,
+ _get_client_requested_model_for_streaming,
+ _restamp_streaming_chunk_model,
+ _serialize_streaming_chunk,
+ async_assistants_data_generator,
+ async_data_generator,
+ data_generator,
+ select_data_generator,
+)
+from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices, Usage
+
+from .conftest import normalize
+
+
+def _user_auth() -> UserAPIKeyAuth:
+ return UserAPIKeyAuth(api_key="sk-test-key", user_id="u")
+
+
+def _simple_chunk(model: str = "gpt-4", content: str = "hi") -> ModelResponseStream:
+ return ModelResponseStream(
+ id="chatcmpl-test",
+ choices=[
+ StreamingChoices(
+ finish_reason=None,
+ index=0,
+ delta=Delta(content=content, role="assistant"),
+ )
+ ],
+ created=0,
+ model=model,
+ object="chat.completion.chunk",
+ )
+
+
+async def _async_iter(items):
+ for it in items:
+ yield it
+
+
+async def _async_iter_raises(exc: Exception):
+ # yield once then raise — exercises the mid-stream failure branch
+ yield _simple_chunk(content="partial")
+ raise exc
+
+
+# ---------------------------------------------------------------------------
+# data_generator
+# ---------------------------------------------------------------------------
+
+
+def test_data_generator_yields_sse_lines_for_dict_chunks():
+ class DictChunk:
+ def __init__(self, payload):
+ self._payload = payload
+
+ def dict(self):
+ return self._payload
+
+ chunks = [
+ DictChunk({"id": "1", "object": "chat.completion.chunk", "model": "gpt-4"}),
+ DictChunk({"id": "2", "object": "chat.completion.chunk", "model": "gpt-4"}),
+ ]
+ out = list(data_generator(chunks))
+
+ assert len(out) == 2
+ payloads = [json.loads(line.removeprefix("data: ").rstrip("\n\n")) for line in out]
+ assert normalize(payloads[0]) == {
+ "id": "",
+ "object": "chat.completion.chunk",
+ "model": "gpt-4",
+ }
+ assert payloads[1]["model"] == "gpt-4"
+
+
+def test_data_generator_fallback_when_dict_raises_exception():
+ class BadChunk:
+ def dict(self):
+ raise RuntimeError("cannot serialize")
+
+ # When .dict() raises, the inner json.dumps(chunk) on a non-JSON-serializable
+ # instance also raises — the generator does not catch the second failure.
+ with pytest.raises((TypeError, RuntimeError)):
+ list(data_generator([BadChunk()]))
+
+
+# ---------------------------------------------------------------------------
+# async_assistants_data_generator
+# ---------------------------------------------------------------------------
+
+
+class _FakeAssistantsStream:
+ """Mimic the async-context-manager + async-iterable shape of the
+ assistants streaming object (e.g. AssistantEventHandler)."""
+
+ def __init__(self, chunks):
+ self._chunks = chunks
+
+ async def __aenter__(self):
+ return self
+
+ async def __aexit__(self, exc_type, exc, tb):
+ return False
+
+ def __aiter__(self):
+ async def _gen():
+ for c in self._chunks:
+ yield c
+
+ return _gen()
+
+
+@pytest.mark.asyncio
+async def test_async_assistants_data_generator_yields_sse_and_done(monkeypatch):
+ chunk = _simple_chunk(content="hello")
+
+ async def _passthrough_hook(*, user_api_key_dict, response, data, **kwargs):
+ return response
+
+ monkeypatch.setattr(
+ ps.proxy_logging_obj,
+ "async_post_call_streaming_hook",
+ _passthrough_hook,
+ )
+
+ stream = _FakeAssistantsStream([chunk])
+ out = []
+ async for line in async_assistants_data_generator(
+ response=stream,
+ user_api_key_dict=_user_auth(),
+ request_data={},
+ ):
+ out.append(line)
+
+ assert out[-1] == "data: [DONE]\n\n"
+ body = json.loads(out[0].removeprefix("data: ").rstrip("\n\n"))
+ assert normalize(body) == {
+ "id": "",
+ "created": "",
+ "model": "gpt-4",
+ "object": "chat.completion.chunk",
+ "choices": [
+ {
+ "index": 0,
+ "delta": {"content": "hello", "role": "assistant"},
+ }
+ ],
+ }
+
+
+@pytest.mark.asyncio
+async def test_async_assistants_data_generator_hook_failure_yields_error_chunk(
+ monkeypatch,
+):
+ async def _boom_hook(*args, **kwargs):
+ raise RuntimeError("hook exploded")
+
+ async def _noop_failure(*args, **kwargs):
+ return None
+
+ monkeypatch.setattr(
+ ps.proxy_logging_obj, "async_post_call_streaming_hook", _boom_hook
+ )
+ monkeypatch.setattr(ps.proxy_logging_obj, "post_call_failure_hook", _noop_failure)
+
+ stream = _FakeAssistantsStream([_simple_chunk()])
+ out = []
+ async for line in async_assistants_data_generator(
+ response=stream,
+ user_api_key_dict=_user_auth(),
+ request_data={},
+ ):
+ out.append(line)
+
+ assert any("error" in line for line in out)
+ assert out[-1].startswith('data: {"error":')
+
+
+# ---------------------------------------------------------------------------
+# _get_client_requested_model_for_streaming
+# ---------------------------------------------------------------------------
+
+
+def test_get_client_requested_model_for_streaming_prefers_client_requested():
+ request_data = {
+ "_litellm_client_requested_model": "gpt-4",
+ "model": "openai/internal-gpt-4",
+ "litellm_call_id": "abc",
+ }
+ result = _get_client_requested_model_for_streaming(request_data)
+ assert result == "gpt-4"
+
+ snapshot = {
+ "result": result,
+ "client_field_preserved": request_data["_litellm_client_requested_model"],
+ "model_field_preserved": request_data["model"],
+ }
+ assert normalize(snapshot) == {
+ "result": "gpt-4",
+ "client_field_preserved": "gpt-4",
+ "model_field_preserved": "openai/internal-gpt-4",
+ }
+
+
+def test_get_client_requested_model_for_streaming_falls_back_to_model_field():
+ result = _get_client_requested_model_for_streaming({"model": "claude-sonnet"})
+ assert result == "claude-sonnet"
+
+
+def test_get_client_requested_model_for_streaming_missing_returns_empty_invalid():
+ """When neither key is set or values are non-strings, the helper returns ""
+ rather than raising — callers depend on this to skip restamping."""
+ assert _get_client_requested_model_for_streaming({}) == ""
+ assert _get_client_requested_model_for_streaming({"model": 123}) == ""
+
+
+# ---------------------------------------------------------------------------
+# _restamp_streaming_chunk_model
+# ---------------------------------------------------------------------------
+
+
+def test_restamp_streaming_chunk_model_overrides_model_on_basemodel():
+ chunk = _simple_chunk(model="openai/internal-x")
+ new_chunk, logged = _restamp_streaming_chunk_model(
+ chunk=chunk,
+ requested_model_from_client="gpt-4",
+ request_data={"litellm_call_id": "id-1"},
+ model_mismatch_logged=False,
+ )
+ snapshot = {
+ "model": new_chunk.model,
+ "logged": logged,
+ "same_object": new_chunk is chunk,
+ }
+ assert snapshot == {"model": "gpt-4", "logged": True, "same_object": True}
+
+
+def test_restamp_streaming_chunk_model_overrides_model_on_dict():
+ chunk = {"model": "internal", "choices": []}
+ new_chunk, logged = _restamp_streaming_chunk_model(
+ chunk=chunk,
+ requested_model_from_client="gpt-4",
+ request_data={},
+ model_mismatch_logged=True,
+ )
+ assert new_chunk["model"] == "gpt-4"
+ assert logged is True
+
+
+def test_restamp_streaming_chunk_model_invalid_chunk_type_unchanged():
+ """For a non-BaseModel, non-dict chunk the helper returns it as-is
+ along with the original ``model_mismatch_logged`` flag."""
+ chunk = "raw string chunk"
+ new_chunk, logged = _restamp_streaming_chunk_model(
+ chunk=chunk,
+ requested_model_from_client="gpt-4",
+ request_data={},
+ model_mismatch_logged=False,
+ )
+ assert new_chunk == "raw string chunk"
+ assert logged is False
+
+
+# ---------------------------------------------------------------------------
+# _fast_serialize_simple_model_response_stream
+# ---------------------------------------------------------------------------
+
+
+def test_fast_serialize_simple_model_response_stream_returns_bytes_payload():
+ chunk = _simple_chunk()
+ result = _fast_serialize_simple_model_response_stream(chunk)
+ assert isinstance(result, bytes)
+ payload = json.loads(result)
+ assert normalize(payload) == {
+ "id": "",
+ "object": "chat.completion.chunk",
+ "created": "",
+ "model": "gpt-4",
+ "choices": [
+ {
+ "index": 0,
+ "delta": {"role": "assistant", "content": "hi"},
+ }
+ ],
+ }
+
+
+def test_fast_serialize_simple_model_response_stream_with_usage_returns_none_invalid():
+ """Fast path bails (returns None) when ``usage`` is populated — the slow
+ path is required to preserve usage fields. Returning None here is the
+ "I cannot handle this" sentinel, not a hard error."""
+ chunk = _simple_chunk()
+ chunk.usage = Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2)
+ assert _fast_serialize_simple_model_response_stream(chunk) is None
+
+
+# ---------------------------------------------------------------------------
+# _serialize_streaming_chunk
+# ---------------------------------------------------------------------------
+
+
+def test_serialize_streaming_chunk_simple_uses_fast_path_bytes():
+ result = _serialize_streaming_chunk(_simple_chunk())
+ assert isinstance(result, bytes)
+ payload = json.loads(result)
+ assert normalize(payload) == {
+ "id": "",
+ "object": "chat.completion.chunk",
+ "created": "",
+ "model": "gpt-4",
+ "choices": [
+ {
+ "index": 0,
+ "delta": {"role": "assistant", "content": "hi"},
+ }
+ ],
+ }
+
+
+def test_serialize_streaming_chunk_invalid_input_raises_attribute_error():
+ """The helper is typed as ``BaseModel`` — handing it a plain dict trips
+ the attribute-access path (no ``model_dump_json``)."""
+ with pytest.raises(AttributeError):
+ _serialize_streaming_chunk({"not": "a model"}) # type: ignore[arg-type]
+
+
+# ---------------------------------------------------------------------------
+# _apply_streaming_chunk_hooks
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_apply_streaming_chunk_hooks_appends_to_str_so_far(monkeypatch):
+ chunk = _simple_chunk(content="abc")
+
+ async def _passthrough(*, user_api_key_dict, response, data, str_so_far=None):
+ return response
+
+ monkeypatch.setattr(
+ ps.proxy_logging_obj, "async_post_call_streaming_hook", _passthrough
+ )
+
+ new_chunk, new_str = await _apply_streaming_chunk_hooks(
+ chunk=chunk,
+ user_api_key_dict=_user_auth(),
+ request_data={},
+ str_so_far="prior:",
+ )
+
+ observed = {
+ "chunk_is_basemodel": isinstance(new_chunk, ModelResponseStream),
+ "str_so_far": new_str,
+ "grew": len(new_str) > len("prior:"),
+ }
+ assert observed == {
+ "chunk_is_basemodel": True,
+ "str_so_far": "prior:abc",
+ "grew": True,
+ }
+
+
+@pytest.mark.asyncio
+async def test_apply_streaming_chunk_hooks_hook_raises_exception(monkeypatch):
+ async def _boom(*args, **kwargs):
+ raise RuntimeError("hook failed")
+
+ monkeypatch.setattr(ps.proxy_logging_obj, "async_post_call_streaming_hook", _boom)
+
+ with pytest.raises(RuntimeError):
+ await _apply_streaming_chunk_hooks(
+ chunk=_simple_chunk(),
+ user_api_key_dict=_user_auth(),
+ request_data={},
+ str_so_far="",
+ )
+
+
+# ---------------------------------------------------------------------------
+# _format_streaming_sse_chunk
+# ---------------------------------------------------------------------------
+
+
+def test_format_streaming_sse_chunk_handles_bytes_and_str_shapes():
+ bytes_out = _format_streaming_sse_chunk(b'{"a":1}')
+ str_out = _format_streaming_sse_chunk('{"a":1}')
+
+ snapshot = {
+ "bytes_out": bytes_out,
+ "str_out": str_out,
+ "bytes_starts_with_data": bytes_out.startswith(b"data: "),
+ }
+ assert snapshot == {
+ "bytes_out": b'data: {"a":1}\n\n',
+ "str_out": 'data: {"a":1}\n\n',
+ "bytes_starts_with_data": True,
+ }
+
+
+def test_format_streaming_sse_chunk_invalid_empty_string_still_wraps():
+ """Edge case: empty string still gets the ``data: \\n\\n`` wrapping
+ — clients expect SSE shape even on empty payloads."""
+ result = _format_streaming_sse_chunk("")
+ assert result == "data: \n\n"
+
+
+# ---------------------------------------------------------------------------
+# async_data_generator
+# ---------------------------------------------------------------------------
+
+
+def _patch_logging_flags(monkeypatch, needs_wrap=False, needs_per_chunk=False):
+ monkeypatch.setattr(
+ ps.proxy_logging_obj,
+ "needs_iterator_wrap",
+ lambda: needs_wrap,
+ )
+ monkeypatch.setattr(
+ ps.proxy_logging_obj,
+ "needs_per_chunk_streaming_hook",
+ lambda: needs_per_chunk,
+ )
+ # ``_fire_deferred_stream_logging`` is a classmethod — patch the
+ # underlying function so the no-wrap branch is a no-op rather than
+ # touching real logging globals.
+ monkeypatch.setattr(
+ ps.ProxyLogging,
+ "_fire_deferred_stream_logging",
+ staticmethod(lambda request_data: None),
+ )
+
+
+@pytest.mark.asyncio
+async def test_async_data_generator_yields_sse_chunks_and_done(monkeypatch):
+ _patch_logging_flags(monkeypatch)
+
+ response = _async_iter([_simple_chunk(content="hello")])
+ out = []
+ async for line in async_data_generator(
+ response=response,
+ user_api_key_dict=_user_auth(),
+ request_data={"model": "gpt-4"},
+ ):
+ out.append(line)
+
+ assert out[-1] == "data: [DONE]\n\n"
+ # First chunk is bytes (fast path) wrapped via _format_streaming_sse_chunk.
+ first = out[0]
+ assert isinstance(first, bytes)
+ payload = json.loads(first.removeprefix(b"data: ").rstrip(b"\n\n"))
+ assert normalize(payload) == {
+ "id": "",
+ "object": "chat.completion.chunk",
+ "created": "",
+ "model": "gpt-4",
+ "choices": [
+ {
+ "index": 0,
+ "delta": {"role": "assistant", "content": "hello"},
+ }
+ ],
+ }
+
+
+@pytest.mark.asyncio
+async def test_async_data_generator_mid_stream_exception_yields_error_payload(
+ monkeypatch,
+):
+ _patch_logging_flags(monkeypatch)
+
+ async def _noop_failure(*args, **kwargs):
+ return None
+
+ monkeypatch.setattr(ps.proxy_logging_obj, "post_call_failure_hook", _noop_failure)
+
+ response = _async_iter_raises(RuntimeError("upstream blew up"))
+ out = []
+ async for line in async_data_generator(
+ response=response,
+ user_api_key_dict=_user_auth(),
+ request_data={},
+ ):
+ out.append(line)
+
+ # First entry is the successful "partial" chunk (bytes), last is the error.
+ assert any(
+ isinstance(item, str) and item.startswith('data: {"error":') for item in out
+ )
+
+
+# ---------------------------------------------------------------------------
+# select_data_generator
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_select_data_generator_returns_async_generator(monkeypatch):
+ _patch_logging_flags(monkeypatch)
+
+ response = _async_iter([_simple_chunk()])
+ gen = select_data_generator(
+ response=response,
+ user_api_key_dict=_user_auth(),
+ request_data={"model": "gpt-4"},
+ )
+
+ # Drain to confirm it really is an async iterator emitting SSE shape.
+ collected = []
+ async for line in gen:
+ collected.append(line)
+
+ snapshot = {
+ "is_async_iterable": hasattr(gen, "__aiter__"),
+ "yielded_at_least_one": len(collected) >= 1,
+ "ends_with_done": collected[-1] == "data: [DONE]\n\n",
+ }
+ assert snapshot == {
+ "is_async_iterable": True,
+ "yielded_at_least_one": True,
+ "ends_with_done": True,
+ }
+
+
+def test_select_data_generator_missing_required_kwarg_raises_type_error():
+ """``select_data_generator`` requires all three keyword args — calling
+ without ``request_data`` raises TypeError at the wrapper, before any
+ streaming starts."""
+ with pytest.raises(TypeError):
+ select_data_generator(response=_async_iter([]), user_api_key_dict=_user_auth()) # type: ignore[call-arg]
From 892838963ca5a11b54327ec293d0f233c4b1ed46 Mon Sep 17 00:00:00 2001
From: ryan-crabbe-berri
Date: Sat, 30 May 2026 00:26:28 -0700
Subject: [PATCH 21/44] test(e2e): cover Internal User create-key flow when in
no teams (#29083)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
* test(e2e): cover Internal User create-key flow when in no teams
The seeded e2e-internal-user is in two teams, so the "no team" branch
of the Create Key modal — where the team dropdown must render empty —
was unreachable. Seeds a noteam@test.local user and adds a spec that
logs in fresh, opens the modal, and asserts the dropdown has zero
options.
* test(e2e): harden no-team dropdown assertion + add with-teams counterpart
Replace the one-shot count() check with a settled-empty assertion: wait for
the dropdown's loaded "No teams found" state before asserting zero options,
so the test can't pass on a transient empty frame while the team-options
request is still in flight.
Add internalUserWithTeams.spec.ts as the differential partner; it logs in as
the seeded e2e-internal-user (two team memberships) and asserts the dropdown
lists exactly those teams. Without it, the no-team spec's zero-options
assertion would still pass against a regression that empties the dropdown for
every user.
---
ui/litellm-dashboard/e2e_tests/constants.ts | 1 +
.../e2e_tests/fixtures/seed.sql | 1 +
.../internal-user/internalUserNoTeam.spec.ts | 41 +++++++++++++++++++
.../internalUserWithTeams.spec.ts | 37 +++++++++++++++++
4 files changed, 80 insertions(+)
create mode 100644 ui/litellm-dashboard/e2e_tests/tests/internal-user/internalUserNoTeam.spec.ts
create mode 100644 ui/litellm-dashboard/e2e_tests/tests/internal-user/internalUserWithTeams.spec.ts
diff --git a/ui/litellm-dashboard/e2e_tests/constants.ts b/ui/litellm-dashboard/e2e_tests/constants.ts
index dbc73432f65..8b1b4313869 100644
--- a/ui/litellm-dashboard/e2e_tests/constants.ts
+++ b/ui/litellm-dashboard/e2e_tests/constants.ts
@@ -18,5 +18,6 @@ export const E2E_TEAM_CRUD_ALIAS = "E2E Team CRUD";
export const E2E_TEAM_DELETE_ID = "e2e-team-delete";
export const E2E_TEAM_DELETE_ALIAS = "E2E Team Delete";
export const E2E_TEAM_ORG_ID = "e2e-team-org";
+export const E2E_TEAM_ORG_ALIAS = "E2E Team In Org";
export const E2E_TEAM_NO_ADMIN_ID = "e2e-team-no-admin";
export const E2E_TEAM_NO_ADMIN_ALIAS = "E2E Team No Admin";
diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/seed.sql b/ui/litellm-dashboard/e2e_tests/fixtures/seed.sql
index 5e5313240e6..a1218633cdb 100644
--- a/ui/litellm-dashboard/e2e_tests/fixtures/seed.sql
+++ b/ui/litellm-dashboard/e2e_tests/fixtures/seed.sql
@@ -33,6 +33,7 @@ VALUES
('e2e-internal-viewer', 'viewer@test.local', 'internal_user_viewer', '{"e2e-team-crud"}', 'scrypt:MU5CcTAi6rVK1HfY1rVPEWq6r4sxg837eq9dG4n5Q6BhDJ44442+seC6LAhLEAYr'),
('e2e-team-admin', 'teamadmin@test.local', 'internal_user', '{"e2e-team-crud","e2e-team-delete"}', 'scrypt:MU5CcTAi6rVK1HfY1rVPEWq6r4sxg837eq9dG4n5Q6BhDJ44442+seC6LAhLEAYr'),
('e2e-invitable-user', 'invitable@test.local', 'internal_user', '{}', 'scrypt:MU5CcTAi6rVK1HfY1rVPEWq6r4sxg837eq9dG4n5Q6BhDJ44442+seC6LAhLEAYr'),
+ ('e2e-internal-noteam', 'noteam@test.local', 'internal_user', '{}', 'scrypt:MU5CcTAi6rVK1HfY1rVPEWq6r4sxg837eq9dG4n5Q6BhDJ44442+seC6LAhLEAYr'),
('e2e-invitable-by-team-admin', 'invitable-team@test.local', 'internal_user', '{}', 'scrypt:MU5CcTAi6rVK1HfY1rVPEWq6r4sxg837eq9dG4n5Q6BhDJ44442+seC6LAhLEAYr'),
('e2e-removable-member', 'removable@test.local', 'internal_user', '{"e2e-team-crud"}', 'scrypt:MU5CcTAi6rVK1HfY1rVPEWq6r4sxg837eq9dG4n5Q6BhDJ44442+seC6LAhLEAYr');
diff --git a/ui/litellm-dashboard/e2e_tests/tests/internal-user/internalUserNoTeam.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/internal-user/internalUserNoTeam.spec.ts
new file mode 100644
index 00000000000..548639d6877
--- /dev/null
+++ b/ui/litellm-dashboard/e2e_tests/tests/internal-user/internalUserNoTeam.spec.ts
@@ -0,0 +1,41 @@
+import { test, expect } from "@playwright/test";
+import { dismissFeedbackPopup } from "../../helpers/navigation";
+
+/**
+ * Logs in fresh inside the test rather than reusing a stored session because
+ * this user (seeded with no team memberships) only exists for this one spec —
+ * extending globalSetup + the Role enum + the storage-path map for a single
+ * assertion isn't worth the maintenance cost.
+ */
+test.describe("Internal User with no team memberships", () => {
+ test.use({ storageState: { cookies: [], origins: [] } });
+
+ test("Create Key team dropdown is empty when the user belongs to no teams", async ({ page }) => {
+ // Log in via the form as the no-team seeded user.
+ await page.goto("/ui/login");
+ await page.getByPlaceholder("Enter your username").fill("noteam@test.local");
+ await page.getByPlaceholder("Enter your password").fill("test");
+ await page.getByRole("button", { name: "Login", exact: true }).click();
+ await expect(page.getByText("Virtual Keys")).toBeVisible({ timeout: 15_000 });
+ await dismissFeedbackPopup(page);
+
+ // Open the Create Key modal.
+ await page.getByRole("button", { name: /Create New Key/i }).click();
+ await expect(page.getByText("Key Ownership")).toBeVisible({ timeout: 10_000 });
+
+ const teamSelect = page.locator(".ant-select", { hasText: "Search or select a team" });
+ await teamSelect.click();
+
+ const dropdown = page.locator(".ant-select-dropdown:visible").first();
+ await expect(dropdown).toBeVisible({ timeout: 5_000 });
+
+ // Wait for the settled-empty state, not a transient one. The dropdown shows
+ // a spinner while teams load and only swaps in "No teams found" once the
+ // request resolves with nothing (team_dropdown.tsx renders the spinner when
+ // isLoading and this copy otherwise). Asserting on it means a regression
+ // where teams DO load for this user fails here instead of racing a one-shot
+ // count() against an in-flight request.
+ await expect(dropdown.getByText("No teams found")).toBeVisible({ timeout: 10_000 });
+ await expect(dropdown.getByRole("option")).toHaveCount(0);
+ });
+});
diff --git a/ui/litellm-dashboard/e2e_tests/tests/internal-user/internalUserWithTeams.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/internal-user/internalUserWithTeams.spec.ts
new file mode 100644
index 00000000000..f6b60f411b4
--- /dev/null
+++ b/ui/litellm-dashboard/e2e_tests/tests/internal-user/internalUserWithTeams.spec.ts
@@ -0,0 +1,37 @@
+import { test, expect } from "@playwright/test";
+import {
+ INTERNAL_USER_STORAGE_PATH,
+ E2E_TEAM_CRUD_ALIAS,
+ E2E_TEAM_ORG_ALIAS,
+} from "../../constants";
+import { Page } from "../../fixtures/pages";
+import { navigateToPage } from "../../helpers/navigation";
+
+/**
+ * Differential partner to internalUserNoTeam.spec.ts: the seeded
+ * e2e-internal-user belongs to exactly two teams, so the Create Key dropdown
+ * must list both. Without this, the no-team spec's "zero options" assertion
+ * would still pass against a bug that empties the dropdown for everyone.
+ */
+test.describe("Internal User with team memberships", () => {
+ test.use({ storageState: INTERNAL_USER_STORAGE_PATH });
+
+ test("Create Key team dropdown lists exactly the teams the user belongs to", async ({ page }) => {
+ await navigateToPage(page, Page.ApiKeys);
+
+ await page.getByRole("button", { name: /Create New Key/i }).click();
+ await expect(page.getByText("Key Ownership")).toBeVisible({ timeout: 10_000 });
+
+ const teamSelect = page.locator(".ant-select", { hasText: "Search or select a team" });
+ await teamSelect.click();
+
+ const dropdown = page.locator(".ant-select-dropdown:visible").first();
+ await expect(dropdown).toBeVisible({ timeout: 5_000 });
+
+ // Both seeded memberships render, and nothing else does — proving the
+ // dropdown is scoped to the user's teams rather than empty or unfiltered.
+ await expect(dropdown.getByText(E2E_TEAM_CRUD_ALIAS, { exact: true })).toBeVisible({ timeout: 10_000 });
+ await expect(dropdown.getByText(E2E_TEAM_ORG_ALIAS, { exact: true })).toBeVisible();
+ await expect(dropdown.getByRole("option")).toHaveCount(2);
+ });
+});
From 37e6e2da1c78e11f4d2edec923f83307334b5e7a Mon Sep 17 00:00:00 2001
From: ryan-crabbe-berri
Date: Sat, 30 May 2026 00:29:30 -0700
Subject: [PATCH 22/44] test(e2e): assert internal-user navbar identity is
scoped to that user (#29077)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
* test(e2e): assert internal-user navbar identity is scoped to that user
The existing login.spec.ts only checks the admin's navbar identity.
This adds the symmetric check for the internal user — verifying the
account button + dropdown surface the internal user's email, id, and
role, and that no admin-scoped values leak through.
* test(e2e): harden navbar identity test per review feedback
Locate the user dropdown panel by a data-testid on the popupRender div
instead of Ant Design internal + Tailwind class names, so styling
refactors no longer risk breaking the identity-scoping assertions.
Source the seeded user emails/ids from shared constants (match seed.sql)
instead of hardcoding them inline.
---
ui/litellm-dashboard/e2e_tests/constants.ts | 6 +++
.../tests/login/internalUserIdentity.spec.ts | 48 +++++++++++++++++++
.../Navbar/UserDropdown/UserDropdown.tsx | 2 +-
3 files changed, 55 insertions(+), 1 deletion(-)
create mode 100644 ui/litellm-dashboard/e2e_tests/tests/login/internalUserIdentity.spec.ts
diff --git a/ui/litellm-dashboard/e2e_tests/constants.ts b/ui/litellm-dashboard/e2e_tests/constants.ts
index 8b1b4313869..236909384b0 100644
--- a/ui/litellm-dashboard/e2e_tests/constants.ts
+++ b/ui/litellm-dashboard/e2e_tests/constants.ts
@@ -5,6 +5,12 @@ export const INTERNAL_USER_STORAGE_PATH = "internalUser.storageState.json";
export const INTERNAL_VIEWER_STORAGE_PATH = "internalViewer.storageState.json";
export const TEAM_ADMIN_STORAGE_PATH = "teamAdmin.storageState.json";
+// Seeded user identities (match seed.sql)
+export const E2E_PROXY_ADMIN_USER_ID = "e2e-proxy-admin";
+export const E2E_PROXY_ADMIN_EMAIL = "admin@test.local";
+export const E2E_INTERNAL_USER_ID = "e2e-internal-user";
+export const E2E_INTERNAL_USER_EMAIL = "internal@test.local";
+
// Key aliases for seeded test keys (match seed.sql)
export const E2E_UPDATE_LIMITS_KEY_ALIAS = "e2eUpdateLimitsKey";
export const E2E_DELETE_KEY_ALIAS = "e2eDeleteKey";
diff --git a/ui/litellm-dashboard/e2e_tests/tests/login/internalUserIdentity.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/login/internalUserIdentity.spec.ts
new file mode 100644
index 00000000000..cbe95276929
--- /dev/null
+++ b/ui/litellm-dashboard/e2e_tests/tests/login/internalUserIdentity.spec.ts
@@ -0,0 +1,48 @@
+import { test, expect } from "@playwright/test";
+import {
+ E2E_INTERNAL_USER_EMAIL,
+ E2E_INTERNAL_USER_ID,
+ E2E_PROXY_ADMIN_EMAIL,
+ E2E_PROXY_ADMIN_USER_ID,
+ INTERNAL_USER_STORAGE_PATH,
+} from "../../constants";
+
+const escapeRegExp = (value: string) => value.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
+
+test.describe("Navbar identity scoping", () => {
+ test.use({ storageState: INTERNAL_USER_STORAGE_PATH });
+
+ test("Internal user navbar dropdown shows their own role and user id, not the admin's", async ({ page }) => {
+ await page.goto("/ui");
+ await expect(page.getByText("Virtual Keys")).toBeVisible({ timeout: 10_000 });
+
+ // The account menu button carries the user's role and email/id in its
+ // aria-label (see UserDropdown.tsx). Match by partial role.
+ const accountButton = page.locator('button[aria-label^="Account menu"]').first();
+ await expect(accountButton).toHaveAttribute("aria-label", /Internal User/, { timeout: 5_000 });
+ await expect(accountButton).toHaveAttribute(
+ "aria-label",
+ new RegExp(
+ `signed in as (${escapeRegExp(E2E_INTERNAL_USER_EMAIL)}|${escapeRegExp(E2E_INTERNAL_USER_ID)})`,
+ ),
+ { timeout: 5_000 },
+ );
+
+ // Open the dropdown (UserDropdown configures trigger=["click"]).
+ await accountButton.click();
+
+ // Locate the panel by its test id (data-testid on the popupRender div in
+ // UserDropdown.tsx) rather than Ant/Tailwind class names, so styling
+ // refactors don't silently break the identity-scoping assertions below.
+ const popup = page.getByTestId("user-dropdown-panel");
+ await expect(popup).toBeVisible({ timeout: 5_000 });
+
+ // The popup must show the internal user's identity — not the seeded
+ // proxy admin's email/id, which would indicate a session/scope leak.
+ await expect(popup.getByText(E2E_INTERNAL_USER_EMAIL)).toBeVisible({ timeout: 5_000 });
+ await expect(popup.getByText(E2E_INTERNAL_USER_ID)).toBeVisible({ timeout: 5_000 });
+ await expect(popup.getByText("Internal User", { exact: true })).toBeVisible({ timeout: 5_000 });
+ await expect(popup.getByText(E2E_PROXY_ADMIN_EMAIL)).toHaveCount(0);
+ await expect(popup.getByText(E2E_PROXY_ADMIN_USER_ID)).toHaveCount(0);
+ });
+});
diff --git a/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.tsx b/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.tsx
index 64a2f1260ba..03312faaa39 100644
--- a/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.tsx
+++ b/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.tsx
@@ -221,7 +221,7 @@ const UserDropdown: React.FC = ({ onLogout }) => {
trigger={["click"]}
menu={{ items: userItems }}
popupRender={(menu) => (
-
+
{renderUserInfoSection()}
Date: Sat, 30 May 2026 00:54:30 -0700
Subject: [PATCH 23/44] feat(otel): add team_metadata, http.route, and model
names to inference spans (#29319)
---
litellm/integrations/opentelemetry.py | 71 ++++++++++++++
.../integrations/test_opentelemetry.py | 92 +++++++++++++++++++
2 files changed, 163 insertions(+)
diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py
index 81fdc5a1e21..814da344f03 100644
--- a/litellm/integrations/opentelemetry.py
+++ b/litellm/integrations/opentelemetry.py
@@ -64,6 +64,9 @@ HTTP_RESPONSE_STATUS_CODE_ATTRIBUTE = "http.response.status_code"
HTTP_ROUTE_ATTRIBUTE = "http.route"
URL_PATH_ATTRIBUTE = "url.path"
PREPROCESSING_DURATION_MS_ATTRIBUTE = "litellm.preprocessing.duration_ms"
+TEAM_METADATA_ATTRIBUTE = "litellm.team.metadata"
+MODEL_GROUP_ATTRIBUTE = "litellm.model_group"
+PROVIDER_MODEL_ATTRIBUTE = "litellm.provider.model"
# Remove the hardcoded LITELLM_RESOURCE dictionary - we'll create it properly later
RAW_REQUEST_SPAN_NAME = "raw_gen_ai_request"
LITELLM_REQUEST_SPAN_NAME = "litellm_request"
@@ -1213,6 +1216,68 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
):
self._set_team_attributes_from_kwargs(proxy_span, kwargs)
+ def _set_inference_identity_attributes(
+ self,
+ span: Span,
+ standard_logging_payload: StandardLoggingPayload,
+ litellm_params: dict,
+ ) -> None:
+ """Stamp request-identity attributes onto an inference span so every
+ LLM-call span is filterable by the route it came in on, the team's
+ metadata, and both the user-facing (model_group alias) and the
+ dispatched (provider) model names. Empty/absent values are skipped.
+ """
+ metadata = standard_logging_payload.get("metadata") or {}
+
+ http_route = metadata.get("user_api_key_request_route")
+ if http_route:
+ self.safe_set_attribute(
+ span=span, key=HTTP_ROUTE_ATTRIBUTE, value=http_route
+ )
+
+ # ``user_api_key_team_metadata`` is dropped from the standard logging
+ # payload metadata, so read it from the raw request metadata in kwargs.
+ # ``metadata`` and ``litellm_metadata`` are alternate names for the same
+ # full metadata dict (the name varies by endpoint), so first-truthy wins.
+ raw_metadata = (
+ litellm_params.get("metadata")
+ or litellm_params.get("litellm_metadata")
+ or {}
+ )
+ team_metadata = self._team_metadata_json(
+ raw_metadata.get("user_api_key_team_metadata")
+ )
+ if team_metadata:
+ self.safe_set_attribute(
+ span=span, key=TEAM_METADATA_ATTRIBUTE, value=team_metadata
+ )
+
+ model_group = standard_logging_payload.get("model_group")
+ if model_group:
+ self.safe_set_attribute(
+ span=span, key=MODEL_GROUP_ATTRIBUTE, value=model_group
+ )
+
+ hidden_params = standard_logging_payload.get("hidden_params") or {}
+ provider_model = hidden_params.get(
+ "litellm_model_name"
+ ) or standard_logging_payload.get("model")
+ if provider_model:
+ self.safe_set_attribute(
+ span=span, key=PROVIDER_MODEL_ATTRIBUTE, value=provider_model
+ )
+
+ @staticmethod
+ def _team_metadata_json(value: Any) -> Optional[str]:
+ """JSON-serialize a team's metadata dict for a single span attribute.
+
+ Returns ``None`` for a missing, non-dict, or empty mapping so the
+ empty case is dropped rather than stamping a useless ``"{}"``.
+ """
+ if not isinstance(value, dict) or not value:
+ return None
+ return safe_dumps(value)
+
def _record_metrics(self, kwargs, response_obj, start_time, end_time):
duration_s = (end_time - start_time).total_seconds()
params = kwargs.get("litellm_params") or {}
@@ -2023,6 +2088,12 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
key="hidden_params",
value=safe_dumps(hidden_params),
)
+
+ self._set_inference_identity_attributes(
+ span=span,
+ standard_logging_payload=standard_logging_payload,
+ litellm_params=litellm_params,
+ )
# Cost breakdown tracking
cost_breakdown: Optional[CostBreakdown] = standard_logging_payload.get(
"cost_breakdown"
diff --git a/tests/test_litellm/integrations/test_opentelemetry.py b/tests/test_litellm/integrations/test_opentelemetry.py
index b65e629c890..461ab39c288 100644
--- a/tests/test_litellm/integrations/test_opentelemetry.py
+++ b/tests/test_litellm/integrations/test_opentelemetry.py
@@ -5142,3 +5142,95 @@ class TestOpenTelemetryPreprocessingDuration(unittest.TestCase):
span, exp = self._span()
otel.set_preprocessing_duration_attribute(span, None)
assert "litellm.preprocessing.duration_ms" not in self._attr(span, exp)
+
+
+class TestOpenTelemetryInferenceIdentityAttributes(unittest.TestCase):
+ """team_metadata, http.route, and both model names (the user-facing
+ model_group alias and the dispatched provider model) must land on the
+ inference span via set_attributes."""
+
+ def _span(self):
+ exporter = InMemorySpanExporter()
+ provider = TracerProvider()
+ provider.add_span_processor(SimpleSpanProcessor(exporter))
+ tracer = provider.get_tracer(__name__)
+ return tracer.start_span("litellm_request"), exporter
+
+ def _attr(self, span, exporter):
+ span.end()
+ return exporter.get_finished_spans()[0].attributes
+
+ def _kwargs(self):
+ return {
+ "model": "gpt-4o",
+ "optional_params": {},
+ "litellm_params": {
+ "custom_llm_provider": "azure",
+ "metadata": {
+ "user_api_key_team_metadata": {
+ "tier": "gold",
+ "cost_center": "42",
+ }
+ },
+ },
+ "standard_logging_object": {
+ "metadata": {
+ "user_api_key_request_route": "/v1/chat/completions",
+ "user_api_key_team_id": "team-1",
+ },
+ "call_type": "completion",
+ "model_group": "gpt-4o",
+ "model": "azure/my-deployment",
+ "hidden_params": {"litellm_model_name": "azure/my-deployment"},
+ "id": "req-1",
+ "litellm_call_id": "call-1",
+ },
+ }
+
+ def test_all_identity_attributes_stamped(self):
+ otel = OpenTelemetry()
+ span, exp = self._span()
+ otel.set_attributes(span, self._kwargs(), {"model": "azure/gpt-4o"})
+ attrs = self._attr(span, exp)
+
+ assert attrs["http.route"] == "/v1/chat/completions"
+ assert json.loads(attrs["litellm.team.metadata"]) == {
+ "tier": "gold",
+ "cost_center": "42",
+ }
+ assert attrs["litellm.model_group"] == "gpt-4o"
+ assert attrs["litellm.provider.model"] == "azure/my-deployment"
+
+ def test_provider_model_falls_back_to_payload_model(self):
+ """Without hidden_params.litellm_model_name the dispatched model is
+ the payload model (the SDK path, where no router renaming happened)."""
+ otel = OpenTelemetry()
+ kwargs = self._kwargs()
+ kwargs["standard_logging_object"]["hidden_params"] = {}
+ span, exp = self._span()
+ otel.set_attributes(span, kwargs, {"model": "azure/gpt-4o"})
+ assert self._attr(span, exp)["litellm.provider.model"] == "azure/my-deployment"
+
+ def test_empty_team_metadata_is_dropped(self):
+ """An empty team_metadata dict must not stamp a useless '{}'."""
+ otel = OpenTelemetry()
+ kwargs = self._kwargs()
+ kwargs["litellm_params"]["metadata"]["user_api_key_team_metadata"] = {}
+ span, exp = self._span()
+ otel.set_attributes(span, kwargs, {"model": "azure/gpt-4o"})
+ assert "litellm.team.metadata" not in self._attr(span, exp)
+
+ def test_missing_route_is_dropped(self):
+ """An SDK request has no route; http.route must be absent, not empty."""
+ otel = OpenTelemetry()
+ kwargs = self._kwargs()
+ del kwargs["standard_logging_object"]["metadata"]["user_api_key_request_route"]
+ span, exp = self._span()
+ otel.set_attributes(span, kwargs, {"model": "azure/gpt-4o"})
+ assert "http.route" not in self._attr(span, exp)
+
+ def test_team_metadata_json_helper_non_dict(self):
+ assert OpenTelemetry._team_metadata_json(None) is None
+ assert OpenTelemetry._team_metadata_json("not-a-dict") is None
+ assert OpenTelemetry._team_metadata_json({}) is None
+ assert json.loads(OpenTelemetry._team_metadata_json({"a": 1})) == {"a": 1}
From 4cc3dd7aad8370244961617d89415020780bd5ef Mon Sep 17 00:00:00 2001
From: Sameer Kankute
Date: Sat, 30 May 2026 21:50:05 +0530
Subject: [PATCH 24/44] feat(context_management): compact_20260112 polyfill for
non-Anthropic providers (#28868)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
* feat(anthropic/messages): in-gateway context_management polyfill for non-Anthropic providers
- Add `context_management/` module with `clear_tool_uses_20250919` editor
dispatched before chat-completions translation on `/v1/messages`
- Hard-protect most-recently completed tool_result from being cleared
- Attach `context_management.applied_edits` to both non-streaming and
streaming (final `message_delta`) responses
- Bedrock Converse: forward `context_management`; filter to
`compact_20260112`-only edits with `compact-2026-01-12` beta header
- token_counter: guard Anthropic-format tools (no `function` key) to
prevent AttributeError during polyfill token counting
- Streaming: handle empty-choices usage-only trailing chunks
- Skip polyfill when `litellm.drop_params = True`
Co-authored-by: Cursor
* fix(bedrock): pop None context_management before sending to Bedrock Converse
If context_management is forwarded as None (e.g. when mapping returns
None for an invalid format), _filter_context_management_for_bedrock_converse
previously returned early without removing the key, leaving
"context_management": null in the request and causing a validation
error. Pop the key when the value is not a dict.
Co-authored-by: Yassin Kortam
* fix(bedrock/converse): pop None context_management; extract helpers to fix PLR0915
Co-authored-by: Cursor
* fix(anthropic/messages): check per-request drop_params alongside global
Co-authored-by: Cursor
* fix(anthropic/messages): preserve drop_params for downstream and respect explicit False
Co-authored-by: Yassin Kortam
* fix: lazy debug logging in clear_tool_uses; remove unused context_management constants
Co-authored-by: Yassin Kortam
* fix(anthropic/messages): guard context_management polyfill with try/except
Wrap apply_context_management() in a try/except so any failure (e.g.
litellm.token_counter raising on an unknown tokenizer or unexpected
message format) is logged but does not crash the underlying LLM
request. The polyfill is a best-effort additive feature; on failure we
forward the original messages without applied edits.
Co-authored-by: Yassin Kortam
* fix(token_counter): guard None input_schema in Anthropic tool fallback
Use `or {}` instead of `.get(..., {})` so explicit null parameters do not
raise AttributeError when formatting function definitions for token counting.
Co-authored-by: Cursor
* fix: minimize context_management polyfill threading
- Use None (not empty list) for polyfill_applied_edits when context
management isn't requested, so semantics of 'feature not requested'
vs 'feature requested but no edits applied' are distinct.
- In the streaming iterator, only pass applied_edits to the per-chunk
translator on the final (finish_reason) chunk; intermediate chunks
ignore it anyway, and this makes intent explicit on both sync and
async paths.
Co-authored-by: Yassin Kortam
* fix(context_management): align tool_use counts and normalize list spec
- _count_tool_uses now requires a string id, matching _collect_tool_use_ids_in_order so the tool_uses trigger can't fire on blocks that aren't clearable.
- apply_context_management dispatcher now accepts the OpenAI list form and normalizes it via AnthropicConfig.map_openai_context_management_to_anthropic, so the polyfill path no longer silently no-ops on list input.
Co-authored-by: Yassin Kortam
* feat(context_management): add compact_20260112 polyfill for non-Anthropic providers
Implements an in-gateway compaction polyfill that summarizes long conversations
using a configurable model when `compact_20260112` is requested for non-Anthropic
targets (e.g. OpenAI, Gemini), matching Anthropic's context management beta
behaviour for those providers.
Co-Authored-By: Claude Sonnet 4.6
* fix(compact): skip tool_result-only user turns; bedrock: elif for context_management
- compact_20260112 Phase D: when keeping the last user turn after a full
summary, skip role=user turns whose content is exclusively tool_result
blocks. Such turns translate to OpenAI tool-role messages with no
preceding assistant tool_calls (those got summarized away), which
non-Anthropic providers reject. Fall back to a synthetic continuation
prompt if no eligible user question exists, so the downstream call
always has a non-empty user message.
- bedrock converse: chain the context_management param as elif so it
follows the same if/elif pattern as the surrounding thinking/
reasoning_effort checks.
Co-authored-by: Yassin Kortam
* fix(anthropic): post-compaction question selection, system type, sync stream merge
- compact.py: select last user question from effective_messages (post-compaction slice) instead of raw messages, so prior summarized turns aren't reintroduced
- handler.py: widen _prepare_completion_kwargs system parameter type to Union[str, List[Dict]] matching PolyfillResult.system
- streaming_iterator.py: mirror async hold-and-merge logic in sync __next__ so context_management is attached to the final merged message_delta when stop_reason and usage arrive in separate chunks
Co-authored-by: Yassin Kortam
* fix(anthropic/messages): apply context_management on sync path; clear held stop_reason chunk in async iterator
- Sync `anthropic_messages_handler` was silently dropping the
`context_management` kwarg via `ANTHROPIC_ONLY_REQUEST_KEYS` after the
polyfill was moved into the async handler. Bridge to the async
dispatcher with `run_async_function` so `litellm.messages.create()`
callers keep working (regressed e.g. `clear_tool_uses_20250919`).
- In the streaming iterator's `__anext__` `StopIteration` handler, clear
`self.holding_stop_reason_chunk` after capturing it (matches `__next__`)
so a subsequent call doesn't re-emit the same chunk.
Co-authored-by: Yassin Kortam
* fix(bugfixes): bedrock None context_mgmt; stream per-instance queue; sync polyfill; trailing-chunk passthrough
Co-authored-by: Yassin Kortam
* fix(anthropic): silently drop trailing chunks after usage; remove dead _polyfill_result key
- streaming_iterator: in sync __next__, after the usage chunk has been
merged and emitted, silently consume any trailing provider events
via 'continue' instead of forwarding them through the queue. Trailing
chunks would translate to content_block_delta or message_delta and
violate Anthropic SSE ordering after the final message_delta. The
async __anext__ already drops these via 'if not self.queued_usage_chunk:'
gating, so this aligns sync and async behavior.
- handler: drop unused '_polyfill_result' from ANTHROPIC_ONLY_REQUEST_KEYS.
PolyfillResult is passed as an explicit arg to the adapter methods, never
through extra_kwargs, so the entry was dead code.
Co-authored-by: Yassin Kortam
* refactor(anthropic): extract usage-merge helper; guard empty slice-only compaction result
- Extract the duplicated hold-and-merge usage logic from the sync __next__ and
async __anext__ paths into a shared _merge_usage_into_held_stop_reason_chunk
helper so the subtle cache-token / context_management attachment lives in
exactly one place.
- In the compact_20260112 slice-only path, fall back to _select_last_user_question
when _strip_compaction_blocks produces an empty list (e.g. messages ending on
an assistant turn whose only content was the compaction block) so the
downstream API never receives an empty messages array.
Co-authored-by: Yassin Kortam
* refactor(anthropic/context_management): streaming iterator compaction fixes and compact polyfill improvements
- Extract usage-merge helper; guard empty slice-only compaction result
- Silently drop trailing chunks after usage; remove dead _polyfill_result key
- Fix bedrock None context_mgmt; stream per-instance queue; sync polyfill; trailing-chunk passthrough
- Apply context_management on sync path; clear held stop_reason chunk in async iterator
- Fix post-compaction question selection, system type, sync stream merge
- Skip tool_result-only user turns; bedrock: elif for context_management
- Add streaming iterator compaction test suite
Co-authored-by: Cursor
* revert(html): restore flat *.html naming in _experimental/out
Reverses the accidental rename from *.html → */index.html introduced in
15ea941fbe. All 35 files moved back to their original flat paths so the
directory structure matches litellm_internal_staging.
Co-authored-by: Cursor
* revert(config): restore proxy_server_config.yaml to litellm_internal_staging
Co-authored-by: Cursor
* Fix: skip client compaction pre-processing when compact_20260112 polyfill will run
The _prepare_context_managed_request helper unconditionally applied
apply_client_compaction_block_history before invoking the polyfill. When
the request also configured a compact_20260112 spec, that pre-processing
consumed the client-sent compaction block and collapsed the message history
to just the latest user question, starving the polyfill of conversation
context. The polyfill's own Phase A (_slice_around_compaction_block)
already handles client compaction blocks correctly and inspects the full
post-compaction tail for the token-threshold check, so the pre-processing
is both redundant and destructive in this case.
Now the pre-processing only runs when no compact_20260112 polyfill spec
will execute (no spec, drop_params on, or only non-compact edits like
clear_tool_uses_20250919).
Co-authored-by: Yassin Kortam
* fix(anthropic): plug compaction-block leak + iteration-usage gaps in streaming adapter
- handler: when polyfill_will_run skipped client-history pre-processing
and the polyfill ultimately returned None (best-effort swallow on
unexpected error), apply the slice-only fallback before returning so
Anthropic-specific 'compaction' content blocks don't leak to non-
Anthropic backends that would reject them.
- streaming_iterator: precompute will_merge_into_held so we don't pass
applied_edits into the translator when the resulting processed_chunk
will be discarded by the held stop-reason merge path.
- streaming_iterator: augment processed_chunk with iterations usage in
the holding_chunk branch (sync and async) for parity with the other
emission branches; ensures usage.iterations is attached on the rare
message_delta-reaches-holding_chunk path.
Co-authored-by: Yassin Kortam
* fix(anthropic): correct streaming usage iteration + translate tools for token counting
- streaming_iterator: skip the trailing "message" iteration entry in the
final message_delta when the held stop_reason chunk carries placeholder
zero usage (no separate usage chunk arrived). Reporting zero tokens was
misleading and inconsistent with the non-streaming path which always
has real usage data.
- streaming_iterator: drop two redundant type checks inside branches
that are already guarded by an outer message_delta type check.
- compact._count_effective_tokens: translate Anthropic-shaped tools
(input_schema) to OpenAI shape before passing to litellm.token_counter
so threshold checks aren't skewed by tokenizer paths that expect the
OpenAI tool wrapper.
Co-authored-by: Yassin Kortam
* Fix lint
* fix(anthropic): plug content drop, compaction SSE shape, and compaction leak
- Sync streaming __next__ no longer drops a buffered holding_chunk when
the usage-merge path has already fired. Restoring the prior unconditional
flush behavior preserves provider-emitted content (the SSE-ordering nit
of a trailing content delta is preferable to silent content loss).
- compaction content_block_start now carries the full block shape
({"type": "compaction", "content": ""}) to match the text-block
pattern and Anthropic's native streaming shape, so clients that key off
content_block_start see the field.
- apply_compact_20260112 now slices around / strips compaction blocks
before the opt-in gate check. Previously, when summary_model was not
configured the editor returned the raw messages, leaking Anthropic-only
compaction content blocks to non-Anthropic providers that reject them.
Co-authored-by: Yassin Kortam
* fix(anthropic): resolve mypy types in context management polyfill
Use AppliedEdit and CompactionBlock consistently in the dispatcher and streaming adapter.
Co-authored-by: Cursor
* fix(anthropic): flush held content chunk in async streaming path
Mirror the sync __next__ behavior: always flush a buffered
holding_chunk after the stream ends, even when usage was already
merged + emitted. Previously the async __anext__ kept the flush
inside the 'if not self.queued_usage_chunk:' guard, silently
dropping the last content delta on the proxy's primary path.
Co-authored-by: Yassin Kortam
* fix(anthropic adapter): correct sync streaming, surface polyfill failures, decouple sync path from proxy router
- translate_completion_output_params_streaming: add is_async flag so the
sync handler returns Iterator[bytes] instead of an unusable
AsyncIterator. Async callers keep the existing behavior via the
default is_async=True.
- _run_polyfill_if_enabled: when the polyfill crashes and the spec
requested non-compact edits (e.g. clear_tool_uses_20250919), raise an
AnthropicContextManagementError instead of silently returning None so
those edits are not dropped without an error surface. The
compaction-block-slicing safety net remains for compact-only specs.
- anthropic_messages_handler (sync): stop auto-attaching the proxy
llm_router. run_async_function bridges to a new thread's event loop;
reusing the proxy's loop-bound httpx clients there causes
'Event loop is closed' errors. The summary editor falls back to
litellm.acompletion when llm_router is None.
Co-authored-by: Yassin Kortam
* fix: address bug detection findings in token counter and streaming iterator
- token_counter: guard against non-dict 'function' field in tool dicts
and skip tools missing a name to avoid emitting 'type None = ...' which
would produce inaccurate token counts.
- streaming_iterator: change sync __next__ generic-error path to raise
StopIteration (was StopAsyncIteration), so sync iteration cleanly stops.
- streaming_iterator: centralize context_management attachment so the
held-stop_reason direct-flush path defensively re-attaches applied_edits
to match the merge path's guarantee.
Co-authored-by: Yassin Kortam
* Fix lint
* fix: correct COMPACT_MIN_TRIGGER_TOKENS to 50_000
Co-authored-by: Yassin Kortam
* Fix lint
* Fix lint
* Fix lint
* fix(compact): reduce to last user question when summary_model not configured but prior compaction block exists
Aligns the summary_model_not_configured path with the under-threshold and
client-compaction-block paths, which both reduce post-compaction messages
to just the latest user question so the downstream provider doesn't get
the summary on system prefix AND the full post-compaction history.
Co-authored-by: Yassin Kortam
* fix(compact): forward caller system prompt to summary model call
The default summarization instructions reference "the initial task above"
and "the raw history above", but the system prompt that holds that task
was not being forwarded to the summary model. The summary call now
prepends an OpenAI-shaped system message translated from the original
Anthropic-shaped system (str or content-block list) so the summarizer
has the agent role and initial task in scope.
* fix(compact_20260112): set default max_tokens and merge prompt when last turn is user
- Set COMPACT_SUMMARY_MAX_TOKENS default for the summary call so providers
like Anthropic (which require max_tokens) don't silently fail and degrade
to summary_call_failed.
- When the trailing translated message is already a user turn, merge the
summarization prompt into it instead of appending a second user turn.
Avoids consecutive role=user messages that strict providers reject.
Co-authored-by: Yassin Kortam
* fix(anthropic adapter): move current_content_block_start to __init__
Move the default TextBlock dict from a class-level attribute to __init__ so
concurrent stream instances don't share the same mutable dict. The class-level
default could be mutated in-place via tool_block['name'] = original_name in
_should_start_new_content_block, leaking state across streams. This mirrors
the existing fix already applied to chunk_queue.
Co-authored-by: Yassin Kortam
* fix(compact_20260112): surface error states + strip tool_result blocks in last user question
applied_edits_for_response() now includes compact_20260112 edits that
carry an error field (summary_model_not_configured, summary_call_failed,
summary_extraction_failed) so clients and operators can see why
compaction was requested but not applied.
_select_last_user_question() now strips tool_result blocks from mixed
[tool_result, text] turns rather than passing them through as-is. After
compaction the paired tool_use assistant turn no longer exists, so
forwarding tool_result blocks translates to orphaned role=tool messages
on non-Anthropic providers and produces a 400.
* fix(compact_20260112): carry prior compaction summary into Phase C summary call
When a request already contains a compaction block, Phase A slices
`effective_messages` to the turns since that block. Previously Phase C
passed the original `system` to the summary model, so multi-round
compaction silently dropped accumulated history each time the polyfill
fired. Pass `augmented_system` (original system + prior summary
prefix) so the summary model can produce a comprehensive summary that
incorporates both the prior round's context and the current slice.
`summarized_system` for the downstream call stays built from the
original `system` + new `summary_text`.
* refactor: delegate handler spec normalization to dispatcher
_normalize_spec_edits in adapters/handler.py duplicated the spec-shape
normalization already implemented by _normalize_spec in
context_management/dispatcher.py. The two could drift: a change in one
(e.g. supporting a new spec shape) without the other would cause the
handler's polyfill_will_run prediction to disagree with the
dispatcher's actual behavior, breaking the client-history pre-processing
skip.
Have the handler delegate to the dispatcher's _normalize_spec while
keeping handler-specific concerns (drop_params short-circuit, swallow
mapping exceptions) at the wrapper level.
Co-authored-by: Yassin Kortam
* fix(compact_20260112): surface warning-only applied edits in response
`applied_edits_for_response()` previously hid `compact_20260112` edits when
they had only warnings (no compaction block, no error). This dropped
diagnostically important warnings such as
`unsupported_trigger_type_X_using_input_tokens` and
`pause_after_compaction_ignored` whenever the conversation was under the
trigger threshold. Operators now see these warnings in the response.
Co-authored-by: Yassin Kortam
* fix: address two low-severity context_management edge cases
- streaming_iterator: keep `sent_content_block_finish` in sync with the
compaction block's emitted start/delta/stop lifecycle and reset it when
the next text block's start is queued.
- bedrock _map_context_management_param: match dispatcher `_normalize_spec`
behavior — only run the OpenAI→Anthropic mapper on list inputs; pass
dict inputs through unchanged so already-Anthropic-format values aren't
silently dropped.
Co-authored-by: Yassin Kortam
* fix(compact_20260112): use beta-header constant; require type discriminator; skip sync bridge when idle
- bedrock: replace hardcoded "compact-2026-01-12" beta string with
ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value in both
Converse (_filter_context_management_for_bedrock_converse) and Invoke
(anthropic_claude3) compact-edit handlers.
- types: mark the "type" discriminator as Required[...] on the new
CompactionBlock and UsageIteration TypedDicts so the discriminator
is not silently optional under total=False.
- adapters/handler: short-circuit the sync /v1/messages adapter path
before spawning the run_async_function worker-thread event loop when
the request has no context_management spec and no client-sent
compaction block in the message history.
Test plan:
- uv run pytest tests/test_litellm/llms/anthropic/experimental_pass_through/ tests/test_litellm/llms/bedrock/test_converse_context_management.py -q
(370 + 10 = 380 passed)
- uv run pytest tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_transformation.py tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py -k compact (3 passed)
* fix(compact_20260112): include system prompt tokens in threshold check
The threshold check in Phase B previously counted only message tokens and
the compaction-block content, omitting the system prompt entirely. When
the system carried a prior compaction summary (via _augment_system_with_summary)
or was otherwise large, the threshold could fire later than intended,
allowing the conversation to exceed the model's context window before
compaction activated.
_count_effective_tokens now also counts the (augmented) system prompt
text. The caller passes compaction_block=None when augmented_system
already includes the prior summary, to avoid double-counting.
Co-authored-by: Yassin Kortam
* Fix SSE ordering and compaction state machine bugs in AnthropicStreamWrapper
- Suppress holding_chunk flush after final message_delta has been emitted
(queued_usage_chunk == True) so a trailing content_block_delta cannot
follow message_delta, which strict Anthropic SDK clients may reject.
When usage has not yet been merged, flush the holding_chunk *before*
the held stop_reason chunk so SSE ordering remains correct.
- Replace _queue_compaction_block_events with _next_compaction_event,
emitting the compaction start/delta/stop events one at a time. The
state machine flags (sent_content_block_finish) and content block
index now advance atomically with the terminal stop event actually
being returned to the caller, eliminating the transient inconsistent
state where flags say the block is finished while its stop event is
still buffered.
Co-authored-by: Yassin Kortam
* fix(compact_20260112): enforce parent key/team allowlist on summary model
The compact_20260112 polyfill summary subrequest used llm_router.acompletion
directly, bypassing the proxy auth checks that gate model access for the
parent key/team. A caller whose key/team was not authorized for the
configured context_management_summary_model could still cause the proxy to
invoke that model and return its output as a compaction block.
Pull the parent's UserAPIKeyAuth out of litellm_metadata in the handler,
thread it through the dispatcher into apply_compact_20260112, and gate the
summary call on _can_object_call_model for both key-level and team-level
allowlists. Failures land as applied_edits[0].error =
summary_model_access_denied without raising. SDK callers (no UserAPIKeyAuth)
remain unaffected.
* fix(compact_20260112): distinguish access-denied from transient errors; greedy summary regex
- _check_summary_model_access now catches ProxyException explicitly for access
denials and logs unexpected exceptions separately. Both still fail closed,
but operators can now tell a denied key/team apart from a router internal
raising during the check.
- _SUMMARY_TAG_RE switches from non-greedy to greedy so a stray
inside the model's summary content no longer silently truncates the
captured text.
* fix(compact_20260112): type object_type as Literal for mypy
* fix(compact_20260112): attribute summary subcall spend to parent key/team
The compact_20260112 polyfill summary subrequest propagated metadata via
the Anthropic-shape `metadata` parameter, which only carries `user_id`.
The proxy auth fields used for spend attribution (`user_api_key`,
`user_api_key_team_id`, `litellm_call_id`, ...) live in
`data["litellm_metadata"]`. As a result, summary subcalls landed on the
router with an empty propagated metadata and the resulting tokens were
not attributed to the caller's key/team budget.
Rename the polyfill chain's spend-propagation parameter to
`litellm_metadata` and pull it from `kwargs["litellm_metadata"]` in
both the async and sync handlers, so the post-call hooks see the parent
key/team and bill the summary tokens accordingly. Add an
`_extract_proxy_litellm_metadata` helper and refactor
`_extract_user_api_key_auth` to use it.
* chore(anthropic adapters): remove unused _extract_user_api_key_auth helper
Co-authored-by: Yassin Kortam
* chore(compact_20260112): non-greedy summary regex; use COMPACT_EDIT_TYPE in bedrock filter
- Make _SUMMARY_TAG_RE non-greedy so a response with multiple
blocks captures only the first complete block.
- Replace the hardcoded 'compact_20260112' literal in
_filter_context_management_for_bedrock_converse with the shared
COMPACT_EDIT_TYPE constant.
* fix: bug fixes from PR review
- streaming_iterator: don't set sent_content_block_finish during compaction
block lifecycle; that flag tracks the regular text/tool_use/thinking block
state machine, conflating the two leaks bad state to introspection paths.
- compact._call_summary_model: send propagated proxy auth/spend-attribution
fields as 'litellm_metadata' instead of 'metadata' so the router's post-call
hooks attribute summary tokens to the caller's key/team budget.
Co-authored-by: Yassin Kortam
* fix(anthropic-streaming): insert content_block_stop between held delta and final message_delta
When the stream exhausts with both `holding_chunk` (a content_block_delta)
and `holding_stop_reason_chunk` (a message_delta) buffered, the after-loop
cleanup previously emitted them back-to-back, producing the invalid
Anthropic SSE sequence `content_block_delta -> message_delta`. Insert a
`content_block_stop` between them in both the sync `__next__` and async
`__anext__` paths so the emitted ordering remains
`content_block_delta -> content_block_stop -> message_delta`.
Co-authored-by: Yassin Kortam
* fix(compact_20260112): propagate allowed_model_region to summary subrequest
The router enforces region restrictions by reading allowed_model_region
from top-level request kwargs (Router._common_checks_available_deployment),
but the compact_20260112 summary subrequest only forwarded litellm_metadata.
A region-restricted caller could trigger compaction and have their conversation
summarized by a deployment outside the permitted region.
Extract allowed_model_region from user_api_key_auth and pass it through
_call_summary_model as a top-level kwarg so the router applies the same
region constraints the parent request would.
* fix(anthropic adapter): emit content_block_stop before held message_delta in drain paths
Co-authored-by: Yassin Kortam
* feat(context_management): configurable summary max_tokens; surface ignored knobs
- compact_20260112: read summary max_tokens from general_settings
(context_management_summary_max_tokens) so operators can fit the
chosen summary model's output budget; falls back to the compiled
default for missing or invalid values.
- clear_tool_uses_20250919: log unsupported knobs at warning level
(was debug, which silently dropped misconfiguration) and surface
them as warnings on the AppliedEdit so clients see what was ignored.
* fix(compact_20260112): bound _call_summary_model with timeout
A slow or unresponsive summary model previously hung the parent
/v1/messages request with no escape hatch. Pass a 60s timeout on the
litellm.acompletion / llm_router.acompletion subrequest; on timeout the
existing summary_call_failed path forwards the request without
compaction rather than blocking indefinitely.
* fix(compact_20260112): preserve post-compaction tail on slice-only path
When a prior compaction block is present and the request is under threshold,
the polyfill was reducing downstream messages to just the latest user
question. The prior summary only covers turns before the compaction block,
so dropping the post-compaction tail silently lost recent context — a
multi-turn conversation that stayed below the threshold would arrive at the
model with no memory of any turn after the prior compaction.
Forward the already-stripped post-compaction tail unchanged on both the
under-threshold path and apply_client_compaction_block_history. Fall
back to _select_last_user_question only when the strip leaves nothing
for the downstream call to answer.
* fix(compact_20260112): enforce user/project/team-member model scopes on summary subrequest
The local gate previously only checked the parent key's and team's
allowed-model lists. A caller restricted by a personal user, project,
or per-team-member allowed_models scope could still trigger the
configured summary model and receive its output as a
compaction block, because llm_router.acompletion bypasses the proxy
common_checks path. Extend _check_summary_model_access to also load
the user_object, project_object, and team_membership and run the
matching allowlist check at each scope before invoking the summary
model.
* fix(compact_20260112): enforce summary model per-model budget and propagate budget metadata
* fix(compact_20260112): forward post-compaction tail when summary model unconfigured
* fix(anthropic endpoints): run failure hook on 500-level context management errors
* fix(compact_20260112): enforce summary model rate limit before summary call
* fix(compact_20260112): propagate end-user/project budget scope to summary call
---------
Co-authored-by: Cursor
Co-authored-by: Yassin Kortam
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Co-authored-by: Claude Sonnet 4.6
---
.../anthropic_interface/messages/__init__.py | 7 +-
litellm/litellm_core_utils/token_counter.py | 22 +-
.../adapters/handler.py | 390 ++-
.../adapters/streaming_iterator.py | 479 +++-
.../adapters/transformation.py | 73 +-
.../context_management/__init__.py | 11 +
.../context_management/constants.py | 45 +
.../context_management/dispatcher.py | 127 +
.../context_management/editors/__init__.py | 4 +
.../editors/clear_tool_uses.py | 210 ++
.../context_management/editors/compact.py | 1206 +++++++++
.../context_management/errors.py | 14 +
.../context_management/placeholders.py | 14 +
.../context_management/result.py | 53 +
.../messages/handler.py | 26 +-
.../bedrock/chat/converse_transformation.py | 76 +-
.../anthropic_claude3_transformation.py | 7 +-
.../proxy/anthropic_endpoints/endpoints.py | 19 +
litellm/types/llms/anthropic.py | 38 +-
.../anthropic_messages/anthropic_response.py | 4 +-
.../test_context_management_polyfill.py | 272 ++
...al_pass_through_adapters_transformation.py | 169 ++
.../test_streaming_iterator_compaction.py | 193 ++
.../context_management/__init__.py | 0
.../test_clear_tool_uses.py | 307 +++
.../context_management/test_compact.py | 2291 +++++++++++++++++
.../context_management/test_dispatcher.py | 131 +
.../test_converse_context_management.py | 114 +
28 files changed, 6201 insertions(+), 101 deletions(-)
create mode 100644 litellm/llms/anthropic/experimental_pass_through/context_management/__init__.py
create mode 100644 litellm/llms/anthropic/experimental_pass_through/context_management/constants.py
create mode 100644 litellm/llms/anthropic/experimental_pass_through/context_management/dispatcher.py
create mode 100644 litellm/llms/anthropic/experimental_pass_through/context_management/editors/__init__.py
create mode 100644 litellm/llms/anthropic/experimental_pass_through/context_management/editors/clear_tool_uses.py
create mode 100644 litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py
create mode 100644 litellm/llms/anthropic/experimental_pass_through/context_management/errors.py
create mode 100644 litellm/llms/anthropic/experimental_pass_through/context_management/placeholders.py
create mode 100644 litellm/llms/anthropic/experimental_pass_through/context_management/result.py
create mode 100644 tests/pass_through_unit_tests/test_context_management_polyfill.py
create mode 100644 tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_compaction.py
create mode 100644 tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/__init__.py
create mode 100644 tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_clear_tool_uses.py
create mode 100644 tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py
create mode 100644 tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_dispatcher.py
create mode 100644 tests/test_litellm/llms/bedrock/test_converse_context_management.py
diff --git a/litellm/anthropic_interface/messages/__init__.py b/litellm/anthropic_interface/messages/__init__.py
index 0996d62c866..f71279b226d 100644
--- a/litellm/anthropic_interface/messages/__init__.py
+++ b/litellm/anthropic_interface/messages/__init__.py
@@ -10,7 +10,7 @@ This is an __init__.py file to allow the following interface
"""
-from typing import Any, AsyncIterator, Coroutine, Dict, List, Optional, Union
+from typing import Any, AsyncIterator, Coroutine, Dict, Iterator, List, Optional, Union
from litellm.llms.anthropic.experimental_pass_through.messages.handler import (
anthropic_messages as _async_anthropic_messages,
@@ -100,8 +100,11 @@ def create(
**kwargs,
) -> Union[
AnthropicMessagesResponse,
+ Iterator[bytes],
AsyncIterator[Any],
- Coroutine[Any, Any, Union[AnthropicMessagesResponse, AsyncIterator[Any]]],
+ Coroutine[
+ Any, Any, Union[AnthropicMessagesResponse, AsyncIterator[Any], Iterator[bytes]]
+ ],
]:
"""
Async wrapper for Anthropic's messages API
diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py
index 7889336f416..74b41062174 100644
--- a/litellm/litellm_core_utils/token_counter.py
+++ b/litellm/litellm_core_utils/token_counter.py
@@ -772,11 +772,29 @@ def _format_function_definitions(tools):
lines.append("namespace functions {")
lines.append("")
for tool in tools:
+ if not isinstance(tool, dict):
+ continue
function = tool.get("function")
+ if not isinstance(function, dict):
+ # Anthropic tool shape → OpenAI function dict for token counting.
+ params = tool.get("input_schema") or tool.get("parameters") or {}
+ if not isinstance(params, dict):
+ params = {}
+ function = {
+ "name": tool.get("name"),
+ "description": tool.get("description"),
+ "parameters": params,
+ }
+ function_name = function.get("name")
+ if not function_name:
+ # Skip malformed tools missing a name to avoid emitting
+ # ``type None = ...`` which would produce inaccurate token counts.
+ continue
if function_description := function.get("description"):
lines.append(f"// {function_description}")
- function_name = function.get("name")
- parameters = function.get("parameters", {})
+ parameters = function.get("parameters") or {}
+ if not isinstance(parameters, dict):
+ parameters = {}
properties = parameters.get("properties")
if properties and properties.keys():
lines.append(f"type {function_name} = (_: {{")
diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py
index 8ed6126d2eb..efb913f709a 100644
--- a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py
+++ b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py
@@ -4,6 +4,7 @@ from typing import (
AsyncIterator,
Coroutine,
Dict,
+ Iterator,
List,
Optional,
Tuple,
@@ -12,9 +13,16 @@ from typing import (
)
import litellm
+from litellm._logging import verbose_logger
+from litellm.litellm_core_utils.asyncify import run_async_function
from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import (
AnthropicAdapter,
)
+from litellm.llms.anthropic.experimental_pass_through.context_management import (
+ AnthropicContextManagementError,
+ PolyfillResult,
+ apply_context_management,
+)
from litellm.llms.anthropic.experimental_pass_through.utils import (
is_reasoning_auto_summary_enabled,
)
@@ -28,15 +36,266 @@ if TYPE_CHECKING:
pass
-# Anthropic-only fields that the translator above already maps into the
-# OpenAI-format completion_kwargs (output_config → reasoning_effort /
-# response_format, etc.). They must be filtered out of the raw
-# extra_kwargs re-merge below or non-Anthropic backends reject the call
-# with 400 "Extra inputs are not permitted". Add new entries here when
-# extending AnthropicMessagesRequestOptionalParams with another Anthropic-
-# specific key.
+# Anthropic-only keys already mapped by the translator; strip on extra_kwargs re-merge.
ANTHROPIC_ONLY_REQUEST_KEYS: frozenset[str] = frozenset({"output_config"})
+
+def _messages_have_compaction_block(messages: List[Dict]) -> bool:
+ """Return True when any message carries a ``compaction`` content block."""
+ for msg in messages:
+ content = msg.get("content")
+ if not isinstance(content, list):
+ continue
+ for block in content:
+ if isinstance(block, dict) and block.get("type") == "compaction":
+ return True
+ return False
+
+
+def _extract_proxy_litellm_metadata(kwargs: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+ """Return ``kwargs["litellm_metadata"]`` when it's a dict; ``None`` otherwise.
+
+ The proxy attaches its auth/spend-attribution fields (``user_api_key``,
+ ``user_api_key_team_id``, ``litellm_call_id``, the full ``UserAPIKeyAuth``
+ object under ``user_api_key_auth``, ...) to ``data["litellm_metadata"]``
+ for ``/v1/messages`` (see
+ ``LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata`` and
+ ``LITELLM_METADATA_ROUTES``). The Anthropic-shape ``metadata`` arg only
+ carries ``user_id`` and must not be conflated. Returns ``None`` for SDK
+ callers that bypass the proxy entirely.
+ """
+ litellm_metadata = kwargs.get("litellm_metadata")
+ if not isinstance(litellm_metadata, dict):
+ return None
+ return litellm_metadata
+
+
+async def _prepare_context_managed_request(
+ *,
+ model: str,
+ messages: List[Dict],
+ tools: Optional[List[Dict]],
+ system: Optional[Any],
+ context_management_spec: Any,
+ litellm_metadata: Optional[Dict],
+ drop_params: Optional[bool],
+ llm_router: Any,
+ user_api_key_auth: Any = None,
+) -> Optional[PolyfillResult]:
+ """Apply client compaction history, then optional context_management polyfill."""
+ from litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact import (
+ apply_client_compaction_block_history,
+ )
+
+ # Skip the client-history pre-processing when a ``compact_20260112``
+ # polyfill spec will run: that editor already slices around any client-sent
+ # compaction block in its Phase A (and uses the full post-compaction tail
+ # for its token-threshold check). Pre-collapsing to just the latest user
+ # question here would starve the polyfill of conversation context and
+ # silently drop intermediate turns.
+ polyfill_will_run = _polyfill_will_run(
+ context_management_spec=context_management_spec,
+ drop_params=drop_params,
+ )
+
+ if polyfill_will_run:
+ history_result: Optional[PolyfillResult] = None
+ working_messages: List[Dict] = messages
+ working_system: Optional[Any] = system
+ else:
+ history_result = apply_client_compaction_block_history(
+ messages=cast(List[Dict[str, Any]], messages),
+ system=system,
+ )
+ working_messages = (
+ history_result.messages if history_result is not None else messages
+ )
+ working_system = history_result.system if history_result is not None else system
+
+ polyfill_result = await _run_polyfill_if_enabled(
+ model=model,
+ messages=working_messages,
+ tools=tools,
+ system=working_system,
+ context_management_spec=context_management_spec,
+ litellm_metadata=litellm_metadata,
+ drop_params=drop_params,
+ llm_router=llm_router,
+ user_api_key_auth=user_api_key_auth,
+ )
+
+ if polyfill_result is not None:
+ return polyfill_result
+
+ # Safety net: if we skipped client-history pre-processing because a
+ # ``compact_20260112`` polyfill was expected to handle the compaction
+ # block itself but the polyfill ultimately did not produce a result
+ # (e.g. it crashed and was best-effort swallowed in
+ # ``_run_polyfill_if_enabled``), apply the slice-only fallback now so
+ # Anthropic-specific ``compaction`` content blocks don't leak through
+ # to non-Anthropic backends that would reject them.
+ if polyfill_will_run and history_result is None:
+ history_result = apply_client_compaction_block_history(
+ messages=cast(List[Dict[str, Any]], messages),
+ system=system,
+ )
+ return history_result
+
+
+def _polyfill_will_run(
+ *,
+ context_management_spec: Any,
+ drop_params: Optional[bool],
+) -> bool:
+ """Return True when ``compact_20260112`` will run via the polyfill dispatcher.
+
+ Mirrors the gating in ``_run_polyfill_if_enabled``: an empty spec or
+ effective ``drop_params`` short-circuits the polyfill. The pre-processing
+ skip only applies when the dispatcher will actually invoke
+ ``apply_compact_20260112`` (which has its own compaction-block slicing).
+ """
+ edits = _normalize_spec_edits(
+ context_management_spec=context_management_spec,
+ drop_params=drop_params,
+ )
+ if edits is None:
+ return False
+
+ from litellm.llms.anthropic.experimental_pass_through.context_management.constants import (
+ COMPACT_EDIT_TYPE,
+ )
+
+ return any(
+ isinstance(edit, dict) and edit.get("type") == COMPACT_EDIT_TYPE
+ for edit in edits
+ )
+
+
+def _spec_has_non_compact_edits(
+ *,
+ context_management_spec: Any,
+ drop_params: Optional[bool],
+) -> bool:
+ """Return True when the spec includes edits other than ``compact_20260112``.
+
+ Used to decide whether a polyfill failure can be silently swallowed
+ (compact-only specs have a safe compaction-block slicing fallback) or
+ must be surfaced (other editors like ``clear_tool_uses_20250919`` have
+ no slice-only fallback and would otherwise be dropped without notice).
+ """
+ edits = _normalize_spec_edits(
+ context_management_spec=context_management_spec,
+ drop_params=drop_params,
+ )
+ if edits is None:
+ return False
+
+ from litellm.llms.anthropic.experimental_pass_through.context_management.constants import (
+ COMPACT_EDIT_TYPE,
+ )
+
+ return any(
+ isinstance(edit, dict)
+ and isinstance(edit.get("type"), str)
+ and edit.get("type") != COMPACT_EDIT_TYPE
+ for edit in edits
+ )
+
+
+def _normalize_spec_edits(
+ *,
+ context_management_spec: Any,
+ drop_params: Optional[bool],
+) -> Optional[List[Dict[str, Any]]]:
+ """Return the normalized ``edits`` list, or ``None`` if the polyfill won't run.
+
+ Delegates spec-shape normalization to the dispatcher's ``_normalize_spec``
+ so the prediction here can't drift from what the dispatcher actually does.
+ """
+ if not context_management_spec:
+ return None
+
+ effective_drop_params = (
+ drop_params if drop_params is not None else litellm.drop_params
+ )
+ if effective_drop_params:
+ return None
+
+ from litellm.llms.anthropic.experimental_pass_through.context_management.dispatcher import (
+ _normalize_spec,
+ )
+
+ try:
+ return _normalize_spec(context_management_spec)
+ except Exception:
+ return None
+
+
+async def _run_polyfill_if_enabled(
+ *,
+ model: str,
+ messages: List[Dict],
+ tools: Optional[List[Dict]],
+ system: Optional[Any],
+ context_management_spec: Any,
+ litellm_metadata: Optional[Dict],
+ drop_params: Optional[bool],
+ llm_router: Any,
+ user_api_key_auth: Any = None,
+) -> Optional[PolyfillResult]:
+ """Run the async context_management polyfill if a spec is present.
+
+ Returns ``None`` when the spec is empty or drop_params is on. Raises
+ ``AnthropicContextManagementError`` so the /v1/messages endpoint can
+ emit an Anthropic-format 400. All other exceptions are best-effort
+ swallowed (matches v0 behavior).
+ """
+ if not context_management_spec:
+ return None
+
+ effective_drop_params = (
+ drop_params if drop_params is not None else litellm.drop_params
+ )
+ if effective_drop_params:
+ return None
+
+ try:
+ return await apply_context_management(
+ model=model,
+ messages=messages,
+ tools=tools,
+ system=system,
+ context_management_spec=context_management_spec,
+ litellm_metadata=litellm_metadata,
+ llm_router=llm_router,
+ user_api_key_auth=user_api_key_auth,
+ )
+ except AnthropicContextManagementError:
+ # Surface validation errors so the endpoint can emit an Anthropic-format
+ # 400. Other exception types fall into the best-effort branch below.
+ raise
+ except Exception as e:
+ verbose_logger.exception(
+ "context_management polyfill: skipping edits due to error: %s", e
+ )
+ # Best-effort swallow is only safe for compact-only specs, where the
+ # caller's compaction-block-slicing safety net produces a correct
+ # (if degraded) result. When the spec also requested non-compact
+ # edits (e.g. ``clear_tool_uses_20250919``), the safety net does
+ # NOT re-run those editors, so silently returning ``None`` would
+ # drop them with no error surface. Raise instead so the endpoint
+ # emits an Anthropic-format error.
+ if _spec_has_non_compact_edits(
+ context_management_spec=context_management_spec,
+ drop_params=drop_params,
+ ):
+ raise AnthropicContextManagementError(
+ status_code=500,
+ message=f"context_management polyfill failed: {e}",
+ ) from e
+ return None
+
+
########################################################
# init adapter
ANTHROPIC_ADAPTER = AnthropicAdapter()
@@ -163,7 +422,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
metadata: Optional[Dict] = None,
stop_sequences: Optional[List[str]] = None,
stream: Optional[bool] = False,
- system: Optional[str] = None,
+ system: Optional[Union[str, List[Dict[str, Any]]]] = None,
temperature: Optional[float] = None,
thinking: Optional[Dict] = None,
tool_choice: Optional[Dict] = None,
@@ -307,19 +566,56 @@ class LiteLLMMessagesToCompletionTransformationHandler:
top_p: Optional[float] = None,
output_format: Optional[Dict] = None,
**kwargs,
- ) -> Union[AnthropicMessagesResponse, AsyncIterator]:
+ ) -> Union[AnthropicMessagesResponse, AsyncIterator[Any], Iterator[bytes]]:
"""Handle non-Anthropic models asynchronously using the adapter"""
+ context_management = kwargs.pop("context_management", None)
+ drop_params: Optional[bool] = kwargs.get("drop_params", None)
+ litellm_router = kwargs.pop("litellm_router", None)
+ if litellm_router is None:
+ try:
+ from litellm.proxy.proxy_server import llm_router as _proxy_router
+
+ litellm_router = _proxy_router
+ except Exception:
+ pass
+
+ proxy_litellm_metadata = _extract_proxy_litellm_metadata(kwargs)
+ user_api_key_auth = (
+ proxy_litellm_metadata.get("user_api_key_auth")
+ if proxy_litellm_metadata is not None
+ else None
+ )
+
+ polyfill_result = await _prepare_context_managed_request(
+ model=model,
+ messages=messages,
+ tools=tools,
+ system=system,
+ context_management_spec=context_management,
+ litellm_metadata=proxy_litellm_metadata,
+ drop_params=drop_params,
+ llm_router=litellm_router,
+ user_api_key_auth=user_api_key_auth,
+ )
+
+ effective_messages = (
+ polyfill_result.messages if polyfill_result is not None else messages
+ )
+ effective_system = (
+ polyfill_result.system if polyfill_result is not None else system
+ )
+
(
completion_kwargs,
tool_name_mapping,
) = LiteLLMMessagesToCompletionTransformationHandler._prepare_completion_kwargs(
max_tokens=max_tokens,
- messages=messages,
+ messages=effective_messages,
model=model,
metadata=metadata,
stop_sequences=stop_sequences,
stream=stream,
- system=system,
+ system=effective_system,
temperature=temperature,
thinking=thinking,
tool_choice=tool_choice,
@@ -338,6 +634,8 @@ class LiteLLMMessagesToCompletionTransformationHandler:
completion_response,
model=model,
tool_name_mapping=tool_name_mapping,
+ polyfill_result=polyfill_result,
+ is_async=True,
)
)
if transformed_stream is not None:
@@ -347,6 +645,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
anthropic_response = ANTHROPIC_ADAPTER.translate_completion_output_params(
cast(ModelResponse, completion_response),
tool_name_mapping=tool_name_mapping,
+ polyfill_result=polyfill_result,
)
if anthropic_response is not None:
return anthropic_response
@@ -372,8 +671,13 @@ class LiteLLMMessagesToCompletionTransformationHandler:
**kwargs,
) -> Union[
AnthropicMessagesResponse,
+ Iterator[bytes],
AsyncIterator[Any],
- Coroutine[Any, Any, Union[AnthropicMessagesResponse, AsyncIterator[Any]]],
+ Coroutine[
+ Any,
+ Any,
+ Union[AnthropicMessagesResponse, AsyncIterator[Any], Iterator[bytes]],
+ ],
]:
"""Handle non-Anthropic models using the adapter."""
if _is_async is True:
@@ -395,17 +699,72 @@ class LiteLLMMessagesToCompletionTransformationHandler:
**kwargs,
)
+ # Run the context_management polyfill on the sync path too so that
+ # ``litellm.messages.create()`` callers don't silently lose edits like
+ # ``clear_tool_uses_20250919``. The dispatcher is async (so the
+ # ``compact_20260112`` editor can ``await`` the summarization model);
+ # bridge to it via ``run_async_function``.
+ context_management = kwargs.pop("context_management", None)
+ drop_params: Optional[bool] = kwargs.get("drop_params", None)
+ # Deliberately do NOT auto-attach the proxy ``llm_router`` here:
+ # ``run_async_function`` spawns a new event loop in a worker thread
+ # to bridge to the async dispatcher, but the proxy router's httpx
+ # ``AsyncClient`` instances are bound to the proxy's main event loop.
+ # Reusing them from the new thread's loop violates httpx's single-loop
+ # invariant and can raise ``RuntimeError: Event loop is closed`` or
+ # produce stalled connections. The summary editor falls back to
+ # ``litellm.acompletion`` (which creates a fresh client per call) when
+ # ``llm_router`` is ``None``, which is safe to call from the bridged
+ # loop. The async ``async_anthropic_messages_handler`` path is
+ # unaffected because it ``await``s within the original event loop.
+ litellm_router = kwargs.pop("litellm_router", None)
+
+ # Skip the async bridge entirely when there is nothing for either the
+ # polyfill or the client-history slice-only fallback to do. The vast
+ # majority of sync ``litellm.messages.create()`` requests carry no
+ # ``context_management`` spec and no client-sent ``compaction`` block,
+ # and bridging through a worker-thread event loop just to discover
+ # there is no work is pure overhead.
+ if context_management is None and not _messages_have_compaction_block(messages):
+ polyfill_result: Optional[PolyfillResult] = None
+ else:
+ proxy_litellm_metadata = _extract_proxy_litellm_metadata(kwargs)
+ user_api_key_auth = (
+ proxy_litellm_metadata.get("user_api_key_auth")
+ if proxy_litellm_metadata is not None
+ else None
+ )
+ polyfill_result = run_async_function(
+ _prepare_context_managed_request,
+ model=model,
+ messages=messages,
+ tools=tools,
+ system=system,
+ context_management_spec=context_management,
+ litellm_metadata=proxy_litellm_metadata,
+ drop_params=drop_params,
+ llm_router=litellm_router,
+ user_api_key_auth=user_api_key_auth,
+ )
+
+ effective_messages = (
+ polyfill_result.messages if polyfill_result is not None else messages
+ )
+ effective_system = (
+ polyfill_result.system if polyfill_result is not None else system
+ )
+
(
completion_kwargs,
tool_name_mapping,
) = LiteLLMMessagesToCompletionTransformationHandler._prepare_completion_kwargs(
max_tokens=max_tokens,
- messages=messages,
+ messages=effective_messages,
model=model,
metadata=metadata,
stop_sequences=stop_sequences,
stream=stream,
- system=system,
+ system=effective_system,
temperature=temperature,
thinking=thinking,
tool_choice=tool_choice,
@@ -424,6 +783,8 @@ class LiteLLMMessagesToCompletionTransformationHandler:
completion_response,
model=model,
tool_name_mapping=tool_name_mapping,
+ polyfill_result=polyfill_result,
+ is_async=False,
)
)
if transformed_stream is not None:
@@ -433,6 +794,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
anthropic_response = ANTHROPIC_ADAPTER.translate_completion_output_params(
cast(ModelResponse, completion_response),
tool_name_mapping=tool_name_mapping,
+ polyfill_result=polyfill_result,
)
if anthropic_response is not None:
return anthropic_response
diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py
index c65dfb22730..bacb9f8ddf6 100644
--- a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py
+++ b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py
@@ -3,11 +3,26 @@
import json
import traceback
from collections import deque
-from typing import TYPE_CHECKING, Any, AsyncIterator, Dict, Iterator, Literal, Optional
+from typing import (
+ TYPE_CHECKING,
+ Any,
+ AsyncIterator,
+ Dict,
+ Iterator,
+ List,
+ Literal,
+ Optional,
+)
-from litellm import verbose_logger
+from litellm._logging import verbose_logger
from litellm._uuid import uuid
-from litellm.types.llms.anthropic import UsageDelta
+from litellm.types.llms.anthropic import (
+ AppliedEdit,
+ CompactionBlock,
+ ContextManagementResponse,
+ UsageDelta,
+ UsageIteration,
+)
from litellm.types.utils import AdapterCompletionStreamWrapper
if TYPE_CHECKING:
@@ -37,22 +52,208 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
holding_stop_reason_chunk: Optional[Any] = None
queued_usage_chunk: bool = False
current_content_block_index: int = 0
- current_content_block_start: ContentBlockContentBlockDict = TextBlock(
- type="text",
- text="",
- )
- chunk_queue: deque = deque() # Queue for buffering multiple chunks
def __init__(
self,
completion_stream: Any,
model: str,
tool_name_mapping: Optional[Dict[str, str]] = None,
+ applied_edits: Optional[List[AppliedEdit]] = None,
+ compaction_block: Optional[CompactionBlock] = None,
+ iterations_usage: Optional[List[UsageIteration]] = None,
):
super().__init__(completion_stream)
self.model = model
# Mapping of truncated tool names to original names (for OpenAI's 64-char limit)
self.tool_name_mapping = tool_name_mapping or {}
+ # Polyfill applied_edits on final message_delta.
+ self.applied_edits: List[AppliedEdit] = list(applied_edits or [])
+ # Synthesized compaction block from compact_20260112 polyfill (streaming).
+ self.compaction_block = compaction_block
+ self.iterations_usage = iterations_usage
+ self.sent_compaction_block: bool = False
+ # Per-phase flags so the compaction block's start/delta/stop events
+ # are emitted (and the public state machine is advanced) in
+ # lock-step with the caller actually consuming each event. Pre-
+ # queuing all three would set ``sent_content_block_finish=True``
+ # before the client received ``content_block_stop``, leaving the
+ # observable state inconsistent during the drain window.
+ self.sent_compaction_block_start: bool = False
+ self.sent_compaction_block_delta: bool = False
+ # Per-instance queue for buffering multiple chunks. Must be initialized
+ # here (not at class level) so concurrent streams don't share the same
+ # deque and corrupt each other's SSE event order.
+ self.chunk_queue: deque = deque()
+ # Per-instance default content block. Must be initialized here (not at
+ # class level) so concurrent streams don't share the same mutable dict
+ # — `_should_start_new_content_block` mutates `tool_block["name"]` in
+ # place, which would otherwise leak across streams.
+ self.current_content_block_start: (
+ "AnthropicStreamWrapper.ContentBlockContentBlockDict"
+ ) = self.TextBlock(
+ type="text",
+ text="",
+ )
+
+ def _merge_usage_into_held_stop_reason_chunk(self, chunk: Any) -> Dict[str, Any]:
+ """Merge usage data from ``chunk`` into the held ``message_delta`` chunk.
+
+ Shared by both the sync ``__next__`` and async ``__anext__`` paths so
+ the subtle hold-and-merge logic (cache tokens, ``context_management``
+ attachment, ``UsageDelta`` shape) lives in exactly one place.
+
+ Caller is responsible for managing ``self.holding_stop_reason_chunk``
+ and ``self.queued_usage_chunk`` state and for queuing the returned
+ merged chunk.
+ """
+ assert self.holding_stop_reason_chunk is not None
+ merged_chunk = self.holding_stop_reason_chunk.copy()
+ if "delta" not in merged_chunk:
+ merged_chunk["delta"] = {}
+
+ uncached_input_tokens = chunk.usage.prompt_tokens or 0
+ if (
+ hasattr(chunk.usage, "prompt_tokens_details")
+ and chunk.usage.prompt_tokens_details
+ ):
+ cached_tokens = (
+ getattr(chunk.usage.prompt_tokens_details, "cached_tokens", 0) or 0
+ )
+ uncached_input_tokens -= cached_tokens
+
+ usage_dict: UsageDelta = {
+ "input_tokens": uncached_input_tokens,
+ "output_tokens": chunk.usage.completion_tokens or 0,
+ }
+ if (
+ hasattr(chunk.usage, "_cache_creation_input_tokens")
+ and chunk.usage._cache_creation_input_tokens > 0
+ ):
+ usage_dict["cache_creation_input_tokens"] = (
+ chunk.usage._cache_creation_input_tokens
+ )
+ if (
+ hasattr(chunk.usage, "_cache_read_input_tokens")
+ and chunk.usage._cache_read_input_tokens > 0
+ ):
+ usage_dict["cache_read_input_tokens"] = chunk.usage._cache_read_input_tokens
+ merged_chunk["usage"] = usage_dict
+ if self.applied_edits and "context_management" not in merged_chunk:
+ merged_chunk["context_management"] = ContextManagementResponse(
+ applied_edits=list(self.applied_edits)
+ )
+ return self._augment_message_delta_usage(merged_chunk)
+
+ def _ensure_context_management_attached(
+ self, message_delta_chunk: Dict[str, Any]
+ ) -> Dict[str, Any]:
+ """Attach ``context_management`` to a ``message_delta`` chunk if
+ ``self.applied_edits`` is non-empty and the chunk does not already
+ carry it. Returns the (possibly new) chunk dict.
+
+ Centralizing this guard ensures every ``message_delta`` emission
+ path (merge-with-usage and direct-flush-of-held) consistently
+ surfaces ``applied_edits`` to the client.
+ """
+ if not self.applied_edits or "context_management" in message_delta_chunk:
+ return message_delta_chunk
+ augmented = message_delta_chunk.copy()
+ augmented["context_management"] = ContextManagementResponse(
+ applied_edits=list(self.applied_edits)
+ )
+ return augmented
+
+ def _augment_message_delta_usage(
+ self, message_delta_chunk: Dict[str, Any]
+ ) -> Dict[str, Any]:
+ """Attach polyfill compaction iteration usage to the final message_delta.
+
+ Also defensively re-attaches ``context_management`` so the direct
+ held-chunk flush path stays in sync with the merge path's guarantee
+ when ``self.applied_edits`` is non-empty.
+ """
+ message_delta_chunk = self._ensure_context_management_attached(
+ message_delta_chunk
+ )
+ if self.iterations_usage is None:
+ return message_delta_chunk
+ usage = message_delta_chunk.get("usage")
+ if not isinstance(usage, dict) or "iterations" in usage:
+ return message_delta_chunk
+
+ input_tokens = usage.get("input_tokens", 0) or 0
+ output_tokens = usage.get("output_tokens", 0) or 0
+ augmented = message_delta_chunk.copy()
+ augmented_usage = dict(usage)
+ iterations: List[UsageIteration] = list(self.iterations_usage)
+ # Only emit a ``message`` iteration when we have real token data.
+ # Without a separate usage chunk (e.g. provider sent finish_reason
+ # alone), the held ``message_delta`` carries placeholder zeros from
+ # the translate step; reporting a zero-token iteration would be
+ # misleading and inconsistent with the non-streaming path.
+ if input_tokens > 0 or output_tokens > 0:
+ message_iteration: UsageIteration = {
+ "type": "message",
+ "input_tokens": input_tokens,
+ "output_tokens": output_tokens,
+ }
+ iterations.append(message_iteration)
+ augmented_usage["iterations"] = iterations # type: ignore[typeddict-unknown-key]
+ augmented["usage"] = augmented_usage
+ return augmented
+
+ def _next_compaction_event(self) -> Optional[Dict[str, Any]]:
+ """Return the next compaction content-block SSE event, or ``None``.
+
+ Anthropic delivers compaction as a single delta (no token-by-token
+ streaming), but we still surface it as a proper
+ start → delta → stop trio. Each call returns exactly one event so
+ the state machine (``sent_content_block_finish``,
+ ``current_content_block_index``) is advanced *only* when the
+ terminal stop event is actually handed back to the caller. This
+ prevents an observable window where the flags claim the block is
+ finished while the stop event is still buffered.
+ """
+ if self.compaction_block is None or self.sent_compaction_block:
+ return None
+
+ compaction_index = self.current_content_block_index
+
+ if not self.sent_compaction_block_start:
+ self.sent_compaction_block_start = True
+ return {
+ "type": "content_block_start",
+ "index": compaction_index,
+ # Mirror the text-block shape ({"type": "text", "text": ""}):
+ # send an empty ``content`` field so clients that introspect
+ # ``content_block_start`` see the full block schema. The
+ # actual summary text arrives via the ``content_block_delta``
+ # below.
+ "content_block": {"type": "compaction", "content": ""},
+ }
+
+ if not self.sent_compaction_block_delta:
+ self.sent_compaction_block_delta = True
+ summary_content = self.compaction_block.get("content") or ""
+ return {
+ "type": "content_block_delta",
+ "index": compaction_index,
+ "delta": {"type": "compaction_delta", "content": summary_content},
+ }
+
+ stop_event = {
+ "type": "content_block_stop",
+ "index": compaction_index,
+ }
+ # Don't touch ``sent_content_block_finish`` here: that flag is the
+ # state machine for the regular text/tool_use/thinking block and is
+ # independent of the synthetic compaction block lifecycle. Conflating
+ # them would let outside observers (subclass overrides, introspection
+ # hooks, exception paths) see ``sent_content_block_finish=True``
+ # without any regular content block ever having started.
+ self._increment_content_block_index()
+ self.sent_compaction_block = True
+ return stop_event
def _create_initial_usage_delta(self) -> UsageDelta:
"""
@@ -75,7 +276,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
cache_read_input_tokens=0,
)
- def __next__(self):
+ def __next__(self): # noqa: PLR0915
from .transformation import LiteLLMAnthropicMessagesAdapter
try:
@@ -103,8 +304,17 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
)
return self.chunk_queue.popleft()
+ if (
+ self.sent_compaction_block is False
+ and self.compaction_block is not None
+ ):
+ compaction_event = self._next_compaction_event()
+ if compaction_event is not None:
+ return compaction_event
+
if self.sent_content_block_start is False:
self.sent_content_block_start = True
+ self.sent_content_block_finish = False
self.chunk_queue.append(
{
"type": "content_block_start",
@@ -122,11 +332,45 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
if should_start_new_block:
self._increment_content_block_index()
+ # applied_edits only needs to flow to the final message_delta
+ # (when finish_reason is set); skip threading it through every
+ # intermediate chunk. For the hold-and-merge path below,
+ # context_management is attached directly to the merged chunk,
+ # so the translated ``processed_chunk`` would be discarded —
+ # skip the applied_edits attachment in that case to avoid
+ # allocating a throwaway ``MessageBlockDelta``.
+ will_merge_into_held = (
+ self.holding_stop_reason_chunk is not None
+ and getattr(chunk, "usage", None) is not None
+ )
+ is_final_chunk = chunk.choices[0].finish_reason is not None
processed_chunk = LiteLLMAnthropicMessagesAdapter().translate_streaming_openai_response_to_anthropic(
response=chunk,
current_content_block_index=self.current_content_block_index,
+ applied_edits=(
+ self.applied_edits
+ if is_final_chunk and not will_merge_into_held
+ else None
+ ),
)
+ # Check if this is a usage chunk and we have a held stop_reason chunk
+ if will_merge_into_held:
+ merged_chunk = self._merge_usage_into_held_stop_reason_chunk(chunk)
+ self.chunk_queue.append(merged_chunk)
+ self.queued_usage_chunk = True
+ self.holding_stop_reason_chunk = None
+ return self.chunk_queue.popleft()
+
+ if self.queued_usage_chunk:
+ # Usage has already been merged + emitted. Any trailing
+ # provider events would violate Anthropic SSE ordering
+ # (no chunks may follow the final ``message_delta``), so
+ # silently drop them — matches the async ``__anext__``
+ # behavior where the block-handling logic is gated on
+ # ``not self.queued_usage_chunk``.
+ continue
+
if should_start_new_block and not self.sent_content_block_finish:
# Queue the sequence: content_block_stop -> content_block_start
# For text blocks the trigger chunk is not emitted as a separate
@@ -178,20 +422,64 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
}
)
self.sent_content_block_finish = True
- self.chunk_queue.append(processed_chunk)
+ if processed_chunk.get("delta", {}).get("stop_reason") is not None:
+ self.holding_stop_reason_chunk = processed_chunk
+ else:
+ processed_chunk = self._augment_message_delta_usage(
+ processed_chunk
+ )
+ self.chunk_queue.append(processed_chunk)
return self.chunk_queue.popleft()
elif self.holding_chunk is not None:
self.chunk_queue.append(self.holding_chunk)
+ if processed_chunk.get("type") == "message_delta":
+ processed_chunk = self._augment_message_delta_usage(
+ processed_chunk
+ )
self.chunk_queue.append(processed_chunk)
self.holding_chunk = None
return self.chunk_queue.popleft()
else:
+ if processed_chunk.get("type") == "message_delta":
+ processed_chunk = self._augment_message_delta_usage(
+ processed_chunk
+ )
self.chunk_queue.append(processed_chunk)
return self.chunk_queue.popleft()
- # Handle any remaining held chunks after stream ends
- if self.holding_chunk is not None:
- self.chunk_queue.append(self.holding_chunk)
+ # Handle any remaining held chunks after stream ends. The
+ # buffered ``holding_chunk`` (a ``content_block_delta``) must
+ # precede the final ``message_delta`` so Anthropic SSE event
+ # ordering is preserved. When ``queued_usage_chunk`` is True,
+ # the final ``message_delta`` has already been emitted; any
+ # buffered content delta is dropped rather than emitted after
+ # ``message_delta`` (which would violate SSE ordering and may
+ # confuse strict Anthropic SDK clients).
+ if not self.queued_usage_chunk:
+ if self.holding_chunk is not None:
+ self.chunk_queue.append(self.holding_chunk)
+ self.holding_chunk = None
+ if self.holding_stop_reason_chunk is not None:
+ # A final ``message_delta`` must be preceded by
+ # ``content_block_stop`` so the emitted SSE stays in
+ # valid Anthropic order (... -> content_block_stop ->
+ # message_delta). Emit ``content_block_stop`` here if
+ # the active content block was not already closed.
+ if not self.sent_content_block_finish:
+ self.chunk_queue.append(
+ {
+ "type": "content_block_stop",
+ "index": self.current_content_block_index,
+ }
+ )
+ self.sent_content_block_finish = True
+ self.chunk_queue.append(
+ self._augment_message_delta_usage(
+ self.holding_stop_reason_chunk
+ )
+ )
+ self.holding_stop_reason_chunk = None
+ else:
self.holding_chunk = None
if not self.sent_last_message:
@@ -205,6 +493,26 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
except StopIteration:
if self.chunk_queue:
return self.chunk_queue.popleft()
+ # Handle any held stop_reason chunk. Emit ``content_block_stop``
+ # first if the active content block was not already closed, so
+ # Anthropic SSE ordering is preserved (content_block_stop ->
+ # message_delta).
+ if self.holding_stop_reason_chunk is not None:
+ if not self.sent_content_block_finish:
+ self.sent_content_block_finish = True
+ self.chunk_queue.append(
+ self._augment_message_delta_usage(
+ self.holding_stop_reason_chunk
+ )
+ )
+ self.holding_stop_reason_chunk = None
+ return {
+ "type": "content_block_stop",
+ "index": self.current_content_block_index,
+ }
+ held = self._augment_message_delta_usage(self.holding_stop_reason_chunk)
+ self.holding_stop_reason_chunk = None
+ return held
if self.sent_last_message is False:
self.sent_last_message = True
return {"type": "message_stop"}
@@ -213,7 +521,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
verbose_logger.error(
"Anthropic Adapter - {}\n{}".format(e, traceback.format_exc())
)
- raise StopAsyncIteration
+ raise StopIteration
async def __anext__(self): # noqa: PLR0915
from .transformation import LiteLLMAnthropicMessagesAdapter
@@ -243,8 +551,17 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
)
return self.chunk_queue.popleft()
+ if (
+ self.sent_compaction_block is False
+ and self.compaction_block is not None
+ ):
+ compaction_event = self._next_compaction_event()
+ if compaction_event is not None:
+ return compaction_event
+
if self.sent_content_block_start is False:
self.sent_content_block_start = True
+ self.sent_content_block_finish = False
self.chunk_queue.append(
{
"type": "content_block_start",
@@ -263,57 +580,31 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
if should_start_new_block:
self._increment_content_block_index()
+ # applied_edits only needs to flow to the final message_delta
+ # (when finish_reason is set); skip threading it through every
+ # intermediate chunk. For the hold-and-merge path below,
+ # context_management is attached directly to the merged chunk,
+ # so the translated ``processed_chunk`` would be discarded —
+ # skip the applied_edits attachment in that case to avoid
+ # allocating a throwaway ``MessageBlockDelta``.
+ will_merge_into_held = (
+ self.holding_stop_reason_chunk is not None
+ and getattr(chunk, "usage", None) is not None
+ )
+ is_final_chunk = chunk.choices[0].finish_reason is not None
processed_chunk = LiteLLMAnthropicMessagesAdapter().translate_streaming_openai_response_to_anthropic(
response=chunk,
current_content_block_index=self.current_content_block_index,
+ applied_edits=(
+ self.applied_edits
+ if is_final_chunk and not will_merge_into_held
+ else None
+ ),
)
# Check if this is a usage chunk and we have a held stop_reason chunk
- if (
- self.holding_stop_reason_chunk is not None
- and getattr(chunk, "usage", None) is not None
- ):
- # Merge usage into the held stop_reason chunk
- merged_chunk = self.holding_stop_reason_chunk.copy()
- if "delta" not in merged_chunk:
- merged_chunk["delta"] = {}
-
- # Add usage to the held chunk
- uncached_input_tokens = chunk.usage.prompt_tokens or 0
- if (
- hasattr(chunk.usage, "prompt_tokens_details")
- and chunk.usage.prompt_tokens_details
- ):
- cached_tokens = (
- getattr(
- chunk.usage.prompt_tokens_details, "cached_tokens", 0
- )
- or 0
- )
- uncached_input_tokens -= cached_tokens
-
- usage_dict: UsageDelta = {
- "input_tokens": uncached_input_tokens,
- "output_tokens": chunk.usage.completion_tokens or 0,
- }
- # Add cache tokens if available (for prompt caching support)
- if (
- hasattr(chunk.usage, "_cache_creation_input_tokens")
- and chunk.usage._cache_creation_input_tokens > 0
- ):
- usage_dict["cache_creation_input_tokens"] = (
- chunk.usage._cache_creation_input_tokens
- )
- if (
- hasattr(chunk.usage, "_cache_read_input_tokens")
- and chunk.usage._cache_read_input_tokens > 0
- ):
- usage_dict["cache_read_input_tokens"] = (
- chunk.usage._cache_read_input_tokens
- )
- merged_chunk["usage"] = usage_dict
-
- # Queue the merged chunk and reset
+ if will_merge_into_held:
+ merged_chunk = self._merge_usage_into_held_stop_reason_chunk(chunk)
self.chunk_queue.append(merged_chunk)
self.queued_usage_chunk = True
self.holding_stop_reason_chunk = None
@@ -379,28 +670,63 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
):
self.holding_stop_reason_chunk = processed_chunk
else:
+ processed_chunk = self._augment_message_delta_usage(
+ processed_chunk
+ )
self.chunk_queue.append(processed_chunk)
return self.chunk_queue.popleft()
elif self.holding_chunk is not None:
# Queue both chunks
self.chunk_queue.append(self.holding_chunk)
+ if processed_chunk.get("type") == "message_delta":
+ processed_chunk = self._augment_message_delta_usage(
+ processed_chunk
+ )
self.chunk_queue.append(processed_chunk)
self.holding_chunk = None
return self.chunk_queue.popleft()
else:
- # Queue the current chunk
+ if processed_chunk.get("type") == "message_delta":
+ processed_chunk = self._augment_message_delta_usage(
+ processed_chunk
+ )
self.chunk_queue.append(processed_chunk)
return self.chunk_queue.popleft()
- # Handle any remaining held chunks after stream ends
+ # Handle any remaining held chunks after stream ends. The
+ # buffered ``holding_chunk`` (a ``content_block_delta``) must
+ # precede the final ``message_delta`` so Anthropic SSE event
+ # ordering is preserved. When ``queued_usage_chunk`` is True,
+ # the final ``message_delta`` has already been emitted; any
+ # buffered content delta is dropped rather than emitted after
+ # ``message_delta`` (which would violate SSE ordering and may
+ # confuse strict Anthropic SDK clients).
if not self.queued_usage_chunk:
- if self.holding_stop_reason_chunk is not None:
- self.chunk_queue.append(self.holding_stop_reason_chunk)
- self.holding_stop_reason_chunk = None
-
if self.holding_chunk is not None:
self.chunk_queue.append(self.holding_chunk)
self.holding_chunk = None
+ if self.holding_stop_reason_chunk is not None:
+ # A final ``message_delta`` must be preceded by
+ # ``content_block_stop`` so the emitted SSE stays in
+ # valid Anthropic order (... -> content_block_stop ->
+ # message_delta). Emit ``content_block_stop`` here if
+ # the active content block was not already closed.
+ if not self.sent_content_block_finish:
+ self.chunk_queue.append(
+ {
+ "type": "content_block_stop",
+ "index": self.current_content_block_index,
+ }
+ )
+ self.sent_content_block_finish = True
+ self.chunk_queue.append(
+ self._augment_message_delta_usage(
+ self.holding_stop_reason_chunk
+ )
+ )
+ self.holding_stop_reason_chunk = None
+ else:
+ self.holding_chunk = None
if not self.sent_last_message:
self.sent_last_message = True
@@ -416,9 +742,28 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
# Handle any remaining queued chunks before stopping
if self.chunk_queue:
return self.chunk_queue.popleft()
- # Handle any held stop_reason chunk
+ # Handle any held stop_reason chunk — clear after capturing so a
+ # subsequent ``__anext__`` call doesn't re-emit the same chunk
+ # (matches the sync ``__next__`` path). Emit ``content_block_stop``
+ # first if the active content block was not already closed, so
+ # Anthropic SSE ordering is preserved (content_block_stop ->
+ # message_delta).
if self.holding_stop_reason_chunk is not None:
- return self.holding_stop_reason_chunk
+ if not self.sent_content_block_finish:
+ self.sent_content_block_finish = True
+ self.chunk_queue.append(
+ self._augment_message_delta_usage(
+ self.holding_stop_reason_chunk
+ )
+ )
+ self.holding_stop_reason_chunk = None
+ return {
+ "type": "content_block_stop",
+ "index": self.current_content_block_index,
+ }
+ held = self._augment_message_delta_usage(self.holding_stop_reason_chunk)
+ self.holding_stop_reason_chunk = None
+ return held
if not self.sent_last_message:
self.sent_last_message = True
return {"type": "message_stop"}
diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py
index 51a1e739a0f..02e0c562654 100644
--- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py
+++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py
@@ -6,6 +6,7 @@ from typing import (
Any,
AsyncIterator,
Dict,
+ Iterator,
List,
Literal,
Optional,
@@ -75,6 +76,9 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
from litellm.litellm_core_utils.prompt_templates.factory import (
THOUGHT_SIGNATURE_SEPARATOR,
)
+from litellm.llms.anthropic.experimental_pass_through.context_management import (
+ PolyfillResult,
+)
from litellm.types.llms.anthropic import (
ANTHROPIC_HOSTED_TOOLS,
AllAnthropicToolsValues,
@@ -87,14 +91,17 @@ from litellm.types.llms.anthropic import (
AnthropicResponseContentBlockText,
AnthropicResponseContentBlockThinking,
AnthropicResponseContentBlockToolUse,
+ AppliedEdit,
ContentBlockDelta,
ContentJsonBlockDelta,
ContentTextBlockDelta,
ContentThinkingBlockDelta,
ContentThinkingSignatureBlockDelta,
+ ContextManagementResponse,
MessageBlockDelta,
MessageDelta,
UsageDelta,
+ UsageIteration,
)
from litellm.types.llms.anthropic_messages.anthropic_response import (
AnthropicMessagesResponse,
@@ -195,6 +202,7 @@ class AnthropicAdapter:
self,
response: ModelResponse,
tool_name_mapping: Optional[Dict[str, str]] = None,
+ polyfill_result: Optional[PolyfillResult] = None,
) -> Optional[AnthropicMessagesResponse]:
"""
Translate OpenAI response to Anthropic format.
@@ -204,10 +212,12 @@ class AnthropicAdapter:
tool_name_mapping: Optional mapping of truncated tool names to original names.
Used to restore original names for tools that exceeded
OpenAI's 64-char limit.
+ polyfill_result: PolyfillResult from context_management polyfill.
"""
return LiteLLMAnthropicMessagesAdapter().translate_openai_response_to_anthropic(
response=response,
tool_name_mapping=tool_name_mapping,
+ polyfill_result=polyfill_result,
)
def translate_completion_output_params_streaming(
@@ -215,7 +225,9 @@ class AnthropicAdapter:
completion_stream: Any,
model: str,
tool_name_mapping: Optional[Dict[str, str]] = None,
- ) -> Union[AsyncIterator[bytes], None]:
+ polyfill_result: Optional[PolyfillResult] = None,
+ is_async: bool = True,
+ ) -> Union[AsyncIterator[bytes], Iterator[bytes], None]:
"""
Translate OpenAI streaming response to Anthropic format.
@@ -223,14 +235,35 @@ class AnthropicAdapter:
completion_stream: The OpenAI streaming response
model: The model name
tool_name_mapping: Optional mapping of truncated tool names to original names.
+ polyfill_result: PolyfillResult from context_management polyfill.
+ is_async: When ``True`` (default, for back-compat with existing
+ async callers) returns an ``AsyncIterator[bytes]``. When
+ ``False`` returns a sync ``Iterator[bytes]`` so sync callers
+ (e.g. ``litellm.anthropic.messages.create(stream=True)`` via
+ the sync handler) don't get back an async iterator they
+ can't iterate without an event loop.
"""
+ applied_edits = (
+ polyfill_result.applied_edits_for_response() if polyfill_result else None
+ )
+ compaction_block = (
+ polyfill_result.compaction_block if polyfill_result is not None else None
+ )
+ iterations_usage = (
+ polyfill_result.iterations_usage if polyfill_result is not None else None
+ )
anthropic_wrapper = AnthropicStreamWrapper(
completion_stream=completion_stream,
model=model,
tool_name_mapping=tool_name_mapping,
+ applied_edits=applied_edits,
+ compaction_block=compaction_block,
+ iterations_usage=iterations_usage,
)
- # Return the SSE-wrapped version for proper event formatting
- return anthropic_wrapper.async_anthropic_sse_wrapper()
+ # Return the SSE-wrapped version for proper event formatting.
+ if is_async:
+ return anthropic_wrapper.async_anthropic_sse_wrapper()
+ return anthropic_wrapper.anthropic_sse_wrapper()
class LiteLLMAnthropicMessagesAdapter:
@@ -1342,6 +1375,7 @@ class LiteLLMAnthropicMessagesAdapter:
self,
response: ModelResponse,
tool_name_mapping: Optional[Dict[str, str]] = None,
+ polyfill_result: Optional[PolyfillResult] = None,
) -> AnthropicMessagesResponse:
"""
Translate OpenAI response to Anthropic format.
@@ -1351,12 +1385,17 @@ class LiteLLMAnthropicMessagesAdapter:
tool_name_mapping: Optional mapping of truncated tool names to original names.
Used to restore original names for tools that exceeded
OpenAI's 64-char limit.
+ polyfill_result: PolyfillResult from context_management polyfill.
"""
## translate content block
anthropic_content = self._translate_openai_content_to_anthropic(
choices=response.choices, # type: ignore
tool_name_mapping=tool_name_mapping,
)
+
+ if polyfill_result is not None and polyfill_result.compaction_block is not None:
+ anthropic_content.insert(0, polyfill_result.compaction_block) # type: ignore[arg-type]
+
## extract finish reason
anthropic_finish_reason = self._translate_openai_finish_reason_to_anthropic(
openai_finish_reason=response.choices[0].finish_reason # type: ignore
@@ -1385,6 +1424,14 @@ class LiteLLMAnthropicMessagesAdapter:
if cached_tokens > 0:
anthropic_usage["cache_read_input_tokens"] = cached_tokens
+ if polyfill_result is not None and polyfill_result.iterations_usage is not None:
+ message_iteration: UsageIteration = {
+ "type": "message",
+ "input_tokens": uncached_input_tokens,
+ "output_tokens": usage.completion_tokens or 0,
+ }
+ anthropic_usage["iterations"] = list(polyfill_result.iterations_usage) + [message_iteration] # type: ignore[typeddict-unknown-key]
+
translated_obj = AnthropicMessagesResponse(
id=response.id,
type="message",
@@ -1396,6 +1443,14 @@ class LiteLLMAnthropicMessagesAdapter:
stop_reason=anthropic_finish_reason,
)
+ applied_edits = (
+ polyfill_result.applied_edits_for_response() if polyfill_result else None
+ )
+ if applied_edits:
+ translated_obj["context_management"] = ContextManagementResponse(
+ applied_edits=list(applied_edits)
+ )
+
return translated_obj
def _translate_streaming_openai_chunk_to_anthropic_content_block(
@@ -1528,7 +1583,10 @@ class LiteLLMAnthropicMessagesAdapter:
return "text_delta", ContentTextBlockDelta(type="text_delta", text=text)
def translate_streaming_openai_response_to_anthropic(
- self, response: ModelResponse, current_content_block_index: int
+ self,
+ response: ModelResponse,
+ current_content_block_index: int,
+ applied_edits: Optional[List[AppliedEdit]] = None,
) -> Union[ContentBlockDelta, MessageBlockDelta]:
## base case - final chunk w/ finish reason
if response.choices[0].finish_reason is not None:
@@ -1578,9 +1636,14 @@ class LiteLLMAnthropicMessagesAdapter:
usage_delta["cache_read_input_tokens"] = cached_tokens
else:
usage_delta = UsageDelta(input_tokens=0, output_tokens=0)
- return MessageBlockDelta(
+ message_block = MessageBlockDelta(
type="message_delta", delta=delta, usage=usage_delta # type: ignore
)
+ if applied_edits:
+ message_block["context_management"] = ContextManagementResponse(
+ applied_edits=list(applied_edits)
+ )
+ return message_block
(
type_of_content,
content_block_delta,
diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/__init__.py b/litellm/llms/anthropic/experimental_pass_through/context_management/__init__.py
new file mode 100644
index 00000000000..729b2864524
--- /dev/null
+++ b/litellm/llms/anthropic/experimental_pass_through/context_management/__init__.py
@@ -0,0 +1,11 @@
+from .constants import CLEARED_TOOL_RESULT_PLACEHOLDER
+from .dispatcher import apply_context_management
+from .errors import AnthropicContextManagementError
+from .result import PolyfillResult
+
+__all__ = [
+ "apply_context_management",
+ "AnthropicContextManagementError",
+ "CLEARED_TOOL_RESULT_PLACEHOLDER",
+ "PolyfillResult",
+]
diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/constants.py b/litellm/llms/anthropic/experimental_pass_through/context_management/constants.py
new file mode 100644
index 00000000000..ebbc182c427
--- /dev/null
+++ b/litellm/llms/anthropic/experimental_pass_through/context_management/constants.py
@@ -0,0 +1,45 @@
+"""Constants for the in-gateway context-management polyfill."""
+
+CLEAR_TOOL_USES_EDIT_TYPE = "clear_tool_uses_20250919"
+
+DEFAULT_INPUT_TOKENS_TRIGGER = 100_000
+DEFAULT_KEEP_TOOL_USES = 3
+
+CLEARED_TOOL_RESULT_PLACEHOLDER = "[Cleared by context management]"
+
+# compact_20260112
+COMPACT_EDIT_TYPE = "compact_20260112"
+COMPACT_DEFAULT_TRIGGER_TOKENS = 150_000
+COMPACT_MIN_TRIGGER_TOKENS = 50_000
+# Default ``max_tokens`` for the summary call. Required by providers like
+# Anthropic that reject requests without it; safely accepted by providers that
+# don't strictly require it. Chosen to comfortably fit a long structured
+# summary. Operators can override via
+# ``general_settings.context_management_summary_max_tokens``.
+COMPACT_SUMMARY_MAX_TOKENS = 4096
+COMPACT_SUMMARY_MAX_TOKENS_SETTING_KEY = "context_management_summary_max_tokens"
+# Wall-clock bound for the summary sub-call. Without this a slow or
+# unresponsive summary model would hang the parent ``/v1/messages`` request
+# with no escape hatch; on timeout the editor falls into the standard
+# ``summary_call_failed`` path and forwards the request without compaction.
+COMPACT_SUMMARY_TIMEOUT_SECONDS = 60.0
+COMPACT_SUMMARY_MODEL_SETTING_KEY = "context_management_summary_model"
+COMPACT_SUMMARY_SYSTEM_PREFIX = "Previous conversation summary: "
+
+# Default summarization prompt from the Anthropic spec.
+COMPACT_DEFAULT_INSTRUCTIONS = (
+ "You have written a partial transcript for the initial task above. Please "
+ "write a summary of the transcript. The purpose of this summary is to "
+ "provide continuity so you can continue to make progress towards solving "
+ "the task in a future context, where the raw history above may not be "
+ "accessible and will be replaced with this summary. Write down anything "
+ "that would be helpful, including the state, next steps, learnings etc. "
+ "You must wrap your summary in a (.*?) ", re.IGNORECASE | re.DOTALL)
+
+
+def _read_summary_model_setting() -> Optional[str]:
+ """Look up the configured summarization model from proxy general_settings."""
+ try:
+ from litellm.proxy.proxy_server import general_settings
+ except Exception:
+ return None
+ value = general_settings.get(COMPACT_SUMMARY_MODEL_SETTING_KEY)
+ return value if isinstance(value, str) and value else None
+
+
+def _read_summary_max_tokens_setting() -> int:
+ """Look up the configured summary ``max_tokens`` from proxy general_settings.
+
+ Falls back to :data:`COMPACT_SUMMARY_MAX_TOKENS` when the setting is
+ missing or invalid (non-positive int, wrong type). Operators tune this
+ when the default doesn't fit their chosen summary model's output budget.
+ """
+ try:
+ from litellm.proxy.proxy_server import general_settings
+ except Exception:
+ return COMPACT_SUMMARY_MAX_TOKENS
+ value = general_settings.get(COMPACT_SUMMARY_MAX_TOKENS_SETTING_KEY)
+ if isinstance(value, int) and value > 0:
+ return value
+ return COMPACT_SUMMARY_MAX_TOKENS
+
+
+async def _check_summary_model_access( # noqa: PLR0915
+ user_api_key_auth: Any,
+ summary_model: str,
+ llm_router: Any,
+) -> bool:
+ """Return True when every model-allowlist scope on the parent request is
+ satisfied for ``summary_model``.
+
+ The summary subrequest does not pass through ``user_api_key_auth`` again,
+ so without this gate a caller whose configured scope at any of these
+ levels excludes ``context_management_summary_model`` could still get the
+ proxy to invoke that model and return its ```` output as a
+ compaction block. Mirrors the model-scope enforcement that
+ ``litellm.proxy.auth.common_checks`` runs for the client-requested model:
+ key, team, user (personal), project, and team-member allowlists.
+
+ Returns True (allow) when ``user_api_key_auth`` is not present — SDK
+ callers and tests run outside the proxy, where no key/team policy exists.
+ Returns False when any of the active allowlists denies the summary model
+ (``ProxyException`` from ``_can_object_call_model`` / ``can_*_model``).
+ Unexpected errors during an access check fail closed but are logged
+ separately so operators can distinguish them from a real access-denied
+ response. DB-lookup failures (object missing from cache or DB) skip the
+ corresponding scope — matching ``common_checks``, which only enforces a
+ scope when its backing object can be loaded.
+ """
+ if user_api_key_auth is None:
+ return True
+ try:
+ from litellm.proxy._types import ProxyException
+ from litellm.proxy.auth.auth_checks import (
+ _can_object_call_model,
+ can_project_access_model,
+ can_user_call_model,
+ get_project_object,
+ get_team_membership,
+ get_user_object,
+ )
+ from litellm.proxy.proxy_server import (
+ prisma_client,
+ proxy_logging_obj,
+ user_api_key_cache,
+ )
+ except Exception:
+ return True
+
+ key_models = list(getattr(user_api_key_auth, "models", None) or [])
+ team_id = getattr(user_api_key_auth, "team_id", None)
+ team_model_aliases = getattr(user_api_key_auth, "team_model_aliases", None)
+ team_models = list(getattr(user_api_key_auth, "team_models", None) or [])
+ user_id = getattr(user_api_key_auth, "user_id", None)
+ project_id = getattr(user_api_key_auth, "project_id", None)
+
+ checks: Tuple[Tuple[Literal["key", "team"], List[str]], ...] = (
+ ("key", key_models),
+ ("team", team_models),
+ )
+ for object_type, models in checks:
+ if not models:
+ continue
+ try:
+ _can_object_call_model(
+ model=summary_model,
+ llm_router=llm_router,
+ models=models,
+ team_model_aliases=team_model_aliases,
+ team_id=team_id,
+ object_type=object_type,
+ )
+ except ProxyException:
+ return False
+ except Exception as e:
+ verbose_logger.warning(
+ "compact_20260112: unexpected error during %s-level access "
+ "check for summary_model=%s; denying access: %s",
+ object_type,
+ summary_model,
+ e,
+ )
+ return False
+
+ if user_id is not None and prisma_client is not None:
+ try:
+ user_obj = await get_user_object(
+ user_id=user_id,
+ prisma_client=prisma_client,
+ user_api_key_cache=user_api_key_cache,
+ user_id_upsert=False,
+ proxy_logging_obj=proxy_logging_obj,
+ )
+ except Exception as e:
+ verbose_logger.debug(
+ "compact_20260112: user object lookup failed for "
+ "summary_model=%s access check; skipping user-level scope: %s",
+ summary_model,
+ e,
+ )
+ user_obj = None
+ if user_obj is not None:
+ try:
+ await can_user_call_model(
+ model=summary_model,
+ llm_router=llm_router,
+ user_object=user_obj,
+ )
+ except ProxyException:
+ return False
+ except Exception as e:
+ verbose_logger.warning(
+ "compact_20260112: unexpected error during user-level "
+ "access check for summary_model=%s; denying access: %s",
+ summary_model,
+ e,
+ )
+ return False
+
+ if project_id is not None and prisma_client is not None:
+ try:
+ project_obj = await get_project_object(
+ project_id=project_id,
+ prisma_client=prisma_client,
+ user_api_key_cache=user_api_key_cache,
+ proxy_logging_obj=proxy_logging_obj,
+ )
+ except Exception as e:
+ verbose_logger.debug(
+ "compact_20260112: project object lookup failed for "
+ "summary_model=%s access check; skipping project-level scope: %s",
+ summary_model,
+ e,
+ )
+ project_obj = None
+ if project_obj is not None and project_obj.models:
+ try:
+ can_project_access_model(
+ model=summary_model,
+ project_object=project_obj,
+ llm_router=llm_router,
+ )
+ except ProxyException:
+ return False
+ except Exception as e:
+ verbose_logger.warning(
+ "compact_20260112: unexpected error during project-level "
+ "access check for summary_model=%s; denying access: %s",
+ summary_model,
+ e,
+ )
+ return False
+
+ if user_id is not None and team_id is not None and prisma_client is not None:
+ try:
+ team_membership = await get_team_membership(
+ user_id=user_id,
+ team_id=team_id,
+ prisma_client=prisma_client,
+ user_api_key_cache=user_api_key_cache,
+ proxy_logging_obj=proxy_logging_obj,
+ )
+ except Exception as e:
+ verbose_logger.debug(
+ "compact_20260112: team membership lookup failed for "
+ "summary_model=%s access check; skipping member-level scope: %s",
+ summary_model,
+ e,
+ )
+ team_membership = None
+ member_allowed_models = (
+ team_membership.litellm_budget_table.allowed_models
+ if team_membership is not None
+ and team_membership.litellm_budget_table is not None
+ else None
+ )
+ if member_allowed_models:
+ try:
+ _can_object_call_model(
+ model=summary_model,
+ llm_router=llm_router,
+ models=list(member_allowed_models),
+ team_model_aliases=team_model_aliases,
+ team_id=team_id,
+ object_type="team",
+ )
+ except ProxyException:
+ return False
+ except Exception as e:
+ verbose_logger.warning(
+ "compact_20260112: unexpected error during member-level "
+ "access check for summary_model=%s; denying access: %s",
+ summary_model,
+ e,
+ )
+ return False
+
+ return True
+
+
+async def _check_summary_model_budget(
+ user_api_key_auth: Any,
+ summary_model: str,
+) -> bool:
+ """Return True when the caller is within their per-model budget for
+ ``summary_model``.
+
+ The summary subrequest never passes back through ``user_api_key_auth``, so
+ without this gate a caller whose ``model_max_budget`` for
+ ``context_management_summary_model`` is exhausted could keep consuming that
+ model via compaction. Mirrors the ``model_max_budget`` /
+ ``end_user_model_max_budget`` enforcement that ``user_api_key_auth`` runs for
+ the client-requested model. Returns True outside the proxy or when no
+ per-model budget is configured.
+ """
+ if user_api_key_auth is None:
+ return True
+ try:
+ from litellm.proxy.proxy_server import model_max_budget_limiter
+ except Exception:
+ return True
+
+ model_max_budget = getattr(user_api_key_auth, "model_max_budget", None)
+ token = getattr(user_api_key_auth, "token", None)
+ if isinstance(model_max_budget, dict) and model_max_budget and token is not None:
+ try:
+ await model_max_budget_limiter.is_key_within_model_budget(
+ user_api_key_dict=user_api_key_auth,
+ model=summary_model,
+ )
+ except litellm.BudgetExceededError:
+ return False
+ except Exception as e:
+ verbose_logger.warning(
+ "compact_20260112: unexpected error during key model-budget "
+ "check for summary_model=%s; denying: %s",
+ summary_model,
+ e,
+ )
+ return False
+
+ end_user_model_max_budget = getattr(
+ user_api_key_auth, "end_user_model_max_budget", None
+ )
+ end_user_id = getattr(user_api_key_auth, "end_user_id", None)
+ if (
+ isinstance(end_user_model_max_budget, dict)
+ and end_user_model_max_budget
+ and end_user_id is not None
+ ):
+ try:
+ await model_max_budget_limiter.is_end_user_within_model_budget(
+ end_user_id=end_user_id,
+ end_user_model_max_budget=end_user_model_max_budget,
+ model=summary_model,
+ )
+ except litellm.BudgetExceededError:
+ return False
+ except Exception as e:
+ verbose_logger.warning(
+ "compact_20260112: unexpected error during end-user model-budget "
+ "check for summary_model=%s; denying: %s",
+ summary_model,
+ e,
+ )
+ return False
+
+ return True
+
+
+async def _check_summary_model_rate_limit(
+ user_api_key_auth: Any,
+ summary_model: str,
+) -> bool:
+ """Return True when the caller is within their configured RPM/TPM limits
+ for ``summary_model``.
+
+ The summary subrequest never passes back through the proxy's pre-call
+ rate limiter, so without this gate a caller already at their key / team /
+ user RPM or TPM could still drive an extra summary-model completion per
+ allowed ``/v1/messages`` request. This mirrors the read side of
+ ``_PROXY_MaxParallelRequestsHandler_v3.async_pre_call_hook`` for the
+ summary model: it builds the same descriptor set and runs the check in
+ ``read_only`` mode so no counter is reserved or incremented — the summary
+ call's actual usage is still charged exactly once by the limiter's
+ post-call success hook (via the propagated ``litellm_metadata``).
+
+ Returns True (allow) outside the proxy, when the active limiter does not
+ expose the read-only descriptor check (legacy limiter), or when the
+ descriptor set cannot be built — the only deny signal is a definitive
+ ``OVER_LIMIT`` response, so an internal error here forwards the request
+ uncompacted rather than blocking every summary.
+ """
+ if user_api_key_auth is None:
+ return True
+ try:
+ from litellm.proxy.proxy_server import proxy_logging_obj
+ except Exception:
+ return True
+
+ limiter = getattr(proxy_logging_obj, "max_parallel_request_limiter", None)
+ if (
+ limiter is None
+ or not hasattr(limiter, "should_rate_limit")
+ or not hasattr(limiter, "_create_rate_limit_descriptors")
+ ):
+ return True
+
+ try:
+ metadata = getattr(user_api_key_auth, "metadata", None) or {}
+ data = {"model": summary_model}
+ descriptors = limiter._create_rate_limit_descriptors(
+ user_api_key_dict=user_api_key_auth,
+ data=data,
+ rpm_limit_type=metadata.get("rpm_limit_type"),
+ tpm_limit_type=metadata.get("tpm_limit_type"),
+ model_has_failures=False,
+ )
+ limiter._add_team_model_rate_limit_descriptor_from_metadata(
+ user_api_key_dict=user_api_key_auth,
+ requested_model=summary_model,
+ descriptors=descriptors,
+ )
+ limiter._add_project_model_rate_limit_descriptor_from_metadata(
+ user_api_key_dict=user_api_key_auth,
+ requested_model=summary_model,
+ descriptors=descriptors,
+ )
+ descriptors.extend(
+ limiter.create_organization_rate_limit_descriptor(
+ user_api_key_auth, summary_model
+ )
+ )
+ if not descriptors:
+ return True
+ response = await limiter.should_rate_limit(
+ descriptors=descriptors,
+ parent_otel_span=getattr(user_api_key_auth, "parent_otel_span", None),
+ read_only=True,
+ )
+ except Exception as e:
+ verbose_logger.warning(
+ "compact_20260112: unexpected error during rate-limit check for "
+ "summary_model=%s; allowing: %s",
+ summary_model,
+ e,
+ )
+ return True
+ return response.get("overall_code") != "OVER_LIMIT"
+
+
+def _find_latest_compaction_index(
+ messages: List[Dict[str, Any]],
+) -> Tuple[Optional[int], Optional[int]]:
+ """Return (message_index, block_index) of the most recent compaction block.
+
+ ``None, None`` if no compaction block is present. Iterates from the end so
+ only the latest one is considered.
+ """
+ for msg_idx in range(len(messages) - 1, -1, -1):
+ content = messages[msg_idx].get("content")
+ if not isinstance(content, list):
+ continue
+ for blk_idx in range(len(content) - 1, -1, -1):
+ block = content[blk_idx]
+ if isinstance(block, dict) and block.get("type") == "compaction":
+ return msg_idx, blk_idx
+ return None, None
+
+
+def _slice_around_compaction_block(
+ messages: List[Dict[str, Any]],
+) -> Tuple[List[Dict[str, Any]], Optional[Dict[str, Any]]]:
+ """Apply Anthropic's "drop everything before the compaction block" rule.
+
+ Returns ``(sliced_messages_with_compaction_block, compaction_block_dict)``
+ if a block was found, else ``(original_messages, None)``. The sliced result
+ keeps the compaction block in the assistant turn that originally carried
+ it (in practice it's the only block in that turn) so callers can still
+ extract the summary text from it.
+ """
+ msg_idx, blk_idx = _find_latest_compaction_index(messages)
+ if msg_idx is None or blk_idx is None:
+ return messages, None
+
+ original_msg = messages[msg_idx]
+ original_content = original_msg["content"]
+ compaction_block = cast(Dict[str, Any], original_content[blk_idx])
+
+ # Per Anthropic's contract everything before the compaction block is
+ # dropped, including earlier blocks within the same assistant message.
+ sliced_content = list(original_content[blk_idx:])
+ sliced_first_msg = {**original_msg, "content": sliced_content}
+
+ sliced_messages: List[Dict[str, Any]] = [sliced_first_msg]
+ sliced_messages.extend(messages[msg_idx + 1 :])
+ return sliced_messages, compaction_block
+
+
+def _strip_compaction_blocks(
+ messages: List[Dict[str, Any]],
+) -> List[Dict[str, Any]]:
+ """Drop any ``compaction`` content blocks from messages.
+
+ Used to build the downstream-bound message list — the adapter has no
+ concept of a compaction block, so it must not see one.
+ """
+ cleaned: List[Dict[str, Any]] = []
+ for msg in messages:
+ content = msg.get("content")
+ if not isinstance(content, list):
+ cleaned.append(msg)
+ continue
+ filtered = [
+ block
+ for block in content
+ if not (isinstance(block, dict) and block.get("type") == "compaction")
+ ]
+ if not filtered:
+ # The compaction block was the only content; drop the whole turn.
+ continue
+ cleaned.append({**msg, "content": filtered})
+ return cleaned
+
+
+def _augment_system_with_summary(
+ system: Optional[Union[str, List[Dict[str, Any]]]],
+ summary_text: str,
+) -> Union[str, List[Dict[str, Any]]]:
+ """Prepend a "Previous conversation summary: ..." block to ``system``."""
+ prefix = f"{COMPACT_SUMMARY_SYSTEM_PREFIX}{summary_text}\n\n"
+ if system is None:
+ return prefix.rstrip()
+ if isinstance(system, str):
+ return f"{prefix}{system}"
+ # List of content blocks: prepend the prefix to the first text block,
+ # otherwise insert a new text block at the head.
+ for idx, block in enumerate(system):
+ if isinstance(block, dict) and block.get("type") == "text":
+ existing = block.get("text", "") or ""
+ new_block = {**block, "text": f"{prefix}{existing}"}
+ return [*system[:idx], new_block, *system[idx + 1 :]]
+ return [{"type": "text", "text": prefix.rstrip()}, *system]
+
+
+def _resolve_trigger_tokens(edit_spec: Dict[str, Any]) -> Tuple[int, List[str]]:
+ """Validate and resolve ``trigger.value``.
+
+ Raises ``AnthropicContextManagementError`` if the explicitly-supplied value
+ is below the 50k minimum. Unknown ``trigger.type`` values fall back to
+ ``input_tokens`` with a warning.
+ """
+ warnings: List[str] = []
+ trigger = edit_spec.get("trigger") or {}
+ if not isinstance(trigger, dict):
+ warnings.append("trigger_not_a_dict_using_default")
+ return COMPACT_DEFAULT_TRIGGER_TOKENS, warnings
+
+ trigger_type = trigger.get("type", "input_tokens")
+ if trigger_type != "input_tokens":
+ warnings.append(f"unsupported_trigger_type_{trigger_type}_using_input_tokens")
+
+ value = trigger.get("value")
+ if value is None:
+ return COMPACT_DEFAULT_TRIGGER_TOKENS, warnings
+ if not isinstance(value, int):
+ warnings.append("trigger_value_not_int_using_default")
+ return COMPACT_DEFAULT_TRIGGER_TOKENS, warnings
+ if value < COMPACT_MIN_TRIGGER_TOKENS:
+ raise AnthropicContextManagementError(
+ status_code=400,
+ message=(
+ f"context_management.compact_20260112.trigger.value must be at "
+ f"least {COMPACT_MIN_TRIGGER_TOKENS} tokens"
+ ),
+ )
+ return value, warnings
+
+
+def _build_summary_prompt(
+ edit_spec: Dict[str, Any], tools: Optional[List[Dict[str, Any]]]
+) -> str:
+ custom = edit_spec.get("instructions")
+ if isinstance(custom, str) and custom.strip():
+ return custom
+ prompt = COMPACT_DEFAULT_INSTRUCTIONS
+ if tools:
+ prompt = f"{prompt}{COMPACT_NO_TOOL_CALLS_SUFFIX}"
+ return prompt
+
+
+def _propagate_metadata(
+ parent_litellm_metadata: Optional[Dict[str, Any]],
+) -> Dict[str, Any]:
+ """Extract the parent request's auth/spend-attribution fields for the summary subcall.
+
+ The proxy attaches ``user_api_key``, ``user_api_key_team_id`` etc. to
+ ``data["litellm_metadata"]`` (see
+ ``LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata``).
+ Without these on the summary subrequest, the router's post-call hooks
+ cannot attribute summary tokens to the caller's key/team budget.
+ """
+ if not parent_litellm_metadata:
+ return {}
+ propagated: Dict[str, Any] = {}
+ for key in _PROPAGATED_METADATA_KEYS:
+ if key in parent_litellm_metadata:
+ propagated[key] = parent_litellm_metadata[key]
+ return propagated
+
+
+def _count_effective_tokens(
+ model: str,
+ effective_messages: List[Dict[str, Any]],
+ compaction_block: Optional[Dict[str, Any]],
+ tools: Optional[List[Dict[str, Any]]],
+ system: Optional[Union[str, List[Dict[str, Any]]]] = None,
+) -> int:
+ """Token-count the conversation as it will appear downstream.
+
+ The compaction block (if any) becomes a system prefix on the downstream
+ call, so its content still counts even though it isn't in ``messages``.
+ The system prompt (which may already include a prior compaction summary
+ prepended via ``_augment_system_with_summary``) is also counted so the
+ threshold check matches the downstream ``input_tokens`` metric.
+ """
+ # Local import to avoid pulling the adapter at module load time.
+ from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import (
+ LiteLLMAnthropicMessagesAdapter,
+ )
+
+ messages_without_compaction = _strip_compaction_blocks(effective_messages)
+ adapter = LiteLLMAnthropicMessagesAdapter()
+ try:
+ openai_shape = adapter.translate_anthropic_messages_to_openai(
+ messages=cast(Any, messages_without_compaction)
+ )
+ except Exception as e:
+ verbose_logger.debug(
+ "compact_20260112: anthropic→openai translation failed during token "
+ "count, falling back to raw messages: %s",
+ e,
+ )
+ openai_shape = cast(Any, messages_without_compaction)
+
+ # Translate Anthropic-shaped tools (``input_schema``) to OpenAI-shaped
+ # tools (``{"type": "function", "function": {...}}``) so ``token_counter``
+ # gets a consistent format regardless of which counting path it uses.
+ # An inaccurate tool token count here could cause the polyfill to skip
+ # needed compaction or trigger unnecessary summarization.
+ openai_tools: Optional[List[Dict[str, Any]]] = None
+ if tools:
+ try:
+ translated_tools, _ = adapter.translate_anthropic_tools_to_openai(
+ tools=cast(Any, tools)
+ )
+ openai_tools = cast(List[Dict[str, Any]], translated_tools)
+ except Exception as e:
+ verbose_logger.debug(
+ "compact_20260112: anthropic→openai tools translation failed "
+ "during token count, falling back to raw tools: %s",
+ e,
+ )
+ openai_tools = tools
+
+ total = litellm.token_counter(
+ model=model,
+ messages=cast(Any, openai_shape),
+ tools=cast(Any, openai_tools),
+ )
+ if compaction_block is not None:
+ content = compaction_block.get("content") or ""
+ if content:
+ total += litellm.token_counter(model=model, text=content)
+ system_text = _system_to_text(system)
+ if system_text:
+ total += litellm.token_counter(model=model, text=system_text)
+ return total
+
+
+def _system_to_text(
+ system: Optional[Union[str, List[Dict[str, Any]]]],
+) -> str:
+ """Flatten an Anthropic-style ``system`` value into a single string for
+ token counting. Returns ``""`` when ``system`` carries no text."""
+ if system is None:
+ return ""
+ if isinstance(system, str):
+ return system
+ parts: List[str] = []
+ for block in system:
+ if isinstance(block, dict) and block.get("type") == "text":
+ text = block.get("text")
+ if isinstance(text, str) and text:
+ parts.append(text)
+ return "\n".join(parts)
+
+
+def _select_last_user_question(
+ messages: List[Dict[str, Any]],
+) -> List[Dict[str, Any]]:
+ """Pick the most recent ``user`` turn that is a real question.
+
+ Returns a one-element message list with any ``tool_result`` blocks
+ stripped: after compaction the paired ``tool_use`` assistant turn no
+ longer exists in the downstream context, so forwarding ``tool_result``
+ blocks would translate to orphaned ``role=tool`` messages on
+ non-Anthropic providers (OpenAI, Gemini, …) and cause a 400 error.
+
+ Falls back to a synthetic continuation prompt if no eligible turn
+ exists (e.g. the conversation only ever contained ``tool_result``
+ turns, or contained no user turns at all). The downstream call always
+ needs a non-empty user message.
+ """
+ for msg in reversed(messages):
+ if msg.get("role") != "user":
+ continue
+ content = msg.get("content")
+ if isinstance(content, list):
+ filtered = [
+ blk
+ for blk in content
+ if not (isinstance(blk, dict) and blk.get("type") == "tool_result")
+ ]
+ if not filtered:
+ # Purely tool_result — skip and look for an earlier turn.
+ continue
+ if len(filtered) < len(content):
+ return [{**msg, "content": filtered}]
+ return [msg]
+ return [
+ {
+ "role": "user",
+ "content": "Please continue based on the conversation summary above.",
+ }
+ ]
+
+
+def _extract_summary_text(raw: Optional[str]) -> Optional[str]:
+ if not raw:
+ return None
+ match = _SUMMARY_TAG_RE.search(raw)
+ if match is None:
+ return None
+ summary = match.group(1).strip()
+ return summary or None
+
+
+def _system_to_openai_message(
+ system: Optional[Union[str, List[Dict[str, Any]]]],
+) -> Optional[Dict[str, Any]]:
+ """Translate Anthropic-shaped ``system`` to an OpenAI system message.
+
+ Accepts a bare string or a list of Anthropic content blocks; returns
+ ``None`` if no usable text is present. Only ``type=="text"`` blocks are
+ carried over — the summary model has no use for ``cache_control`` or
+ other non-text metadata.
+ """
+ if isinstance(system, str):
+ return {"role": "system", "content": system} if system else None
+ if isinstance(system, list):
+ parts = [
+ block.get("text", "")
+ for block in system
+ if isinstance(block, dict) and block.get("type") == "text"
+ ]
+ joined = "\n\n".join(part for part in parts if part)
+ return {"role": "system", "content": joined} if joined else None
+ return None
+
+
+def _build_summary_messages(
+ effective_messages: List[Dict[str, Any]],
+ prompt: str,
+ system: Optional[Union[str, List[Dict[str, Any]]]] = None,
+) -> List[Dict[str, Any]]:
+ """Build the OpenAI-shape message list for the summary call.
+
+ The caller's ``system`` prompt is prepended (the default summarization
+ instructions reference "the initial task above", which lives in that
+ system prompt); the conversation history is translated to OpenAI shape;
+ the summarization prompt is appended as a final user turn.
+ """
+ from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import (
+ LiteLLMAnthropicMessagesAdapter,
+ )
+
+ stripped = _strip_compaction_blocks(effective_messages)
+ try:
+ openai_messages = (
+ LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai(
+ messages=cast(Any, stripped)
+ )
+ )
+ except Exception as e:
+ verbose_logger.warning(
+ "compact_20260112: anthropic→openai translation failed when "
+ "building summary call; falling back to raw shape: %s",
+ e,
+ )
+ openai_messages = cast(Any, stripped)
+
+ summary_messages: List[Dict[str, Any]] = []
+ system_message = _system_to_openai_message(system)
+ if system_message is not None:
+ summary_messages.append(system_message)
+ summary_messages.extend(openai_messages)
+ # If the last turn is already a user message, merge the summarization
+ # prompt into it. Some providers (and strict OpenAI-compatible endpoints)
+ # reject two consecutive ``role=user`` messages, which would otherwise
+ # silently fall into the ``summary_call_failed`` error path.
+ if summary_messages and _is_user_message(summary_messages[-1]):
+ last_msg = summary_messages[-1]
+ summary_messages[-1] = {
+ **last_msg,
+ "content": _append_text_to_content(last_msg.get("content"), prompt),
+ }
+ else:
+ summary_messages.append({"role": "user", "content": prompt})
+ return summary_messages
+
+
+def _is_user_message(msg: Any) -> bool:
+ return isinstance(msg, dict) and msg.get("role") == "user"
+
+
+def _append_text_to_content(content: Any, extra_text: str) -> Any:
+ """Append ``extra_text`` to an OpenAI-shape message ``content`` field.
+
+ Handles the two common shapes: ``str`` and ``list`` of content parts.
+ For unexpected/empty shapes, fall back so the caller gets a usable value.
+ """
+ if content is None or content == "":
+ return extra_text
+ if isinstance(content, str):
+ return f"{content}\n\n{extra_text}"
+ if isinstance(content, list):
+ return [*content, {"type": "text", "text": extra_text}]
+ return [content, {"type": "text", "text": extra_text}]
+
+
+async def _call_summary_model(
+ *,
+ summary_model: str,
+ summary_messages: List[Dict[str, Any]],
+ metadata: Dict[str, Any],
+ llm_router: Any,
+ allowed_model_region: Optional[str] = None,
+ max_tokens: int = COMPACT_SUMMARY_MAX_TOKENS,
+) -> Any:
+ """Invoke the configured summary model.
+
+ Prefers ``llm_router.acompletion`` so the model alias resolves against the
+ proxy's ``model_list``; falls back to ``litellm.acompletion`` if no router
+ is available (e.g. SDK usage outside the proxy).
+ """
+ # ``max_tokens`` is required by providers like Anthropic and silently
+ # accepted by providers that don't strictly require it (OpenAI etc.).
+ # Setting a sensible default here means the feature works regardless of
+ # which model an admin configures as ``context_management_summary_model``;
+ # operators can override via ``context_management_summary_max_tokens`` in
+ # ``general_settings`` when the default doesn't fit the chosen model's
+ # output budget.
+ # The propagated proxy auth/spend-attribution fields (``user_api_key`` etc.)
+ # must travel as ``litellm_metadata`` — that is the parameter the proxy's
+ # post-call spend hooks read for budget attribution. The provider-level
+ # ``metadata`` kwarg corresponds to the upstream API request body and would
+ # not flow into spend tracking.
+ # ``allowed_model_region`` must travel as a top-level kwarg because the
+ # router enforces region restrictions by reading ``request_kwargs`` directly
+ # (see ``Router._common_checks_available_deployment``); without this the
+ # summary subrequest could be routed to a deployment outside the caller's
+ # permitted region.
+ # ``timeout`` bounds how long a slow/unresponsive summary model can stall
+ # the parent ``/v1/messages`` request. On timeout the caller catches the
+ # exception and surfaces ``applied_edits[0].error = "summary_call_failed"``,
+ # forwarding the request without compaction rather than hanging.
+ call_kwargs: Dict[str, Any] = {
+ "model": summary_model,
+ "messages": summary_messages,
+ "max_tokens": max_tokens,
+ "timeout": COMPACT_SUMMARY_TIMEOUT_SECONDS,
+ "litellm_metadata": metadata,
+ }
+ # The end-user id must also travel as the top-level ``user`` kwarg: legacy
+ # limiter hooks and prometheus end-user tracking read it from there rather
+ # than from ``litellm_metadata``, so without it the summary tokens would not
+ # debit the caller's end-user counters.
+ end_user_id = metadata.get("user_api_key_end_user_id")
+ if end_user_id:
+ call_kwargs["user"] = end_user_id
+ if allowed_model_region is not None:
+ call_kwargs["allowed_model_region"] = allowed_model_region
+ if llm_router is not None and hasattr(llm_router, "acompletion"):
+ return await llm_router.acompletion(**call_kwargs)
+ return await litellm.acompletion(**call_kwargs)
+
+
+def _extract_response_text(response: Any) -> Optional[str]:
+ try:
+ choice = response.choices[0]
+ message = choice.message
+ content = getattr(message, "content", None)
+ if isinstance(content, str):
+ return content
+ # Some providers return a list of content parts.
+ if isinstance(content, list):
+ text_parts = [
+ part.get("text", "")
+ for part in content
+ if isinstance(part, dict) and part.get("type") == "text"
+ ]
+ return "".join(text_parts) or None
+ except (AttributeError, IndexError, KeyError):
+ return None
+ return None
+
+
+def _extract_usage(response: Any) -> Tuple[int, int]:
+ usage = getattr(response, "usage", None)
+ if usage is None:
+ return 0, 0
+ return (
+ int(getattr(usage, "prompt_tokens", 0) or 0),
+ int(getattr(usage, "completion_tokens", 0) or 0),
+ )
+
+
+def apply_client_compaction_block_history(
+ *,
+ messages: List[Dict[str, Any]],
+ system: Optional[Union[str, List[Dict[str, Any]]]],
+) -> Optional[PolyfillResult]:
+ """Honor client-sent compaction blocks without a ``compact_20260112`` edit.
+
+ When the request omits ``context_management`` but the message history already
+ contains a ``compaction`` content block (e.g. Claude Code client-side
+ compaction), apply the same slice-only forwarding as the under-threshold
+ path: the prior summary is prepended to ``system`` and the post-compaction
+ tail is forwarded unchanged (with compaction blocks stripped) so recent
+ turns the summary does not cover are preserved.
+ """
+ effective_messages, prior_compaction_block = _slice_around_compaction_block(
+ messages
+ )
+ if prior_compaction_block is None:
+ return None
+
+ verbose_logger.info(
+ "compact_20260112: client compaction block in message history; "
+ "applying slice-only forwarding (no context_management edit)"
+ )
+
+ prior_summary_text = prior_compaction_block.get("content") or ""
+ augmented_system: Union[str, List[Dict[str, Any]], None] = system
+ if isinstance(prior_summary_text, str) and prior_summary_text:
+ augmented_system = _augment_system_with_summary(system, prior_summary_text)
+ verbose_logger.info(
+ "compact_20260112: compaction summary added to main call system prefix (%s chars)",
+ len(prior_summary_text),
+ )
+
+ # Post-compaction turns are recent context the prior summary does not cover,
+ # so forward them unchanged. Only fall back to the last user question if the
+ # strip leaves the downstream call with nothing to answer.
+ downstream_messages = _strip_compaction_blocks(effective_messages)
+ if not downstream_messages:
+ downstream_messages = _select_last_user_question(effective_messages)
+
+ return PolyfillResult(
+ messages=downstream_messages,
+ system=augmented_system,
+ applied_edits=[],
+ )
+
+
+async def apply_compact_20260112( # noqa: PLR0915
+ *,
+ model: str,
+ messages: List[Dict[str, Any]],
+ tools: Optional[List[Dict[str, Any]]],
+ system: Optional[Union[str, List[Dict[str, Any]]]],
+ edit_spec: Dict[str, Any],
+ litellm_metadata: Optional[Dict[str, Any]] = None,
+ llm_router: Any = None,
+ user_api_key_auth: Any = None,
+) -> PolyfillResult:
+ """Apply ``compact_20260112``; return a ``PolyfillResult``.
+
+ See module docstring for the algorithm. Errors are best-effort: when the
+ summary call fails or the response is malformed, the editor returns the
+ pre-summary state (with ``applied_edits[0].error`` populated) so the
+ original request still proceeds.
+ """
+ # Validation runs first. Raising AnthropicContextManagementError here is
+ # the only path on which the polyfill aborts the request.
+ trigger_tokens, warnings = _resolve_trigger_tokens(edit_spec)
+ verbose_logger.info(
+ "compact_20260112: request has compaction trigger (input_tokens threshold=%s)",
+ trigger_tokens,
+ )
+ if edit_spec.get("pause_after_compaction"):
+ warnings.append("pause_after_compaction_ignored")
+
+ applied: AppliedEdit = {"type": COMPACT_EDIT_TYPE}
+ if warnings:
+ applied["warnings"] = warnings
+
+ # Phase A: slice around any existing compaction block. Runs before the
+ # opt-in gate below so that even when summarization is disabled we still
+ # strip Anthropic-only ``compaction`` blocks from messages going to
+ # non-Anthropic backends (which would reject them).
+ effective_messages, prior_compaction_block = _slice_around_compaction_block(
+ messages
+ )
+ prior_summary_text = (
+ prior_compaction_block.get("content") if prior_compaction_block else None
+ )
+ augmented_system: Union[str, List[Dict[str, Any]], None] = system
+ if isinstance(prior_summary_text, str) and prior_summary_text:
+ augmented_system = _augment_system_with_summary(system, prior_summary_text)
+ verbose_logger.info(
+ "compact_20260112: compaction summary added to main call system prefix (%s chars)",
+ len(prior_summary_text),
+ )
+
+ downstream_messages = _strip_compaction_blocks(effective_messages)
+
+ # Opt-in gate: no summary model configured → no-op (but still return the
+ # Phase A-sliced/stripped messages so compaction blocks don't leak).
+ summary_model = _read_summary_model_setting()
+ if summary_model is None:
+ applied["error"] = "summary_model_not_configured"
+ # Slice-only forwarding: ``augmented_system`` already carries any prior
+ # compaction summary, and the post-compaction tail in
+ # ``downstream_messages`` is recent context the summary does not cover,
+ # so forward it unchanged. Only fall back to the last user question when
+ # the strip leaves nothing for the downstream call to answer.
+ if not downstream_messages:
+ downstream_messages = _select_last_user_question(effective_messages)
+ return PolyfillResult(
+ messages=downstream_messages,
+ system=augmented_system,
+ applied_edits=[applied],
+ )
+
+ # Phase B: threshold check.
+ try:
+ current_tokens = _count_effective_tokens(
+ model=model,
+ effective_messages=effective_messages,
+ # ``augmented_system`` already carries the prior compaction summary
+ # (prepended via ``_augment_system_with_summary``); pass ``None``
+ # here so we don't double-count the summary text.
+ compaction_block=None,
+ tools=tools,
+ system=augmented_system,
+ )
+ except Exception as e:
+ verbose_logger.warning(
+ "compact_20260112: token_counter failed; assuming under threshold: %s", e
+ )
+ current_tokens = 0
+
+ verbose_logger.debug(
+ "compact_20260112: current_tokens=%s trigger=%s", current_tokens, trigger_tokens
+ )
+
+ if current_tokens <= trigger_tokens:
+ # Slice-only path: the prior compaction summary already lives in
+ # ``augmented_system``. Post-compaction turns are recent context the
+ # summary does not cover, so forward ``downstream_messages`` (the
+ # post-compaction tail with compaction blocks stripped) unchanged.
+ # Only fall back to the last user question when the strip leaves
+ # nothing for the downstream call to answer.
+ if not downstream_messages:
+ downstream_messages = _select_last_user_question(effective_messages)
+ return PolyfillResult(
+ messages=downstream_messages,
+ system=augmented_system,
+ applied_edits=[applied],
+ )
+
+ # Phase C: summarize. ``augmented_system`` carries any prior compaction
+ # summary so multi-round compaction does not lose accumulated history —
+ # ``effective_messages`` only contains turns since the last compaction.
+ if not await _check_summary_model_access(
+ user_api_key_auth=user_api_key_auth,
+ summary_model=summary_model,
+ llm_router=llm_router,
+ ):
+ verbose_logger.warning(
+ "compact_20260112: caller not authorized for summary_model=%s; "
+ "skipping summary call",
+ summary_model,
+ )
+ applied["error"] = "summary_model_access_denied"
+ return PolyfillResult(
+ messages=downstream_messages,
+ system=augmented_system,
+ applied_edits=[applied],
+ )
+
+ if not await _check_summary_model_budget(
+ user_api_key_auth=user_api_key_auth,
+ summary_model=summary_model,
+ ):
+ verbose_logger.warning(
+ "compact_20260112: caller over model budget for summary_model=%s; "
+ "skipping summary call",
+ summary_model,
+ )
+ applied["error"] = "summary_model_budget_exceeded"
+ return PolyfillResult(
+ messages=downstream_messages,
+ system=augmented_system,
+ applied_edits=[applied],
+ )
+
+ if not await _check_summary_model_rate_limit(
+ user_api_key_auth=user_api_key_auth,
+ summary_model=summary_model,
+ ):
+ verbose_logger.warning(
+ "compact_20260112: caller over rate limit for summary_model=%s; "
+ "skipping summary call",
+ summary_model,
+ )
+ applied["error"] = "summary_model_rate_limit_exceeded"
+ return PolyfillResult(
+ messages=downstream_messages,
+ system=augmented_system,
+ applied_edits=[applied],
+ )
+
+ prompt = _build_summary_prompt(edit_spec, tools)
+ summary_messages = _build_summary_messages(
+ effective_messages, prompt, system=augmented_system
+ )
+ propagated_metadata = _propagate_metadata(litellm_metadata)
+ allowed_model_region = getattr(user_api_key_auth, "allowed_model_region", None)
+
+ try:
+ response = await _call_summary_model(
+ summary_model=summary_model,
+ summary_messages=summary_messages,
+ metadata=propagated_metadata,
+ llm_router=llm_router,
+ allowed_model_region=allowed_model_region,
+ max_tokens=_read_summary_max_tokens_setting(),
+ )
+ except Exception as e:
+ verbose_logger.warning("compact_20260112: summary call failed: %s", e)
+ applied["error"] = "summary_call_failed"
+ return PolyfillResult(
+ messages=downstream_messages,
+ system=augmented_system,
+ applied_edits=[applied],
+ )
+
+ summary_text = _extract_summary_text(_extract_response_text(response))
+ if summary_text is None:
+ applied["error"] = "summary_extraction_failed"
+ return PolyfillResult(
+ messages=downstream_messages,
+ system=augmented_system,
+ applied_edits=[applied],
+ )
+
+ summary_input_tokens, summary_output_tokens = _extract_usage(response)
+ applied["summary_input_tokens"] = summary_input_tokens
+ applied["summary_output_tokens"] = summary_output_tokens
+
+ compaction_block: CompactionBlock = {
+ "type": "compaction",
+ "content": summary_text,
+ }
+ iterations_usage: List[UsageIteration] = [
+ {
+ "type": "compaction",
+ "input_tokens": summary_input_tokens,
+ "output_tokens": summary_output_tokens,
+ }
+ ]
+
+ # Per Anthropic's contract, everything before the compaction block is
+ # dropped. Phase D: the user/assistant log goes empty; the summary lives
+ # on the system message instead. Anthropic requires a non-empty messages
+ # array, so keep the most recent original user *question* turn so the
+ # model has something to answer. Skip ``tool_result``-only user turns:
+ # in Anthropic's format those are role=user but represent the response
+ # from a tool, and surfacing one as the sole downstream message would
+ # produce an orphaned ``tool``-role message on non-Anthropic providers
+ # with no matching ``tool_calls`` in the prior assistant history. If no
+ # eligible turn exists, fall back to a synthetic continuation prompt so
+ # the downstream call still has a non-empty user message.
+ summarized_system = _augment_system_with_summary(system, summary_text)
+ verbose_logger.info(
+ "compact_20260112: compaction summary added to main call system prefix (%s chars)",
+ len(summary_text),
+ )
+ downstream_messages_after_summary = _select_last_user_question(effective_messages)
+
+ return PolyfillResult(
+ messages=downstream_messages_after_summary,
+ system=summarized_system,
+ applied_edits=[applied],
+ compaction_block=compaction_block,
+ iterations_usage=iterations_usage,
+ )
diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/errors.py b/litellm/llms/anthropic/experimental_pass_through/context_management/errors.py
new file mode 100644
index 00000000000..1b14089a451
--- /dev/null
+++ b/litellm/llms/anthropic/experimental_pass_through/context_management/errors.py
@@ -0,0 +1,14 @@
+"""Exceptions raised by the context_management polyfill."""
+
+
+class AnthropicContextManagementError(Exception):
+ """Validation error from the polyfill, surfaced as an Anthropic-format 4xx.
+
+ The `/v1/messages` endpoint catches this in its exception handler and
+ emits an Anthropic-shaped error body instead of the default OpenAI shape.
+ """
+
+ def __init__(self, *, status_code: int, message: str) -> None:
+ super().__init__(message)
+ self.status_code = status_code
+ self.message = message
diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/placeholders.py b/litellm/llms/anthropic/experimental_pass_through/context_management/placeholders.py
new file mode 100644
index 00000000000..f684d970df4
--- /dev/null
+++ b/litellm/llms/anthropic/experimental_pass_through/context_management/placeholders.py
@@ -0,0 +1,14 @@
+"""Placeholder content for cleared ``tool_result`` blocks (string or block list)."""
+
+from typing import Any, List, Union
+
+from .constants import CLEARED_TOOL_RESULT_PLACEHOLDER
+
+
+def build_cleared_tool_result_content(
+ original_content: Any,
+) -> Union[str, List[dict]]:
+ """Return a string or single text block list, matching ``original_content`` shape."""
+ if isinstance(original_content, list):
+ return [{"type": "text", "text": CLEARED_TOOL_RESULT_PLACEHOLDER}]
+ return CLEARED_TOOL_RESULT_PLACEHOLDER
diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/result.py b/litellm/llms/anthropic/experimental_pass_through/context_management/result.py
new file mode 100644
index 00000000000..36bcde98d0c
--- /dev/null
+++ b/litellm/llms/anthropic/experimental_pass_through/context_management/result.py
@@ -0,0 +1,53 @@
+"""``PolyfillResult`` — the shape returned by the context-management dispatcher.
+
+Threaded from the dispatcher through ``async_anthropic_messages_handler`` into
+the adapter so it can prepend the ``compaction`` block to the response and
+attach ``iterations`` to ``usage``.
+"""
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Union
+
+from litellm.types.llms.anthropic import (
+ AppliedEdit,
+ CompactionBlock,
+ UsageIteration,
+)
+
+from .constants import COMPACT_EDIT_TYPE
+
+
+@dataclass
+class PolyfillResult:
+ messages: List[Dict[str, Any]]
+ system: Optional[Union[str, List[Dict[str, Any]]]]
+ applied_edits: List[AppliedEdit] = field(default_factory=list)
+ compaction_block: Optional[CompactionBlock] = None
+ iterations_usage: Optional[List[UsageIteration]] = None
+
+ def applied_edits_for_response(self) -> Optional[List[AppliedEdit]]:
+ """``applied_edits`` to attach on the client-visible response.
+
+ ``compact_20260112`` is included when a new compaction block was
+ synthesized (success), when the edit carries an ``error`` field
+ (``summary_model_not_configured``, ``summary_call_failed``,
+ ``summary_extraction_failed``), or when the edit carries
+ ``warnings`` (e.g. ``unsupported_trigger_type_X_using_input_tokens``,
+ ``pause_after_compaction_ignored``) — operators and clients need to
+ see why compaction was requested but not applied as expected.
+ Slice-only / under-threshold paths that produced no edit at all
+ (no block, no error, no warnings) are omitted. Other edit types are
+ included when the editor returned an ``AppliedEdit``.
+ """
+ visible: List[AppliedEdit] = []
+ for edit in self.applied_edits:
+ if edit.get("type") == COMPACT_EDIT_TYPE:
+ if (
+ self.compaction_block is not None
+ or edit.get("error")
+ or edit.get("warnings")
+ ):
+ visible.append(edit)
+ else:
+ visible.append(edit)
+ return visible or None
diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py
index 14e06e047ea..62eced8e6f0 100644
--- a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py
+++ b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py
@@ -8,7 +8,17 @@
import asyncio
import contextvars
from functools import partial
-from typing import Any, AsyncIterator, Coroutine, Dict, List, Optional, Union, cast
+from typing import (
+ Any,
+ AsyncIterator,
+ Coroutine,
+ Dict,
+ Iterator,
+ List,
+ Optional,
+ Union,
+ cast,
+)
import litellm
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
@@ -189,7 +199,7 @@ async def anthropic_messages(
client: Optional[AsyncHTTPHandler] = None,
custom_llm_provider: Optional[str] = None,
**kwargs,
-) -> Union[AnthropicMessagesResponse, AsyncIterator]:
+) -> Union[AnthropicMessagesResponse, Iterator[bytes], AsyncIterator[Any]]:
"""
Async: Make llm api request in Anthropic /messages API spec.
@@ -346,8 +356,11 @@ def anthropic_messages_handler(
**kwargs,
) -> Union[
AnthropicMessagesResponse,
+ Iterator[bytes],
AsyncIterator[Any],
- Coroutine[Any, Any, Union[AnthropicMessagesResponse, AsyncIterator[Any]]],
+ Coroutine[
+ Any, Any, Union[AnthropicMessagesResponse, AsyncIterator[Any], Iterator[bytes]]
+ ],
]:
"""
Makes Anthropic `/v1/messages` API calls In the Anthropic API Spec
@@ -456,9 +469,14 @@ def anthropic_messages_handler(
return LiteLLMMessagesToResponsesAPIHandler.anthropic_messages_handler(
**_shared_kwargs
)
+
+ # The in-gateway context_management polyfill runs inside
+ # ``async_anthropic_messages_handler`` so it can ``await`` the
+ # summarization model for ``compact_20260112``. ``context_management``
+ # is passed through as a regular kwarg.
return (
LiteLLMMessagesToCompletionTransformationHandler.anthropic_messages_handler(
- **_shared_kwargs
+ **_shared_kwargs,
)
)
diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py
index 1a8e59ba8db..90dfa13e938 100644
--- a/litellm/llms/bedrock/chat/converse_transformation.py
+++ b/litellm/llms/bedrock/chat/converse_transformation.py
@@ -586,6 +586,9 @@ class AmazonConverseConfig(BaseConfig):
):
supported_params.append("thinking")
supported_params.append("reasoning_effort")
+
+ if base_model.startswith("anthropic"):
+ supported_params.append("context_management")
return supported_params
def map_tool_choice_values(
@@ -947,10 +950,10 @@ class AmazonConverseConfig(BaseConfig):
self._handle_reasoning_effort_parameter(
model=model, reasoning_effort=value, optional_params=optional_params
)
+ elif param == "context_management" and isinstance(value, (dict, list)):
+ self._map_context_management_param(value, optional_params)
if param == "requestMetadata":
- if value is not None and isinstance(value, dict):
- self._validate_request_metadata(value) # type: ignore
- optional_params["requestMetadata"] = value
+ self._map_request_metadata_param(value, optional_params)
if param == "service_tier" and isinstance(value, str):
self._map_service_tier_param(value, optional_params)
@@ -983,6 +986,32 @@ class AmazonConverseConfig(BaseConfig):
return optional_params
+ def _map_request_metadata_param(self, value: Any, optional_params: dict) -> None:
+ if value is not None and isinstance(value, dict):
+ self._validate_request_metadata(value) # type: ignore
+ optional_params["requestMetadata"] = value
+
+ def _map_context_management_param(
+ self, value: Union[dict, list], optional_params: dict
+ ) -> None:
+ # Match the dispatcher's ``_normalize_spec`` behavior: only run the
+ # OpenAI→Anthropic mapper for list inputs. Dict inputs are already in
+ # Anthropic-native shape (``{"edits": [...]}``) and should pass
+ # through unchanged so an Anthropic-format ``context_management``
+ # value isn't silently dropped when the mapper can't classify it.
+ if isinstance(value, list):
+ mapped = AnthropicConfig.map_openai_context_management_to_anthropic(
+ cast(Union[dict, list], value)
+ )
+ else:
+ mapped = value
+ # Skip when the mapper returned None for malformed input — leaving the
+ # key out is safer than passing `context_management: null` downstream,
+ # which Bedrock would reject and which can confuse intermediate checks
+ # before the final _filter_context_management_for_bedrock_converse step.
+ if mapped is not None:
+ optional_params["context_management"] = mapped
+
def _map_service_tier_param(self, value: str, optional_params: dict) -> None:
"""Map OpenAI service_tier (string) to Bedrock serviceTier (object).
@@ -1488,6 +1517,11 @@ class AmazonConverseConfig(BaseConfig):
if ANTHROPIC_EFFORT_BETA_HEADER not in anthropic_beta_list:
anthropic_beta_list.append(ANTHROPIC_EFFORT_BETA_HEADER)
+ # Bedrock Converse: compact_20260112 edits only (+ beta header).
+ AmazonConverseConfig._filter_context_management_for_bedrock_converse(
+ additional_request_params, anthropic_beta_list
+ )
+
# Set anthropic_beta in additional_request_params if we have any beta features
# ONLY apply to Anthropic/Claude models - other models (e.g., Qwen, Llama) don't support this field
if anthropic_beta_list and base_model.startswith("anthropic"):
@@ -1495,6 +1529,42 @@ class AmazonConverseConfig(BaseConfig):
return bedrock_tools, anthropic_beta_list
+ @staticmethod
+ def _filter_context_management_for_bedrock_converse(
+ additional_request_params: dict,
+ anthropic_beta_list: list,
+ ) -> None:
+ """Keep only compact_20260112 edits for Bedrock; add beta header or drop field."""
+ from litellm.llms.anthropic.experimental_pass_through.context_management.constants import (
+ COMPACT_EDIT_TYPE,
+ )
+ from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES
+
+ cm = additional_request_params.get("context_management")
+ if not isinstance(cm, dict):
+ additional_request_params.pop("context_management", None)
+ return
+ edits = cm.get("edits")
+ if not isinstance(edits, list):
+ additional_request_params.pop("context_management", None)
+ return
+
+ compact_edits = [
+ e
+ for e in edits
+ if isinstance(e, dict) and e.get("type") == COMPACT_EDIT_TYPE
+ ]
+ if compact_edits:
+ compact_beta = ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value
+ if compact_beta not in anthropic_beta_list:
+ anthropic_beta_list.append(compact_beta)
+ additional_request_params["context_management"] = {
+ **cm,
+ "edits": compact_edits,
+ }
+ else:
+ additional_request_params.pop("context_management", None)
+
def _transform_request_helper(
self,
model: str,
diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py
index b223f4534fa..42c3bd517a9 100644
--- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py
+++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py
@@ -41,7 +41,10 @@ from litellm.llms.bedrock.common_utils import (
pop_bedrock_invoke_output_config_format,
remove_custom_field_from_tools,
)
-from litellm.types.llms.anthropic import ANTHROPIC_TOOL_SEARCH_BETA_HEADER
+from litellm.types.llms.anthropic import (
+ ANTHROPIC_BETA_HEADER_VALUES,
+ ANTHROPIC_TOOL_SEARCH_BETA_HEADER,
+)
from litellm.types.llms.bedrock import BedrockInvokeAnthropicMessagesRequest
from litellm.types.llms.openai import AllMessageValues
from litellm.types.router import GenericLiteLLMParams
@@ -445,7 +448,7 @@ class AmazonAnthropicClaudeMessagesConfig(
if isinstance(e, dict) and e.get("type") == "compact_20260112"
]
if compact_edits:
- beta_set.add("compact-2026-01-12")
+ beta_set.add(ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value)
anthropic_messages_request["context_management"] = {
**cm,
"edits": compact_edits,
diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py
index 69d69354fd1..900386f3d7b 100644
--- a/litellm/proxy/anthropic_endpoints/endpoints.py
+++ b/litellm/proxy/anthropic_endpoints/endpoints.py
@@ -3,10 +3,14 @@ Unified /v1/messages endpoint - (Anthropic Spec)
"""
from fastapi import APIRouter, Depends, HTTPException, Request, Response
+from fastapi.responses import JSONResponse
from litellm._logging import verbose_proxy_logger
from litellm.anthropic_interface.exceptions import AnthropicExceptionMapping
from litellm.integrations.custom_guardrail import ModifyResponseException
+from litellm.llms.anthropic.experimental_pass_through.context_management import (
+ AnthropicContextManagementError,
+)
from litellm.proxy._types import *
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_request_processing import (
@@ -114,6 +118,21 @@ async def anthropic_response( # noqa: PLR0915
)
return _anthropic_response
+ except AnthropicContextManagementError as e:
+ if e.status_code >= 500:
+ # Server-side polyfill failures hit the failure hook for spend/alert
+ # parity with the generic handler; 4xx validation errors do not.
+ await proxy_logging_obj.post_call_failure_hook(
+ user_api_key_dict=user_api_key_dict,
+ original_exception=e,
+ request_data=data,
+ )
+ body = AnthropicExceptionMapping.transform_to_anthropic_error(
+ status_code=e.status_code,
+ raw_message=e.message,
+ request_id=request.headers.get("x-request-id"),
+ )
+ return JSONResponse(status_code=e.status_code, content=body)
except Exception as e:
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data
diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py
index bbb892a0276..a4a059dc88a 100644
--- a/litellm/types/llms/anthropic.py
+++ b/litellm/types/llms/anthropic.py
@@ -2,7 +2,7 @@ from enum import Enum
from typing import Any, Dict, Iterable, List, Optional, Union
from pydantic import BaseModel, ConfigDict
-from typing_extensions import Literal, Required, TypedDict
+from typing_extensions import Literal, NotRequired, Required, TypedDict
from .openai import (
ChatCompletionCachedContent,
@@ -515,6 +515,41 @@ class UsageDelta(TypedDict, total=False):
cache_read_input_tokens: int
+class AppliedEdit(TypedDict, total=False):
+ """One applied context_management edit (Anthropic response shape)."""
+
+ type: str
+ cleared_input_tokens: int
+ cleared_tool_uses: int
+ cleared_thinking_turns: int
+ # compact_20260112 fields
+ summary_input_tokens: int
+ summary_output_tokens: int
+ error: str
+ warnings: List[str]
+
+
+class ContextManagementResponse(TypedDict, total=False):
+ """Response ``context_management`` with ``applied_edits``."""
+
+ applied_edits: List[AppliedEdit]
+
+
+class CompactionBlock(TypedDict, total=False):
+ """Synthesized ``compaction`` content block (compact_20260112)."""
+
+ type: Required[Literal["compaction"]]
+ content: Optional[str]
+
+
+class UsageIteration(TypedDict, total=False):
+ """One sampling iteration's token usage (compact_20260112)."""
+
+ type: Required[Literal["compaction", "message"]]
+ input_tokens: int
+ output_tokens: int
+
+
class MessageBlockDelta(TypedDict):
"""
Anthropic
@@ -524,6 +559,7 @@ class MessageBlockDelta(TypedDict):
type: Literal["message_delta"]
delta: MessageDelta
usage: UsageDelta
+ context_management: NotRequired[ContextManagementResponse]
class MessageChunk(TypedDict, total=False):
diff --git a/litellm/types/llms/anthropic_messages/anthropic_response.py b/litellm/types/llms/anthropic_messages/anthropic_response.py
index 1eab1b37e06..85a2b3fee7c 100644
--- a/litellm/types/llms/anthropic_messages/anthropic_response.py
+++ b/litellm/types/llms/anthropic_messages/anthropic_response.py
@@ -1,10 +1,11 @@
from typing import Any, Dict, List, Literal, Optional, Union
-from typing_extensions import TypeAlias, TypedDict
+from typing_extensions import NotRequired, TypeAlias, TypedDict
from litellm.types.llms.anthropic import (
AnthropicResponseContentBlockText,
AnthropicResponseContentBlockToolUse,
+ ContextManagementResponse,
)
@@ -94,3 +95,4 @@ class AnthropicMessagesResponse(TypedDict, total=False):
stop_sequence: Optional[str]
type: Optional[Literal["message"]]
usage: Optional[AnthropicUsage]
+ context_management: NotRequired[ContextManagementResponse]
diff --git a/tests/pass_through_unit_tests/test_context_management_polyfill.py b/tests/pass_through_unit_tests/test_context_management_polyfill.py
new file mode 100644
index 00000000000..564dbe36f66
--- /dev/null
+++ b/tests/pass_through_unit_tests/test_context_management_polyfill.py
@@ -0,0 +1,272 @@
+"""Integration tests for context_management polyfill on /v1/messages adapter path."""
+
+import json
+from unittest.mock import patch
+
+import pytest
+
+import litellm
+from litellm.llms.anthropic.experimental_pass_through.context_management.constants import (
+ CLEARED_TOOL_RESULT_PLACEHOLDER,
+)
+from litellm.types.utils import (
+ Choices,
+ Message,
+ ModelResponse,
+ ModelResponseStream,
+ StreamingChoices,
+ Delta,
+ Usage,
+)
+
+MODEL = "xai/grok-4"
+
+
+def _make_history(n_pairs: int, result_filler: str = "x" * 50):
+ messages = [{"role": "user", "content": "Compare weather across cities."}]
+ for i in range(n_pairs):
+ messages.append(
+ {
+ "role": "assistant",
+ "content": [
+ {
+ "type": "tool_use",
+ "id": f"toolu_{i:02d}",
+ "name": "get_weather",
+ "input": {"location": f"City{i}"},
+ }
+ ],
+ }
+ )
+ messages.append(
+ {
+ "role": "user",
+ "content": [
+ {
+ "type": "tool_result",
+ "tool_use_id": f"toolu_{i:02d}",
+ "content": f"Result {i}: {result_filler}",
+ }
+ ],
+ }
+ )
+ return messages
+
+
+def _mock_completion_response() -> ModelResponse:
+ return ModelResponse(
+ id="chatcmpl-test",
+ choices=[
+ Choices(
+ finish_reason="stop",
+ index=0,
+ message=Message(role="assistant", content="ok"),
+ )
+ ],
+ created=0,
+ model="grok-4",
+ object="chat.completion",
+ usage=Usage(prompt_tokens=10, completion_tokens=2, total_tokens=12),
+ )
+
+
+async def _mock_streaming_chunks():
+ yield ModelResponseStream(
+ id="chatcmpl-test",
+ created=0,
+ model="grok-4",
+ object="chat.completion.chunk",
+ choices=[
+ StreamingChoices(
+ finish_reason=None,
+ index=0,
+ delta=Delta(role="assistant", content="ok"),
+ )
+ ],
+ )
+ yield ModelResponseStream(
+ id="chatcmpl-test",
+ created=0,
+ model="grok-4",
+ object="chat.completion.chunk",
+ choices=[
+ StreamingChoices(
+ finish_reason="stop",
+ index=0,
+ delta=Delta(),
+ )
+ ],
+ usage=Usage(prompt_tokens=10, completion_tokens=2, total_tokens=12),
+ )
+
+
+@pytest.mark.asyncio
+async def test_polyfill_round_trip_non_streaming():
+ captured = {}
+
+ async def fake_acompletion(**kwargs):
+ captured.update(kwargs)
+ return _mock_completion_response()
+
+ with patch("litellm.acompletion", side_effect=fake_acompletion):
+ response = await litellm.anthropic.messages.acreate(
+ model=MODEL,
+ messages=_make_history(n_pairs=5),
+ max_tokens=128,
+ api_key="sk-test",
+ context_management={
+ "edits": [
+ {
+ "type": "clear_tool_uses_20250919",
+ "trigger": {"type": "tool_uses", "value": 1},
+ "keep": {"type": "tool_uses", "value": 2},
+ }
+ ]
+ },
+ )
+
+ # 1. Downstream got the edited messages — older tool_result.content cleared.
+ downstream_messages = captured.get("messages")
+ assert downstream_messages is not None
+ cleared_ids = {"toolu_00", "toolu_01", "toolu_02"}
+ kept_ids = {"toolu_03", "toolu_04"}
+ found_cleared = 0
+ for msg in downstream_messages:
+ # The adapter may have translated the messages out of Anthropic shape;
+ # we accept either Anthropic-shape (tool_result block) or OpenAI-shape
+ # (tool-role message whose content is the placeholder).
+ if isinstance(msg, dict) and msg.get("role") == "tool":
+ if msg.get("tool_call_id") in cleared_ids:
+ content = msg.get("content")
+ if isinstance(content, str):
+ if CLEARED_TOOL_RESULT_PLACEHOLDER in content:
+ found_cleared += 1
+ elif isinstance(content, list):
+ text = "".join(
+ b.get("text", "") for b in content if isinstance(b, dict)
+ )
+ if CLEARED_TOOL_RESULT_PLACEHOLDER in text:
+ found_cleared += 1
+ elif msg.get("tool_call_id") in kept_ids:
+ content = msg.get("content")
+ if isinstance(content, str):
+ assert CLEARED_TOOL_RESULT_PLACEHOLDER not in content
+ assert found_cleared == 3
+
+ # 2. context_management must not leak into downstream kwargs.
+ assert "context_management" not in captured
+
+ # 3. Response carries the applied_edits in Anthropic's documented shape.
+ assert isinstance(response, dict)
+ cm = response.get("context_management")
+ assert cm is not None, f"context_management missing from response: {response}"
+ edits = cm.get("applied_edits")
+ assert isinstance(edits, list) and len(edits) == 1
+ edit = edits[0]
+ assert edit["type"] == "clear_tool_uses_20250919"
+ assert edit["cleared_tool_uses"] == 3
+ assert "cleared_input_tokens" in edit
+
+
+@pytest.mark.asyncio
+async def test_polyfill_trigger_not_met_passes_through_unchanged():
+ captured = {}
+
+ async def fake_acompletion(**kwargs):
+ captured.update(kwargs)
+ return _mock_completion_response()
+
+ with patch("litellm.acompletion", side_effect=fake_acompletion):
+ response = await litellm.anthropic.messages.acreate(
+ model=MODEL,
+ messages=_make_history(n_pairs=2),
+ max_tokens=128,
+ api_key="sk-test",
+ context_management={
+ "edits": [
+ {
+ "type": "clear_tool_uses_20250919",
+ "trigger": {"type": "input_tokens", "value": 10_000_000},
+ "keep": {"type": "tool_uses", "value": 1},
+ }
+ ]
+ },
+ )
+
+ # Downstream still got the request, but no edits applied.
+ assert captured.get("messages") is not None
+ assert "context_management" not in captured
+
+ # Response shouldn't carry context_management when nothing fired.
+ assert isinstance(response, dict)
+ assert (
+ response.get("context_management") is None
+ or response.get("context_management") == {"applied_edits": []}
+ or "context_management" not in response
+ )
+
+
+@pytest.mark.asyncio
+async def test_polyfill_streaming_attaches_to_message_delta():
+ async def fake_acompletion(**kwargs):
+ return _mock_streaming_chunks()
+
+ with patch("litellm.acompletion", side_effect=fake_acompletion):
+ response = await litellm.anthropic.messages.acreate(
+ model=MODEL,
+ messages=_make_history(n_pairs=5),
+ max_tokens=128,
+ api_key="sk-test",
+ stream=True,
+ context_management={
+ "edits": [
+ {
+ "type": "clear_tool_uses_20250919",
+ "trigger": {"type": "tool_uses", "value": 1},
+ "keep": {"type": "tool_uses", "value": 2},
+ }
+ ]
+ },
+ )
+
+ # Collect all SSE bytes.
+ collected = []
+ async for chunk in response: # type: ignore[union-attr]
+ if isinstance(chunk, (bytes, bytearray)):
+ collected.append(chunk.decode("utf-8"))
+ else:
+ collected.append(str(chunk))
+ sse_text = "".join(collected)
+
+ # Find the message_delta event payload and check it carries context_management
+ # as a sibling of `usage` per Anthropic's spec.
+ found_delta_with_cm = False
+ for block in sse_text.split("\n\n"):
+ if "message_delta" not in block:
+ continue
+ data_line = next(
+ (
+ line[len("data:") :].strip()
+ for line in block.splitlines()
+ if line.startswith("data:")
+ ),
+ None,
+ )
+ if data_line is None:
+ continue
+ payload = json.loads(data_line)
+ if payload.get("type") != "message_delta":
+ continue
+ cm = payload.get("context_management")
+ if cm is None:
+ continue
+ assert "applied_edits" in cm
+ assert len(cm["applied_edits"]) == 1
+ assert cm["applied_edits"][0]["type"] == "clear_tool_uses_20250919"
+ assert cm["applied_edits"][0]["cleared_tool_uses"] == 3
+ found_delta_with_cm = True
+ break
+ assert found_delta_with_cm, (
+ "Expected `context_management` on the message_delta SSE event. "
+ f"SSE text was: {sse_text!r}"
+ )
diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py
index 44530fecebd..74e1e17e6d7 100644
--- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py
+++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py
@@ -2472,3 +2472,172 @@ def test_translate_anthropic_tool_choice_none():
result = adapter.translate_anthropic_tool_choice_to_openai({"type": "none"})
assert result == "none"
+
+
+# ---------------------------------------------------------------------------
+# PolyfillResult integration tests
+# ---------------------------------------------------------------------------
+
+
+def _make_simple_openai_response(
+ text: str = "Hello", prompt_tokens: int = 10, completion_tokens: int = 5
+) -> ModelResponse:
+ return ModelResponse(
+ id="resp_polyfill_test",
+ model="gpt-4o",
+ choices=[
+ Choices(
+ finish_reason="stop",
+ message=Message(role="assistant", content=text),
+ )
+ ],
+ usage=Usage(prompt_tokens=prompt_tokens, completion_tokens=completion_tokens),
+ )
+
+
+def test_translate_openai_response_to_anthropic_with_polyfill_compaction_block():
+ """compaction_block from PolyfillResult must be prepended to content at index 0."""
+ from litellm.llms.anthropic.experimental_pass_through.context_management.result import (
+ PolyfillResult,
+ )
+
+ compaction_block = {"type": "compaction", "content": "Summary of prior turns."}
+ polyfill = PolyfillResult(
+ messages=[],
+ system=None,
+ applied_edits=[{"type": "compact_20260112"}],
+ compaction_block=compaction_block,
+ iterations_usage=None,
+ )
+ response = _make_simple_openai_response(text="Hello after compaction.")
+ adapter = LiteLLMAnthropicMessagesAdapter()
+ result = adapter.translate_openai_response_to_anthropic(
+ response=response, polyfill_result=polyfill
+ )
+
+ content = result.get("content")
+ assert content is not None
+ assert content[0]["type"] == "compaction"
+ assert content[0]["content"] == "Summary of prior turns."
+ assert content[1]["type"] == "text"
+ assert content[1]["text"] == "Hello after compaction."
+
+ # applied_edits must surface on context_management
+ cm = result.get("context_management")
+ assert cm is not None
+ assert cm["applied_edits"][0]["type"] == "compact_20260112"
+
+
+def test_translate_openai_response_to_anthropic_with_polyfill_iterations_usage():
+ """iterations_usage from PolyfillResult must produce usage['iterations'] with a message entry."""
+ from litellm.llms.anthropic.experimental_pass_through.context_management.result import (
+ PolyfillResult,
+ )
+
+ polyfill = PolyfillResult(
+ messages=[],
+ system=None,
+ applied_edits=[{"type": "compact_20260112"}],
+ compaction_block=None,
+ iterations_usage=[
+ {"type": "compaction", "input_tokens": 200, "output_tokens": 50},
+ ],
+ )
+ response = _make_simple_openai_response(prompt_tokens=100, completion_tokens=30)
+ adapter = LiteLLMAnthropicMessagesAdapter()
+ result = adapter.translate_openai_response_to_anthropic(
+ response=response, polyfill_result=polyfill
+ )
+
+ usage = result.get("usage")
+ assert usage is not None
+ iterations = usage.get("iterations")
+ assert iterations is not None
+ assert len(iterations) == 2
+ assert iterations[0] == {
+ "type": "compaction",
+ "input_tokens": 200,
+ "output_tokens": 50,
+ }
+ assert iterations[1]["type"] == "message"
+ assert iterations[1]["input_tokens"] == 100
+ assert iterations[1]["output_tokens"] == 30
+
+ # Top-level tokens must still reflect the message iteration
+ assert usage["input_tokens"] == 100
+ assert usage["output_tokens"] == 30
+
+
+def test_translate_openai_response_to_anthropic_no_polyfill_no_change():
+ """Without a PolyfillResult the response must be unchanged (no compaction, no iterations)."""
+ response = _make_simple_openai_response()
+ adapter = LiteLLMAnthropicMessagesAdapter()
+ result = adapter.translate_openai_response_to_anthropic(response=response)
+
+ content = result.get("content")
+ assert content is not None
+ assert content[0]["type"] == "text"
+
+ usage = result.get("usage")
+ assert usage is not None
+ assert "iterations" not in usage
+
+
+def test_translate_openai_response_to_anthropic_with_polyfill_both_compaction_and_iterations():
+ """Full summary path: compaction_block and iterations_usage both present simultaneously."""
+ from litellm.llms.anthropic.experimental_pass_through.context_management.result import (
+ PolyfillResult,
+ )
+
+ compaction_block = {
+ "type": "compaction",
+ "content": "Summary of a long conversation.",
+ }
+ polyfill = PolyfillResult(
+ messages=[],
+ system=None,
+ applied_edits=[{"type": "compact_20260112"}],
+ compaction_block=compaction_block,
+ iterations_usage=[
+ {"type": "compaction", "input_tokens": 300, "output_tokens": 75},
+ ],
+ )
+ response = _make_simple_openai_response(
+ text="After compaction.", prompt_tokens=120, completion_tokens=40
+ )
+ adapter = LiteLLMAnthropicMessagesAdapter()
+ result = adapter.translate_openai_response_to_anthropic(
+ response=response, polyfill_result=polyfill
+ )
+
+ # compaction block must come first
+ content = result.get("content")
+ assert content is not None
+ assert content[0]["type"] == "compaction"
+ assert content[0]["content"] == "Summary of a long conversation."
+ assert content[1]["type"] == "text"
+ assert content[1]["text"] == "After compaction."
+
+ # iterations: compaction entry + message entry
+ usage = result.get("usage")
+ assert usage is not None
+ iterations = usage.get("iterations")
+ assert iterations is not None
+ assert len(iterations) == 2
+ assert iterations[0] == {
+ "type": "compaction",
+ "input_tokens": 300,
+ "output_tokens": 75,
+ }
+ assert iterations[1]["type"] == "message"
+ assert iterations[1]["input_tokens"] == 120
+ assert iterations[1]["output_tokens"] == 40
+
+ # top-level tokens match the message iteration
+ assert usage["input_tokens"] == 120
+ assert usage["output_tokens"] == 40
+
+ # context_management applied_edits must surface
+ cm = result.get("context_management")
+ assert cm is not None
+ assert cm["applied_edits"][0]["type"] == "compact_20260112"
diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_compaction.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_compaction.py
new file mode 100644
index 00000000000..076d4392f05
--- /dev/null
+++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_compaction.py
@@ -0,0 +1,193 @@
+"""Compaction block SSE events from AnthropicStreamWrapper (compact_20260112 polyfill)."""
+
+import os
+import sys
+from typing import List
+from unittest.mock import MagicMock
+
+import pytest
+
+sys.path.insert(0, os.path.abspath("../../../../.."))
+
+from litellm.llms.anthropic.experimental_pass_through.adapters.streaming_iterator import (
+ AnthropicStreamWrapper,
+)
+from litellm.types.utils import Delta, StreamingChoices, Usage
+
+
+def _make_text_chunk(
+ text: str,
+ finish_reason: str = None,
+ usage: "Usage | None" = None,
+) -> MagicMock:
+ chunk = MagicMock()
+ chunk.choices = [
+ StreamingChoices(
+ finish_reason=finish_reason,
+ index=0,
+ delta=Delta(
+ content=text, role="assistant" if text else None, tool_calls=None
+ ),
+ logprobs=None,
+ )
+ ]
+ chunk.usage = usage
+ chunk._hidden_params = {}
+ return chunk
+
+
+async def _collect_events_async(wrapper: AnthropicStreamWrapper) -> List[dict]:
+ events = []
+ async for event in wrapper:
+ events.append(event)
+ return events
+
+
+@pytest.mark.asyncio
+async def test_stream_emits_compaction_block_before_text():
+ """Polyfill compaction_block must surface as compaction SSE events at index 0."""
+
+ async def mock_stream():
+ yield _make_text_chunk("Hi")
+ yield _make_text_chunk(
+ "",
+ finish_reason="stop",
+ usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
+ )
+
+ compaction_block = {
+ "type": "compaction",
+ "content": "Summary of prior conversation turns.",
+ }
+ iterations_usage = [
+ {"type": "compaction", "input_tokens": 100, "output_tokens": 50},
+ ]
+
+ wrapper = AnthropicStreamWrapper(
+ completion_stream=mock_stream(),
+ model="claude-sonnet-4-6",
+ compaction_block=compaction_block,
+ iterations_usage=iterations_usage,
+ applied_edits=[{"type": "compact_20260112"}],
+ )
+
+ events = await _collect_events_async(wrapper)
+
+ compaction_start = next(
+ e
+ for e in events
+ if e.get("type") == "content_block_start"
+ and e.get("content_block", {}).get("type") == "compaction"
+ )
+ assert compaction_start["index"] == 0
+
+ compaction_delta = next(
+ e
+ for e in events
+ if e.get("type") == "content_block_delta"
+ and e.get("delta", {}).get("type") == "compaction_delta"
+ )
+ assert compaction_delta["index"] == 0
+ assert (
+ compaction_delta["delta"]["content"] == "Summary of prior conversation turns."
+ )
+
+ compaction_stop = next(
+ e
+ for e in events
+ if e.get("type") == "content_block_stop" and e.get("index") == 0
+ )
+ assert compaction_stop is not None
+
+ text_start = next(
+ e
+ for e in events
+ if e.get("type") == "content_block_start"
+ and e.get("content_block", {}).get("type") == "text"
+ )
+ assert text_start["index"] == 1
+
+ message_delta = next(e for e in events if e.get("type") == "message_delta")
+ iterations = message_delta.get("usage", {}).get("iterations")
+ assert iterations is not None
+ assert iterations[0]["type"] == "compaction"
+ assert iterations[1]["type"] == "message"
+ assert iterations[1]["input_tokens"] == 10
+ assert iterations[1]["output_tokens"] == 5
+
+
+@pytest.mark.asyncio
+async def test_stream_omits_message_iteration_when_no_usage_chunk():
+ """When provider sends finish_reason without usage, the held message_delta
+ carries placeholder zeros — we must not emit a misleading zero-token
+ ``message`` iteration entry."""
+
+ async def mock_stream():
+ yield _make_text_chunk("Hi")
+ yield _make_text_chunk("", finish_reason="stop")
+
+ iterations_usage = [
+ {"type": "compaction", "input_tokens": 100, "output_tokens": 50},
+ ]
+
+ wrapper = AnthropicStreamWrapper(
+ completion_stream=mock_stream(),
+ model="claude-sonnet-4-6",
+ iterations_usage=iterations_usage,
+ )
+
+ events = await _collect_events_async(wrapper)
+ message_delta = next(e for e in events if e.get("type") == "message_delta")
+ iterations = message_delta.get("usage", {}).get("iterations")
+ assert iterations is not None
+ assert len(iterations) == 1
+ assert iterations[0]["type"] == "compaction"
+
+
+@pytest.mark.asyncio
+async def test_stream_omits_context_management_when_no_compaction_applied():
+ """applied_edits without a compaction block must not emit context_management."""
+
+ async def mock_stream():
+ yield _make_text_chunk("Hello")
+ yield _make_text_chunk("", finish_reason="stop")
+
+ wrapper = AnthropicStreamWrapper(
+ completion_stream=mock_stream(),
+ model="claude-sonnet-4-6",
+ applied_edits=None,
+ )
+
+ events = await _collect_events_async(wrapper)
+ message_deltas = [e for e in events if e.get("type") == "message_delta"]
+ assert message_deltas
+ assert "context_management" not in message_deltas[-1]
+
+
+@pytest.mark.asyncio
+async def test_stream_without_compaction_block_unchanged():
+ """No compaction_block means no compaction SSE events."""
+
+ async def mock_stream():
+ yield _make_text_chunk("Hello")
+ yield _make_text_chunk("", finish_reason="stop")
+
+ wrapper = AnthropicStreamWrapper(
+ completion_stream=mock_stream(),
+ model="claude-sonnet-4-6",
+ )
+
+ events = await _collect_events_async(wrapper)
+
+ assert not any(
+ e.get("content_block", {}).get("type") == "compaction"
+ for e in events
+ if e.get("type") == "content_block_start"
+ )
+ text_start = next(
+ e
+ for e in events
+ if e.get("type") == "content_block_start"
+ and e.get("content_block", {}).get("type") == "text"
+ )
+ assert text_start["index"] == 0
diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/__init__.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_clear_tool_uses.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_clear_tool_uses.py
new file mode 100644
index 00000000000..09ac95ab16e
--- /dev/null
+++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_clear_tool_uses.py
@@ -0,0 +1,307 @@
+"""
+Unit tests for the in-gateway `clear_tool_uses_20250919` polyfill editor.
+"""
+
+from copy import deepcopy
+
+from litellm.llms.anthropic.experimental_pass_through.context_management.constants import (
+ CLEARED_TOOL_RESULT_PLACEHOLDER,
+)
+from litellm.llms.anthropic.experimental_pass_through.context_management.editors.clear_tool_uses import (
+ apply_clear_tool_uses_20250919,
+)
+
+MODEL = "xai/grok-4"
+
+
+def _make_pair(tool_use_id: str, result_text: str, location: str = "Mumbai"):
+ """Return an (assistant, user) message pair with one tool_use + tool_result."""
+ assistant_msg = {
+ "role": "assistant",
+ "content": [
+ {
+ "type": "tool_use",
+ "id": tool_use_id,
+ "name": "get_weather",
+ "input": {"location": location},
+ }
+ ],
+ }
+ user_msg = {
+ "role": "user",
+ "content": [
+ {
+ "type": "tool_result",
+ "tool_use_id": tool_use_id,
+ "content": result_text,
+ }
+ ],
+ }
+ return assistant_msg, user_msg
+
+
+def _make_history(n_pairs: int, result_filler: str = "x" * 200):
+ messages = [{"role": "user", "content": "Compare weather across cities."}]
+ for i in range(n_pairs):
+ assistant_msg, user_msg = _make_pair(
+ tool_use_id=f"toolu_{i:02d}",
+ result_text=f"Result {i}: {result_filler}",
+ location=f"City{i}",
+ )
+ messages.append(assistant_msg)
+ messages.append(user_msg)
+ return messages
+
+
+def test_below_trigger_returns_unchanged():
+ """If trigger threshold isn't exceeded, editor is a no-op."""
+ messages = _make_history(n_pairs=2)
+ original = deepcopy(messages)
+ new_messages, applied = apply_clear_tool_uses_20250919(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ edit_spec={
+ "type": "clear_tool_uses_20250919",
+ "trigger": {"type": "input_tokens", "value": 10_000_000},
+ "keep": {"type": "tool_uses", "value": 1},
+ },
+ )
+ assert applied is None
+ assert new_messages == original
+
+
+def test_keep_preserves_most_recent_pairs():
+ """With keep=2 and 5 pairs, the 3 oldest pairs are cleared."""
+ messages = _make_history(n_pairs=5)
+ new_messages, applied = apply_clear_tool_uses_20250919(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ edit_spec={
+ "type": "clear_tool_uses_20250919",
+ "trigger": {"type": "tool_uses", "value": 1},
+ "keep": {"type": "tool_uses", "value": 2},
+ },
+ )
+ assert applied is not None
+ assert applied["type"] == "clear_tool_uses_20250919"
+ assert applied["cleared_tool_uses"] == 3
+
+ # Tool results for the first 3 pairs should be the placeholder, last 2 untouched.
+ cleared_ids = {"toolu_00", "toolu_01", "toolu_02"}
+ kept_ids = {"toolu_03", "toolu_04"}
+ for msg in new_messages:
+ if msg.get("role") != "user":
+ continue
+ content = msg.get("content")
+ if not isinstance(content, list):
+ continue
+ for block in content:
+ if block.get("type") != "tool_result":
+ continue
+ if block["tool_use_id"] in cleared_ids:
+ assert block["content"] == CLEARED_TOOL_RESULT_PLACEHOLDER
+ elif block["tool_use_id"] in kept_ids:
+ assert "Result" in block["content"]
+
+
+def test_tool_use_input_is_not_cleared():
+ """clear_tool_inputs defaults to false — tool_use.input must remain intact."""
+ messages = _make_history(n_pairs=3)
+ new_messages, applied = apply_clear_tool_uses_20250919(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ edit_spec={
+ "type": "clear_tool_uses_20250919",
+ "trigger": {"type": "tool_uses", "value": 0},
+ "keep": {"type": "tool_uses", "value": 1},
+ },
+ )
+ assert applied is not None
+ # Every tool_use block still has its original `input`.
+ for msg in new_messages:
+ if msg.get("role") != "assistant":
+ continue
+ for block in msg.get("content", []):
+ if block.get("type") == "tool_use":
+ assert block["input"] == {"location": block["input"]["location"]}
+ assert block["input"]["location"].startswith("City")
+
+
+def test_message_array_length_and_roles_preserved():
+ messages = _make_history(n_pairs=4)
+ original_roles = [m["role"] for m in messages]
+ new_messages, applied = apply_clear_tool_uses_20250919(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ edit_spec={
+ "type": "clear_tool_uses_20250919",
+ "trigger": {"type": "tool_uses", "value": 0},
+ "keep": {"type": "tool_uses", "value": 1},
+ },
+ )
+ assert applied is not None
+ assert len(new_messages) == len(messages)
+ assert [m["role"] for m in new_messages] == original_roles
+
+
+def test_defaults_applied_when_knobs_omitted():
+ """No trigger/keep specified — defaults are 100k input_tokens / 3 tool_uses."""
+ messages = _make_history(n_pairs=2)
+ # Below 100k tokens; should not fire.
+ new_messages, applied = apply_clear_tool_uses_20250919(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ edit_spec={"type": "clear_tool_uses_20250919"},
+ )
+ assert applied is None
+ assert new_messages == messages
+
+
+def test_tool_uses_trigger_variant():
+ """Trigger by raw count of tool_use blocks, not tokens."""
+ messages = _make_history(n_pairs=4)
+ _, applied = apply_clear_tool_uses_20250919(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ edit_spec={
+ "type": "clear_tool_uses_20250919",
+ "trigger": {"type": "tool_uses", "value": 2},
+ "keep": {"type": "tool_uses", "value": 1},
+ },
+ )
+ assert applied is not None
+ # 4 total - 1 kept = 3 cleared
+ assert applied["cleared_tool_uses"] == 3
+
+
+def test_cleared_input_tokens_is_nonnegative():
+ messages = _make_history(n_pairs=4)
+ _, applied = apply_clear_tool_uses_20250919(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ edit_spec={
+ "type": "clear_tool_uses_20250919",
+ "trigger": {"type": "tool_uses", "value": 1},
+ "keep": {"type": "tool_uses", "value": 1},
+ },
+ )
+ assert applied is not None
+ assert applied["cleared_input_tokens"] >= 0
+
+
+def test_ignored_knobs_do_not_alter_behavior():
+ """clear_at_least / exclude_tools / clear_tool_inputs are accepted but ignored in v0."""
+ messages = _make_history(n_pairs=3)
+ _, applied = apply_clear_tool_uses_20250919(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ edit_spec={
+ "type": "clear_tool_uses_20250919",
+ "trigger": {"type": "tool_uses", "value": 0},
+ "keep": {"type": "tool_uses", "value": 1},
+ "clear_at_least": {"type": "input_tokens", "value": 999_999_999},
+ "exclude_tools": ["get_weather"],
+ "clear_tool_inputs": True,
+ },
+ )
+ # Despite clear_at_least being huge, polyfill still applies (knob ignored).
+ # Despite clear_tool_inputs=True, inputs are NOT cleared (knob ignored).
+ assert applied is not None
+ assert applied["cleared_tool_uses"] == 2
+ # Ignored knobs surface as warnings on the AppliedEdit so operators can
+ # see what was dropped (the v0 polyfill silently dropping them at debug
+ # log level made misconfiguration invisible from the response).
+ assert set(applied.get("warnings", [])) == {
+ "clear_at_least_ignored",
+ "exclude_tools_ignored",
+ "clear_tool_inputs_ignored",
+ }
+
+
+def test_no_ignored_knobs_omits_warnings_field():
+ """When the caller doesn't pass any unsupported knobs, no ``warnings`` are added."""
+ messages = _make_history(n_pairs=3)
+ _, applied = apply_clear_tool_uses_20250919(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ edit_spec={
+ "type": "clear_tool_uses_20250919",
+ "trigger": {"type": "tool_uses", "value": 0},
+ "keep": {"type": "tool_uses", "value": 1},
+ },
+ )
+ assert applied is not None
+ assert "warnings" not in applied
+
+
+def test_tool_result_list_content_shape_preserved():
+ """When tool_result.content is a list of blocks, replacement returns a list shape."""
+ messages = [
+ {"role": "user", "content": "Hi"},
+ {
+ "role": "assistant",
+ "content": [
+ {"type": "tool_use", "id": "toolu_a", "name": "f", "input": {}}
+ ],
+ },
+ {
+ "role": "user",
+ "content": [
+ {
+ "type": "tool_result",
+ "tool_use_id": "toolu_a",
+ "content": [{"type": "text", "text": "huge result"}],
+ }
+ ],
+ },
+ {
+ "role": "assistant",
+ "content": [
+ {"type": "tool_use", "id": "toolu_b", "name": "f", "input": {}}
+ ],
+ },
+ {
+ "role": "user",
+ "content": [
+ {
+ "type": "tool_result",
+ "tool_use_id": "toolu_b",
+ "content": [{"type": "text", "text": "keep me"}],
+ }
+ ],
+ },
+ ]
+ new_messages, applied = apply_clear_tool_uses_20250919(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ edit_spec={
+ "type": "clear_tool_uses_20250919",
+ "trigger": {"type": "tool_uses", "value": 0},
+ "keep": {"type": "tool_uses", "value": 1},
+ },
+ )
+ assert applied is not None
+ cleared_block = new_messages[2]["content"][0]
+ assert isinstance(cleared_block["content"], list)
+ assert cleared_block["content"][0]["type"] == "text"
+ assert cleared_block["content"][0]["text"] == CLEARED_TOOL_RESULT_PLACEHOLDER
diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py
new file mode 100644
index 00000000000..be430db9eed
--- /dev/null
+++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py
@@ -0,0 +1,2291 @@
+"""
+Unit tests for the compact_20260112 polyfill editor.
+
+Coverage:
+- trigger.value < 50k → AnthropicContextManagementError(400)
+- opt-in gate (no summary model) → summary_model_not_configured
+- slice-only path (existing compaction block, under threshold)
+- full summary path (over threshold, summary fires)
+- summary call raises → summary_call_failed
+- summary response missing tags → summary_extraction_failed
+- pause_after_compaction: true → pause_after_compaction_ignored warning, proceeds
+- custom instructions → default prompt is not used even when tools present
+"""
+
+from typing import Any, Dict, List
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+from litellm.llms.anthropic.experimental_pass_through.context_management import (
+ AnthropicContextManagementError,
+ apply_context_management,
+)
+from litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact import (
+ _augment_system_with_summary,
+ _extract_summary_text,
+ _select_last_user_question,
+ _slice_around_compaction_block,
+ _strip_compaction_blocks,
+ apply_client_compaction_block_history,
+ apply_compact_20260112,
+)
+from litellm.llms.anthropic.experimental_pass_through.context_management.result import (
+ PolyfillResult,
+)
+
+MODEL = "openai/gpt-4o"
+
+_EDIT_SPEC_DEFAULT: Dict[str, Any] = {"type": "compact_20260112"}
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _simple_messages() -> List[Dict[str, Any]]:
+ return [
+ {"role": "user", "content": "Hello"},
+ {"role": "assistant", "content": [{"type": "text", "text": "Hi there"}]},
+ {"role": "user", "content": "What is 2+2?"},
+ ]
+
+
+def _messages_with_compaction(summary: str = "prev summary") -> List[Dict[str, Any]]:
+ """History that already has a compaction block in an assistant turn."""
+ return [
+ {"role": "user", "content": "older question"},
+ {
+ "role": "assistant",
+ "content": [{"type": "compaction", "content": summary}],
+ },
+ {"role": "user", "content": "newer question"},
+ {"role": "assistant", "content": [{"type": "text", "text": "newer reply"}]},
+ {"role": "user", "content": "latest question"},
+ ]
+
+
+def _make_mock_response(
+ content: str,
+ prompt_tokens: int = 50,
+ completion_tokens: int = 100,
+) -> MagicMock:
+ response = MagicMock()
+ choice = MagicMock()
+ message = MagicMock()
+ message.content = content
+ choice.message = message
+ response.choices = [choice]
+ usage = MagicMock()
+ usage.prompt_tokens = prompt_tokens
+ usage.completion_tokens = completion_tokens
+ response.usage = usage
+ return response
+
+
+# ---------------------------------------------------------------------------
+# Unit: helper functions
+# ---------------------------------------------------------------------------
+
+
+def test_applied_edits_for_response_omits_compact_without_block_or_error():
+ """No compaction block and no error: omit the compact_20260112 edit."""
+ result = PolyfillResult(
+ messages=[],
+ system="summary on system",
+ applied_edits=[{"type": "compact_20260112"}],
+ compaction_block=None,
+ )
+ assert result.applied_edits_for_response() is None
+
+
+def test_applied_edits_for_response_includes_compact_when_error_present():
+ """Error states must surface to the client so operators can debug."""
+ for error in (
+ "summary_model_not_configured",
+ "summary_call_failed",
+ "summary_extraction_failed",
+ ):
+ result = PolyfillResult(
+ messages=[],
+ system=None,
+ applied_edits=[{"type": "compact_20260112", "error": error}],
+ compaction_block=None,
+ )
+ visible = result.applied_edits_for_response()
+ assert visible is not None, error
+ assert visible[0]["error"] == error
+
+
+def test_applied_edits_for_response_includes_compact_when_block_present():
+ result = PolyfillResult(
+ messages=[],
+ system=None,
+ applied_edits=[
+ {
+ "type": "compact_20260112",
+ "summary_input_tokens": 10,
+ "summary_output_tokens": 5,
+ }
+ ],
+ compaction_block={"type": "compaction", "content": "summary"},
+ )
+ visible = result.applied_edits_for_response()
+ assert visible is not None
+ assert visible[0]["type"] == "compact_20260112"
+ assert visible[0]["summary_input_tokens"] == 10
+
+
+def test_slice_around_compaction_block_found():
+ messages = _messages_with_compaction("my summary")
+ sliced, block = _slice_around_compaction_block(messages)
+ assert block is not None
+ assert block["type"] == "compaction"
+ assert block["content"] == "my summary"
+ # Sliced list starts at the assistant turn containing the compaction block
+ assert sliced[0]["role"] == "assistant"
+ assert len(sliced) == 4 # assistant(compaction), user, assistant, user
+
+
+def test_slice_around_compaction_block_not_found():
+ messages = _simple_messages()
+ sliced, block = _slice_around_compaction_block(messages)
+ assert block is None
+ assert sliced is messages # same object, no copy
+
+
+def test_strip_compaction_blocks_removes_block():
+ messages = [
+ {
+ "role": "assistant",
+ "content": [
+ {"type": "compaction", "content": "summary"},
+ {"type": "text", "text": "hello"},
+ ],
+ }
+ ]
+ stripped = _strip_compaction_blocks(messages)
+ assert len(stripped) == 1
+ content = stripped[0]["content"]
+ assert all(b["type"] != "compaction" for b in content)
+ assert len(content) == 1
+ assert content[0]["type"] == "text"
+
+
+def test_select_last_user_question_strips_tool_result_from_mixed_turn():
+ """Mixed [tool_result, text] turn: keep text, drop tool_result blocks."""
+ messages = [
+ {"role": "user", "content": "earlier"},
+ {
+ "role": "user",
+ "content": [
+ {"type": "tool_result", "tool_use_id": "a", "content": "res"},
+ {"type": "text", "text": "follow-up question"},
+ ],
+ },
+ ]
+ selected = _select_last_user_question(messages)
+ assert len(selected) == 1
+ assert selected[0]["role"] == "user"
+ content = selected[0]["content"]
+ assert isinstance(content, list)
+ assert all(b.get("type") != "tool_result" for b in content)
+ assert any(
+ b.get("type") == "text" and b.get("text") == "follow-up question"
+ for b in content
+ )
+
+
+def test_select_last_user_question_skips_pure_tool_result_turn():
+ """Pure tool_result turn: skip and walk back to a real user turn."""
+ messages = [
+ {"role": "user", "content": "real question"},
+ {
+ "role": "assistant",
+ "content": [{"type": "tool_use", "id": "a", "name": "x", "input": {}}],
+ },
+ {
+ "role": "user",
+ "content": [{"type": "tool_result", "tool_use_id": "a", "content": "res"}],
+ },
+ ]
+ selected = _select_last_user_question(messages)
+ assert len(selected) == 1
+ assert selected[0]["content"] == "real question"
+
+
+def test_select_last_user_question_falls_back_when_no_eligible_turn():
+ """Only tool_result-only user turns: emit a synthetic continuation prompt."""
+ messages = [
+ {
+ "role": "user",
+ "content": [{"type": "tool_result", "tool_use_id": "a", "content": "res"}],
+ },
+ ]
+ selected = _select_last_user_question(messages)
+ assert len(selected) == 1
+ assert selected[0]["role"] == "user"
+ assert isinstance(selected[0]["content"], str)
+
+
+def test_strip_compaction_blocks_drops_compaction_only_turn():
+ messages = [
+ {"role": "user", "content": "hi"},
+ {
+ "role": "assistant",
+ "content": [{"type": "compaction", "content": "summary"}],
+ },
+ {"role": "user", "content": "bye"},
+ ]
+ stripped = _strip_compaction_blocks(messages)
+ assert len(stripped) == 2
+ assert stripped[0]["role"] == "user"
+ assert stripped[1]["role"] == "user"
+
+
+def test_augment_system_with_summary_none_system():
+ result = _augment_system_with_summary(None, "my summary")
+ assert isinstance(result, str)
+ assert "my summary" in result
+
+
+def test_augment_system_with_summary_string_system():
+ result = _augment_system_with_summary("You are helpful.", "my summary")
+ assert isinstance(result, str)
+ assert result.startswith("Previous conversation summary:")
+ assert "my summary" in result
+ assert "You are helpful." in result
+
+
+def test_augment_system_with_summary_list_system():
+ system = [{"type": "text", "text": "existing system"}]
+ result = _augment_system_with_summary(system, "my summary")
+ assert isinstance(result, list)
+ assert result[0]["type"] == "text"
+ text = result[0]["text"]
+ assert "my summary" in text
+ assert "existing system" in text
+
+
+def test_extract_summary_text_found():
+ raw = "Here is the summary:\nKey points from chat \nDone."
+ assert _extract_summary_text(raw) == "Key points from chat"
+
+
+def test_extract_summary_text_missing_tags():
+ assert _extract_summary_text("No tags here") is None
+
+
+def test_extract_summary_text_none():
+ assert _extract_summary_text(None) is None
+
+
+def test_extract_summary_text_case_insensitive():
+ raw = "uppercase tags "
+ assert _extract_summary_text(raw) == "uppercase tags"
+
+
+# ---------------------------------------------------------------------------
+# Editor: validation
+# ---------------------------------------------------------------------------
+
+
+async def test_trigger_below_minimum_raises():
+ with pytest.raises(AnthropicContextManagementError) as exc_info:
+ await apply_compact_20260112(
+ model=MODEL,
+ messages=_simple_messages(),
+ tools=None,
+ system=None,
+ edit_spec={
+ "type": "compact_20260112",
+ "trigger": {"type": "input_tokens", "value": 10_000},
+ },
+ )
+ assert exc_info.value.status_code == 400
+ assert "50000" in exc_info.value.message
+
+
+async def test_trigger_at_minimum_does_not_raise():
+ """Exactly 50 000 is allowed — only strictly less than 50k is rejected."""
+ with patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
+ return_value=None,
+ ):
+ result = await apply_compact_20260112(
+ model=MODEL,
+ messages=_simple_messages(),
+ tools=None,
+ system=None,
+ edit_spec={
+ "type": "compact_20260112",
+ "trigger": {"type": "input_tokens", "value": 50_000},
+ },
+ )
+ # Reached opt-in gate (no summary model); no error raised from trigger check
+ assert result.applied_edits[0]["error"] == "summary_model_not_configured"
+
+
+# ---------------------------------------------------------------------------
+# Editor: opt-in gate
+# ---------------------------------------------------------------------------
+
+
+async def test_opt_in_gating_no_summary_model_configured():
+ messages = _simple_messages()
+ with patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
+ return_value=None,
+ ):
+ result = await apply_compact_20260112(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system="system prompt",
+ edit_spec=_EDIT_SPEC_DEFAULT,
+ )
+ assert result.applied_edits[0]["error"] == "summary_model_not_configured"
+ assert result.messages == messages
+ assert result.system == "system prompt"
+ assert result.compaction_block is None
+ assert result.iterations_usage is None
+
+
+async def test_opt_in_gating_no_summary_model_keeps_post_compaction_tail():
+ """No summary model + prior compaction block forwards the full tail.
+
+ The prior summary lives on the system prefix; the post-compaction turns it
+ does not cover must be forwarded unchanged rather than collapsed to the
+ latest user question (which would strip intermediate turns the model needs).
+ """
+ messages = _messages_with_compaction("prior summary text")
+
+ with patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
+ return_value=None,
+ ):
+ result = await apply_compact_20260112(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ edit_spec=_EDIT_SPEC_DEFAULT,
+ )
+
+ assert result.applied_edits[0]["error"] == "summary_model_not_configured"
+ assert result.system is not None
+ assert "prior summary text" in str(result.system)
+ assert result.compaction_block is None
+ assert result.iterations_usage is None
+ # Post-compaction tail forwarded unchanged (compaction blocks stripped).
+ assert [m["role"] for m in result.messages] == ["user", "assistant", "user"]
+ assert result.messages[0]["content"] == "newer question"
+ assert result.messages[-1]["content"] == "latest question"
+ for msg in result.messages:
+ content = msg.get("content")
+ if isinstance(content, list):
+ for block in content:
+ assert block.get("type") != "compaction"
+
+
+# ---------------------------------------------------------------------------
+# Client compaction block without context_management
+# ---------------------------------------------------------------------------
+
+
+def test_client_compaction_block_history_without_context_management():
+ """Compaction in messages alone triggers slice-only forwarding.
+
+ The prior summary is prepended to ``system``; the post-compaction tail is
+ forwarded unchanged so the model sees the recent turns the summary does
+ not cover. Compaction blocks themselves are stripped from messages so
+ non-Anthropic backends don't reject them.
+ """
+ messages = _messages_with_compaction("prior summary text")
+
+ result = apply_client_compaction_block_history(messages=messages, system=None)
+
+ assert result is not None
+ assert result.system is not None
+ assert "prior summary text" in str(result.system)
+ assert result.compaction_block is None
+ assert result.applied_edits == []
+ # Post-compaction tail: newer question, newer reply, latest question.
+ assert [m["role"] for m in result.messages] == ["user", "assistant", "user"]
+ assert result.messages[0]["content"] == "newer question"
+ assert result.messages[-1]["content"] == "latest question"
+ for msg in result.messages:
+ content = msg.get("content")
+ if isinstance(content, list):
+ for block in content:
+ assert block.get("type") != "compaction"
+
+
+def test_client_compaction_block_history_no_compaction_returns_none():
+ result = apply_client_compaction_block_history(
+ messages=_simple_messages(), system="base"
+ )
+ assert result is None
+
+
+# ---------------------------------------------------------------------------
+# Editor: slice-only path
+# ---------------------------------------------------------------------------
+
+
+async def test_slice_only_path_with_existing_compaction_block():
+ """Phase A slices; Phase B token count is below threshold; no summary call.
+
+ The prior compaction summary lives on the system prefix; the
+ post-compaction tail is forwarded unchanged so the model retains the
+ recent turns the summary does not cover. Compaction blocks themselves
+ are stripped from messages.
+ """
+ messages = _messages_with_compaction("prior summary text")
+
+ with (
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
+ return_value="claude-haiku-4-5",
+ ),
+ patch("litellm.token_counter", return_value=500), # well under threshold
+ ):
+ result = await apply_compact_20260112(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ edit_spec=_EDIT_SPEC_DEFAULT,
+ )
+
+ # System should have the prior summary prefixed
+ assert result.system is not None
+ assert "prior summary text" in str(result.system)
+
+ # No new compaction block; no iterations_usage
+ assert result.compaction_block is None
+ assert result.iterations_usage is None
+
+ # Main call: summary on system + full post-compaction tail (no compaction blocks).
+ assert [m["role"] for m in result.messages] == ["user", "assistant", "user"]
+ assert result.messages[0]["content"] == "newer question"
+ assert result.messages[-1]["content"] == "latest question"
+ for msg in result.messages:
+ content = msg.get("content")
+ if isinstance(content, list):
+ for block in content:
+ assert block.get("type") != "compaction"
+
+
+async def test_slice_only_no_compaction_block_under_threshold():
+ """No prior compaction block, and token count is below threshold — pure pass-through."""
+ messages = _simple_messages()
+ with (
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
+ return_value="claude-haiku-4-5",
+ ),
+ patch("litellm.token_counter", return_value=500),
+ ):
+ result = await apply_compact_20260112(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ edit_spec=_EDIT_SPEC_DEFAULT,
+ )
+
+ assert result.messages == messages
+ assert result.compaction_block is None
+ assert result.iterations_usage is None
+ assert not result.applied_edits[0].get("error")
+
+
+# ---------------------------------------------------------------------------
+# Editor: full summary path
+# ---------------------------------------------------------------------------
+
+
+async def test_full_summary_path():
+ """Over threshold: summary call fires, compaction_block and iterations_usage returned."""
+ messages = _simple_messages()
+ mock_response = _make_mock_response(
+ "Condensed history ", prompt_tokens=200, completion_tokens=50
+ )
+
+ with (
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
+ return_value="claude-haiku-4-5",
+ ),
+ patch("litellm.token_counter", return_value=200_000), # over 150k threshold
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._call_summary_model",
+ new_callable=AsyncMock,
+ return_value=mock_response,
+ ),
+ ):
+ result = await apply_compact_20260112(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ edit_spec=_EDIT_SPEC_DEFAULT,
+ )
+
+ assert result.compaction_block is not None
+ assert result.compaction_block["type"] == "compaction"
+ assert result.compaction_block["content"] == "Condensed history"
+
+ assert result.iterations_usage is not None
+ assert len(result.iterations_usage) == 1
+ assert result.iterations_usage[0]["type"] == "compaction"
+ assert result.iterations_usage[0]["input_tokens"] == 200
+ assert result.iterations_usage[0]["output_tokens"] == 50
+
+ # System must have summary prefixed
+ assert "Condensed history" in str(result.system)
+
+ # applied_edits should have usage fields
+ edit = result.applied_edits[0]
+ assert edit["type"] == "compact_20260112"
+ assert edit.get("summary_input_tokens") == 200
+ assert edit.get("summary_output_tokens") == 50
+
+ # Downstream messages must not contain a compaction block
+ for msg in result.messages:
+ content = msg.get("content")
+ if isinstance(content, list):
+ for block in content:
+ assert block.get("type") != "compaction"
+
+
+async def test_full_summary_path_uses_router_when_available():
+ """When llm_router is provided, its acompletion method is called instead of litellm."""
+ messages = _simple_messages()
+ mock_response = _make_mock_response("Router summary ")
+ mock_router = MagicMock()
+ mock_router.acompletion = AsyncMock(return_value=mock_response)
+
+ with (
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
+ return_value="my-summary-model",
+ ),
+ patch("litellm.token_counter", return_value=200_000),
+ ):
+ result = await apply_compact_20260112(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ edit_spec=_EDIT_SPEC_DEFAULT,
+ llm_router=mock_router,
+ )
+
+ mock_router.acompletion.assert_called_once()
+ call_kwargs = mock_router.acompletion.call_args.kwargs
+ assert call_kwargs["model"] == "my-summary-model"
+
+ assert result.compaction_block is not None
+ assert result.compaction_block["content"] == "Router summary"
+
+
+async def test_litellm_metadata_propagated_to_summary_call():
+ """Auth fields from the proxy ``litellm_metadata`` are forwarded to the summary call."""
+ messages = _simple_messages()
+ mock_response = _make_mock_response("Summary ")
+ parent_litellm_metadata = {
+ "user_api_key": "sk-test",
+ "user_api_key_team_id": "team-123",
+ "user_api_key_user_id": "user-456",
+ "litellm_call_id": "call-789",
+ "should_not_propagate": "secret",
+ }
+
+ with (
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
+ return_value="claude-haiku-4-5",
+ ),
+ patch("litellm.token_counter", return_value=200_000),
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._call_summary_model",
+ new_callable=AsyncMock,
+ return_value=mock_response,
+ ) as mock_call,
+ ):
+ await apply_compact_20260112(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ edit_spec=_EDIT_SPEC_DEFAULT,
+ litellm_metadata=parent_litellm_metadata,
+ )
+
+ call_kwargs = mock_call.call_args.kwargs
+ propagated = call_kwargs["metadata"]
+ assert propagated["user_api_key"] == "sk-test"
+ assert propagated["user_api_key_team_id"] == "team-123"
+ assert "should_not_propagate" not in propagated
+
+
+# ---------------------------------------------------------------------------
+# Editor: error paths
+# ---------------------------------------------------------------------------
+
+
+async def test_summary_call_failed():
+ """When the summary model raises, applied_edits[0].error == 'summary_call_failed'."""
+ messages = _simple_messages()
+
+ with (
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
+ return_value="claude-haiku-4-5",
+ ),
+ patch("litellm.token_counter", return_value=200_000),
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._call_summary_model",
+ new_callable=AsyncMock,
+ side_effect=RuntimeError("network error"),
+ ),
+ ):
+ result = await apply_compact_20260112(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ edit_spec=_EDIT_SPEC_DEFAULT,
+ )
+
+ assert result.applied_edits[0]["error"] == "summary_call_failed"
+ assert result.compaction_block is None
+ assert result.iterations_usage is None
+ # Messages passed through (at minimum sliced, no compaction blocks)
+ for msg in result.messages:
+ content = msg.get("content")
+ if isinstance(content, list):
+ for block in content:
+ assert block.get("type") != "compaction"
+
+
+async def test_summary_extraction_failed_no_tags():
+ """When summary response has no tags, applied_edits[0].error == 'summary_extraction_failed'."""
+ messages = _simple_messages()
+ mock_response = _make_mock_response("I cannot summarize that.")
+
+ with (
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
+ return_value="claude-haiku-4-5",
+ ),
+ patch("litellm.token_counter", return_value=200_000),
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._call_summary_model",
+ new_callable=AsyncMock,
+ return_value=mock_response,
+ ),
+ ):
+ result = await apply_compact_20260112(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ edit_spec=_EDIT_SPEC_DEFAULT,
+ )
+
+ assert result.applied_edits[0]["error"] == "summary_extraction_failed"
+ assert result.compaction_block is None
+ assert result.iterations_usage is None
+
+
+# ---------------------------------------------------------------------------
+# Editor: warnings
+# ---------------------------------------------------------------------------
+
+
+async def test_pause_after_compaction_ignored_warning():
+ """pause_after_compaction: true → warning recorded, request proceeds normally."""
+ messages = _simple_messages()
+ with patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
+ return_value=None,
+ ):
+ result = await apply_compact_20260112(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ edit_spec={
+ "type": "compact_20260112",
+ "pause_after_compaction": True,
+ },
+ )
+
+ edit = result.applied_edits[0]
+ assert "pause_after_compaction_ignored" in (edit.get("warnings") or [])
+ # Request still proceeds (here it hits opt-in gate because no model configured)
+ assert edit.get("error") == "summary_model_not_configured"
+
+
+async def test_unsupported_trigger_type_falls_back_to_default():
+ messages = _simple_messages()
+ with patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
+ return_value=None,
+ ):
+ result = await apply_compact_20260112(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ edit_spec={
+ "type": "compact_20260112",
+ "trigger": {"type": "output_tokens", "value": 200_000},
+ },
+ )
+
+ edit = result.applied_edits[0]
+ warnings = edit.get("warnings") or []
+ assert any("unsupported_trigger_type" in w for w in warnings)
+
+
+# ---------------------------------------------------------------------------
+# Editor: custom instructions
+# ---------------------------------------------------------------------------
+
+
+async def test_custom_instructions_used_verbatim():
+ """Custom instructions are used as-is; the default prompt is NOT appended."""
+ messages = _simple_messages()
+ tools = [{"name": "search", "description": "Search tool"}]
+ mock_response = _make_mock_response("Custom summary ")
+
+ captured_calls: list = []
+
+ async def _fake_call_summary_model(**kwargs):
+ captured_calls.append(kwargs)
+ return mock_response
+
+ with (
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
+ return_value="claude-haiku-4-5",
+ ),
+ patch("litellm.token_counter", return_value=200_000),
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._call_summary_model",
+ side_effect=_fake_call_summary_model,
+ ),
+ ):
+ await apply_compact_20260112(
+ model=MODEL,
+ messages=messages,
+ tools=tools,
+ system=None,
+ edit_spec={
+ "type": "compact_20260112",
+ "instructions": "Summarize everything briefly.",
+ },
+ )
+
+ assert len(captured_calls) == 1
+ summary_messages = captured_calls[0]["summary_messages"]
+ # The custom instruction prompt is appended to the trailing user turn so
+ # we don't end up with two consecutive ``role=user`` messages (some
+ # providers reject that).
+ last_msg = summary_messages[-1]
+ assert last_msg["role"] == "user"
+ assert "Summarize everything briefly." in last_msg["content"]
+ # The "do not call tools" suffix should NOT be in the prompt since custom was set
+ assert "do not call" not in last_msg["content"].lower()
+
+
+async def test_default_instructions_appended_with_no_tool_suffix_when_no_tools():
+ """Without tools, default prompt is used but the no-tool-calls suffix is absent."""
+ messages = _simple_messages()
+ mock_response = _make_mock_response("Default summary ")
+
+ captured_calls: list = []
+
+ async def _fake_call_summary_model(**kwargs):
+ captured_calls.append(kwargs)
+ return mock_response
+
+ with (
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
+ return_value="claude-haiku-4-5",
+ ),
+ patch("litellm.token_counter", return_value=200_000),
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._call_summary_model",
+ side_effect=_fake_call_summary_model,
+ ),
+ ):
+ await apply_compact_20260112(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ edit_spec=_EDIT_SPEC_DEFAULT,
+ )
+
+ prompt = captured_calls[0]["summary_messages"][-1]["content"]
+ # Should not contain the no-tool-calls guidance
+ assert "do not call" not in prompt.lower()
+
+
+async def test_default_instructions_with_tools_appends_no_tool_suffix():
+ """With tools and no custom instructions, the no-tool-calls suffix is appended."""
+ messages = _simple_messages()
+ tools = [{"name": "search"}]
+ mock_response = _make_mock_response("Tool-aware summary ")
+
+ captured_calls: list = []
+
+ async def _fake_call_summary_model(**kwargs):
+ captured_calls.append(kwargs)
+ return mock_response
+
+ with (
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
+ return_value="claude-haiku-4-5",
+ ),
+ patch("litellm.token_counter", return_value=200_000),
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._call_summary_model",
+ side_effect=_fake_call_summary_model,
+ ),
+ ):
+ await apply_compact_20260112(
+ model=MODEL,
+ messages=messages,
+ tools=tools,
+ system=None,
+ edit_spec=_EDIT_SPEC_DEFAULT,
+ )
+
+ prompt = captured_calls[0]["summary_messages"][-1]["content"]
+ assert "tool" in prompt.lower()
+
+
+async def test_system_prompt_forwarded_to_summary_call_as_string():
+ """A bare-string ``system`` is prepended as a system message to the summary call."""
+ messages = _simple_messages()
+ mock_response = _make_mock_response("With system ")
+
+ captured_calls: list = []
+
+ async def _fake_call_summary_model(**kwargs):
+ captured_calls.append(kwargs)
+ return mock_response
+
+ with (
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
+ return_value="claude-haiku-4-5",
+ ),
+ patch("litellm.token_counter", return_value=200_000),
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._call_summary_model",
+ side_effect=_fake_call_summary_model,
+ ),
+ ):
+ await apply_compact_20260112(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system="You are a helpful coding agent. The initial task is to fix bug X.",
+ edit_spec=_EDIT_SPEC_DEFAULT,
+ )
+
+ summary_messages = captured_calls[0]["summary_messages"]
+ assert summary_messages[0]["role"] == "system"
+ assert "initial task is to fix bug X" in summary_messages[0]["content"]
+
+
+async def test_system_prompt_forwarded_to_summary_call_as_content_blocks():
+ """An Anthropic-shaped list ``system`` is flattened to text and prepended."""
+ messages = _simple_messages()
+ mock_response = _make_mock_response("With list system ")
+
+ captured_calls: list = []
+
+ async def _fake_call_summary_model(**kwargs):
+ captured_calls.append(kwargs)
+ return mock_response
+
+ system_blocks = [
+ {"type": "text", "text": "Agent role: code reviewer."},
+ {"type": "text", "text": "Initial task: review PR #123."},
+ ]
+
+ with (
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
+ return_value="claude-haiku-4-5",
+ ),
+ patch("litellm.token_counter", return_value=200_000),
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._call_summary_model",
+ side_effect=_fake_call_summary_model,
+ ),
+ ):
+ await apply_compact_20260112(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=system_blocks,
+ edit_spec=_EDIT_SPEC_DEFAULT,
+ )
+
+ summary_messages = captured_calls[0]["summary_messages"]
+ assert summary_messages[0]["role"] == "system"
+ content = summary_messages[0]["content"]
+ assert "Agent role: code reviewer." in content
+ assert "Initial task: review PR #123." in content
+
+
+async def test_summary_call_carries_prior_compaction_summary_into_system():
+ """Multi-round: when a prior compaction block is present, the summary
+ model receives the augmented system (with ``Previous conversation
+ summary: ``) so it can produce a comprehensive summary that
+ incorporates both the prior round's context and the current slice.
+ Without this, multi-round compaction would silently drop accumulated
+ history each time the polyfill fires.
+ """
+ messages = _messages_with_compaction(summary="ROUND_ONE_SUMMARY_TEXT")
+ mock_response = _make_mock_response("Round two ")
+
+ captured_calls: list = []
+
+ async def _fake_call_summary_model(**kwargs):
+ captured_calls.append(kwargs)
+ return mock_response
+
+ with (
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
+ return_value="claude-haiku-4-5",
+ ),
+ patch("litellm.token_counter", return_value=200_000),
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._call_summary_model",
+ side_effect=_fake_call_summary_model,
+ ),
+ ):
+ await apply_compact_20260112(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system="Original agent role.",
+ edit_spec=_EDIT_SPEC_DEFAULT,
+ )
+
+ summary_messages = captured_calls[0]["summary_messages"]
+ assert summary_messages[0]["role"] == "system"
+ system_content = summary_messages[0]["content"]
+ assert "ROUND_ONE_SUMMARY_TEXT" in system_content
+ assert "Original agent role." in system_content
+
+
+async def test_summary_call_omits_system_message_when_system_is_none():
+ """No system message is prepended when the caller did not provide one."""
+ messages = _simple_messages()
+ mock_response = _make_mock_response("No system ")
+
+ captured_calls: list = []
+
+ async def _fake_call_summary_model(**kwargs):
+ captured_calls.append(kwargs)
+ return mock_response
+
+ with (
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
+ return_value="claude-haiku-4-5",
+ ),
+ patch("litellm.token_counter", return_value=200_000),
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._call_summary_model",
+ side_effect=_fake_call_summary_model,
+ ),
+ ):
+ await apply_compact_20260112(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ edit_spec=_EDIT_SPEC_DEFAULT,
+ )
+
+ summary_messages = captured_calls[0]["summary_messages"]
+ assert all(msg.get("role") != "system" for msg in summary_messages)
+
+
+async def test_summary_call_does_not_emit_consecutive_user_turns():
+ """When the trailing message is already a user turn, the summarization
+ prompt is merged into it instead of appended as a second user message.
+
+ Some providers (and strict OpenAI-compatible endpoints) reject two
+ consecutive ``role=user`` messages, which would silently fall into the
+ ``summary_call_failed`` error path.
+ """
+ messages = _simple_messages()
+ assert messages[-1]["role"] == "user"
+ mock_response = _make_mock_response("x ")
+
+ captured_calls: list = []
+
+ async def _fake_call_summary_model(**kwargs):
+ captured_calls.append(kwargs)
+ return mock_response
+
+ with (
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
+ return_value="claude-haiku-4-5",
+ ),
+ patch("litellm.token_counter", return_value=200_000),
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._call_summary_model",
+ side_effect=_fake_call_summary_model,
+ ),
+ ):
+ await apply_compact_20260112(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ edit_spec=_EDIT_SPEC_DEFAULT,
+ )
+
+ summary_messages = captured_calls[0]["summary_messages"]
+ user_indices = [
+ idx for idx, msg in enumerate(summary_messages) if msg.get("role") == "user"
+ ]
+ # No two adjacent indices.
+ assert all(
+ b - a > 1 for a, b in zip(user_indices, user_indices[1:])
+ ), f"two consecutive user turns produced: {summary_messages}"
+
+
+async def test_summary_call_sends_default_max_tokens():
+ """``max_tokens`` is set on the summary call so providers like Anthropic
+ (which require it) don't reject the request and silently fall back to
+ ``summary_call_failed``.
+ """
+ from litellm.llms.anthropic.experimental_pass_through.context_management.constants import (
+ COMPACT_SUMMARY_MAX_TOKENS,
+ )
+ from litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact import (
+ _call_summary_model,
+ )
+
+ captured_kwargs: dict = {}
+
+ class _FakeRouter:
+ async def acompletion(self, **kwargs):
+ captured_kwargs.update(kwargs)
+ return _make_mock_response("x ")
+
+ await _call_summary_model(
+ summary_model="claude-haiku-4-5",
+ summary_messages=[{"role": "user", "content": "hi"}],
+ metadata={},
+ llm_router=_FakeRouter(),
+ )
+
+ assert captured_kwargs.get("max_tokens") == COMPACT_SUMMARY_MAX_TOKENS
+
+
+async def test_summary_call_honors_max_tokens_override():
+ """Operators can override the default summary ``max_tokens`` via
+ ``general_settings.context_management_summary_max_tokens``."""
+ from litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact import (
+ _read_summary_max_tokens_setting,
+ )
+
+ captured_kwargs: dict = {}
+
+ class _FakeRouter:
+ async def acompletion(self, **kwargs):
+ captured_kwargs.update(kwargs)
+ return _make_mock_response("x ")
+
+ with patch(
+ "litellm.proxy.proxy_server.general_settings",
+ {"context_management_summary_max_tokens": 8192},
+ ):
+ assert _read_summary_max_tokens_setting() == 8192
+
+ from litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact import (
+ _call_summary_model,
+ )
+
+ await _call_summary_model(
+ summary_model="claude-haiku-4-5",
+ summary_messages=[{"role": "user", "content": "hi"}],
+ metadata={},
+ llm_router=_FakeRouter(),
+ max_tokens=_read_summary_max_tokens_setting(),
+ )
+
+ assert captured_kwargs.get("max_tokens") == 8192
+
+
+def test_summary_max_tokens_setting_falls_back_for_invalid_values():
+ """Invalid override values (non-int, non-positive, missing) fall back to
+ the compiled default so a typo in ``general_settings`` doesn't break the
+ summary call."""
+ from litellm.llms.anthropic.experimental_pass_through.context_management.constants import (
+ COMPACT_SUMMARY_MAX_TOKENS,
+ )
+ from litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact import (
+ _read_summary_max_tokens_setting,
+ )
+
+ for bad in ("4096", 0, -1, None, {"value": 1024}):
+ with patch(
+ "litellm.proxy.proxy_server.general_settings",
+ {"context_management_summary_max_tokens": bad},
+ ):
+ assert (
+ _read_summary_max_tokens_setting() == COMPACT_SUMMARY_MAX_TOKENS
+ ), f"expected default for invalid override {bad!r}"
+
+
+async def test_summary_call_sends_default_timeout():
+ """``timeout`` is set on the summary call so a slow or unresponsive summary
+ model cannot hang the parent ``/v1/messages`` request indefinitely."""
+ from litellm.llms.anthropic.experimental_pass_through.context_management.constants import (
+ COMPACT_SUMMARY_TIMEOUT_SECONDS,
+ )
+ from litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact import (
+ _call_summary_model,
+ )
+
+ captured_kwargs: dict = {}
+
+ class _FakeRouter:
+ async def acompletion(self, **kwargs):
+ captured_kwargs.update(kwargs)
+ return _make_mock_response("x ")
+
+ await _call_summary_model(
+ summary_model="claude-haiku-4-5",
+ summary_messages=[{"role": "user", "content": "hi"}],
+ metadata={},
+ llm_router=_FakeRouter(),
+ )
+
+ assert captured_kwargs.get("timeout") == COMPACT_SUMMARY_TIMEOUT_SECONDS
+
+
+# ---------------------------------------------------------------------------
+# Editor: summary model key/team access gate
+# ---------------------------------------------------------------------------
+
+
+def _fake_user_api_key_auth(
+ *,
+ key_models=None,
+ team_models=None,
+ team_id=None,
+ model_max_budget=None,
+ end_user_model_max_budget=None,
+ end_user_id=None,
+ token=None,
+):
+ """Build a minimal stand-in for ``UserAPIKeyAuth`` with just the fields
+ consulted by ``_check_summary_model_access`` and
+ ``_check_summary_model_budget``. Avoids pulling the proxy deps into this
+ unit test."""
+
+ class _Auth:
+ pass
+
+ auth = _Auth()
+ auth.models = list(key_models) if key_models is not None else []
+ auth.team_models = list(team_models) if team_models is not None else []
+ auth.team_id = team_id
+ auth.team_model_aliases = None
+ auth.model_max_budget = model_max_budget
+ auth.end_user_model_max_budget = end_user_model_max_budget
+ auth.end_user_id = end_user_id
+ auth.token = token
+ return auth
+
+
+async def test_summary_model_denied_when_key_not_in_allowlist():
+ """Caller key restricted to specific models cannot trigger an unauthorized summary model."""
+ messages = _simple_messages()
+ mock_call = AsyncMock(return_value=_make_mock_response("x "))
+
+ with (
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
+ return_value="claude-haiku-4-5",
+ ),
+ patch("litellm.token_counter", return_value=200_000),
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._call_summary_model",
+ mock_call,
+ ),
+ ):
+ result = await apply_compact_20260112(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ edit_spec=_EDIT_SPEC_DEFAULT,
+ user_api_key_auth=_fake_user_api_key_auth(key_models=["gpt-4o"]),
+ )
+
+ mock_call.assert_not_awaited()
+ assert result.compaction_block is None
+ assert result.iterations_usage is None
+ assert result.applied_edits[0]["type"] == "compact_20260112"
+ assert result.applied_edits[0].get("error") == "summary_model_access_denied"
+
+
+async def test_summary_model_denied_when_team_not_in_allowlist():
+ """Team-level model allowlist is enforced even if the key allows all models."""
+ messages = _simple_messages()
+ mock_call = AsyncMock(return_value=_make_mock_response("x "))
+
+ with (
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
+ return_value="claude-haiku-4-5",
+ ),
+ patch("litellm.token_counter", return_value=200_000),
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._call_summary_model",
+ mock_call,
+ ),
+ ):
+ result = await apply_compact_20260112(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ edit_spec=_EDIT_SPEC_DEFAULT,
+ user_api_key_auth=_fake_user_api_key_auth(
+ key_models=["all-proxy-models"], team_models=["gpt-4o"]
+ ),
+ )
+
+ mock_call.assert_not_awaited()
+ assert result.applied_edits[0].get("error") == "summary_model_access_denied"
+
+
+async def test_summary_model_allowed_when_in_key_allowlist():
+ """Caller key that explicitly allows the summary model is permitted to use it."""
+ messages = _simple_messages()
+ mock_call = AsyncMock(return_value=_make_mock_response("ok "))
+
+ with (
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
+ return_value="claude-haiku-4-5",
+ ),
+ patch("litellm.token_counter", return_value=200_000),
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._call_summary_model",
+ mock_call,
+ ),
+ ):
+ result = await apply_compact_20260112(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ edit_spec=_EDIT_SPEC_DEFAULT,
+ user_api_key_auth=_fake_user_api_key_auth(
+ key_models=["claude-haiku-4-5", "gpt-4o"]
+ ),
+ )
+
+ mock_call.assert_awaited_once()
+ assert result.compaction_block is not None
+ assert result.compaction_block["content"] == "ok"
+ assert not result.applied_edits[0].get("error")
+
+
+async def test_summary_model_allowed_when_no_user_api_key_auth():
+ """SDK callers (no proxy auth object) are not gated."""
+ messages = _simple_messages()
+ mock_call = AsyncMock(return_value=_make_mock_response("ok "))
+
+ with (
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
+ return_value="claude-haiku-4-5",
+ ),
+ patch("litellm.token_counter", return_value=200_000),
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._call_summary_model",
+ mock_call,
+ ),
+ ):
+ result = await apply_compact_20260112(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ edit_spec=_EDIT_SPEC_DEFAULT,
+ )
+
+ mock_call.assert_awaited_once()
+ assert result.compaction_block is not None
+
+
+async def test_summary_model_denied_when_user_scope_excludes_it():
+ """Personal user allowed-models scope denies the summary model even when
+ key/team allowlists permit it."""
+ messages = _simple_messages()
+ mock_call = AsyncMock(return_value=_make_mock_response("x "))
+
+ auth = _fake_user_api_key_auth(key_models=["all-proxy-models"])
+ auth.user_id = "user-123"
+
+ class _User:
+ user_id = "user-123"
+ models = ["gpt-3.5-turbo"]
+ organization_memberships = []
+
+ with (
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
+ return_value="claude-haiku-4-5",
+ ),
+ patch("litellm.token_counter", return_value=200_000),
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._call_summary_model",
+ mock_call,
+ ),
+ patch(
+ "litellm.proxy.auth.auth_checks.get_user_object",
+ AsyncMock(return_value=_User()),
+ ),
+ patch(
+ "litellm.proxy.auth.auth_checks.get_team_membership",
+ AsyncMock(return_value=None),
+ ),
+ patch(
+ "litellm.proxy.auth.auth_checks.get_project_object",
+ AsyncMock(return_value=None),
+ ),
+ patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
+ ):
+ result = await apply_compact_20260112(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ edit_spec=_EDIT_SPEC_DEFAULT,
+ user_api_key_auth=auth,
+ )
+
+ mock_call.assert_not_awaited()
+ assert result.applied_edits[0].get("error") == "summary_model_access_denied"
+
+
+async def test_summary_model_denied_when_project_scope_excludes_it():
+ """Project allowed-models scope denies the summary model even when
+ key/team allowlists permit it."""
+ messages = _simple_messages()
+ mock_call = AsyncMock(return_value=_make_mock_response("x "))
+
+ auth = _fake_user_api_key_auth(key_models=["all-proxy-models"])
+ auth.project_id = "project-1"
+
+ class _Project:
+ project_id = "project-1"
+ models = ["gpt-3.5-turbo"]
+
+ with (
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
+ return_value="claude-haiku-4-5",
+ ),
+ patch("litellm.token_counter", return_value=200_000),
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._call_summary_model",
+ mock_call,
+ ),
+ patch(
+ "litellm.proxy.auth.auth_checks.get_user_object",
+ AsyncMock(return_value=None),
+ ),
+ patch(
+ "litellm.proxy.auth.auth_checks.get_team_membership",
+ AsyncMock(return_value=None),
+ ),
+ patch(
+ "litellm.proxy.auth.auth_checks.get_project_object",
+ AsyncMock(return_value=_Project()),
+ ),
+ patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
+ ):
+ result = await apply_compact_20260112(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ edit_spec=_EDIT_SPEC_DEFAULT,
+ user_api_key_auth=auth,
+ )
+
+ mock_call.assert_not_awaited()
+ assert result.applied_edits[0].get("error") == "summary_model_access_denied"
+
+
+async def test_summary_model_denied_when_team_member_scope_excludes_it():
+ """Per-team-member allowed-models scope denies the summary model even
+ when key/team allowlists permit it."""
+ messages = _simple_messages()
+ mock_call = AsyncMock(return_value=_make_mock_response("x "))
+
+ auth = _fake_user_api_key_auth(key_models=["all-proxy-models"], team_id="team-1")
+ auth.user_id = "user-123"
+
+ class _Budget:
+ allowed_models = ["gpt-3.5-turbo"]
+
+ class _Membership:
+ litellm_budget_table = _Budget()
+
+ with (
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
+ return_value="claude-haiku-4-5",
+ ),
+ patch("litellm.token_counter", return_value=200_000),
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._call_summary_model",
+ mock_call,
+ ),
+ patch(
+ "litellm.proxy.auth.auth_checks.get_user_object",
+ AsyncMock(return_value=None),
+ ),
+ patch(
+ "litellm.proxy.auth.auth_checks.get_team_membership",
+ AsyncMock(return_value=_Membership()),
+ ),
+ patch(
+ "litellm.proxy.auth.auth_checks.get_project_object",
+ AsyncMock(return_value=None),
+ ),
+ patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
+ ):
+ result = await apply_compact_20260112(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ edit_spec=_EDIT_SPEC_DEFAULT,
+ user_api_key_auth=auth,
+ )
+
+ mock_call.assert_not_awaited()
+ assert result.applied_edits[0].get("error") == "summary_model_access_denied"
+
+
+async def test_summary_model_denied_when_key_over_model_budget():
+ """A caller whose per-model budget for the summary model is exhausted cannot
+ trigger the summary call via compaction."""
+ import litellm
+
+ messages = _simple_messages()
+ mock_call = AsyncMock(return_value=_make_mock_response("x "))
+
+ auth = _fake_user_api_key_auth(
+ key_models=["all-proxy-models"],
+ model_max_budget={"claude-haiku-4-5": {"budget_limit": 5}},
+ token="hashed-token",
+ )
+
+ limiter = MagicMock()
+ limiter.is_key_within_model_budget = AsyncMock(
+ side_effect=litellm.BudgetExceededError(
+ message="over budget", current_cost=10, max_budget=5
+ )
+ )
+
+ with (
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
+ return_value="claude-haiku-4-5",
+ ),
+ patch("litellm.token_counter", return_value=200_000),
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._call_summary_model",
+ mock_call,
+ ),
+ patch("litellm.proxy.proxy_server.model_max_budget_limiter", limiter),
+ ):
+ result = await apply_compact_20260112(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ edit_spec=_EDIT_SPEC_DEFAULT,
+ user_api_key_auth=auth,
+ )
+
+ mock_call.assert_not_awaited()
+ limiter.is_key_within_model_budget.assert_awaited_once()
+ assert result.applied_edits[0].get("error") == "summary_model_budget_exceeded"
+
+
+async def test_summary_model_denied_when_end_user_over_model_budget():
+ """End-user per-model budget is enforced for the summary subrequest too."""
+ import litellm
+
+ messages = _simple_messages()
+ mock_call = AsyncMock(return_value=_make_mock_response("x "))
+
+ auth = _fake_user_api_key_auth(
+ key_models=["all-proxy-models"],
+ end_user_model_max_budget={"claude-haiku-4-5": {"budget_limit": 5}},
+ end_user_id="end-user-1",
+ token="hashed-token",
+ )
+
+ limiter = MagicMock()
+ limiter.is_key_within_model_budget = AsyncMock(return_value=True)
+ limiter.is_end_user_within_model_budget = AsyncMock(
+ side_effect=litellm.BudgetExceededError(
+ message="over budget", current_cost=10, max_budget=5
+ )
+ )
+
+ with (
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
+ return_value="claude-haiku-4-5",
+ ),
+ patch("litellm.token_counter", return_value=200_000),
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._call_summary_model",
+ mock_call,
+ ),
+ patch("litellm.proxy.proxy_server.model_max_budget_limiter", limiter),
+ ):
+ result = await apply_compact_20260112(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ edit_spec=_EDIT_SPEC_DEFAULT,
+ user_api_key_auth=auth,
+ )
+
+ mock_call.assert_not_awaited()
+ limiter.is_end_user_within_model_budget.assert_awaited_once()
+ assert result.applied_edits[0].get("error") == "summary_model_budget_exceeded"
+
+
+async def test_summary_model_allowed_when_within_model_budget():
+ """When the per-model budget check passes, the summary call proceeds."""
+ messages = _simple_messages()
+ mock_call = AsyncMock(return_value=_make_mock_response("ok "))
+
+ auth = _fake_user_api_key_auth(
+ key_models=["all-proxy-models"],
+ model_max_budget={"claude-haiku-4-5": {"budget_limit": 5}},
+ token="hashed-token",
+ )
+
+ limiter = MagicMock()
+ limiter.is_key_within_model_budget = AsyncMock(return_value=True)
+ limiter.is_end_user_within_model_budget = AsyncMock(return_value=True)
+
+ with (
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
+ return_value="claude-haiku-4-5",
+ ),
+ patch("litellm.token_counter", return_value=200_000),
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._call_summary_model",
+ mock_call,
+ ),
+ patch("litellm.proxy.proxy_server.model_max_budget_limiter", limiter),
+ ):
+ result = await apply_compact_20260112(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ edit_spec=_EDIT_SPEC_DEFAULT,
+ user_api_key_auth=auth,
+ )
+
+ mock_call.assert_awaited_once()
+ limiter.is_key_within_model_budget.assert_awaited_once()
+ assert not result.applied_edits[0].get("error")
+
+
+class _FakeRateLimiter:
+ """Minimal stand-in for ``_PROXY_MaxParallelRequestsHandler_v3`` exposing
+ just the descriptor-build + read-only check surface the editor consults."""
+
+ def __init__(self, overall_code: str):
+ self._overall_code = overall_code
+ self.read_only_checked = False
+
+ def _create_rate_limit_descriptors(self, **kwargs):
+ return [
+ {
+ "key": "api_key",
+ "value": "hashed-token",
+ "rate_limit": {"requests_per_unit": 10},
+ }
+ ]
+
+ def _add_team_model_rate_limit_descriptor_from_metadata(self, **kwargs):
+ return None
+
+ def _add_project_model_rate_limit_descriptor_from_metadata(self, **kwargs):
+ return None
+
+ def create_organization_rate_limit_descriptor(self, *args, **kwargs):
+ return []
+
+ async def should_rate_limit(self, **kwargs):
+ self.read_only_checked = kwargs.get("read_only") is True
+ return {"overall_code": self._overall_code}
+
+
+async def test_summary_model_denied_when_over_rate_limit():
+ """A caller already at their configured RPM/TPM for the summary model cannot
+ drive an extra summary completion via compaction."""
+ messages = _simple_messages()
+ mock_call = AsyncMock(return_value=_make_mock_response("x "))
+
+ auth = _fake_user_api_key_auth(key_models=["all-proxy-models"])
+ limiter = _FakeRateLimiter("OVER_LIMIT")
+ proxy_logging = MagicMock()
+ proxy_logging.max_parallel_request_limiter = limiter
+
+ with (
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
+ return_value="claude-haiku-4-5",
+ ),
+ patch("litellm.token_counter", return_value=200_000),
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._call_summary_model",
+ mock_call,
+ ),
+ patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging),
+ ):
+ result = await apply_compact_20260112(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ edit_spec=_EDIT_SPEC_DEFAULT,
+ user_api_key_auth=auth,
+ )
+
+ mock_call.assert_not_awaited()
+ assert limiter.read_only_checked is True
+ assert result.compaction_block is None
+ assert result.applied_edits[0].get("error") == "summary_model_rate_limit_exceeded"
+
+
+async def test_summary_model_allowed_when_within_rate_limit():
+ """When the read-only rate-limit check is under limit, the summary call proceeds."""
+ messages = _simple_messages()
+ mock_call = AsyncMock(return_value=_make_mock_response("ok "))
+
+ auth = _fake_user_api_key_auth(key_models=["all-proxy-models"])
+ limiter = _FakeRateLimiter("OK")
+ proxy_logging = MagicMock()
+ proxy_logging.max_parallel_request_limiter = limiter
+
+ with (
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
+ return_value="claude-haiku-4-5",
+ ),
+ patch("litellm.token_counter", return_value=200_000),
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._call_summary_model",
+ mock_call,
+ ),
+ patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging),
+ ):
+ result = await apply_compact_20260112(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ edit_spec=_EDIT_SPEC_DEFAULT,
+ user_api_key_auth=auth,
+ )
+
+ mock_call.assert_awaited_once()
+ assert limiter.read_only_checked is True
+ assert result.compaction_block is not None
+ assert not result.applied_edits[0].get("error")
+
+
+async def test_summary_model_rate_limit_skipped_for_legacy_limiter():
+ """A limiter without the v3 read-only check surface fails open so the summary
+ call still proceeds (its usage is still charged post-call)."""
+ messages = _simple_messages()
+ mock_call = AsyncMock(return_value=_make_mock_response("ok "))
+
+ auth = _fake_user_api_key_auth(key_models=["all-proxy-models"])
+
+ class _LegacyLimiter:
+ async def async_pre_call_hook(self, **kwargs):
+ return None
+
+ proxy_logging = MagicMock()
+ proxy_logging.max_parallel_request_limiter = _LegacyLimiter()
+
+ with (
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
+ return_value="claude-haiku-4-5",
+ ),
+ patch("litellm.token_counter", return_value=200_000),
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._call_summary_model",
+ mock_call,
+ ),
+ patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging),
+ ):
+ result = await apply_compact_20260112(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ edit_spec=_EDIT_SPEC_DEFAULT,
+ user_api_key_auth=auth,
+ )
+
+ mock_call.assert_awaited_once()
+ assert result.compaction_block is not None
+ assert not result.applied_edits[0].get("error")
+
+
+async def test_scoped_budget_metadata_propagated_to_summary_call():
+ """The end-user/project scope identifiers and the end-user budget the post-call
+ spend and rate-limit hooks key on are forwarded to the summary subrequest, and
+ the end-user id is also passed as the top-level ``user`` kwarg the legacy
+ limiter hooks read, so the summary tokens debit those scoped budgets/counters."""
+ messages = _simple_messages()
+ mock_response = _make_mock_response("Summary ")
+ parent_litellm_metadata = {
+ "user_api_key": "sk-test",
+ "user_api_key_end_user_id": "customer-1",
+ "user_api_end_user_max_budget": 10,
+ "user_api_key_project_id": "project-9",
+ }
+
+ with (
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
+ return_value="claude-haiku-4-5",
+ ),
+ patch("litellm.token_counter", return_value=200_000),
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._call_summary_model",
+ new_callable=AsyncMock,
+ return_value=mock_response,
+ ) as mock_call,
+ ):
+ await apply_compact_20260112(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ edit_spec=_EDIT_SPEC_DEFAULT,
+ litellm_metadata=parent_litellm_metadata,
+ )
+
+ propagated = mock_call.call_args.kwargs["metadata"]
+ assert propagated["user_api_key_end_user_id"] == "customer-1"
+ assert propagated["user_api_end_user_max_budget"] == 10
+ assert propagated["user_api_key_project_id"] == "project-9"
+
+
+async def test_summary_call_passes_end_user_id_as_top_level_user():
+ """``_call_summary_model`` forwards the propagated end-user id as the top-level
+ ``user`` kwarg that legacy limiter / prometheus end-user tracking reads."""
+ from litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact import (
+ _call_summary_model,
+ )
+
+ captured_kwargs: dict = {}
+
+ class _FakeRouter:
+ async def acompletion(self, **kwargs):
+ captured_kwargs.update(kwargs)
+ return _make_mock_response("x ")
+
+ await _call_summary_model(
+ summary_model="claude-haiku-4-5",
+ summary_messages=[{"role": "user", "content": "hi"}],
+ metadata={"user_api_key_end_user_id": "customer-1"},
+ llm_router=_FakeRouter(),
+ )
+
+ assert captured_kwargs.get("user") == "customer-1"
+
+
+async def test_summary_call_omits_user_when_no_end_user_id():
+ """No end-user id on the parent request means no ``user`` kwarg is sent."""
+ from litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact import (
+ _call_summary_model,
+ )
+
+ captured_kwargs: dict = {}
+
+ class _FakeRouter:
+ async def acompletion(self, **kwargs):
+ captured_kwargs.update(kwargs)
+ return _make_mock_response("x ")
+
+ await _call_summary_model(
+ summary_model="claude-haiku-4-5",
+ summary_messages=[{"role": "user", "content": "hi"}],
+ metadata={},
+ llm_router=_FakeRouter(),
+ )
+
+ assert "user" not in captured_kwargs
+
+
+async def test_model_budget_metadata_propagated_to_summary_call():
+ """The per-model budget metadata the spend caches rely on is forwarded to the
+ summary subrequest so its spend counts against the caller's model budget."""
+ messages = _simple_messages()
+ mock_response = _make_mock_response("Summary ")
+ parent_litellm_metadata = {
+ "user_api_key": "sk-test",
+ "user_api_key_model_max_budget": {"claude-haiku-4-5": {"budget_limit": 5}},
+ "user_api_key_end_user_model_max_budget": {
+ "claude-haiku-4-5": {"budget_limit": 2}
+ },
+ }
+
+ with (
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
+ return_value="claude-haiku-4-5",
+ ),
+ patch("litellm.token_counter", return_value=200_000),
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._call_summary_model",
+ new_callable=AsyncMock,
+ return_value=mock_response,
+ ) as mock_call,
+ ):
+ await apply_compact_20260112(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ edit_spec=_EDIT_SPEC_DEFAULT,
+ litellm_metadata=parent_litellm_metadata,
+ )
+
+ propagated = mock_call.call_args.kwargs["metadata"]
+ assert propagated["user_api_key_model_max_budget"] == {
+ "claude-haiku-4-5": {"budget_limit": 5}
+ }
+ assert propagated["user_api_key_end_user_model_max_budget"] == {
+ "claude-haiku-4-5": {"budget_limit": 2}
+ }
+
+
+async def test_summary_call_propagates_allowed_model_region():
+ """``allowed_model_region`` from ``user_api_key_auth`` is propagated to the
+ summary subrequest as a top-level kwarg so the router applies the same
+ region restriction the parent request would.
+ """
+ messages = _simple_messages()
+ mock_call = AsyncMock(return_value=_make_mock_response("ok "))
+
+ auth = _fake_user_api_key_auth(key_models=["all-proxy-models"])
+ auth.allowed_model_region = "eu"
+
+ with (
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
+ return_value="claude-haiku-4-5",
+ ),
+ patch("litellm.token_counter", return_value=200_000),
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._call_summary_model",
+ mock_call,
+ ),
+ ):
+ await apply_compact_20260112(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ edit_spec=_EDIT_SPEC_DEFAULT,
+ user_api_key_auth=auth,
+ )
+
+ mock_call.assert_awaited_once()
+ assert mock_call.await_args.kwargs.get("allowed_model_region") == "eu"
+
+
+async def test_summary_call_omits_allowed_model_region_when_unset():
+ """Callers without a region restriction must not get an ``allowed_model_region=None``
+ kwarg, which would otherwise force the router to evaluate region filtering.
+ """
+ from litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact import (
+ _call_summary_model,
+ )
+
+ captured_kwargs: dict = {}
+
+ class _FakeRouter:
+ async def acompletion(self, **kwargs):
+ captured_kwargs.update(kwargs)
+ return _make_mock_response("x ")
+
+ await _call_summary_model(
+ summary_model="claude-haiku-4-5",
+ summary_messages=[{"role": "user", "content": "hi"}],
+ metadata={},
+ llm_router=_FakeRouter(),
+ )
+
+ assert "allowed_model_region" not in captured_kwargs
+
+
+async def test_summary_call_forwards_allowed_model_region_when_set():
+ """When the caller is region-restricted, the kwarg reaches the router."""
+ from litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact import (
+ _call_summary_model,
+ )
+
+ captured_kwargs: dict = {}
+
+ class _FakeRouter:
+ async def acompletion(self, **kwargs):
+ captured_kwargs.update(kwargs)
+ return _make_mock_response("x ")
+
+ await _call_summary_model(
+ summary_model="claude-haiku-4-5",
+ summary_messages=[{"role": "user", "content": "hi"}],
+ metadata={},
+ llm_router=_FakeRouter(),
+ allowed_model_region="eu",
+ )
+
+ assert captured_kwargs.get("allowed_model_region") == "eu"
+
+
+# ---------------------------------------------------------------------------
+# Dispatcher integration: compact_20260112 via apply_context_management
+# ---------------------------------------------------------------------------
+
+
+async def test_dispatcher_routes_compact_edit():
+ """compact_20260112 in the dispatcher resolves to opt-in gate when no model set."""
+ messages = _simple_messages()
+ with patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
+ return_value=None,
+ ):
+ result = await apply_context_management(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ context_management_spec={"edits": [{"type": "compact_20260112"}]},
+ )
+
+ assert len(result.applied_edits) == 1
+ assert result.applied_edits[0]["type"] == "compact_20260112"
+ assert result.applied_edits[0].get("error") == "summary_model_not_configured"
+
+
+async def test_dispatcher_trigger_below_minimum_raises_through():
+ """AnthropicContextManagementError from the editor bubbles up through the dispatcher."""
+ with pytest.raises(AnthropicContextManagementError):
+ await apply_context_management(
+ model=MODEL,
+ messages=_simple_messages(),
+ tools=None,
+ system=None,
+ context_management_spec={
+ "edits": [
+ {
+ "type": "compact_20260112",
+ "trigger": {"type": "input_tokens", "value": 1_000},
+ }
+ ]
+ },
+ )
+
+
+# ---------------------------------------------------------------------------
+# _run_polyfill_if_enabled: drop_params gate
+# ---------------------------------------------------------------------------
+
+
+async def test_run_polyfill_skipped_when_drop_params_true():
+ """When drop_params=True the polyfill must be skipped (returns None)."""
+ from litellm.llms.anthropic.experimental_pass_through.adapters.handler import (
+ _run_polyfill_if_enabled,
+ )
+
+ result = await _run_polyfill_if_enabled(
+ model=MODEL,
+ messages=_simple_messages(),
+ tools=None,
+ system=None,
+ context_management_spec={"edits": [{"type": "compact_20260112"}]},
+ litellm_metadata={},
+ drop_params=True,
+ llm_router=None,
+ )
+ assert result is None
+
+
+async def test_run_polyfill_skipped_when_spec_empty():
+ """Empty context_management_spec must also return None (no polyfill work)."""
+ from litellm.llms.anthropic.experimental_pass_through.adapters.handler import (
+ _run_polyfill_if_enabled,
+ )
+
+ result = await _run_polyfill_if_enabled(
+ model=MODEL,
+ messages=_simple_messages(),
+ tools=None,
+ system=None,
+ context_management_spec=None,
+ litellm_metadata={},
+ drop_params=False,
+ llm_router=None,
+ )
+ assert result is None
+
+
+async def test_prepare_context_managed_request_forwards_proxy_litellm_metadata():
+ """The handler must hand the polyfill the proxy ``litellm_metadata`` (which
+ carries ``user_api_key`` / ``user_api_key_team_id`` / ...), not the
+ Anthropic-shape ``metadata`` arg (which only carries ``user_id``). Otherwise
+ the summary subcall lands on the router with no parent attribution, and
+ those tokens go unbilled to the caller's key/team."""
+ from litellm.llms.anthropic.experimental_pass_through.adapters.handler import (
+ _prepare_context_managed_request,
+ )
+
+ captured_summary_metadata: Dict[str, Any] = {}
+
+ class _RouterStub:
+ async def acompletion(self, **kwargs):
+ captured_summary_metadata.update(kwargs.get("litellm_metadata", {}))
+ return _make_mock_response("s ")
+
+ with (
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
+ return_value="claude-haiku-4-5",
+ ),
+ patch("litellm.token_counter", return_value=200_000),
+ ):
+ result = await _prepare_context_managed_request(
+ model=MODEL,
+ messages=_simple_messages(),
+ tools=None,
+ system=None,
+ context_management_spec={"edits": [_EDIT_SPEC_DEFAULT]},
+ litellm_metadata={
+ "user_api_key": "sk-parent",
+ "user_api_key_team_id": "team-abc",
+ "user_api_key_user_id": "user-xyz",
+ "litellm_call_id": "call-1",
+ },
+ drop_params=False,
+ llm_router=_RouterStub(),
+ )
+
+ assert result is not None
+ assert captured_summary_metadata.get("user_api_key") == "sk-parent"
+ assert captured_summary_metadata.get("user_api_key_team_id") == "team-abc"
+ assert captured_summary_metadata.get("user_api_key_user_id") == "user-xyz"
+ assert captured_summary_metadata.get("litellm_call_id") == "call-1"
+ # Anthropic-shape ``metadata.user_id`` must not leak in as a propagated field.
+ assert "user_id" not in captured_summary_metadata
+
+
+# ---------------------------------------------------------------------------
+# Endpoint error format: AnthropicContextManagementError → Anthropic 400 body
+# ---------------------------------------------------------------------------
+
+
+def test_anthropic_context_management_error_format():
+ """AnthropicContextManagementError must produce an Anthropic-format body via
+ AnthropicExceptionMapping.transform_to_anthropic_error — the same path the
+ /v1/messages endpoint takes when it catches this exception."""
+ from litellm.anthropic_interface.exceptions import AnthropicExceptionMapping
+
+ body = AnthropicExceptionMapping.transform_to_anthropic_error(
+ status_code=400,
+ raw_message="trigger.value must be at least 50000 tokens",
+ request_id=None,
+ )
+
+ assert body["type"] == "error"
+ assert body["error"]["type"] == "invalid_request_error"
+ assert "50000" in body["error"]["message"]
+
+
+def test_anthropic_context_management_error_attrs():
+ """AnthropicContextManagementError carries status_code and message correctly."""
+ err = AnthropicContextManagementError(
+ status_code=400,
+ message="trigger.value must be at least 50000 tokens",
+ )
+
+ assert err.status_code == 400
+ assert "50000" in err.message
+
+
+# ---------------------------------------------------------------------------
+# Endpoint integration: /v1/messages → Anthropic 400 on context management error
+# ---------------------------------------------------------------------------
+
+
+def test_endpoint_returns_anthropic_400_on_context_management_error():
+ """The /v1/messages endpoint must catch AnthropicContextManagementError and
+ return an Anthropic-format 400 JSONResponse — not a 500 ProxyException."""
+ import sys
+ from unittest.mock import AsyncMock, MagicMock, patch
+
+ from fastapi import FastAPI
+ from fastapi.testclient import TestClient
+
+ from litellm.proxy.anthropic_endpoints.endpoints import router
+ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+
+ # Stub proxy_server to avoid apscheduler/heavy proxy deps imported lazily
+ # inside the route handler at request time.
+ mock_proxy_server = MagicMock()
+ mock_proxy_server.general_settings = {}
+ mock_proxy_server.llm_router = None
+ mock_proxy_server.proxy_config = MagicMock()
+ mock_proxy_server.proxy_logging_obj = MagicMock()
+ mock_proxy_server.user_api_base = None
+ mock_proxy_server.user_max_tokens = None
+ mock_proxy_server.user_model = None
+ mock_proxy_server.user_request_timeout = None
+ mock_proxy_server.user_temperature = None
+ mock_proxy_server.version = "test"
+
+ with patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}):
+ with patch(
+ "litellm.proxy.anthropic_endpoints.endpoints.ProxyBaseLLMRequestProcessing"
+ ) as mock_cls:
+ mock_instance = MagicMock()
+ mock_instance.base_process_llm_request = AsyncMock(
+ side_effect=AnthropicContextManagementError(
+ status_code=400,
+ message="trigger.value must be at least 50000 tokens",
+ )
+ )
+ mock_cls.return_value = mock_instance
+
+ app = FastAPI()
+ app.include_router(router)
+ app.dependency_overrides[user_api_key_auth] = lambda: MagicMock()
+
+ client = TestClient(app, raise_server_exceptions=False)
+ response = client.post(
+ "/v1/messages",
+ json={
+ "model": "gpt-4o",
+ "messages": [{"role": "user", "content": "hi"}],
+ },
+ headers={"Authorization": "Bearer test-key"},
+ )
+
+ assert response.status_code == 400
+ body = response.json()
+ assert body["type"] == "error"
+ assert body["error"]["type"] == "invalid_request_error"
+ assert "50000" in body["error"]["message"]
+
+
+def test_endpoint_runs_failure_hook_on_500_context_management_error():
+ """A 500-level AnthropicContextManagementError (internal polyfill failure)
+ must invoke post_call_failure_hook for spend/alerting parity, while still
+ returning the Anthropic-format error body."""
+ import sys
+ from unittest.mock import AsyncMock, MagicMock, patch
+
+ from fastapi import FastAPI
+ from fastapi.testclient import TestClient
+
+ from litellm.proxy.anthropic_endpoints.endpoints import router
+ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+
+ failure_hook = AsyncMock()
+ mock_proxy_server = MagicMock()
+ mock_proxy_server.general_settings = {}
+ mock_proxy_server.llm_router = None
+ mock_proxy_server.proxy_config = MagicMock()
+ mock_proxy_server.proxy_logging_obj = MagicMock()
+ mock_proxy_server.proxy_logging_obj.post_call_failure_hook = failure_hook
+ mock_proxy_server.user_api_base = None
+ mock_proxy_server.user_max_tokens = None
+ mock_proxy_server.user_model = None
+ mock_proxy_server.user_request_timeout = None
+ mock_proxy_server.user_temperature = None
+ mock_proxy_server.version = "test"
+
+ with patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}):
+ with patch(
+ "litellm.proxy.anthropic_endpoints.endpoints.ProxyBaseLLMRequestProcessing"
+ ) as mock_cls:
+ mock_instance = MagicMock()
+ mock_instance.base_process_llm_request = AsyncMock(
+ side_effect=AnthropicContextManagementError(
+ status_code=500,
+ message="context_management polyfill failed: boom",
+ )
+ )
+ mock_cls.return_value = mock_instance
+
+ app = FastAPI()
+ app.include_router(router)
+ app.dependency_overrides[user_api_key_auth] = lambda: MagicMock()
+
+ client = TestClient(app, raise_server_exceptions=False)
+ response = client.post(
+ "/v1/messages",
+ json={
+ "model": "gpt-4o",
+ "messages": [{"role": "user", "content": "hi"}],
+ },
+ headers={"Authorization": "Bearer test-key"},
+ )
+
+ assert response.status_code == 500
+ body = response.json()
+ assert body["type"] == "error"
+ failure_hook.assert_awaited_once()
diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_dispatcher.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_dispatcher.py
new file mode 100644
index 00000000000..50c72cfe8d0
--- /dev/null
+++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_dispatcher.py
@@ -0,0 +1,131 @@
+"""
+Unit tests for the context_management polyfill dispatcher.
+"""
+
+from litellm.llms.anthropic.experimental_pass_through.context_management import (
+ apply_context_management,
+)
+
+MODEL = "xai/grok-4"
+
+
+def _history_with_two_tool_pairs():
+ return [
+ {"role": "user", "content": "Hi"},
+ {
+ "role": "assistant",
+ "content": [{"type": "tool_use", "id": "t1", "name": "f", "input": {}}],
+ },
+ {
+ "role": "user",
+ "content": [
+ {
+ "type": "tool_result",
+ "tool_use_id": "t1",
+ "content": "first result",
+ }
+ ],
+ },
+ {
+ "role": "assistant",
+ "content": [{"type": "tool_use", "id": "t2", "name": "f", "input": {}}],
+ },
+ {
+ "role": "user",
+ "content": [
+ {
+ "type": "tool_result",
+ "tool_use_id": "t2",
+ "content": "second result",
+ }
+ ],
+ },
+ ]
+
+
+async def test_unknown_edit_type_is_noop():
+ messages = _history_with_two_tool_pairs()
+ result = await apply_context_management(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ context_management_spec={
+ "edits": [{"type": "totally_not_a_real_edit_20999999"}]
+ },
+ )
+ assert result.applied_edits == []
+ assert result.messages == messages
+
+
+async def test_known_edit_is_applied():
+ messages = _history_with_two_tool_pairs()
+ result = await apply_context_management(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ context_management_spec={
+ "edits": [
+ {
+ "type": "clear_tool_uses_20250919",
+ "trigger": {"type": "tool_uses", "value": 1},
+ "keep": {"type": "tool_uses", "value": 1},
+ }
+ ]
+ },
+ )
+ assert len(result.applied_edits) == 1
+ assert result.applied_edits[0]["type"] == "clear_tool_uses_20250919"
+ assert result.applied_edits[0]["cleared_tool_uses"] == 1
+
+
+async def test_mixed_known_unknown_only_known_applied():
+ messages = _history_with_two_tool_pairs()
+ result = await apply_context_management(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ context_management_spec={
+ "edits": [
+ {"type": "unknown_foo"},
+ {
+ "type": "clear_tool_uses_20250919",
+ "trigger": {"type": "tool_uses", "value": 0},
+ "keep": {"type": "tool_uses", "value": 1},
+ },
+ {"type": "another_unknown"},
+ ]
+ },
+ )
+ assert len(result.applied_edits) == 1
+ assert result.applied_edits[0]["type"] == "clear_tool_uses_20250919"
+
+
+async def test_empty_or_missing_edits_list():
+ messages = _history_with_two_tool_pairs()
+ for spec in [{}, {"edits": None}, {"edits": []}, None]:
+ result = await apply_context_management(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ context_management_spec=spec, # type: ignore[arg-type]
+ )
+ assert result.applied_edits == []
+ assert result.messages == messages
+
+
+async def test_malformed_edit_entries_are_skipped():
+ """Non-dict entries in `edits` list should be silently skipped."""
+ messages = _history_with_two_tool_pairs()
+ result = await apply_context_management(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ context_management_spec={"edits": ["not a dict", 42, None, {"type": None}]},
+ )
+ assert result.applied_edits == []
+ assert result.messages == messages
diff --git a/tests/test_litellm/llms/bedrock/test_converse_context_management.py b/tests/test_litellm/llms/bedrock/test_converse_context_management.py
new file mode 100644
index 00000000000..709fc4e8b39
--- /dev/null
+++ b/tests/test_litellm/llms/bedrock/test_converse_context_management.py
@@ -0,0 +1,114 @@
+"""Bedrock Converse context_management forwarding (compact_20260112 only)."""
+
+from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig
+
+CLAUDE_MODEL = "anthropic.claude-opus-4-7-20250115-v1:0"
+
+
+def test_supported_params_include_context_management_for_anthropic():
+ cfg = AmazonConverseConfig()
+ params = cfg.get_supported_openai_params(CLAUDE_MODEL)
+ assert "context_management" in params
+
+
+def test_supported_params_exclude_context_management_for_non_anthropic():
+ cfg = AmazonConverseConfig()
+ params = cfg.get_supported_openai_params("meta.llama3-70b-instruct-v1:0")
+ assert "context_management" not in params
+
+
+def test_map_openai_params_forwards_anthropic_shape():
+ cfg = AmazonConverseConfig()
+ optional_params: dict = {}
+ cfg.map_openai_params(
+ non_default_params={
+ "context_management": {"edits": [{"type": "compact_20260112"}]}
+ },
+ optional_params=optional_params,
+ model=CLAUDE_MODEL,
+ drop_params=False,
+ )
+ assert optional_params.get("context_management") == {
+ "edits": [{"type": "compact_20260112"}]
+ }
+
+
+def test_map_openai_params_normalizes_openai_list_shape():
+ """OpenAI Responses-API style list of {type: "compaction"} normalizes to Anthropic dict."""
+ cfg = AmazonConverseConfig()
+ optional_params: dict = {}
+ cfg.map_openai_params(
+ non_default_params={"context_management": [{"type": "compaction"}]},
+ optional_params=optional_params,
+ model=CLAUDE_MODEL,
+ drop_params=False,
+ )
+ forwarded = optional_params.get("context_management")
+ assert isinstance(forwarded, dict)
+ edits = forwarded.get("edits")
+ assert isinstance(edits, list) and len(edits) == 1
+ assert edits[0].get("type") == "compact_20260112"
+
+
+def test_filter_keeps_only_compact_edits_and_adds_beta_header():
+ additional = {
+ "context_management": {
+ "edits": [
+ {"type": "clear_tool_uses_20250919"},
+ {"type": "compact_20260112"},
+ {"type": "clear_thinking_20251015"},
+ ]
+ }
+ }
+ betas: list = []
+ AmazonConverseConfig._filter_context_management_for_bedrock_converse(
+ additional, betas
+ )
+ assert additional["context_management"]["edits"] == [{"type": "compact_20260112"}]
+ assert "compact-2026-01-12" in betas
+
+
+def test_filter_drops_field_when_no_compact_edit_remains():
+ additional = {
+ "context_management": {
+ "edits": [
+ {"type": "clear_tool_uses_20250919"},
+ {"type": "clear_thinking_20251015"},
+ ]
+ }
+ }
+ betas: list = []
+ AmazonConverseConfig._filter_context_management_for_bedrock_converse(
+ additional, betas
+ )
+ assert "context_management" not in additional
+ assert betas == []
+
+
+def test_filter_is_noop_when_field_absent():
+ additional: dict = {}
+ betas: list = []
+ AmazonConverseConfig._filter_context_management_for_bedrock_converse(
+ additional, betas
+ )
+ assert additional == {}
+ assert betas == []
+
+
+def test_filter_drops_malformed_edits_list():
+ additional = {"context_management": {"edits": "not a list"}}
+ betas: list = []
+ AmazonConverseConfig._filter_context_management_for_bedrock_converse(
+ additional, betas
+ )
+ assert "context_management" not in additional
+ assert betas == []
+
+
+def test_filter_does_not_duplicate_beta_header():
+ additional = {"context_management": {"edits": [{"type": "compact_20260112"}]}}
+ betas: list = ["compact-2026-01-12"]
+ AmazonConverseConfig._filter_context_management_for_bedrock_converse(
+ additional, betas
+ )
+ assert betas.count("compact-2026-01-12") == 1
From 967fed1fa1e03ae476be83d7c3f484973151dbbf Mon Sep 17 00:00:00 2001
From: Shivam Rawat
Date: Sat, 30 May 2026 10:09:05 -0700
Subject: [PATCH 25/44] feat(enterprise): add RESEND_FROM_EMAIL for self-hosted
Resend sends (#28830)
Allow self-hosted installs to override the default LiteLLM sender address
via RESEND_FROM_EMAIL, matching SendGrid's SENDGRID_SENDER_EMAIL pattern.
Co-authored-by: Cursor
---
.../send_emails/resend_email.py | 19 ++++-
.../send_emails/test_resend_email.py | 76 +++++++++++++++++--
2 files changed, 88 insertions(+), 7 deletions(-)
diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/resend_email.py b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/resend_email.py
index 7593e66aa47..3fad5601f52 100644
--- a/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/resend_email.py
+++ b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/resend_email.py
@@ -19,12 +19,26 @@ RESEND_API_ENDPOINT = "https://api.resend.com/emails"
class ResendEmailLogger(BaseEmailLogger):
+ """
+ Send emails using Resend's API.
+
+ Required env vars:
+ - RESEND_API_KEY
+
+ Optional env vars:
+ - RESEND_FROM_EMAIL: Override the default sender address. Must be on a
+ domain verified in your Resend account. When unset, falls back to the
+ `from_email` argument passed by the caller (which defaults to
+ `notifications@alerts.litellm.ai` and only works on LiteLLM Cloud).
+ """
+
def __init__(self, internal_usage_cache=None, **kwargs):
super().__init__(internal_usage_cache=internal_usage_cache, **kwargs)
self.async_httpx_client = get_async_httpx_client(
llm_provider=httpxSpecialProvider.LoggingCallback
)
self.resend_api_key = os.getenv("RESEND_API_KEY")
+ self.resend_from_email = os.getenv("RESEND_FROM_EMAIL")
async def send_email(
self,
@@ -33,13 +47,14 @@ class ResendEmailLogger(BaseEmailLogger):
subject: str,
html_body: str,
):
+ sender_email = self.resend_from_email or from_email
verbose_logger.debug(
- f"Sending email from {from_email} to {to_email} with subject {subject}"
+ f"Sending email from {sender_email} to {to_email} with subject {subject}"
)
response = await self.async_httpx_client.post(
url=RESEND_API_ENDPOINT,
json={
- "from": from_email,
+ "from": sender_email,
"to": to_email,
"subject": subject,
"html": html_body,
diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_resend_email.py b/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_resend_email.py
index b07216921eb..88cc2275ae2 100644
--- a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_resend_email.py
+++ b/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_resend_email.py
@@ -32,7 +32,11 @@ def clear_client_cache():
@pytest.fixture
def mock_env_vars():
- with mock.patch.dict(os.environ, {"RESEND_API_KEY": "test_api_key"}):
+ # Set test API key and ensure RESEND_FROM_EMAIL is unset for isolation
+ # so tests can verify the default `from_email` argument is used.
+ patched = {"RESEND_API_KEY": "test_api_key"}
+ with mock.patch.dict(os.environ, patched):
+ os.environ.pop("RESEND_FROM_EMAIL", None)
yield
@@ -87,7 +91,7 @@ async def test_send_email_success(mock_env_vars):
async def test_send_email_missing_api_key():
# Remove the API key from environment before initializing logger
original_key = os.environ.pop("RESEND_API_KEY", None)
-
+
try:
# Initialize the logger after removing the API key
logger = ResendEmailLogger()
@@ -104,16 +108,19 @@ async def test_send_email_missing_api_key():
mock_response.raise_for_status.return_value = None
mock_response.status_code = 200
mock_response.json.return_value = {"id": "test_email_id"}
-
+
mock_async_client = mock.AsyncMock()
mock_async_client.post.return_value = mock_response
-
+
# Directly inject the mock client to bypass any caching
logger.async_httpx_client = mock_async_client
# Send email
await logger.send_email(
- from_email=from_email, to_email=to_email, subject=subject, html_body=html_body
+ from_email=from_email,
+ to_email=to_email,
+ subject=subject,
+ html_body=html_body,
)
# Verify the HTTP client was called with None as the API key
@@ -159,3 +166,62 @@ async def test_send_email_multiple_recipients(mock_env_vars):
call_args = mock_async_client.post.call_args
request_body = call_args[1]["json"]
assert request_body["to"] == to_email
+
+
+@pytest.mark.asyncio
+async def test_send_email_uses_resend_from_email_override():
+ """RESEND_FROM_EMAIL overrides the caller-supplied from_email."""
+ with mock.patch.dict(
+ os.environ,
+ {
+ "RESEND_API_KEY": "test_api_key",
+ "RESEND_FROM_EMAIL": "alerts@my-verified-domain.com",
+ },
+ ):
+ logger = ResendEmailLogger()
+
+ mock_response = mock.Mock(spec=Response)
+ mock_response.status_code = 200
+ mock_response.json.return_value = {"id": "test_email_id"}
+ mock_response.raise_for_status.return_value = None
+
+ mock_async_client = mock.AsyncMock()
+ mock_async_client.post.return_value = mock_response
+ logger.async_httpx_client = mock_async_client
+
+ await logger.send_email(
+ from_email="notifications@alerts.litellm.ai",
+ to_email=["recipient@example.com"],
+ subject="Test Subject",
+ html_body="
Date: Sat, 30 May 2026 12:01:22 -0700
Subject: [PATCH 27/44] fix(mcp): preserve source_url in GET /v1/mcp/server
list responses (#29249)
* fix(mcp): preserve source_url in GET /v1/mcp/server list responses
The list endpoint builds responses from the in-memory registry, but
source_url was dropped during the DB-to-registry roundtrip even though
GET /v1/mcp/server/{id} returned it correctly from the database.
Co-authored-by: Cursor
* fix(tests/mcp): set source_url on MagicMock table records
MagicMock auto-creates source_url as a mock object, which fails MCPServer
Pydantic validation after source_url was wired through build_mcp_server_from_table.
Co-authored-by: Cursor
---------
Co-authored-by: Cursor
---
.../mcp_server/mcp_server_manager.py | 2 ++
.../types/mcp_server/mcp_server_manager.py | 1 +
tests/mcp_tests/test_mcp_server.py | 3 +++
.../mcp_server/test_mcp_server_manager.py | 25 +++++++++++++++++++
.../mcp_server/test_mcp_sigv4_auth.py | 2 ++
5 files changed, 33 insertions(+)
diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
index d0e9ad7b2a4..f35aa30a7c9 100644
--- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
+++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
@@ -892,6 +892,7 @@ class MCPServerManager:
is_byok=bool(getattr(mcp_server, "is_byok", False)),
byok_description=getattr(mcp_server, "byok_description", None) or [],
byok_api_key_help_url=getattr(mcp_server, "byok_api_key_help_url", None),
+ source_url=getattr(mcp_server, "source_url", None),
# AWS SigV4 fields
aws_access_key_id=aws_creds.get("aws_access_key_id"),
aws_secret_access_key=aws_creds.get("aws_secret_access_key"),
@@ -3750,6 +3751,7 @@ class MCPServerManager:
is_byok=server.is_byok,
byok_description=server.byok_description,
byok_api_key_help_url=server.byok_api_key_help_url,
+ source_url=server.source_url,
instructions=server.instructions,
)
diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py
index 776c7fa67a6..13e325838dc 100644
--- a/litellm/types/mcp_server/mcp_server_manager.py
+++ b/litellm/types/mcp_server/mcp_server_manager.py
@@ -77,6 +77,7 @@ class MCPServer(BaseModel):
is_byok: bool = False
byok_description: List[str] = []
byok_api_key_help_url: Optional[str] = None
+ source_url: Optional[str] = None
created_at: Optional[datetime] = None
updated_at: Optional[datetime] = None
# OAuth2 flow type. Defaults to None (interactive / authorization_code).
diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py
index 27df05225e5..d76ebb0072f 100644
--- a/tests/mcp_tests/test_mcp_server.py
+++ b/tests/mcp_tests/test_mcp_server.py
@@ -1505,6 +1505,7 @@ async def test_add_update_server_with_alias():
mock_mcp_server.created_at = None
mock_mcp_server.updated_at = None
mock_mcp_server.instructions = None
+ mock_mcp_server.source_url = None
mock_mcp_server.approval_status = "active"
# Add server to manager
@@ -1563,6 +1564,7 @@ async def test_add_update_server_without_alias():
mock_mcp_server.created_at = None
mock_mcp_server.updated_at = None
mock_mcp_server.instructions = None
+ mock_mcp_server.source_url = None
mock_mcp_server.approval_status = "active"
# Add server to manager
@@ -1622,6 +1624,7 @@ async def test_add_update_server_fallback_to_server_id():
mock_mcp_server.created_at = None
mock_mcp_server.updated_at = None
mock_mcp_server.instructions = None
+ mock_mcp_server.source_url = None
mock_mcp_server.approval_status = "active"
# Add server to manager
await test_manager.add_server(mock_mcp_server)
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py
index d7078412a44..2db9845c765 100644
--- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py
@@ -2888,6 +2888,31 @@ class TestMCPServerTimestamps:
assert rebuilt_table.created_at == created
assert rebuilt_table.updated_at == updated
+ @pytest.mark.asyncio
+ async def test_round_trip_source_url_preserved(self):
+ """source_url survives the full round-trip: LiteLLM_MCPServerTable -> MCPServer -> LiteLLM_MCPServerTable.
+
+ Regression test: the list endpoint (GET /v1/mcp/server) builds its
+ response from the registry via this round-trip, so a dropped field
+ here surfaces as a null source_url in the list response even though
+ the value is stored in the DB.
+ """
+ manager = MCPServerManager()
+
+ table_record = LiteLLM_MCPServerTable(
+ server_id="src-url-server",
+ server_name="src_url_server",
+ url="https://example.com/mcp",
+ transport=MCPTransport.http,
+ source_url="https://github.com/org/mcp-server",
+ )
+
+ mcp_server = await manager.build_mcp_server_from_table(table_record)
+ assert mcp_server.source_url == "https://github.com/org/mcp-server"
+
+ rebuilt_table = manager._build_mcp_server_table(mcp_server)
+ assert rebuilt_table.source_url == "https://github.com/org/mcp-server"
+
class TestInternalDelegatePkceWarningLog:
@pytest.mark.asyncio
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py
index 32b988ddb22..0c2a8bb8087 100644
--- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py
@@ -857,6 +857,7 @@ class TestSigV4BuildFromTable:
table_record.byok_api_key_help_url = None
table_record.oauth2_flow = None
table_record.instructions = None
+ table_record.source_url = None
manager = MCPServerManager()
@@ -915,6 +916,7 @@ class TestSigV4BuildFromTable:
table_record.byok_api_key_help_url = None
table_record.oauth2_flow = None
table_record.instructions = None
+ table_record.source_url = None
manager = MCPServerManager()
From 2d4c13c00fa95c4eeab95d1636bc7954161c326f Mon Sep 17 00:00:00 2001
From: Shivam Rawat
Date: Sat, 30 May 2026 12:03:22 -0700
Subject: [PATCH 28/44] fix(mcp): preserve omitted fields on PUT /v1/mcp/server
partial updates (#29253)
* fix(mcp): preserve omitted fields on PUT /v1/mcp/server partial updates
Use model_dump(exclude_unset=True) for updates so schema defaults (transport=sse,
allow_all_keys=false, etc.) are not written when callers omit them. Serialize
JSON fields from the filtered dict and only force is_byok on create.
Co-authored-by: Cursor
* fix(mcp): allow explicit alias=None on MCP server partial updates
Snapshot caller-provided fields before normalization so omitted alias is not written while an intentional alias=None still clears the stored value.
Co-authored-by: Cursor
---------
Co-authored-by: Cursor
---
litellm/proxy/_experimental/mcp_server/db.py | 82 +++++---
.../mcp_management_endpoints.py | 3 +
.../mcp_server/test_mcp_partial_update.py | 177 ++++++++++++++++++
3 files changed, 237 insertions(+), 25 deletions(-)
create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py
diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py
index a6f0d145e9b..e30667776c1 100644
--- a/litellm/proxy/_experimental/mcp_server/db.py
+++ b/litellm/proxy/_experimental/mcp_server/db.py
@@ -30,6 +30,8 @@ from litellm.types.mcp import MCPCredentials
def _prepare_mcp_server_data(
data: Union[NewMCPServerRequest, UpdateMCPServerRequest],
+ exclude_unset: bool = False,
+ fields_set: Optional[Set[str]] = None,
) -> Dict[str, Any]:
"""
Helper function to prepare MCP server data for database operations.
@@ -37,17 +39,39 @@ def _prepare_mcp_server_data(
Args:
data: NewMCPServerRequest or UpdateMCPServerRequest object
+ exclude_unset: When True, only fields the caller explicitly provided are
+ included. Used for partial updates (PUT /v1/mcp/server) so omitted
+ fields keep their existing DB value instead of being silently reset
+ to a Pydantic schema default. ``exclude_none`` is not enough here:
+ non-Optional fields (e.g. ``transport=MCPTransport.sse``,
+ ``mcp_access_groups=[]``, ``allow_all_keys=False``) are backfilled
+ with their default when omitted, and a non-None default survives the
+ ``exclude_none`` filter and overwrites the row.
Returns:
Dict with properly serialized JSON fields
"""
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
- # Convert model to dict
- data_dict = data.model_dump(exclude_none=True)
- # Ensure alias is always present in the dict (even if None)
- if "alias" not in data_dict:
- data_dict["alias"] = getattr(data, "alias", None)
+ # Convert model to dict.
+ # - Partial update (exclude_unset): only caller-provided keys are emitted, so
+ # omitted fields are never written and keep their existing DB value.
+ # - Create (exclude_none): drop None-valued fields and let DB defaults apply.
+ if exclude_unset:
+ if fields_set is None:
+ fields_set = data.fields_set()
+ data_dict = data.model_dump(exclude_unset=True)
+ # ``validate_and_normalize_mcp_server_payload`` always assigns ``alias``
+ # on the payload, which marks it as set even when the caller omitted it.
+ # Drop it only when the original request omitted alias; an explicit
+ # ``alias=None`` is a valid request to clear the stored alias.
+ if data_dict.get("alias") is None and "alias" not in fields_set:
+ data_dict.pop("alias", None)
+ else:
+ data_dict = data.model_dump(exclude_none=True)
+ # Ensure alias is always present in the dict (even if None)
+ if "alias" not in data_dict:
+ data_dict["alias"] = getattr(data, "alias", None)
# Handle credentials serialization
credentials = data_dict.get("credentials")
@@ -57,33 +81,33 @@ def _prepare_mcp_server_data(
)
data_dict["credentials"] = safe_dumps(data_dict["credentials"])
- # Handle static_headers serialization
- if data.static_headers is not None:
- data_dict["static_headers"] = safe_dumps(data.static_headers)
+ # Serialize JSON fields from ``data_dict`` (not ``data``) so the
+ # exclude_unset filter is respected. Reading back from ``data`` would
+ # reintroduce defaults (e.g. ``env={}``) for fields the caller never set.
+ if data_dict.get("static_headers") is not None:
+ data_dict["static_headers"] = safe_dumps(data_dict["static_headers"])
- # Handle mcp_info serialization
- if data.mcp_info is not None:
- data_dict["mcp_info"] = safe_dumps(data.mcp_info)
+ if data_dict.get("mcp_info") is not None:
+ data_dict["mcp_info"] = safe_dumps(data_dict["mcp_info"])
- # Handle env serialization
- if data.env is not None:
- data_dict["env"] = safe_dumps(data.env)
+ if data_dict.get("env") is not None:
+ data_dict["env"] = safe_dumps(data_dict["env"])
- # Handle tool name override serialization
- if data.tool_name_to_display_name is not None:
+ if data_dict.get("tool_name_to_display_name") is not None:
data_dict["tool_name_to_display_name"] = safe_dumps(
- data.tool_name_to_display_name
+ data_dict["tool_name_to_display_name"]
)
- if data.tool_name_to_description is not None:
+ if data_dict.get("tool_name_to_description") is not None:
data_dict["tool_name_to_description"] = safe_dumps(
- data.tool_name_to_description
+ data_dict["tool_name_to_description"]
)
# mcp_access_groups is already List[str], no serialization needed
- # Force include is_byok even when False (exclude_none=True would not drop it,
- # but be explicit to ensure a False value is always written to the DB).
- data_dict["is_byok"] = getattr(data, "is_byok", False)
+ # On create, force is_byok so a False value is always written to the DB. On
+ # partial update, only write it when the caller explicitly provided it.
+ if not exclude_unset:
+ data_dict["is_byok"] = getattr(data, "is_byok", False)
return data_dict
@@ -398,7 +422,10 @@ async def create_mcp_server(
async def update_mcp_server(
- prisma_client: PrismaClient, data: UpdateMCPServerRequest, touched_by: str
+ prisma_client: PrismaClient,
+ data: UpdateMCPServerRequest,
+ touched_by: str,
+ fields_set: Optional[Set[str]] = None,
) -> LiteLLM_MCPServerTable:
"""
Update a new mcp server record in the db
@@ -407,8 +434,13 @@ async def update_mcp_server(
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
- # Use helper to prepare data with proper JSON serialization
- data_dict = _prepare_mcp_server_data(data)
+ # Use helper to prepare data with proper JSON serialization.
+ # exclude_unset=True makes this a true partial update: fields the caller did
+ # not provide are not written, so they keep their existing DB value instead
+ # of being reset to a schema default (transport=sse, allow_all_keys=False...).
+ data_dict = _prepare_mcp_server_data(
+ data, exclude_unset=True, fields_set=fields_set
+ )
# Pre-fetch existing record once if we need it for auth_type or credential logic
existing = None
diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py
index 431ff49c7ce..b35e2b6e3fd 100644
--- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py
+++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py
@@ -2135,6 +2135,8 @@ if MCP_AVAILABLE:
"Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys"
)
+ payload_fields_set = set(payload.fields_set())
+
# Validate and normalize payload fields
validate_and_normalize_mcp_server_payload(payload)
@@ -2154,6 +2156,7 @@ if MCP_AVAILABLE:
prisma_client,
payload,
touched_by=user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME,
+ fields_set=payload_fields_set,
)
if mcp_server_record_updated is None:
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py
new file mode 100644
index 00000000000..b5e0f20f660
--- /dev/null
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py
@@ -0,0 +1,177 @@
+"""
+Tests for partial-update semantics of PUT /v1/mcp/server.
+
+A partial update must only write the fields the caller explicitly provided.
+Omitting a field must NOT reset it to its Pydantic schema default (e.g.
+``transport=sse``, ``mcp_access_groups=[]``, ``allow_all_keys=False``), which
+would silently overwrite the existing DB row.
+"""
+
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from litellm.proxy._experimental.mcp_server.db import (
+ create_mcp_server,
+ update_mcp_server,
+)
+from litellm.proxy._types import NewMCPServerRequest, UpdateMCPServerRequest
+
+
+def _mock_prisma():
+ mock_prisma = MagicMock()
+ mock_prisma.db.litellm_mcpservertable = AsyncMock()
+ mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock())
+ mock_prisma.db.litellm_mcpservertable.create = AsyncMock(return_value=MagicMock())
+ return mock_prisma
+
+
+async def _run_update(data: UpdateMCPServerRequest, fields_set=None) -> dict:
+ mock_prisma = _mock_prisma()
+ await update_mcp_server(mock_prisma, data, "test-user", fields_set=fields_set)
+ return mock_prisma.db.litellm_mcpservertable.update.call_args[1]["data"]
+
+
+@pytest.mark.asyncio
+async def test_partial_update_omits_unset_defaultful_fields():
+ """
+ A PUT touching only allowed_tools must not write transport,
+ mcp_access_groups, allow_all_keys, available_on_public_internet,
+ delegate_auth_to_upstream, is_byok, args, env or byok_description.
+ """
+ data = UpdateMCPServerRequest(
+ server_id="my-test-server",
+ allowed_tools=["foo"],
+ )
+
+ data_dict = await _run_update(data)
+
+ # The intended change is present.
+ assert data_dict["allowed_tools"] == ["foo"]
+
+ # Fields the caller did not provide must not be in the write payload, so the
+ # existing DB value is preserved.
+ for trapped_field in (
+ "transport",
+ "mcp_access_groups",
+ "allow_all_keys",
+ "available_on_public_internet",
+ "delegate_auth_to_upstream",
+ "is_byok",
+ "args",
+ "env",
+ "byok_description",
+ ):
+ assert trapped_field not in data_dict, (
+ f"{trapped_field} should not be written on a partial update that "
+ f"omitted it (would reset the row to a schema default)"
+ )
+
+
+@pytest.mark.asyncio
+async def test_partial_update_preserves_http_transport():
+ """The reported prod incident: a PUT without transport must not flip http->sse."""
+ data = UpdateMCPServerRequest(
+ server_id="atlassian_url",
+ allowed_tools=[],
+ )
+
+ data_dict = await _run_update(data)
+
+ assert "transport" not in data_dict
+ assert data_dict["allowed_tools"] == []
+
+
+@pytest.mark.asyncio
+async def test_partial_update_writes_explicitly_provided_fields():
+ """Explicitly provided fields are written, including falsy/default-equal values."""
+ data = UpdateMCPServerRequest(
+ server_id="my-test-server",
+ url="https://example.com/mcp",
+ transport="http",
+ allow_all_keys=False,
+ mcp_access_groups=["mcp-dev-sandbox"],
+ available_on_public_internet=True,
+ )
+
+ data_dict = await _run_update(data)
+
+ assert data_dict["transport"] == "http"
+ # Explicitly provided False must still be written.
+ assert data_dict["allow_all_keys"] is False
+ assert data_dict["mcp_access_groups"] == ["mcp-dev-sandbox"]
+ assert data_dict["available_on_public_internet"] is True
+
+
+@pytest.mark.asyncio
+async def test_partial_update_can_explicitly_reset_allow_all_keys():
+ """Caller can still reset a field to its default by sending it explicitly."""
+ enabled = await _run_update(
+ UpdateMCPServerRequest(server_id="s", allow_all_keys=True)
+ )
+ assert enabled["allow_all_keys"] is True
+
+ disabled = await _run_update(
+ UpdateMCPServerRequest(server_id="s", allow_all_keys=False)
+ )
+ assert disabled["allow_all_keys"] is False
+
+
+@pytest.mark.asyncio
+async def test_partial_update_does_not_clear_alias_when_unset():
+ """alias is force-normalized on the payload; an unset/None alias must not be written."""
+ data = UpdateMCPServerRequest(
+ server_id="my-test-server",
+ allowed_tools=["foo"],
+ )
+ fields_set = set(data.fields_set())
+ # Simulate validate_and_normalize_mcp_server_payload assigning alias=None.
+ data.alias = None
+
+ data_dict = await _run_update(data, fields_set=fields_set)
+
+ assert "alias" not in data_dict
+
+
+@pytest.mark.asyncio
+async def test_partial_update_can_explicitly_clear_alias():
+ """Caller can clear an existing alias by explicitly sending alias=None."""
+ data = UpdateMCPServerRequest(
+ server_id="my-test-server",
+ alias=None,
+ )
+ fields_set = set(data.fields_set())
+ # Simulate validate_and_normalize_mcp_server_payload preserving alias=None.
+ data.alias = None
+
+ data_dict = await _run_update(data, fields_set=fields_set)
+
+ assert "alias" in data_dict
+ assert data_dict["alias"] is None
+
+
+@pytest.mark.asyncio
+async def test_create_still_writes_defaults():
+ """
+ Regression guard: create (POST) must keep writing defaults so DB columns
+ without a default get populated. exclude_unset is update-only.
+ """
+ mock_prisma = _mock_prisma()
+ data = NewMCPServerRequest(
+ server_id="new-server",
+ url="https://example.com/mcp",
+ transport="http",
+ )
+
+ await create_mcp_server(mock_prisma, data, "test-user")
+
+ data_dict = mock_prisma.db.litellm_mcpservertable.create.call_args[1]["data"]
+
+ assert data_dict["transport"] == "http"
+ # is_byok is force-written on create.
+ assert data_dict["is_byok"] is False
+ # alias key is always present on create (even if None).
+ assert "alias" in data_dict
+ # audit fields set by create_mcp_server.
+ assert data_dict["created_by"] == "test-user"
+ assert data_dict["updated_by"] == "test-user"
From bfbb5d23756dcabc0b992982c7d8d185f76fa817 Mon Sep 17 00:00:00 2001
From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 30 May 2026 13:57:18 -0700
Subject: [PATCH 29/44] fix(ci): make litellm_internal_staging green (logging
test + Bedrock Opus 4.7 self-heal) (#29344)
* test(logging): align DB metrics event_metadata assertions with safe redaction
PR #28909 hardened log_db_metrics to emit a minimal, non-sensitive
event_metadata (only table_name when present, otherwise None) instead of
dumping function_name, function_kwargs, and function_args onto the span. The
test in test_log_db_redis_services was not updated and still asserted
"function_name" in event_metadata, which raised TypeError (argument of type
'NoneType' is not iterable) and turned the logging_testing CI job red on
litellm_internal_staging.
Update test_log_db_metrics_success to assert event_metadata is None when no
table_name is passed, and add test_log_db_metrics_event_metadata_is_safe as a
regression guard verifying that only the table name surfaces and that sensitive
kwargs (tokens, prisma client) are never dumped.
* test(bedrock): self-heal opus-4-7 grid cells when unentitled on CI
The bedrock-claude-opus-4-7 converse cells are unentitled on the Bedrock CI
account, so they were marked xfail. xfail keeps reporting them as expected
failures even after access is granted, so the wire translation never gets
verified again. Now the cell makes the call and skips only when Bedrock
replies "is not available for this account"; the moment the model is
entitled the same cells run their full assertions with no edit.
A focused unit test pins the tolerance predicate so any other failure still
surfaces loudly and the available path still runs the assertions.
---
.../reasoning_effort_grid/grid_spec.py | 3 +-
.../test_reasoning_effort_grid.py | 37 ++++++++++++++++++-
.../test_log_db_redis_services.py | 34 ++++++++++++++++-
3 files changed, 70 insertions(+), 4 deletions(-)
diff --git a/tests/llm_translation/reasoning_effort_grid/grid_spec.py b/tests/llm_translation/reasoning_effort_grid/grid_spec.py
index 0d709584bde..3eaa9bc2950 100644
--- a/tests/llm_translation/reasoning_effort_grid/grid_spec.py
+++ b/tests/llm_translation/reasoning_effort_grid/grid_spec.py
@@ -1,7 +1,6 @@
from dataclasses import dataclass, field
from typing import Dict, FrozenSet, List, Optional, Tuple
-
OMIT = object()
@@ -22,6 +21,7 @@ class ModelEntry:
extra_params: Tuple[Tuple[str, str], ...] = field(default_factory=tuple)
required_env: FrozenSet[str] = field(default_factory=frozenset)
caps: FrozenSet[str] = field(default_factory=frozenset)
+ unavailable_error: Optional[str] = None
bedrock_effort_ceiling: Optional[str] = None
def params(self) -> Dict[str, str]:
@@ -234,6 +234,7 @@ BEDROCK_CONVERSE_MODELS: Tuple[ModelEntry, ...] = (
extra_params=(("aws_region_name", "us-east-1"),),
required_env=_BEDROCK_REQ,
caps=_CAPS_OPUS_4_7,
+ unavailable_error="is not available for this account",
),
ModelEntry(
alias="bedrock-claude-opus-4-6",
diff --git a/tests/llm_translation/reasoning_effort_grid/test_reasoning_effort_grid.py b/tests/llm_translation/reasoning_effort_grid/test_reasoning_effort_grid.py
index 28e2e402d67..b4382e16fed 100644
--- a/tests/llm_translation/reasoning_effort_grid/test_reasoning_effort_grid.py
+++ b/tests/llm_translation/reasoning_effort_grid/test_reasoning_effort_grid.py
@@ -15,7 +15,6 @@ from .grid_spec import (
all_cells,
)
-
_PROMPT_MESSAGES: List[Dict[str, str]] = [
{"role": "user", "content": "Step by step, calculate 47 * 53. Show your work."}
]
@@ -133,6 +132,12 @@ def _classify_status(exc: Exception) -> int:
return 500
+def _model_unavailable(model: ModelEntry, exc: Optional[Exception]) -> bool:
+ if not model.unavailable_error or exc is None:
+ return False
+ return model.unavailable_error in str(exc)
+
+
async def _call_chat(model: ModelEntry, effort: str) -> Tuple[int, Optional[Exception]]:
kwargs = _build_completion_kwargs(model, effort)
try:
@@ -173,6 +178,9 @@ async def test_reasoning_effort_grid(
else:
status, exc = await _call_chat(model, effort)
+ if _model_unavailable(model, exc):
+ pytest.skip(f"{model.alias}: {model.unavailable_error}")
+
record = wire_capture.latest()
body = record["body"] if record else None
if route_name == "bedrock_converse" and isinstance(body, str):
@@ -205,3 +213,30 @@ def test_grid_route_coverage() -> None:
"bedrock_invoke_chat",
"bedrock_invoke_messages",
}
+
+
+def test_model_unavailable_tolerates_only_the_declared_error() -> None:
+ gated = ModelEntry(
+ alias="bedrock-claude-opus-4-7",
+ model="bedrock/converse/us.anthropic.claude-opus-4-7",
+ mode="adaptive",
+ unavailable_error="is not available for this account",
+ )
+ entitlement_error = Exception(
+ "litellm.APIConnectionError: BedrockException - "
+ '{"message":"anthropic.claude-opus-4-7 is not available for this account."}'
+ )
+
+ assert _model_unavailable(gated, entitlement_error) is True
+ assert (
+ _model_unavailable(gated, Exception("ThrottlingException: rate exceeded"))
+ is False
+ )
+ assert _model_unavailable(gated, None) is False
+
+ ungated = ModelEntry(
+ alias="bedrock-claude-opus-4-6",
+ model="bedrock/converse/us.anthropic.claude-opus-4-6-v1",
+ mode="adaptive",
+ )
+ assert _model_unavailable(ungated, entitlement_error) is False
diff --git a/tests/logging_callback_tests/test_log_db_redis_services.py b/tests/logging_callback_tests/test_log_db_redis_services.py
index fa0c3b595a0..a8c3929be16 100644
--- a/tests/logging_callback_tests/test_log_db_redis_services.py
+++ b/tests/logging_callback_tests/test_log_db_redis_services.py
@@ -2,7 +2,6 @@ import io
import os
import sys
-
sys.path.insert(0, os.path.abspath("../.."))
import asyncio
@@ -59,7 +58,38 @@ async def test_log_db_metrics_success():
assert isinstance(call_args["duration"], float)
assert isinstance(call_args["start_time"], datetime)
assert isinstance(call_args["end_time"], datetime)
- assert "function_name" in call_args["event_metadata"]
+ assert call_args["event_metadata"] is None
+
+
+@pytest.mark.asyncio
+async def test_log_db_metrics_event_metadata_is_safe():
+ """event_metadata must surface only the table name, never the raw
+ kwargs/args which carry live clients (Prisma, OTel spans) and secrets.
+
+ Regression guard for #28909: a previous version dumped function_kwargs and
+ function_args onto the span.
+ """
+ with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging:
+ mock_proxy_logging.service_logging_obj.async_service_success_hook = AsyncMock()
+
+ @log_db_metrics
+ async def db_call(**kwargs):
+ return "success"
+
+ await db_call(
+ parent_otel_span="test_span",
+ table_name="LiteLLM_SpendLogs",
+ token="sk-secret-should-not-leak",
+ prisma_client=object(),
+ )
+ await asyncio.sleep(0)
+
+ call_args = (
+ mock_proxy_logging.service_logging_obj.async_service_success_hook.call_args[
+ 1
+ ]
+ )
+ assert call_args["event_metadata"] == {"table_name": "LiteLLM_SpendLogs"}
@pytest.mark.asyncio
From 94a043efb226c5ccdbfc028fbb930ce45fb965eb Mon Sep 17 00:00:00 2001
From: yuneng-jiang
Date: Sat, 30 May 2026 14:04:22 -0700
Subject: [PATCH 30/44] refactor(proxy/auth): normalize Bearer prefix in
safe-hash helper (#29343)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
* refactor(proxy/auth): normalize Bearer prefix in safe-hash helper
UserAPIKeyAuth._safe_hash_litellm_api_key now strips a leading
"Bearer "/"bearer " prefix before its existing sk-/JWT classification, so
the helper produces the same hashed output regardless of whether the
caller stripped the Authorization header prefix or passed the header
value through unchanged.
* refactor(proxy/auth): make Bearer-prefix strip case-insensitive
Per RFC 7235 the HTTP authorization scheme token is case-insensitive.
Replace the two-prefix loop with a single case-insensitive check so the
helper normalizes "Bearer ", "bearer ", "BEARER ", and any mixed-case
variant before classifying the remainder as sk- or JWT. The contract
test gains coverage of "BEARER " and "BeArEr ".
* test(mcp): align auth-handler test expectations with safe-hash helper
The two MCP auth tests asserted that UserAPIKeyAuth(api_key="Bearer ...")
retained the raw header bytes on the api_key field. _safe_hash_litellm_api_key
now normalizes that input — stripping the Bearer prefix and hashing the
resulting sk- key — so the expectations move to the normalized form:
the bare token in the parametrize case, and hash_token("sk-...") in the
backward-compat assertion. This matches what the real auth flow produces
(the builder strips Bearer and the DB stores the hashed token), so the
mocks now line up with production rather than with the un-normalized
validator output.
---
litellm/proxy/_types.py | 13 ++++++++-----
.../auth/test_user_api_key_auth_mcp.py | 6 ++++--
tests/test_litellm/proxy/test_proxy_types.py | 18 ++++++++++++++++++
3 files changed, 30 insertions(+), 7 deletions(-)
diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
index 522e85632dc..95294f48386 100644
--- a/litellm/proxy/_types.py
+++ b/litellm/proxy/_types.py
@@ -2747,13 +2747,16 @@ class UserAPIKeyAuth(
1. Regular API keys from LiteLLM DB
2. JWT tokens used for connecting to LiteLLM API
"""
- if api_key.startswith("sk-"):
- return hash_token(api_key)
+ normalized = api_key
+ if normalized[:7].lower() == "bearer ":
+ normalized = normalized[7:]
+ if normalized.startswith("sk-"):
+ return hash_token(normalized)
from litellm.proxy.auth.handle_jwt import JWTHandler
- if JWTHandler.is_jwt(token=api_key):
- return f"hashed-jwt-{hash_token(token=api_key)}"
- return api_key
+ if JWTHandler.is_jwt(token=normalized):
+ return f"hashed-jwt-{hash_token(token=normalized)}"
+ return normalized
@classmethod
def get_litellm_internal_health_check_user_api_key_auth(cls) -> "UserAPIKeyAuth":
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py
index a7d9ce64f8d..b3508e13127 100644
--- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py
@@ -213,7 +213,7 @@ class TestMCPRequestHandler:
# Test case 2: Authorization header present (fallback)
(
[(b"authorization", b"Bearer test-auth-token")],
- "Bearer test-auth-token",
+ "test-auth-token",
None,
{},
),
@@ -674,7 +674,9 @@ class TestMCPOAuth2AuthFlow:
) = await MCPRequestHandler.process_mcp_request(scope)
# Should succeed with the LiteLLM key from Authorization header
- assert auth_result.api_key == "Bearer sk-litellm-valid-key"
+ from litellm.proxy.utils import hash_token
+
+ assert auth_result.api_key == hash_token("sk-litellm-valid-key")
mock_auth.assert_called_once()
async def test_non_auth_http_exception_still_raises(self):
diff --git a/tests/test_litellm/proxy/test_proxy_types.py b/tests/test_litellm/proxy/test_proxy_types.py
index 0fa86798999..dbb952968ed 100644
--- a/tests/test_litellm/proxy/test_proxy_types.py
+++ b/tests/test_litellm/proxy/test_proxy_types.py
@@ -69,3 +69,21 @@ def test_internal_jobs_user_has_proxy_admin_role():
assert system_user.user_id == "system"
assert system_user.team_id == "system"
assert system_user.team_alias == "system"
+
+
+def test_user_api_key_auth_hashes_authorization_header_form_of_key():
+ from litellm.proxy._types import UserAPIKeyAuth
+
+ raw_key = "sk-AbCdEfGhIjKlMnOpQrStUvWxYz0123456789"
+ baseline = UserAPIKeyAuth(api_key=raw_key)
+
+ for header_form in (
+ f"Bearer {raw_key}",
+ f"bearer {raw_key}",
+ f"BEARER {raw_key}",
+ f"BeArEr {raw_key}",
+ ):
+ from_header = UserAPIKeyAuth(api_key=header_form)
+ assert from_header.api_key == baseline.api_key
+ assert from_header.token == baseline.token
+ assert not from_header.api_key.lower().startswith("bearer")
From 152b1177e57387f1ef4e2a6c1876c723bdac8139 Mon Sep 17 00:00:00 2001
From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 30 May 2026 14:12:57 -0700
Subject: [PATCH 31/44] test(reasoning-effort-grid): cover Claude Opus 4.8
across provider routes (#29327)
* test(logging): align DB metrics event_metadata assertions with safe redaction
PR #28909 hardened log_db_metrics to emit a minimal, non-sensitive
event_metadata (only table_name when present, otherwise None) instead of
dumping function_name, function_kwargs, and function_args onto the span. The
test in test_log_db_redis_services was not updated and still asserted
"function_name" in event_metadata, which raised TypeError (argument of type
'NoneType' is not iterable) and turned the logging_testing CI job red on
litellm_internal_staging.
Update test_log_db_metrics_success to assert event_metadata is None when no
table_name is passed, and add test_log_db_metrics_event_metadata_is_safe as a
regression guard verifying that only the table name surfaces and that sensitive
kwargs (tokens, prisma client) are never dumped.
* test(bedrock): self-heal opus-4-7 grid cells when unentitled on CI
The bedrock-claude-opus-4-7 converse cells are unentitled on the Bedrock CI
account, so they were marked xfail. xfail keeps reporting them as expected
failures even after access is granted, so the wire translation never gets
verified again. Now the cell makes the call and skips only when Bedrock
replies "is not available for this account"; the moment the model is
entitled the same cells run their full assertions with no edit.
A focused unit test pins the tolerance predicate so any other failure still
surfaces loudly and the available path still runs the assertions.
* test(reasoning-effort-grid): add claude-opus-4-8 across provider routes
Adds claude-opus-4-8 to the anthropic, azure, vertex and bedrock-converse
routes (275 cells total) so the reasoning-effort wire translation is
covered for the new model. The bedrock opus-4-8 and opus-4-7 cells reuse
the self-heal path: they run the call and skip only on Bedrock's "is not
available for this account" reply, then assert in full once the model is
entitled. The azure and vertex opus-4-8 cells stay xfail until a Foundry
deployment exists and Vertex availability is confirmed. The shared
xhigh+max capability set is renamed to _CAPS_XHIGH_MAX now that more than
one model uses it.
---
.../reasoning_effort_grid/grid_spec.py | 56 +++++++++++++++++--
.../test_reasoning_effort_grid.py | 7 ++-
2 files changed, 56 insertions(+), 7 deletions(-)
diff --git a/tests/llm_translation/reasoning_effort_grid/grid_spec.py b/tests/llm_translation/reasoning_effort_grid/grid_spec.py
index 3eaa9bc2950..51a027ec101 100644
--- a/tests/llm_translation/reasoning_effort_grid/grid_spec.py
+++ b/tests/llm_translation/reasoning_effort_grid/grid_spec.py
@@ -22,6 +22,7 @@ class ModelEntry:
required_env: FrozenSet[str] = field(default_factory=frozenset)
caps: FrozenSet[str] = field(default_factory=frozenset)
unavailable_error: Optional[str] = None
+ fail_reason: Optional[str] = None
bedrock_effort_ceiling: Optional[str] = None
def params(self) -> Dict[str, str]:
@@ -126,7 +127,7 @@ _VERTEX_REQ = frozenset({"VERTEX_PROJECT"})
_BEDROCK_REQ = frozenset({"AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY"})
-_CAPS_OPUS_4_7: FrozenSet[str] = frozenset(
+_CAPS_XHIGH_MAX: FrozenSet[str] = frozenset(
{"supports_xhigh_reasoning_effort", "supports_max_reasoning_effort"}
)
_CAPS_4_6: FrozenSet[str] = frozenset({"supports_max_reasoning_effort"})
@@ -134,12 +135,19 @@ _CAPS_NONE: FrozenSet[str] = frozenset()
ANTHROPIC_DIRECT_MODELS: Tuple[ModelEntry, ...] = (
+ ModelEntry(
+ alias="claude-opus-4-8",
+ model="anthropic/claude-opus-4-8",
+ mode="adaptive",
+ required_env=_ANTHROPIC_REQ,
+ caps=_CAPS_XHIGH_MAX,
+ ),
ModelEntry(
alias="claude-opus-4-7",
model="anthropic/claude-opus-4-7",
mode="adaptive",
required_env=_ANTHROPIC_REQ,
- caps=_CAPS_OPUS_4_7,
+ caps=_CAPS_XHIGH_MAX,
),
ModelEntry(
alias="claude-sonnet-4-6",
@@ -159,12 +167,25 @@ ANTHROPIC_DIRECT_MODELS: Tuple[ModelEntry, ...] = (
AZURE_AI_MODELS: Tuple[ModelEntry, ...] = (
+ ModelEntry(
+ alias="azure-claude-opus-4-8",
+ model="azure_ai/claude-opus-4-8",
+ mode="adaptive",
+ required_env=_AZURE_FOUNDRY_REQ,
+ caps=_CAPS_XHIGH_MAX,
+ fail_reason=(
+ "claude-opus-4-8 has no deployment on the CI Microsoft Foundry "
+ "resource yet; Foundry returns DeploymentNotFound until someone "
+ "creates the opus-4-8 deployment, so this cell stays loud in CI. "
+ "Remove this fail_reason once the deployment exists."
+ ),
+ ),
ModelEntry(
alias="azure-claude-opus-4-7",
model="azure_ai/claude-opus-4-7",
mode="adaptive",
required_env=_AZURE_FOUNDRY_REQ,
- caps=_CAPS_OPUS_4_7,
+ caps=_CAPS_XHIGH_MAX,
),
ModelEntry(
alias="azure-claude-opus-4-6",
@@ -191,13 +212,27 @@ AZURE_AI_MODELS: Tuple[ModelEntry, ...] = (
VERTEX_AI_MODELS: Tuple[ModelEntry, ...] = (
+ ModelEntry(
+ alias="vertex-claude-opus-4-8",
+ model="vertex_ai/claude-opus-4-8",
+ mode="adaptive",
+ extra_params=(("vertex_location", "global"),),
+ required_env=_VERTEX_REQ,
+ caps=_CAPS_XHIGH_MAX,
+ fail_reason=(
+ "claude-opus-4-8 availability on the CI Vertex project is not yet "
+ "confirmed for this brand-new release, so this cell stays loud in "
+ "CI until verified. Remove this fail_reason once the model is "
+ "confirmed available on the global Vertex endpoint."
+ ),
+ ),
ModelEntry(
alias="vertex-claude-opus-4-7",
model="vertex_ai/claude-opus-4-7",
mode="adaptive",
extra_params=(("vertex_location", "global"),),
required_env=_VERTEX_REQ,
- caps=_CAPS_OPUS_4_7,
+ caps=_CAPS_XHIGH_MAX,
),
ModelEntry(
alias="vertex-claude-opus-4-6",
@@ -227,13 +262,24 @@ VERTEX_AI_MODELS: Tuple[ModelEntry, ...] = (
BEDROCK_CONVERSE_MODELS: Tuple[ModelEntry, ...] = (
+ ModelEntry(
+ alias="bedrock-claude-opus-4-8",
+ model="bedrock/converse/us.anthropic.claude-opus-4-8",
+ mode="adaptive",
+ extra_params=(("aws_region_name", "us-east-1"),),
+ required_env=_BEDROCK_REQ,
+ caps=_CAPS_XHIGH_MAX,
+ bedrock_effort_ceiling="xhigh",
+ unavailable_error="is not available for this account",
+ ),
ModelEntry(
alias="bedrock-claude-opus-4-7",
model="bedrock/converse/us.anthropic.claude-opus-4-7",
mode="adaptive",
extra_params=(("aws_region_name", "us-east-1"),),
required_env=_BEDROCK_REQ,
- caps=_CAPS_OPUS_4_7,
+ caps=_CAPS_XHIGH_MAX,
+ bedrock_effort_ceiling="xhigh",
unavailable_error="is not available for this account",
),
ModelEntry(
diff --git a/tests/llm_translation/reasoning_effort_grid/test_reasoning_effort_grid.py b/tests/llm_translation/reasoning_effort_grid/test_reasoning_effort_grid.py
index b4382e16fed..a5ef1b7ea3c 100644
--- a/tests/llm_translation/reasoning_effort_grid/test_reasoning_effort_grid.py
+++ b/tests/llm_translation/reasoning_effort_grid/test_reasoning_effort_grid.py
@@ -173,6 +173,9 @@ async def test_reasoning_effort_grid(
if skip_reason:
pytest.skip(skip_reason)
+ if model.fail_reason:
+ pytest.xfail(model.fail_reason)
+
if route_name == "bedrock_invoke_messages":
status, exc = await _call_messages(model, effort)
else:
@@ -197,8 +200,8 @@ async def test_reasoning_effort_grid(
def test_grid_cell_count() -> None:
- assert len(_PARAMS) == 21 * 11, (
- f"expected 231 cells (21 provider x model combos x 11 efforts), "
+ assert len(_PARAMS) == 25 * 11, (
+ f"expected 275 cells (25 provider x model combos x 11 efforts), "
f"got {len(_PARAMS)}"
)
From 4c3efe9c7cfe98ad4a2f4aa919624eff88c303e5 Mon Sep 17 00:00:00 2001
From: Shivam Rawat
Date: Sat, 30 May 2026 14:28:46 -0700
Subject: [PATCH 32/44] fix(guardrails): return HTTP 400 for litellm content
filter blocks (#28418)
* fix(guardrails): return HTTP 400 for litellm content filter blocks
Align litellm_content_filter hard rejects with the standard guardrail block status code so clients receive 400 instead of 403.
Co-authored-by: Cursor
* fix(guardrails): return HTTP 400 for custom code guardrail blocks
Pre-call custom code guardrail blocks now raise HTTPException(400) instead of using the passthrough ModifyResponseException path that returned a synthetic 200 response.
Co-authored-by: Cursor
* fix(guardrails): preserve custom code passthrough blocks
Keep standalone custom code guardrail blocks on the passthrough contract while covering policy pipeline block handling for passthrough-style guardrail interventions.
Co-authored-by: Cursor
---------
Co-authored-by: Cursor
---
.../custom_code/custom_code_guardrail.py | 4 +
.../litellm_content_filter/content_filter.py | 10 +-
.../guardrail_benchmarks/test_eval.py | 4 +-
.../content_filter/test_competitor_intent.py | 2 +-
.../content_filter/test_content_filter.py | 44 +++---
.../guardrails/test_custom_code_security.py | 46 ++++++-
.../policy_engine/test_pipeline_executor.py | 125 ++++++++++++++++++
7 files changed, 204 insertions(+), 31 deletions(-)
diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py
index 58502e309ef..22d4548aa99 100644
--- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py
+++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py
@@ -41,6 +41,7 @@ from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Type, cast
from fastapi import HTTPException
from litellm._logging import verbose_proxy_logger
+from litellm.exceptions import ModifyResponseException
from litellm.integrations.custom_guardrail import (
CustomGuardrail,
log_guardrail_information,
@@ -253,6 +254,9 @@ class CustomCodeGuardrail(CustomGuardrail):
except HTTPException:
# Re-raise HTTP exceptions (from block action)
raise
+ except ModifyResponseException:
+ # Pre-call block uses passthrough; must not wrap as execution error (500)
+ raise
except Exception as e:
verbose_proxy_logger.error(
f"Custom code guardrail '{self.guardrail_name}' execution error: {e}"
diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py
index d6d2e014948..d6065ef73f5 100644
--- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py
+++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py
@@ -1202,7 +1202,7 @@ class ContentFilterGuardrail(CustomGuardrail):
)
verbose_proxy_logger.warning(error_msg)
raise HTTPException(
- status_code=403,
+ status_code=400,
detail={
"error": error_msg,
"category": category_name,
@@ -1242,7 +1242,7 @@ class ContentFilterGuardrail(CustomGuardrail):
)
verbose_proxy_logger.warning(error_msg)
raise HTTPException(
- status_code=403,
+ status_code=400,
detail={
"error": error_msg,
"category": category_name,
@@ -1285,7 +1285,7 @@ class ContentFilterGuardrail(CustomGuardrail):
error_msg = f"Content blocked: {pattern_name} pattern detected"
verbose_proxy_logger.warning(error_msg)
raise HTTPException(
- status_code=403,
+ status_code=400,
detail={"error": error_msg, "pattern": pattern_name},
)
elif action == ContentFilterAction.MASK:
@@ -1325,7 +1325,7 @@ class ContentFilterGuardrail(CustomGuardrail):
error_msg += f" ({description})"
verbose_proxy_logger.warning(error_msg)
raise HTTPException(
- status_code=403,
+ status_code=400,
detail={
"error": error_msg,
"keyword": keyword,
@@ -1677,7 +1677,7 @@ class ContentFilterGuardrail(CustomGuardrail):
"ContentFilterGuardrail: competitor intent refuse - %s", intent_val
)
raise HTTPException(
- status_code=403,
+ status_code=400,
detail={
"error": msg,
"intent": intent_val,
diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/guardrail_benchmarks/test_eval.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/guardrail_benchmarks/test_eval.py
index 56398739b9b..aedc6acc810 100644
--- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/guardrail_benchmarks/test_eval.py
+++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/guardrail_benchmarks/test_eval.py
@@ -59,7 +59,7 @@ def _run(checker, text: str) -> dict:
checker.check(text)
return {"decision": "ALLOW", "score": 0.0, "matched_topic": None}
except HTTPException as e:
- if e.status_code == 403:
+ if e.status_code == 400:
detail: Dict[str, Any] = e.detail if isinstance(e.detail, dict) else {}
return {
"decision": "BLOCK",
@@ -542,7 +542,7 @@ class _LlmJudgeChecker:
if "BLOCK" in decision:
raise HTTPException(
- status_code=403,
+ status_code=400,
detail={
"error": "Content blocked by LLM judge",
"topic": "financial_advice",
diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_competitor_intent.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_competitor_intent.py
index 545b75fa06b..723d4b7db75 100644
--- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_competitor_intent.py
+++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_competitor_intent.py
@@ -226,7 +226,7 @@ class TestContentFilterWithCompetitorIntent:
await guardrail.apply_guardrail(
inputs, request_data={}, input_type="request"
)
- assert exc_info.value.status_code == 403
+ assert exc_info.value.status_code == 400
# Exact config from litellm/proxy/_new_secret_config.yaml (lines 27-53).
diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py
index fb952d4b18b..bb079ea6580 100644
--- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py
+++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py
@@ -198,7 +198,7 @@ class TestContentFilterGuardrail:
input_type="request",
)
- assert exc_info.value.status_code == 403
+ assert exc_info.value.status_code == 400
assert "us_ssn" in str(exc_info.value.detail)
@pytest.mark.asyncio
@@ -563,7 +563,7 @@ class TestContentFilterGuardrail:
):
pass
- assert exc_info.value.status_code == 403
+ assert exc_info.value.status_code == 400
assert "us_ssn" in str(exc_info.value.detail)
@pytest.mark.asyncio
@@ -1010,7 +1010,7 @@ class TestContentFilterGuardrail:
input_type="request",
)
- assert exc_info.value.status_code == 403
+ assert exc_info.value.status_code == 400
assert "danger_word" in str(exc_info.value.detail)
@pytest.mark.asyncio
@@ -1298,7 +1298,7 @@ class TestContentFilterGuardrail:
input_type="request",
)
- assert exc_info.value.status_code == 403
+ assert exc_info.value.status_code == 400
detail = exc_info.value.detail
if isinstance(detail, dict):
assert detail.get("category") == "harm_toxic_abuse"
@@ -1327,7 +1327,7 @@ class TestContentFilterGuardrail:
input_type="request",
)
- assert exc_info.value.status_code == 403
+ assert exc_info.value.status_code == 400
detail = exc_info.value.detail
if isinstance(detail, dict):
assert detail.get("category") == "harm_toxic_abuse"
@@ -1375,7 +1375,7 @@ class TestContentFilterGuardrail:
input_type="request",
)
- assert exc_info.value.status_code == 403, f"Failed to block: '{test_input}'"
+ assert exc_info.value.status_code == 400, f"Failed to block: '{test_input}'"
detail = exc_info.value.detail
if isinstance(detail, dict):
assert detail.get("category") == "harm_toxic_abuse"
@@ -1443,7 +1443,7 @@ class TestContentFilterGuardrail:
input_type="request",
)
- assert exc_info.value.status_code == 403
+ assert exc_info.value.status_code == 400
assert "te*st" in str(exc_info.value.detail)
def test_check_category_keywords_asterisk_pattern_matching(self):
@@ -1510,7 +1510,7 @@ class TestContentFilterGuardrail:
input_type="request",
)
- assert exc_info.value.status_code == 403, f"Failed to block: '{test_input}'"
+ assert exc_info.value.status_code == 400, f"Failed to block: '{test_input}'"
detail = exc_info.value.detail
if isinstance(detail, dict):
assert detail.get("category") == "harm_toxic_abuse"
@@ -1560,7 +1560,7 @@ class TestContentFilterGuardrail:
input_type="request",
)
- assert exc_info.value.status_code == 403, f"Failed to block: '{test_input}'"
+ assert exc_info.value.status_code == 400, f"Failed to block: '{test_input}'"
detail = exc_info.value.detail
if isinstance(detail, dict):
assert detail.get("category") == "harm_toxic_abuse"
@@ -1646,7 +1646,7 @@ class TestContentFilterGuardrail:
)
assert (
- exc_info.value.status_code == 403
+ exc_info.value.status_code == 400
), f"Failed to block Spanish: '{test_input}'"
@pytest.mark.asyncio
@@ -1683,7 +1683,7 @@ class TestContentFilterGuardrail:
)
assert (
- exc_info.value.status_code == 403
+ exc_info.value.status_code == 400
), f"Failed to block French: '{test_input}'"
@pytest.mark.asyncio
@@ -1720,7 +1720,7 @@ class TestContentFilterGuardrail:
)
assert (
- exc_info.value.status_code == 403
+ exc_info.value.status_code == 400
), f"Failed to block German: '{test_input}'"
@pytest.mark.asyncio
@@ -1766,7 +1766,7 @@ class TestContentFilterGuardrail:
)
assert (
- exc_info.value.status_code == 403
+ exc_info.value.status_code == 400
), f"Failed to block Australian: '{test_input}'"
async def test_html_tags_in_messages_not_blocked(self):
@@ -1942,7 +1942,7 @@ class TestContentFilterGuardrail:
request_data={},
input_type="request",
)
- assert exc_info.value.status_code == 403
+ assert exc_info.value.status_code == 400
assert "harmful_child_safety" in str(exc_info.value.detail)
# Test case 2: Should BLOCK - identifier + block word combination
@@ -1956,7 +1956,7 @@ class TestContentFilterGuardrail:
request_data={},
input_type="request",
)
- assert exc_info.value.status_code == 403
+ assert exc_info.value.status_code == 400
# Test case 3: Should BLOCK - explicit content + minors
with pytest.raises(HTTPException) as exc_info:
@@ -1967,7 +1967,7 @@ class TestContentFilterGuardrail:
request_data={},
input_type="request",
)
- assert exc_info.value.status_code == 403
+ assert exc_info.value.status_code == 400
# Test case 4: Should NOT block - identifier word alone (no block word)
result = await guardrail.apply_guardrail(
@@ -2009,7 +2009,7 @@ class TestContentFilterGuardrail:
request_data={},
input_type="request",
)
- assert exc_info.value.status_code == 403
+ assert exc_info.value.status_code == 400
@pytest.mark.asyncio
async def test_conditional_category_sentence_boundaries(self):
@@ -2093,7 +2093,7 @@ class TestContentFilterGuardrail:
request_data={},
input_type="request",
)
- assert exc_info.value.status_code == 403
+ assert exc_info.value.status_code == 400
assert "bias_racial" in str(exc_info.value.detail)
# Test case 2: Should BLOCK - identifier + dehumanizing language
@@ -2107,7 +2107,7 @@ class TestContentFilterGuardrail:
request_data={},
input_type="request",
)
- assert exc_info.value.status_code == 403
+ assert exc_info.value.status_code == 400
# Test case 3: Should BLOCK - supremacist content
with pytest.raises(HTTPException) as exc_info:
@@ -2120,7 +2120,7 @@ class TestContentFilterGuardrail:
request_data={},
input_type="request",
)
- assert exc_info.value.status_code == 403
+ assert exc_info.value.status_code == 400
# Test case 4: Should BLOCK - elimination rhetoric
with pytest.raises(HTTPException) as exc_info:
@@ -2133,7 +2133,7 @@ class TestContentFilterGuardrail:
request_data={},
input_type="request",
)
- assert exc_info.value.status_code == 403
+ assert exc_info.value.status_code == 400
# Test case 5: Should NOT block - identifier word alone (no block word)
result = await guardrail.apply_guardrail(
@@ -2171,7 +2171,7 @@ class TestContentFilterGuardrail:
request_data={},
input_type="request",
)
- assert exc_info.value.status_code == 403
+ assert exc_info.value.status_code == 400
# Test case 9: Should NOT block - block word alone (no identifier)
result = await guardrail.apply_guardrail(
diff --git a/tests/test_litellm/proxy/guardrails/test_custom_code_security.py b/tests/test_litellm/proxy/guardrails/test_custom_code_security.py
index 00cf3f317c9..f93ecfc3010 100644
--- a/tests/test_litellm/proxy/guardrails/test_custom_code_security.py
+++ b/tests/test_litellm/proxy/guardrails/test_custom_code_security.py
@@ -1,11 +1,12 @@
import pytest
+from fastapi import HTTPException
+from litellm.exceptions import ModifyResponseException
from litellm.proxy.guardrails.guardrail_hooks.custom_code.custom_code_guardrail import (
CustomCodeCompilationError,
CustomCodeGuardrail,
)
-
# str.mro() + generator gi_code + code.replace(co_names=...) + __setattr__
# to swap a function's bytecode and read http_get's real builtins dict.
BYTECODE_REWRITE_PAYLOAD = (
@@ -153,6 +154,49 @@ async def test_async_guardrail_compiles_and_runs():
assert result["texts"][0] == "test"
+@pytest.mark.asyncio
+async def test_custom_code_pre_call_block_uses_passthrough():
+ code = (
+ "def apply_guardrail(inputs, request_data, input_type):\n"
+ ' return block("blocked by test")\n'
+ )
+ guardrail = _compile(code)
+
+ with pytest.raises(ModifyResponseException) as exc_info:
+ await guardrail.apply_guardrail(
+ inputs={"texts": ["test"]},
+ request_data={"model": "test-model"},
+ input_type="request",
+ )
+
+ assert exc_info.value.message == "blocked by test"
+ assert exc_info.value.model == "test-model"
+ assert exc_info.value.guardrail_name == "t"
+
+
+@pytest.mark.asyncio
+async def test_custom_code_post_call_block_raises_http_400():
+ code = (
+ "def apply_guardrail(inputs, request_data, input_type):\n"
+ ' return block("blocked by test")\n'
+ )
+ guardrail = _compile(code)
+
+ with pytest.raises(HTTPException) as exc_info:
+ await guardrail.apply_guardrail(
+ inputs={"texts": ["test"]},
+ request_data={"model": "test-model"},
+ input_type="response",
+ )
+
+ assert exc_info.value.status_code == 400
+ assert exc_info.value.detail == {
+ "error": "blocked by test",
+ "guardrail": "t",
+ "detection_info": {},
+ }
+
+
def test_typical_sync_guardrail_still_works():
code = (
"def apply_guardrail(inputs, request_data, input_type):\n"
diff --git a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py
index 16c3c696519..058d5b0283b 100644
--- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py
+++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py
@@ -10,6 +10,9 @@ import pytest
import litellm
from litellm.integrations.custom_guardrail import CustomGuardrail
+from litellm.proxy.guardrails.guardrail_hooks.custom_code.custom_code_guardrail import (
+ CustomCodeGuardrail,
+)
from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor
from litellm.types.proxy.policy_engine.pipeline_types import (
GuardrailPipeline,
@@ -85,6 +88,29 @@ class AlwaysPassGuardrail(CustomGuardrail):
return None
+class PassthroughBlockGuardrail(CustomGuardrail):
+ """Mock guardrail that blocks using the legacy passthrough contract."""
+
+ def __init__(self, guardrail_name: str):
+ super().__init__(
+ guardrail_name=guardrail_name,
+ event_hook="pre_call",
+ default_on=True,
+ )
+ self.calls = 0
+
+ def should_run_guardrail(self, data, event_type) -> bool:
+ return True
+
+ async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
+ self.calls += 1
+ self.raise_passthrough_exception(
+ violation_message="Content policy violation",
+ request_data=data,
+ detection_info={"source": "passthrough"},
+ )
+
+
class PiiMaskingGuardrail(CustomGuardrail):
"""Mock guardrail that masks PII in messages and returns modified data."""
@@ -183,6 +209,105 @@ async def test_escalation_step1_fails_step2_blocks():
litellm.callbacks = original_callbacks
+@pytest.mark.asyncio
+async def test_passthrough_guardrail_failure_can_pipeline_block():
+ """
+ Pipeline: passthrough guardrail (on_fail: block)
+ Expected: passthrough ModifyResponseException is treated as policy fail,
+ and the pipeline terminal action is block.
+ """
+ passthrough_guard = PassthroughBlockGuardrail(guardrail_name="passthrough-filter")
+
+ pipeline = GuardrailPipeline(
+ mode="pre_call",
+ steps=[
+ PipelineStep(
+ guardrail="passthrough-filter",
+ on_fail="block",
+ on_pass="allow",
+ ),
+ ],
+ )
+
+ original_callbacks = litellm.callbacks.copy()
+ litellm.callbacks = [passthrough_guard]
+
+ try:
+ result = await PipelineExecutor.execute_steps(
+ steps=pipeline.steps,
+ mode=pipeline.mode,
+ data={
+ "model": "fake-model",
+ "messages": [{"role": "user", "content": "bad content"}],
+ },
+ user_api_key_dict=MagicMock(),
+ call_type="completion",
+ policy_name="content-safety",
+ )
+
+ assert passthrough_guard.calls == 1
+ assert result.terminal_action == "block"
+ assert len(result.step_results) == 1
+ assert result.step_results[0].guardrail_name == "passthrough-filter"
+ assert result.step_results[0].outcome == "fail"
+ assert result.step_results[0].action_taken == "block"
+ assert result.error_message == "Content policy violation"
+ finally:
+ litellm.callbacks = original_callbacks
+
+
+@pytest.mark.asyncio
+async def test_custom_code_guardrail_failure_can_pipeline_block():
+ """
+ Pipeline: custom code guardrail (on_fail: block)
+ Expected: custom code keeps its standalone passthrough block behavior, and
+ the pipeline converts that guardrail intervention into a block action.
+ """
+ custom_guard = CustomCodeGuardrail(
+ guardrail_name="custom-code-filter",
+ custom_code=(
+ "def apply_guardrail(inputs, request_data, input_type):\n"
+ ' return block("SSN detected")\n'
+ ),
+ )
+
+ pipeline = GuardrailPipeline(
+ mode="pre_call",
+ steps=[
+ PipelineStep(
+ guardrail="custom-code-filter",
+ on_fail="block",
+ on_pass="allow",
+ ),
+ ],
+ )
+
+ original_callbacks = litellm.callbacks.copy()
+ litellm.callbacks = [custom_guard]
+
+ try:
+ result = await PipelineExecutor.execute_steps(
+ steps=pipeline.steps,
+ mode=pipeline.mode,
+ data={
+ "model": "fake-model",
+ "messages": [{"role": "user", "content": "123-45-6789"}],
+ },
+ user_api_key_dict=MagicMock(),
+ call_type="completion",
+ policy_name="content-safety",
+ )
+
+ assert result.terminal_action == "block"
+ assert len(result.step_results) == 1
+ assert result.step_results[0].guardrail_name == "custom-code-filter"
+ assert result.step_results[0].outcome == "fail"
+ assert result.step_results[0].action_taken == "block"
+ assert result.error_message == "SSN detected"
+ finally:
+ litellm.callbacks = original_callbacks
+
+
@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed")
@pytest.mark.asyncio
async def test_early_allow_step1_passes_step2_skipped():
From 7ca796beb10fd1f84d673f2a84c41ebeffece06c Mon Sep 17 00:00:00 2001
From: Shivam Rawat
Date: Sat, 30 May 2026 15:10:21 -0700
Subject: [PATCH 33/44] fix(proxy): restrict vector store index create/delete
to proxy admins (#29202)
* fix(proxy): restrict vector store index create/delete to proxy admins
Prevent non-admin API keys from registering indexes via POST /v1/indexes or deleting Azure AI Search indexes through managed pass-through routes.
Co-authored-by: Cursor
* fix(proxy): tighten vector store index lifecycle checks
Co-authored-by: Cursor
---------
Co-authored-by: Cursor
---
.../proxy/vector_store_endpoints/endpoints.py | 6 +
litellm/proxy/vector_store_endpoints/utils.py | 84 ++++++-
.../test_vector_store_endpoints.py | 216 +++++++++++++++++-
3 files changed, 302 insertions(+), 4 deletions(-)
diff --git a/litellm/proxy/vector_store_endpoints/endpoints.py b/litellm/proxy/vector_store_endpoints/endpoints.py
index ccf15c206b0..b3bdfecbe55 100644
--- a/litellm/proxy/vector_store_endpoints/endpoints.py
+++ b/litellm/proxy/vector_store_endpoints/endpoints.py
@@ -12,6 +12,7 @@ from litellm.proxy.vector_store_endpoints.management_endpoints import (
_resolve_embedding_config,
)
from litellm.proxy.vector_store_endpoints.utils import (
+ assert_proxy_admin_for_vector_store_index_management,
assert_user_can_access_vector_store,
get_litellm_managed_vector_store,
)
@@ -575,6 +576,11 @@ async def index_create(
"""
from litellm.proxy.proxy_server import prisma_client
+ assert_proxy_admin_for_vector_store_index_management(
+ user_api_key_dict,
+ operation="create",
+ )
+
if prisma_client is None:
raise HTTPException(
status_code=500,
diff --git a/litellm/proxy/vector_store_endpoints/utils.py b/litellm/proxy/vector_store_endpoints/utils.py
index 1221ccf119f..d4afc547031 100644
--- a/litellm/proxy/vector_store_endpoints/utils.py
+++ b/litellm/proxy/vector_store_endpoints/utils.py
@@ -1,4 +1,5 @@
import json
+import re
from typing import Any, Dict, Literal, Optional
from fastapi import HTTPException, Request
@@ -37,6 +38,64 @@ def _is_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> bool:
)
+def assert_proxy_admin_for_vector_store_index_management(
+ user_api_key_dict: UserAPIKeyAuth,
+ *,
+ operation: Literal["create", "delete", "update"] = "create",
+) -> None:
+ """Raise 403 unless the caller is a proxy admin."""
+ if _is_proxy_admin(user_api_key_dict):
+ return
+ raise HTTPException(
+ status_code=403,
+ detail=(
+ f"Only proxy admins can {operation} vector store indexes. "
+ "Contact your LiteLLM administrator."
+ ),
+ )
+
+
+def _suffix_after_index_name(request_path: str, index_name: str) -> Optional[str]:
+ """Return the path suffix after ``/indexes/{index_name}``, or None if absent."""
+ match = re.search(rf"/indexes/{re.escape(index_name)}(?=$|[/?])", request_path)
+ if match is None:
+ return None
+ return request_path[match.end() :]
+
+
+def _is_vector_store_index_lifecycle_request(
+ request_method: str,
+ request_path: str,
+ index_name: str,
+) -> bool:
+ """
+ True when the request creates or deletes a search index itself (not documents).
+
+ Examples (admin-only):
+ - DELETE /azure_ai/indexes/my-index
+ - PUT /azure_ai/indexes/my-index
+ - POST /azure_ai/indexes
+ """
+ if request_method not in ("POST", "PUT", "DELETE", "PATCH"):
+ return False
+
+ suffix = _suffix_after_index_name(request_path, index_name)
+ if suffix is not None:
+ # Document operations live under /indexes/{name}/docs/...
+ if suffix.startswith("/docs"):
+ return False
+ # DELETE/PUT/PATCH on /indexes/{name} itself is index lifecycle.
+ if suffix == "" or suffix.startswith("?"):
+ return True
+
+ # POST /indexes (create index at service level; no index name in path).
+ normalized = request_path.rstrip("/")
+ if request_method == "POST" and normalized.endswith("/indexes"):
+ return True
+
+ return False
+
+
def _object_permission_allows_vector_store(
object_permission: Optional[LiteLLM_ObjectPermissionTable],
vector_store_id: str,
@@ -335,6 +394,22 @@ def is_allowed_to_call_vector_store_endpoint(
request_route = get_request_route(request)
+ if _is_vector_store_index_lifecycle_request(
+ request_method=request.method,
+ request_path=request_route,
+ index_name=index_name,
+ ):
+ operation_label: Literal["create", "delete", "update"] = "create"
+ if request.method == "DELETE":
+ operation_label = "delete"
+ elif request.method in ("PUT", "PATCH"):
+ operation_label = "update"
+ assert_proxy_admin_for_vector_store_index_management(
+ user_api_key_dict,
+ operation=operation_label,
+ )
+ return True
+
# Determine the permission type based on the request
permission_type = None
for endpoint in provider_vector_store_endpoints["read"]:
@@ -353,7 +428,14 @@ def is_allowed_to_call_vector_store_endpoint(
break
if permission_type is None:
- return None
+ raise HTTPException(
+ status_code=403,
+ detail=(
+ f"User does not have permission to call vector store endpoint "
+ f"{index_name}. Ask your administrator to add the necessary "
+ "permissions to your API key/Team."
+ ),
+ )
# Check if key has specific permission for allowed_vector_store_indexes
has_permission = check_vector_store_permission(
diff --git a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py
index 81b67e8bc50..882081b33c3 100644
--- a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py
+++ b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py
@@ -16,9 +16,10 @@ import litellm
from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import (
LiteLLM_ManagedVectorStore,
)
-from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.vector_store_endpoints.endpoints import (
_update_request_data_with_litellm_managed_vector_store_registry,
+ index_create,
)
from litellm.proxy.vector_store_endpoints.management_endpoints import (
_check_vector_store_access,
@@ -33,6 +34,7 @@ from litellm.proxy.vector_store_endpoints.utils import (
is_allowed_to_call_vector_store_endpoint,
is_allowed_to_call_vector_store_files_endpoint,
)
+from litellm.types.vector_stores import IndexCreateRequest
from litellm.types.utils import LlmProviders
@@ -674,18 +676,151 @@ class TestIsAllowedToCallVectorStoreEndpoint:
"write": [("POST", "/create")],
}
+ with patch(
+ "litellm.proxy.vector_store_endpoints.utils.ProviderConfigManager.get_provider_vector_stores_config",
+ return_value=mock_provider_config,
+ ):
+ with pytest.raises(HTTPException) as exc_info:
+ is_allowed_to_call_vector_store_endpoint(
+ provider=LlmProviders.OPENAI,
+ index_name="my-index",
+ request=mock_request,
+ user_api_key_dict=mock_user_api_key,
+ )
+
+ assert exc_info.value.status_code == 403
+
+ def test_delete_index_requires_admin(self):
+ """Non-admin users must not delete managed search indexes via pass-through."""
+ mock_request = MagicMock(spec=Request)
+ mock_request.method = "DELETE"
+ mock_request.url.path = "/azure_ai/indexes/my-index"
+
+ mock_user_api_key = MagicMock(spec=UserAPIKeyAuth)
+ mock_user_api_key.user_role = None
+ mock_user_api_key.metadata = {
+ "allowed_vector_store_indexes": [
+ {"index_name": "my-index", "index_permissions": ["read", "write"]}
+ ]
+ }
+ mock_user_api_key.team_metadata = None
+
+ mock_provider_config = MagicMock()
+ mock_provider_config.get_vector_store_endpoints_by_type.return_value = {
+ "read": [("GET", "/docs/search"), ("POST", "/docs/search")],
+ "write": [("PUT", "/docs")],
+ }
+
+ with patch(
+ "litellm.proxy.vector_store_endpoints.utils.ProviderConfigManager.get_provider_vector_stores_config",
+ return_value=mock_provider_config,
+ ):
+ with pytest.raises(HTTPException) as exc_info:
+ is_allowed_to_call_vector_store_endpoint(
+ provider=LlmProviders.AZURE_AI,
+ index_name="my-index",
+ request=mock_request,
+ user_api_key_dict=mock_user_api_key,
+ )
+
+ assert exc_info.value.status_code == 403
+ assert "Only proxy admins can delete" in exc_info.value.detail
+
+ def test_delete_index_allowed_for_admin(self):
+ """Proxy admins can delete managed search indexes via pass-through."""
+ mock_request = MagicMock(spec=Request)
+ mock_request.method = "DELETE"
+ mock_request.url.path = "/azure_ai/indexes/my-index"
+
+ mock_user_api_key = MagicMock(spec=UserAPIKeyAuth)
+ mock_user_api_key.user_role = LitellmUserRoles.PROXY_ADMIN
+
+ mock_provider_config = MagicMock()
+ mock_provider_config.get_vector_store_endpoints_by_type.return_value = {
+ "read": [("GET", "/docs/search"), ("POST", "/docs/search")],
+ "write": [("PUT", "/docs")],
+ }
+
with patch(
"litellm.proxy.vector_store_endpoints.utils.ProviderConfigManager.get_provider_vector_stores_config",
return_value=mock_provider_config,
):
result = is_allowed_to_call_vector_store_endpoint(
- provider=LlmProviders.OPENAI,
+ provider=LlmProviders.AZURE_AI,
index_name="my-index",
request=mock_request,
user_api_key_dict=mock_user_api_key,
)
- assert result is None
+ assert result is True
+
+ def test_update_index_requires_admin_with_update_message(self):
+ """Non-admin users get an update-specific message for index replacement."""
+ mock_request = MagicMock(spec=Request)
+ mock_request.method = "PUT"
+ mock_request.url.path = "/azure_ai/indexes/my-index"
+
+ mock_user_api_key = MagicMock(spec=UserAPIKeyAuth)
+ mock_user_api_key.user_role = None
+ mock_user_api_key.metadata = {
+ "allowed_vector_store_indexes": [
+ {"index_name": "my-index", "index_permissions": ["read", "write"]}
+ ]
+ }
+ mock_user_api_key.team_metadata = None
+
+ mock_provider_config = MagicMock()
+ mock_provider_config.get_vector_store_endpoints_by_type.return_value = {
+ "read": [("GET", "/docs/search"), ("POST", "/docs/search")],
+ "write": [("PUT", "/docs")],
+ }
+
+ with patch(
+ "litellm.proxy.vector_store_endpoints.utils.ProviderConfigManager.get_provider_vector_stores_config",
+ return_value=mock_provider_config,
+ ):
+ with pytest.raises(HTTPException) as exc_info:
+ is_allowed_to_call_vector_store_endpoint(
+ provider=LlmProviders.AZURE_AI,
+ index_name="my-index",
+ request=mock_request,
+ user_api_key_dict=mock_user_api_key,
+ )
+
+ assert exc_info.value.status_code == 403
+ assert "Only proxy admins can update" in exc_info.value.detail
+
+ def test_index_name_prefix_does_not_match_lifecycle_request(self):
+ """An index name that is only a path prefix must not trigger lifecycle checks."""
+ mock_request = MagicMock(spec=Request)
+ mock_request.method = "DELETE"
+ mock_request.url.path = "/azure_ai/indexes/my-index-archive"
+
+ mock_user_api_key = MagicMock(spec=UserAPIKeyAuth)
+ mock_user_api_key.user_role = None
+ mock_user_api_key.metadata = None
+ mock_user_api_key.team_metadata = None
+
+ mock_provider_config = MagicMock()
+ mock_provider_config.get_vector_store_endpoints_by_type.return_value = {
+ "read": [],
+ "write": [],
+ }
+
+ with patch(
+ "litellm.proxy.vector_store_endpoints.utils.ProviderConfigManager.get_provider_vector_stores_config",
+ return_value=mock_provider_config,
+ ):
+ with pytest.raises(HTTPException) as exc_info:
+ is_allowed_to_call_vector_store_endpoint(
+ provider=LlmProviders.AZURE_AI,
+ index_name="my-index",
+ request=mock_request,
+ user_api_key_dict=mock_user_api_key,
+ )
+
+ assert exc_info.value.status_code == 403
+ assert "Only proxy admins" not in exc_info.value.detail
def test_team_metadata_permissions(self):
"""Test that team metadata permissions work."""
@@ -800,6 +935,81 @@ class TestIsAllowedToCallVectorStoreEndpoint:
assert exc_info.value.status_code == 403
+class TestIndexCreate:
+ @pytest.mark.asyncio
+ async def test_index_create_requires_admin(self):
+ """Non-admin users must not register managed vector store indexes."""
+ request = IndexCreateRequest(
+ index_name="test-index",
+ litellm_params={
+ "vector_store_index": "real-index",
+ "vector_store_name": "azure-ai-search",
+ },
+ )
+ mock_request = MagicMock(spec=Request)
+ mock_response = MagicMock()
+
+ with pytest.raises(HTTPException) as exc_info:
+ await index_create(
+ request=mock_request,
+ index_create_request=request,
+ fastapi_response=mock_response,
+ user_api_key_dict=UserAPIKeyAuth(
+ token="sk-test",
+ key_name="sk-...test",
+ user_role=LitellmUserRoles.INTERNAL_USER,
+ ),
+ )
+
+ assert exc_info.value.status_code == 403
+ assert "Only proxy admins can create" in exc_info.value.detail
+
+ @pytest.mark.asyncio
+ async def test_index_create_allowed_for_admin(self):
+ """Proxy admins can register managed vector store indexes."""
+ create_request = IndexCreateRequest(
+ index_name="test-index",
+ litellm_params={
+ "vector_store_index": "real-index",
+ "vector_store_name": "azure-ai-search",
+ },
+ )
+ mock_request = MagicMock(spec=Request)
+ mock_response = MagicMock()
+ mock_row = MagicMock()
+ mock_row.model_dump.return_value = {
+ "index_name": "test-index",
+ "litellm_params": create_request.litellm_params.model_dump(),
+ }
+
+ mock_prisma = MagicMock()
+ mock_prisma.db.litellm_managedvectorstoreindextable.find_unique = AsyncMock(
+ return_value=None
+ )
+ mock_prisma.db.litellm_managedvectorstoreindextable.create = AsyncMock(
+ return_value=mock_row
+ )
+
+ with patch(
+ "litellm.proxy.proxy_server.prisma_client",
+ mock_prisma,
+ ):
+ result = await index_create(
+ request=mock_request,
+ index_create_request=create_request,
+ fastapi_response=mock_response,
+ user_api_key_dict=UserAPIKeyAuth(
+ token="sk-test",
+ key_name="sk-...test",
+ user_role=LitellmUserRoles.PROXY_ADMIN,
+ user_id="admin-user",
+ ),
+ )
+
+ assert result["index_name"] == "test-index"
+ mock_prisma.db.litellm_managedvectorstoreindextable.create.assert_awaited_once()
+
+
class TestIsAllowedToCallVectorStoreFilesEndpoint:
def _mock_provider_config(self):
provider_config = MagicMock()
From ba2699740cefe1322f3e54f74028459ec0fbb368 Mon Sep 17 00:00:00 2001
From: Sameer Kankute
Date: Sun, 31 May 2026 05:00:10 +0530
Subject: [PATCH 34/44] feat(pass_through): extend
passthrough_managed_object_ids to Azure (#29160)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
* feat(pass_through): extend passthrough_managed_object_ids to Azure
Adds managed ID minting/resolution for Azure passthrough endpoints
(/azure/...) alongside the existing OpenAI passthrough support.
Key changes:
- pass_through_endpoints.py: detect azure/azure_ai custom_llm_provider
(string or enum) to set _is_managed_id_provider and _managed_id_provider;
both INPUT and OUTPUT rewrite blocks now fire for Azure.
- llm_passthrough_endpoints.py: forward custom_llm_provider into
create_pass_through_route so it reaches pass_through_request (was None).
- managed_id_rewriter.py: extend _PASSTHROUGH_PREFIX_RE and _canonical_path
to strip /azure/openai prefix and add /v1/ for Azure paths that omit it;
add ("azure", method, path) entries to BUILTIN_OUTPUT_ID_FIELD_MAP for
files and batches endpoints.
- managed_id_codec.py / types/utils.py: supporting codec and enum constant.
- proxy_server.py: register llm_passthrough_router before batches_router to
prevent route collision for /openai_passthrough/* paths.
Co-authored-by: Cursor
* fix(pass_through): remove unused imports for ruff F401
Co-authored-by: Cursor
* fix(pass_through): satisfy mypy for optional parsed_body
Co-authored-by: Cursor
* fix(pass_through): compute query params string after managed-ID rewrite
Move requested_query_params_str computation to after the managed-ID
input rewrite block so logging_url reflects the rewritten raw-provider
query params actually sent upstream, instead of the original
managed IDs.
Co-authored-by: Yassin Kortam
* Add support for managed ids for passthrough responses api
* Add support for list batches and list files
* style: run Black on passthrough managed ID files
Fix CI formatting for managed_id_rewriter.py and pass_through_endpoints.py.
Co-authored-by: Cursor
* fix(passthrough): parse json file_object and implement before-cursor pagination
- Parse row.file_object via json.loads when Prisma returns it as a string;
mirrors openai_files_endpoints/common_utils.py so list responses keep all
stored detail fields (status, timestamps, etc.).
- Implement the previously-parsed-but-unused 'before' cursor for list
pagination by flipping fetch order to ascending with a 'gt' bound on
created_at, then reversing rows so the response stays newest-first.
Co-authored-by: Yassin Kortam
* Remove logger
* refactor: split list_passthrough_ids_from_db to fix PLR0915
Extract pagination, fetch, and serialization helpers so the main list
function stays under the statement limit without changing behavior.
Co-authored-by: Cursor
* fix: scope passthrough managed ID dedup and list by provider
Validate embedded provider before reusing deduped file/object rows so
OpenAI and Azure cannot share the same managed ID for an identical raw ID.
Filter list responses to rows whose managed IDs decode to the current
provider, with over-fetch scanning when needed.
Co-authored-by: Cursor
* fix(managed_id_rewriter): cap pagination trim at effective limit
When raw_limit > 100, fetch_limit is capped at 101 (one extra row to
detect has_more), but trimming with rows[:raw_limit] failed to drop
the sentinel row. Use the capped effective limit instead.
Co-authored-by: Yassin Kortam
* fix: cross-provider object collision and fail-closed list error handling
Greptile P1: move provider check before access check in _mint_or_reuse_object
so a cross-provider raw ID collision (OpenAI and Azure share the same batch_
ID) falls through to mint a new provider-scoped row instead of raising 404.
Veria-ai medium: _fetch_provider_scoped_list_rows now always returns
(page, has_more) — DB errors break out of the scan loop and return matched
rows so far. list_passthrough_ids_from_db never returns None for a recognised
list route, so the caller can never fall through to the upstream provider.
Co-authored-by: Cursor
* fix: namespace passthrough model_object_id by provider to prevent unique violation
Store model_object_id as 'passthrough:{provider}:{raw_id}' instead of the
bare raw ID so OpenAI and Azure can each own a row for the same raw batch ID
without hitting the @unique constraint. Dedup lookup uses the same namespaced
key so it is implicitly provider-scoped and the _managed_id_matches_provider
check is no longer needed on the object path.
Co-authored-by: Cursor
* fix: gate list interception on managed_files hook like input/output rewrites
Without the hook no managed IDs are minted so the DB is empty. Intercepting
GET /v1/files without the hook returned an empty list and hid the caller's
real upstream files/batches. Matches the guard used by the input and output
rewrite blocks.
Co-authored-by: Cursor
* fix: set has_more=True when scan cap is hit with a full final DB batch
When max_scans (20) is exhausted and the last DB page was full-sized,
there are almost certainly more rows beyond the scan window. Track
last_batch_full across iterations so the scan_cap_hit condition sets
has_more=True in that case, preventing silent pagination truncation in
high-mixed-provider pools.
Co-authored-by: Cursor
* fix(managed_id_rewriter): scope pagination cursor lookup to caller-owned rows
Prevent a cross-tenant timing oracle by constraining the after/before
cursor row lookup to the caller's owner_filter, and cover the real Azure
responses path form (no /v1/) in tests.
* fix(managed_id_rewriter): align passthrough list metadata with direct GET
Persist upstream file metadata when minting a managed file ID and rewrite
nested batch file IDs before snapshotting the object, so DB-served file/batch
list responses return the same fields and managed IDs as a direct endpoint
GET.
* fix(managed_id_rewriter): degrade to raw id on cross-owner object collision
The OUTPUT (mint) path of _mint_or_reuse_object raised HTTPException(404)
when a dedup hit on the namespaced model_object_id belonged to a different
owner, converting a successful upstream batch/response creation into a 404
for the caller. Two upstream accounts under one provider name can issue the
same raw id, so this is reachable in multi-tenant deployments.
Return the caller's raw id unmanaged instead: the upstream create already
succeeded, a new managed row can't be minted (model_object_id is @unique),
and reusing the other owner's managed id would later fail the access check.
* fix(managed_id_rewriter): scope list cursor by provider and cap body-rewrite recursion depth
* perf(managed_id_rewriter): push batch list provider scope to DB and anchor canonical-path prefix
Object (batch) list rows store model_object_id as passthrough:{provider}:{raw},
so the provider filter is now applied at the indexed DB column, collapsing the
application-layer multi-scan to a single query for that table. File rows keep
the decode-based scan since they have no provider column.
Anchor the canonical-path prefix regex at a path boundary so routes such as
/openai_realtime/... are no longer mis-stripped.
* fix(managed_id_rewriter): refresh stored batch snapshot on reuse
The dedup-reuse path in _mint_or_reuse_object returned the existing managed id
without updating the stored file_object, so DB-served list responses kept the
creation-time snapshot and showed null output_file_id/error_file_id even after
the batch completed. Refresh the snapshot when an owned row is reused so the
list reflects the batch's latest state.
* fix(managed_id_rewriter): deny cross-owner object access on retrieve/cancel/delete
Returning the raw id when can_access_resource fails only made sense for create
responses, where the caller's own upstream create succeeded under a raw id that a
different owner already holds. On retrieve/cancel/delete the caller reaches that
branch only by supplying another tenant's raw id (which bypasses the managed-id
input gate), so echoing the upstream object back leaked it cross-tenant. Restrict
the raw fallback to create routes and return 404 otherwise.
* fix(managed_id_rewriter): deny cross-owner file access on retrieve/delete
_mint_or_reuse_file scoped the raw file dedup lookup to the current caller, so a
raw file-... id belonging to another tenant was never found and the OUTPUT path
minted a fresh managed id for that same upstream file under the caller. A raw id
only reaches this path by skipping the managed-id input gate (raw provider ids
are opt-out), so a different-owner row means the caller is touching someone
else's file. Look up flat_model_file_ids globally and run can_access_resource;
deny with 404 on retrieve/delete and leave the raw id unmanaged on create, which
mirrors the cross-owner handling already in _mint_or_reuse_object.
* fix(managed_id_rewriter): deterministic provider-scoped file dedup
Replace the unscoped find_first in _mint_or_reuse_file with a find_many
ordered by created_at and an application-layer provider filter. The file
table has no provider column, so a raw file id shared across OpenAI and
Azure could map to one row per provider; find_first then picked a row
non-deterministically and, on a provider mismatch, minted a fresh managed
row on every call, accumulating duplicates. Selecting the oldest matching
same-provider row the caller can access keeps reuse stable and prevents
duplicate rows while preserving the cross-tenant deny/leave-raw behaviour.
* refactor(pass_through): scope passthrough managed IDs on the explicit provider
Move the openai/azure detection out of pass_through_request into
resolve_passthrough_managed_id_provider in llms/base_llm/managed_resources,
and key managed-ID rewriting on the forwarded custom_llm_provider rather than
the upstream URL's EndpointType. The helper documents why azure and azure_ai
collapse to one "azure" scope (they share the same Azure OpenAI files/batches
surface, so an ID minted on one must resolve on the other) and returns None for
any other provider so a third-party OpenAI-compatible endpoint never triggers
managed-ID minting.
Add TestManagedIdProviderScope covering the azure_ai -> azure collapse and the
non-openai/azure exclusion.
* test(log_db_metrics): assert sanitized event_metadata contract
test_log_db_metrics_success still asserted the legacy event_metadata
shape (function_name/function_kwargs/function_args), which #28909
intentionally removed so that live Prisma clients, OTel spans, and
secrets never land on a service-log span. The decorator now emits only
a sanitized payload: None when no table_name is present, and
{"table_name": ...} when it is. Update the test to verify both branches
of that contract.
* fix(managed_id_rewriter): page provider-scoped file list by offset
The file list scan advanced its cursor with a strict created_at boundary.
When several rows shared a created_at timestamp and a non-matching provider
row sat on the page boundary, the next query skipped the remaining rows at
that timestamp, dropping matching files from the response. Page by a stable
offset over a total order (created_at plus the unique id column) so tied
rows are never skipped or repeated.
* fix(managed_id_rewriter): push file-list provider scope to the DB
The file-list helper had no provider column to query, so it scanned the
table and filtered by decoding each managed ID in the application layer,
capped at 20 pages. For an admin with a large mixed-provider file pool
that cap could truncate a page.
Mint now writes a _passthrough_provider:{provider} marker into
flat_model_file_ids, giving the file table the same DB-queryable provider
scope object rows already get from the namespaced model_object_id. The
list helper pushes the scope into the query so a single round-trip serves
the page. The scan loop, the offset paging, and the cap are gone, so pages
can no longer truncate, leak the other provider, or skip rows that share a
created_at timestamp.
* fix(managed_id_rewriter): deny raw provider IDs that map to another tenant's managed resource
Clients only ever receive managed IDs on passthrough, so a raw file/batch/response ID for another tenant's managed object can only be recovered by decoding that tenant's managed ID. Raw IDs were forwarded upstream untouched (deliberate opt-out), which on a retrieve/cancel/delete executed upstream before the response-side ownership check ran, leaking a cross-tenant action.
Guard raw provider IDs on the input path: when a raw file-/batch_/resp_ ID resolves to a managed row the caller cannot access, return 404 before forwarding. Genuinely unmanaged raw IDs (no DB row) and IDs the caller owns are left untouched, preserving the opt-out.
* test(managed_id_rewriter): cover azure_ai and pre-versioned azure passthrough paths
* fix(managed_id_rewriter): fall back to raw id when persistence fails
A DB persistence failure after a successful upstream create left the
client holding a minted managed ID with no backing row, so every later
resolve returned 404 and the resource was permanently unreachable. Mint
the managed ID only when the row is stored; on persistence failure return
the raw provider id, matching the no-persistence-available fallback, so
the freshly-created resource stays reachable.
* fix(managed_id_rewriter): log only rewritten query param keys
* fix(managed_id_rewriter): use compound (created_at, id) list cursor boundary
A timestamp-only lt/gt cursor boundary skips list rows that share the
cursor row's created_at across a page boundary, silently dropping them.
Compare the unique id (the secondary sort key) alongside created_at so
the page walk stays complete when timestamps tie.
* fix(managed_id_rewriter): converge concurrent object creates on one managed id
* fix(managed_id_rewriter): bound raw-id guard DB lookups per request
The INPUT guard fired one DB lookup for every file-/batch_/resp_ prefixed
string in the path, query, and body. The file-id guard is an unindexed
array-containment scan over LiteLLM_ManagedFileTable, so an authenticated
caller could amplify a single passthrough request into thousands of
full-table scans by packing a body with id-shaped strings.
De-dupe raw ids within a request and cap the distinct guard lookups,
failing closed with 400 instead of skipping the guard. Legitimate callers
hold managed ids (resolved via an indexed unified_*_id lookup, not the
guard), so the cap only trips under abuse.
---------
Co-authored-by: Cursor
Co-authored-by: Yassin Kortam
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
---
.../base_llm/managed_resources/__init__.py | 2 +
.../llms/base_llm/managed_resources/utils.py | 35 +-
.../llm_passthrough_endpoints.py | 5 +
.../managed_id_codec.py | 97 +
.../managed_id_rewriter.py | 1234 ++++
.../pass_through_endpoints.py | 6317 +++++++++--------
litellm/proxy/proxy_server.py | 2 +-
litellm/types/utils.py | 4 +
.../test_passthrough_managed_ids.py | 2087 ++++++
9 files changed, 6716 insertions(+), 3067 deletions(-)
create mode 100644 litellm/proxy/pass_through_endpoints/managed_id_codec.py
create mode 100644 litellm/proxy/pass_through_endpoints/managed_id_rewriter.py
create mode 100644 tests/pass_through_unit_tests/test_passthrough_managed_ids.py
diff --git a/litellm/llms/base_llm/managed_resources/__init__.py b/litellm/llms/base_llm/managed_resources/__init__.py
index 5eb9b46f89f..a5543e631c0 100644
--- a/litellm/llms/base_llm/managed_resources/__init__.py
+++ b/litellm/llms/base_llm/managed_resources/__init__.py
@@ -24,10 +24,12 @@ from .utils import (
generate_unified_id_string,
is_base64_encoded_unified_id,
parse_unified_id,
+ resolve_passthrough_managed_id_provider,
)
__all__ = [
"BaseManagedResource",
+ "resolve_passthrough_managed_id_provider",
"is_base64_encoded_unified_id",
"extract_target_model_names_from_unified_id",
"extract_resource_type_from_unified_id",
diff --git a/litellm/llms/base_llm/managed_resources/utils.py b/litellm/llms/base_llm/managed_resources/utils.py
index 6e30b6cb252..e9a6aef689e 100644
--- a/litellm/llms/base_llm/managed_resources/utils.py
+++ b/litellm/llms/base_llm/managed_resources/utils.py
@@ -7,7 +7,40 @@ different managed resource types (files, vector stores, etc.).
import base64
import re
-from typing import List, Optional, Union, Literal
+from typing import Any, List, Literal, Optional, Union
+
+PASSTHROUGH_MANAGED_ID_AZURE_PROVIDERS = ("azure", "azure_ai")
+
+
+def resolve_passthrough_managed_id_provider(
+ custom_llm_provider: Any,
+) -> Optional[str]:
+ """Map a pass-through ``custom_llm_provider`` to the provider scope that
+ namespaces passthrough managed object IDs, or ``None`` when the route is not
+ an OpenAI/Azure pass-through and managed IDs must not apply.
+
+ Scoping is keyed on the explicit provider that the pass-through route
+ forwards (``openai``, ``azure``, ``azure_ai``), not on the upstream URL, so
+ a third-party OpenAI-compatible endpoint never triggers managed-ID minting.
+
+ ``azure`` and ``azure_ai`` deliberately collapse to one ``"azure"`` scope:
+ they expose the same Azure OpenAI files/batches surface, so an ID minted
+ while routing as one must still resolve while routing as the other.
+ Splitting them would make a managed ID minted on ``azure`` fail to resolve
+ when replayed on ``azure_ai`` and vice versa.
+ """
+ provider = str(
+ getattr(custom_llm_provider, "value", custom_llm_provider) or ""
+ ).lower()
+ if not provider:
+ return None
+ if provider in PASSTHROUGH_MANAGED_ID_AZURE_PROVIDERS or provider.endswith(
+ (".azure", ".azure_ai")
+ ):
+ return "azure"
+ if provider == "openai" or provider.endswith(".openai"):
+ return "openai"
+ return None
def is_base64_encoded_unified_id(
diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
index 7ca28a5d4ac..e94f56302a8 100644
--- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
+++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
@@ -2091,6 +2091,11 @@ class BaseOpenAIPassThroughHandler:
api_key=api_key, request=request, extra_headers=extra_headers
),
is_streaming_request=is_streaming_request, # type: ignore
+ custom_llm_provider=(
+ custom_llm_provider.value
+ if hasattr(custom_llm_provider, "value")
+ else str(custom_llm_provider) if custom_llm_provider else None
+ ),
) # dynamically construct pass-through endpoint based on incoming path
received_value = await endpoint_func(
request,
diff --git a/litellm/proxy/pass_through_endpoints/managed_id_codec.py b/litellm/proxy/pass_through_endpoints/managed_id_codec.py
new file mode 100644
index 00000000000..f0c24bbaf39
--- /dev/null
+++ b/litellm/proxy/pass_through_endpoints/managed_id_codec.py
@@ -0,0 +1,97 @@
+"""
+Codec for LiteLLM passthrough-managed object IDs.
+
+Plaintext format (before urlsafe-base64 encoding):
+ litellm_proxy:passthrough;provider:{p};unified_id,{u};raw_id,{r}
+
+Uses the same base64.urlsafe_b64encode / padding-restore convention as
+``_is_base64_encoded_unified_file_id`` in
+``openai_files_endpoints/common_utils.py``.
+
+The ``passthrough;`` discriminator distinguishes these rows from
+unified-endpoint rows that share the same LiteLLM_ManagedFileTable /
+LiteLLM_ManagedObjectTable. ``_resolve_one`` in the rewriter module rejects
+any row whose decoded plaintext lacks this discriminator, making cross-system
+replay safe.
+"""
+
+from __future__ import annotations
+
+import base64
+import uuid as _uuid_mod
+from dataclasses import dataclass
+from typing import Optional
+
+from litellm.types.utils import SpecialEnums
+
+_PREFIX = SpecialEnums.LITELM_MANAGED_FILE_ID_PREFIX.value # "litellm_proxy"
+_DISCRIMINATOR = "passthrough"
+
+
+@dataclass(frozen=True)
+class ManagedIdPayload:
+ """Decoded contents of a passthrough managed ID."""
+
+ provider: str
+ unified_uuid: str
+ raw_provider_id: str
+
+
+def encode(provider: str, unified_uuid: str, raw_provider_id: str) -> str:
+ """Return a urlsafe-base64 managed ID string (trailing ``=`` stripped)."""
+ plaintext = SpecialEnums.LITELLM_PASSTHROUGH_MANAGED_ID_COMPLETE_STR.value.format(
+ provider, unified_uuid, raw_provider_id
+ )
+ return base64.urlsafe_b64encode(plaintext.encode()).decode().rstrip("=")
+
+
+def decode(managed_id: str) -> Optional[ManagedIdPayload]:
+ """
+ Decode *managed_id*.
+
+ Returns ``None`` for anything that is not a passthrough managed ID — raw
+ OpenAI IDs, unified-endpoint IDs, garbage, wrong types. Never raises.
+ """
+ if not isinstance(managed_id, str):
+ return None
+ # Restore stripped padding before decoding
+ padded = managed_id + "=" * (-len(managed_id) % 4)
+ try:
+ plaintext = base64.urlsafe_b64decode(padded).decode()
+ except Exception:
+ return None
+
+ # Must start with "litellm_proxy:passthrough;"
+ expected_head = f"{_PREFIX}:{_DISCRIMINATOR};"
+ if not plaintext.startswith(expected_head):
+ return None
+
+ rest = plaintext[len(expected_head) :]
+ try:
+ # Split only on first two ';' so a raw_id containing ';' cannot
+ # break parsing (OpenAI IDs don't use ';', but defensive).
+ provider_part, rest2 = rest.split(";", 1)
+ unified_part, raw_id_part = rest2.split(";", 1)
+ if not (
+ provider_part.startswith("provider:")
+ and unified_part.startswith("unified_id,")
+ and raw_id_part.startswith("raw_id,")
+ ):
+ return None
+ return ManagedIdPayload(
+ provider=provider_part[len("provider:") :],
+ unified_uuid=unified_part[len("unified_id,") :],
+ raw_provider_id=raw_id_part[len("raw_id,") :],
+ )
+ except Exception:
+ return None
+
+
+def is_managed(value: str) -> bool:
+ """Return ``True`` iff *value* decodes to a passthrough managed ID."""
+ return decode(value) is not None
+
+
+def new_managed_id(provider: str, raw_provider_id: str) -> str:
+ """Mint a fresh managed ID for a given raw provider ID."""
+ return encode(provider, str(_uuid_mod.uuid4()), raw_provider_id)
diff --git a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py
new file mode 100644
index 00000000000..a267c97c0e8
--- /dev/null
+++ b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py
@@ -0,0 +1,1234 @@
+"""
+Rewrite passthrough-managed IDs in pass-through endpoint requests and responses.
+
+OUTPUT (response) path
+----------------------
+``rewrite_response_ids()`` is called after the upstream response is received.
+It looks up the (provider, method, path) combination in ``BUILTIN_OUTPUT_ID_FIELD_MAP``,
+mints a managed ID for each listed field whose raw provider value is present,
+stores / reuses a DB row (dedup), and swaps the value in the body before the
+response is returned to the client.
+
+INPUT (request) path
+--------------------
+``rewrite_path_ids()``, ``rewrite_query_ids()``, and ``rewrite_body_ids()``
+are called just before the request is forwarded upstream. Each one walks its
+respective location (URL path, query params, JSON body) and calls
+``_resolve_one()`` for every string that looks like a passthrough managed ID
+(decode-first detection). ``_resolve_one()`` enforces:
+
+ 1. Cross-route check: the provider embedded in the ID must match the current
+ route's provider, else HTTPException(404).
+ 2. DB existence check: unknown / forged IDs raise HTTPException(404); the
+ raw string is NEVER forwarded to upstream.
+ 3. Access check: ``can_access_resource()`` raises HTTPException(403) on
+ mismatch.
+
+When a value does not decode as a passthrough managed ID it is passed through
+untouched (deliberate opt-out for raw OpenAI IDs).
+"""
+
+from __future__ import annotations
+
+import json
+import re
+from typing import Any, Dict, FrozenSet, List, Optional, Tuple
+from urllib.parse import quote, unquote
+
+from fastapi import HTTPException
+
+from litellm._logging import verbose_proxy_logger
+from litellm.llms.base_llm.managed_resources.isolation import (
+ build_owner_filter,
+ can_access_resource,
+)
+from litellm.proxy._types import UserAPIKeyAuth
+from litellm.types.llms.openai import OpenAIFileObject
+
+from .managed_id_codec import ManagedIdPayload, decode, is_managed, new_managed_id
+
+# ---------------------------------------------------------------------------
+# Field map
+# ---------------------------------------------------------------------------
+
+_FieldSpec = Tuple[str, str] # (field_name, expected_raw_id_prefix)
+_MapKey = Tuple[str, str, str] # (provider, HTTP_METHOD, canonical_path)
+
+# ``canonical_path`` uses ``/v1/...`` form without any ``/openai/`` prefix.
+# Both ``/openai/...`` and ``/openai_passthrough/...`` are normalised by
+# ``_canonical_path()`` before the lookup so only one set of entries is needed.
+BUILTIN_OUTPUT_ID_FIELD_MAP: Dict[_MapKey, List[_FieldSpec]] = {
+ # ------------------------------------------------------------------ files
+ ("openai", "POST", "/v1/files"): [
+ ("id", "file-"),
+ ],
+ ("openai", "GET", "/v1/files/{file_id}"): [
+ ("id", "file-"),
+ ],
+ ("openai", "DELETE", "/v1/files/{file_id}"): [
+ ("id", "file-"),
+ ],
+ # ----------------------------------------------------------------- batches
+ ("openai", "POST", "/v1/batches"): [
+ ("id", "batch_"),
+ ("input_file_id", "file-"),
+ ("output_file_id", "file-"),
+ ("error_file_id", "file-"),
+ ],
+ ("openai", "GET", "/v1/batches/{batch_id}"): [
+ ("id", "batch_"),
+ ("input_file_id", "file-"),
+ ("output_file_id", "file-"),
+ ("error_file_id", "file-"),
+ ],
+ ("openai", "POST", "/v1/batches/{batch_id}/cancel"): [
+ ("id", "batch_"),
+ ("input_file_id", "file-"),
+ ("output_file_id", "file-"),
+ ("error_file_id", "file-"),
+ ],
+ # --------------------------------------------------------------- responses
+ ("openai", "POST", "/v1/responses"): [
+ ("id", "resp_"),
+ ],
+ ("openai", "GET", "/v1/responses/{response_id}"): [
+ ("id", "resp_"),
+ ],
+ ("openai", "DELETE", "/v1/responses/{response_id}"): [
+ ("id", "resp_"),
+ ],
+ # ================================================================ azure
+ # Azure OpenAI exposes the same files/batches surface as OpenAI.
+ # IDs are scoped to "azure" so they are never confused with "openai" ones.
+ # ------------------------------------------------------------------ files
+ ("azure", "POST", "/v1/files"): [
+ ("id", "file-"),
+ ],
+ ("azure", "GET", "/v1/files/{file_id}"): [
+ ("id", "file-"),
+ ],
+ ("azure", "DELETE", "/v1/files/{file_id}"): [
+ ("id", "file-"),
+ ],
+ # ----------------------------------------------------------------- batches
+ ("azure", "POST", "/v1/batches"): [
+ ("id", "batch_"),
+ ("input_file_id", "file-"),
+ ("output_file_id", "file-"),
+ ("error_file_id", "file-"),
+ ],
+ ("azure", "GET", "/v1/batches/{batch_id}"): [
+ ("id", "batch_"),
+ ("input_file_id", "file-"),
+ ("output_file_id", "file-"),
+ ("error_file_id", "file-"),
+ ],
+ ("azure", "POST", "/v1/batches/{batch_id}/cancel"): [
+ ("id", "batch_"),
+ ("input_file_id", "file-"),
+ ("output_file_id", "file-"),
+ ("error_file_id", "file-"),
+ ],
+ # --------------------------------------------------------------- responses
+ ("azure", "POST", "/v1/responses"): [
+ ("id", "resp_"),
+ ],
+ ("azure", "GET", "/v1/responses/{response_id}"): [
+ ("id", "resp_"),
+ ],
+ ("azure", "DELETE", "/v1/responses/{response_id}"): [
+ ("id", "resp_"),
+ ],
+}
+
+# Prefixes that live in the *file* table rather than the object table.
+_FILE_PREFIXES: FrozenSet[str] = frozenset({"file-"})
+
+# Raw provider-ID prefixes that live in the object table (batches, responses).
+_OBJECT_PREFIXES: FrozenSet[str] = frozenset({"batch_", "resp_"})
+
+# Guards request-body rewriting against stack exhaustion from adversarially
+# deep payloads. Real OpenAI files/batches bodies nest only a few levels.
+_MAX_BODY_REWRITE_DEPTH = 64
+
+# Caps the distinct raw-provider-id guard lookups issued per request. A raw
+# file-id guard is an unindexed array-containment scan over
+# LiteLLM_ManagedFileTable (flat_model_file_ids has no index), so a body packed
+# with id-shaped strings could otherwise amplify one request into thousands of
+# full-table scans. Legitimate callers reference managed IDs (resolved via an
+# indexed lookup, never the guard), so guarding more raw ids than this only
+# happens under abuse; the request is rejected rather than skipping the guard.
+_MAX_RAW_ID_GUARD_LOOKUPS = 100
+
+
+class _RawIdGuardBudget:
+ """Per-request de-dupe + cap for raw-provider-id guard DB lookups."""
+
+ __slots__ = ("_remaining", "_seen")
+
+ def __init__(self, limit: int = _MAX_RAW_ID_GUARD_LOOKUPS) -> None:
+ self._remaining = limit
+ self._seen: set = set()
+
+ def reserve(self, raw_id: str) -> bool:
+ """Return True when a guard lookup for *raw_id* should run. Returns
+ False for a raw id already checked this request (de-dupe). Raises
+ ``HTTPException(400)`` once the per-request lookup budget is exhausted."""
+ if raw_id in self._seen:
+ return False
+ if self._remaining <= 0:
+ raise HTTPException(
+ status_code=400,
+ detail="Too many resource identifiers in request.",
+ )
+ self._remaining -= 1
+ self._seen.add(raw_id)
+ return True
+
+
+# ---------------------------------------------------------------------------
+# List routes — GET requests that return a paginated {object:"list", data:[…]}
+# These are intercepted and served entirely from the DB rather than forwarded
+# to the upstream provider, so each caller only sees IDs they own.
+# ---------------------------------------------------------------------------
+
+# Maps (provider, canonical_path) -> "files" | "batches"
+_LIST_ROUTE_TABLE: Dict[Tuple[str, str], str] = {
+ ("openai", "/v1/files"): "files",
+ ("openai", "/v1/batches"): "batches",
+ ("azure", "/v1/files"): "files",
+ ("azure", "/v1/batches"): "batches",
+}
+
+
+# Sentinel model_id written to model_mappings for passthrough-created rows.
+# Prevents the unified-endpoint deployment-resolution path from ever finding a
+# real deployment, so a passthrough ID replayed on a unified endpoint fails
+# cleanly (no silent raw-ID leak).
+def _passthrough_sentinel_model_id(provider: str) -> str:
+ return f"_passthrough_{provider}"
+
+
+# Key under which the provider marker is stored in a file row's model_mappings.
+# Its value lands in flat_model_file_ids (built from model_mappings.values()),
+# giving the file table a DB-queryable provider scope it otherwise lacks.
+_PASSTHROUGH_PROVIDER_MARKER_KEY = "_passthrough_provider_marker"
+
+
+def _passthrough_provider_marker(provider: str) -> str:
+ return f"_passthrough_provider:{provider}"
+
+
+def _managed_id_matches_provider(unified_id: str, provider: str) -> bool:
+ payload = decode(unified_id)
+ return payload is not None and payload.provider == provider
+
+
+# Strip /openai or /openai_passthrough prefix to produce canonical /v1/... path.
+# Strips provider-specific passthrough prefixes before the /v1/... path:
+# /openai_passthrough/v1/files -> /v1/files
+# /openai/v1/files -> /v1/files
+# /azure/openai/files -> /files (_canonical_path then prepends /v1/)
+# /azure_ai/openai/files -> /files
+_PASSTHROUGH_PREFIX_RE = re.compile(
+ r"^/(?:azure(?:_ai)?/)?openai(?:_passthrough)?(?=/|$)"
+)
+
+
+def _canonical_path(route: str) -> str:
+ """
+ Normalise a passthrough route to a bare /v1/... path for map lookup.
+
+ Examples:
+ /openai_passthrough/v1/files -> /v1/files
+ /openai/v1/files -> /v1/files
+ /azure/openai/files -> /v1/files (Azure omits /v1/)
+ /azure/openai/batches/batch_x -> /v1/batches/batch_x
+ """
+ stripped = _PASSTHROUGH_PREFIX_RE.sub("", route) or "/"
+ # Azure API paths don't include /v1/ — add it so they match the map keys.
+ if not stripped.startswith("/v1/") and stripped != "/":
+ stripped = "/v1" + stripped
+ return stripped
+
+
+# ---------------------------------------------------------------------------
+# Shared resolver — used by all INPUT path extractors
+# ---------------------------------------------------------------------------
+
+
+async def _resolve_one(
+ managed_id: str,
+ provider: str,
+ user_api_key_dict: UserAPIKeyAuth,
+ prisma_client: Any,
+ managed_files_hook: Any,
+) -> str:
+ """
+ Resolve a single value that may be a passthrough managed ID.
+
+ Returns the raw provider ID on success.
+ Returns *managed_id* unchanged when it is NOT a managed ID so callers
+ need not pre-filter.
+ Raises ``HTTPException(403)`` on access denial.
+ Raises ``HTTPException(404)`` on unknown / forged managed IDs — never
+ forwarded upstream as a literal string.
+ """
+ payload: Optional[ManagedIdPayload] = decode(managed_id)
+ if payload is None:
+ return managed_id # not a passthrough managed ID; pass through
+ verbose_proxy_logger.debug(
+ "managed_id_rewriter: resolving managed id provider=%s raw_prefix=%s",
+ provider,
+ (
+ payload.raw_provider_id.split("_", 1)[0]
+ if "_" in payload.raw_provider_id
+ else payload.raw_provider_id.split("-", 1)[0]
+ ),
+ )
+
+ # 1. Cross-route (cross-provider) check
+ if payload.provider != provider:
+ raise HTTPException(
+ status_code=404,
+ detail=(
+ f"Managed ID was minted for provider '{payload.provider}', "
+ f"not '{provider}'."
+ ),
+ )
+
+ row_created_by: Optional[str] = None
+ row_team_id: Optional[str] = None
+ found = False
+
+ raw_id = payload.raw_provider_id
+
+ # 2. DB lookup — pick table based on raw ID prefix
+ if any(raw_id.startswith(p) for p in _FILE_PREFIXES):
+ # File table — use hook's internal cache for speed when available
+ if managed_files_hook is not None:
+ try:
+ file_row = await managed_files_hook.get_unified_file_id(
+ managed_id,
+ litellm_parent_otel_span=None,
+ )
+ if file_row is not None:
+ row_created_by = file_row.created_by
+ row_team_id = file_row.team_id
+ found = True
+ except Exception:
+ verbose_proxy_logger.debug(
+ "managed_id_rewriter._resolve_one: file hook lookup failed",
+ exc_info=True,
+ )
+ if not found and prisma_client is not None:
+ try:
+ db_row = await prisma_client.db.litellm_managedfiletable.find_first(
+ where={"unified_file_id": managed_id}
+ )
+ if db_row is not None:
+ row_created_by = db_row.created_by
+ row_team_id = db_row.team_id
+ found = True
+ except Exception:
+ verbose_proxy_logger.debug(
+ "managed_id_rewriter._resolve_one: file DB lookup failed",
+ exc_info=True,
+ )
+ else:
+ # Object table (batches, responses)
+ if prisma_client is not None:
+ try:
+ obj_row = await prisma_client.db.litellm_managedobjecttable.find_first(
+ where={"unified_object_id": managed_id}
+ )
+ if obj_row is not None:
+ row_created_by = obj_row.created_by
+ row_team_id = obj_row.team_id
+ found = True
+ except Exception:
+ verbose_proxy_logger.debug(
+ "managed_id_rewriter._resolve_one: object DB lookup failed",
+ exc_info=True,
+ )
+
+ # 3. Hard 404 for unknown / forged IDs — NEVER forward to upstream
+ if not found:
+ raise HTTPException(
+ status_code=404,
+ detail="Managed resource not found.",
+ )
+
+ # 4. Access check
+ if not can_access_resource(user_api_key_dict, row_created_by, row_team_id):
+ raise HTTPException(
+ status_code=403,
+ detail="Access denied to managed resource.",
+ )
+
+ return payload.raw_provider_id
+
+
+async def _guard_raw_provider_id(
+ raw_id: str,
+ provider: str,
+ user_api_key_dict: UserAPIKeyAuth,
+ prisma_client: Any,
+ budget: Optional[_RawIdGuardBudget] = None,
+) -> None:
+ """Deny a raw provider ID that maps to a managed resource the caller does
+ not own, before it is forwarded upstream.
+
+ Clients only ever receive managed IDs (response bodies are rewritten), so a
+ raw provider ID for another tenant's managed resource can only have been
+ recovered by decoding that tenant's managed ID. Raw IDs are otherwise
+ forwarded untouched (deliberate opt-out), which on a retrieve / cancel /
+ delete would execute upstream before the response-side ownership check ever
+ runs. Resolving the access check here, on input, keeps the raw fallback
+ from becoming a cross-tenant bypass. Genuinely unmanaged raw IDs (no DB
+ row) are left untouched; ``HTTPException(404)`` mirrors the managed-ID
+ resolver so callers cannot probe which raw IDs exist.
+ """
+ if prisma_client is None:
+ return
+
+ if any(raw_id.startswith(p) for p in _FILE_PREFIXES):
+ if budget is not None and not budget.reserve(raw_id):
+ return
+ # File rows have no provider column, so fetch every row holding this raw
+ # id and scope to the current provider in the application layer (same as
+ # _mint_or_reuse_file's dedup).
+ try:
+ candidates = await prisma_client.db.litellm_managedfiletable.find_many(
+ where={"flat_model_file_ids": {"has": raw_id}},
+ )
+ except Exception:
+ verbose_proxy_logger.debug(
+ "managed_id_rewriter: raw file-id guard lookup failed", exc_info=True
+ )
+ return
+ provider_rows = [
+ row
+ for row in (candidates or [])
+ if _managed_id_matches_provider(row.unified_file_id, provider)
+ ]
+ if provider_rows and not any(
+ can_access_resource(user_api_key_dict, row.created_by, row.team_id)
+ for row in provider_rows
+ ):
+ raise HTTPException(status_code=404, detail="Managed resource not found.")
+ return
+
+ if any(raw_id.startswith(p) for p in _OBJECT_PREFIXES):
+ if budget is not None and not budget.reserve(raw_id):
+ return
+ # Object rows store model_object_id as "passthrough:{provider}:{raw}", so
+ # the lookup is exact and already provider-scoped.
+ try:
+ existing = await prisma_client.db.litellm_managedobjecttable.find_first(
+ where={"model_object_id": f"passthrough:{provider}:{raw_id}"}
+ )
+ except Exception:
+ verbose_proxy_logger.debug(
+ "managed_id_rewriter: raw object-id guard lookup failed", exc_info=True
+ )
+ return
+ if existing is not None and not can_access_resource(
+ user_api_key_dict, existing.created_by, existing.team_id
+ ):
+ raise HTTPException(status_code=404, detail="Managed resource not found.")
+
+
+# ---------------------------------------------------------------------------
+# OUTPUT path — helpers for minting and storing managed IDs
+# ---------------------------------------------------------------------------
+
+
+def _build_managed_file_object(
+ snapshot: Optional[Dict[str, Any]], managed_id: str
+) -> Optional[OpenAIFileObject]:
+ """Build an ``OpenAIFileObject`` (with the managed ID swapped in) from an
+ upstream file response so the DB-served list returns the same metadata as a
+ direct file GET. Returns ``None`` when no usable snapshot is available, in
+ which case the row is stored without metadata (previous behaviour)."""
+ if not snapshot:
+ return None
+ try:
+ return OpenAIFileObject(**{**snapshot, "id": managed_id})
+ except Exception:
+ verbose_proxy_logger.debug(
+ "managed_id_rewriter: file object snapshot incomplete; "
+ "storing file row without list metadata",
+ exc_info=True,
+ )
+ return None
+
+
+async def _mint_or_reuse_file(
+ raw_id: str,
+ provider: str,
+ user_api_key_dict: UserAPIKeyAuth,
+ prisma_client: Any,
+ managed_files_hook: Any,
+ file_object_snapshot: Optional[Dict[str, Any]] = None,
+ is_create_route: bool = True,
+) -> str:
+ """Return an existing managed file ID or mint + store a new one."""
+ if prisma_client is None and managed_files_hook is None:
+ return raw_id # no persistence available; leave raw
+
+ # Dedup + cross-tenant guard. Look up existing passthrough rows for this
+ # raw id WITHOUT scoping to the caller, so a raw file id that belongs to a
+ # different tenant is denied rather than re-minted under the caller. A raw
+ # id only reaches this OUTPUT path by skipping the managed-id input gate (raw
+ # provider ids are opt-out), so a row owned by someone else means the caller
+ # is touching another tenant's upstream file. flat_model_file_ids uses array
+ # containment (no index, acceptable at the scale managed-file features run).
+ #
+ # The file table has no provider column, so the same raw id can map to one
+ # row per provider (OpenAI and Azure both use the ``file-`` format). Fetch
+ # all matches and filter to this provider in the application layer, picking
+ # the oldest match deterministically so two providers issuing the same raw id
+ # reuse a stable row instead of minting duplicate rows on every call.
+ if prisma_client is not None:
+ try:
+ candidates = await prisma_client.db.litellm_managedfiletable.find_many(
+ where={"flat_model_file_ids": {"has": raw_id}},
+ order={"created_at": "asc"},
+ )
+ except Exception:
+ candidates = []
+ verbose_proxy_logger.debug(
+ "managed_id_rewriter: file dedup lookup failed", exc_info=True
+ )
+ provider_rows = [
+ row
+ for row in (candidates or [])
+ if _managed_id_matches_provider(row.unified_file_id, provider)
+ ]
+ owned_row = next(
+ (
+ row
+ for row in provider_rows
+ if can_access_resource(user_api_key_dict, row.created_by, row.team_id)
+ ),
+ None,
+ )
+ if owned_row is not None:
+ verbose_proxy_logger.debug(
+ "managed_id_rewriter: reusing existing managed file id for raw prefix=%s",
+ raw_id.split("-", 1)[0],
+ )
+ return owned_row.unified_file_id
+ if provider_rows:
+ if not is_create_route:
+ # Retrieve / delete: the caller supplied another owner's raw file
+ # id, so deny instead of minting a fresh managed id that would
+ # grant them cross-tenant access.
+ raise HTTPException(
+ status_code=404,
+ detail="Managed resource not found.",
+ )
+ # Create only: the caller's own upstream upload reused a raw id a
+ # different owner already holds (two upstream accounts under one
+ # provider name); the file is the caller's, so leave it unmanaged.
+ verbose_proxy_logger.debug(
+ "managed_id_rewriter: file dedup hit different owner on create; "
+ "leaving raw id unmanaged for prefix=%s",
+ raw_id.split("-", 1)[0],
+ )
+ return raw_id
+
+ # No existing row — mint a new managed ID and store it.
+ managed_id = new_managed_id(provider, raw_id)
+ verbose_proxy_logger.debug(
+ "managed_id_rewriter: minted new managed file id for raw prefix=%s",
+ raw_id.split("-", 1)[0],
+ )
+ if managed_files_hook is not None:
+ try:
+ await managed_files_hook.store_unified_file_id(
+ file_id=managed_id,
+ file_object=_build_managed_file_object(
+ file_object_snapshot, managed_id
+ ),
+ litellm_parent_otel_span=None,
+ model_mappings={
+ _passthrough_sentinel_model_id(provider): raw_id,
+ _PASSTHROUGH_PROVIDER_MARKER_KEY: _passthrough_provider_marker(
+ provider
+ ),
+ },
+ user_api_key_dict=user_api_key_dict,
+ )
+ except Exception:
+ # No row backs the minted ID, so every later resolve would 404. Fall
+ # back to the raw id (as when no persistence is available) to keep the
+ # caller's freshly-created resource reachable rather than orphaned.
+ verbose_proxy_logger.warning(
+ "managed_id_rewriter: could not persist file row; "
+ "leaving raw id unmanaged",
+ exc_info=True,
+ )
+ return raw_id
+ return managed_id
+
+
+async def _mint_or_reuse_object(
+ raw_id: str,
+ provider: str,
+ file_purpose: str,
+ body_snapshot: dict,
+ user_api_key_dict: UserAPIKeyAuth,
+ prisma_client: Any,
+ is_create_route: bool,
+) -> str:
+ """Return an existing managed object ID (batch/response) or mint + store one."""
+ if prisma_client is None:
+ return raw_id
+
+ # Namespace raw_id with provider so two providers that happen to issue
+ # the same raw batch/response ID get distinct rows. The @unique constraint
+ # on model_object_id would otherwise cause a UniqueConstraintViolation when
+ # the second provider tries to insert, silently losing the persisted mapping
+ # and causing every subsequent _resolve_one for that ID to return 404.
+ # This mirrors the pattern in container_endpoints/ownership.py which uses
+ # f"{purpose}:{provider}:{raw_id}" for the same reason.
+ namespaced_model_object_id = f"passthrough:{provider}:{raw_id}"
+
+ async def _reuse_existing(existing: Any, refresh_snapshot: bool) -> str:
+ """Resolve an already-persisted namespaced row: enforce the access
+ check, optionally refresh the snapshot, and return its managed ID."""
+ if not can_access_resource(
+ user_api_key_dict, existing.created_by, existing.team_id
+ ):
+ if not is_create_route:
+ # Retrieve / cancel / delete: the caller supplied a raw ID whose
+ # managed row belongs to someone else. A raw ID only reaches the
+ # upstream by bypassing the managed-ID input gate, so deny here
+ # instead of echoing another owner's object back to the caller.
+ raise HTTPException(
+ status_code=404,
+ detail="Managed resource not found.",
+ )
+ # Create only: the caller's upstream create just succeeded under a
+ # raw id a different owner already holds (two upstream accounts under
+ # one provider name). The object is the caller's own, so leave the raw
+ # id unmanaged rather than 404 a successful create; a new row can't be
+ # minted because model_object_id is @unique.
+ verbose_proxy_logger.debug(
+ "managed_id_rewriter: object dedup hit different owner on create; "
+ "leaving raw id unmanaged for prefix=%s",
+ raw_id.split("_", 1)[0],
+ )
+ return raw_id
+ if refresh_snapshot:
+ # Refresh the stored snapshot so DB-served list responses reflect
+ # the batch's latest state (e.g. output_file_id / error_file_id that
+ # were null at creation but populated once the batch completed).
+ try:
+ await prisma_client.db.litellm_managedobjecttable.update(
+ where={"unified_object_id": existing.unified_object_id},
+ data={
+ "file_object": json.dumps(body_snapshot),
+ "updated_by": user_api_key_dict.user_id,
+ },
+ )
+ except Exception:
+ verbose_proxy_logger.debug(
+ "managed_id_rewriter: object snapshot refresh failed",
+ exc_info=True,
+ )
+ verbose_proxy_logger.debug(
+ "managed_id_rewriter: reusing existing managed object id for raw prefix=%s",
+ raw_id.split("_", 1)[0],
+ )
+ return existing.unified_object_id
+
+ # Dedup: look up by the namespaced key — guaranteed unique per provider.
+ try:
+ existing = await prisma_client.db.litellm_managedobjecttable.find_first(
+ where={"model_object_id": namespaced_model_object_id}
+ )
+ except Exception:
+ verbose_proxy_logger.debug(
+ "managed_id_rewriter: object dedup lookup failed", exc_info=True
+ )
+ existing = None
+
+ if existing is not None:
+ return await _reuse_existing(existing, refresh_snapshot=True)
+
+ # No existing row — mint and upsert.
+ managed_id = new_managed_id(provider, raw_id)
+ verbose_proxy_logger.debug(
+ "managed_id_rewriter: minted new managed object id for raw prefix=%s",
+ raw_id.split("_", 1)[0],
+ )
+ try:
+ await prisma_client.db.litellm_managedobjecttable.upsert(
+ where={"unified_object_id": managed_id},
+ data={
+ "create": {
+ "unified_object_id": managed_id,
+ "file_object": json.dumps(body_snapshot),
+ "model_object_id": namespaced_model_object_id,
+ "file_purpose": file_purpose,
+ "created_by": user_api_key_dict.user_id,
+ "team_id": user_api_key_dict.team_id,
+ "updated_by": user_api_key_dict.user_id,
+ },
+ "update": {
+ "updated_by": user_api_key_dict.user_id,
+ },
+ },
+ )
+ except Exception:
+ # A concurrent caller may have inserted the same namespaced row between
+ # our dedup lookup and this insert (model_object_id is @unique, so the
+ # loser's create hits a UniqueConstraintViolation). Re-read it and reuse
+ # the winner's managed ID so both callers converge on one ID instead of
+ # the loser silently keeping the raw id.
+ try:
+ raced = await prisma_client.db.litellm_managedobjecttable.find_first(
+ where={"model_object_id": namespaced_model_object_id}
+ )
+ except Exception:
+ raced = None
+ if raced is not None:
+ return await _reuse_existing(raced, refresh_snapshot=False)
+ # No row backs the minted ID, so every later resolve would 404. Fall
+ # back to the raw id (as when no persistence is available) to keep the
+ # caller's freshly-created resource reachable rather than orphaned.
+ verbose_proxy_logger.warning(
+ "managed_id_rewriter: could not persist object row; "
+ "leaving raw id unmanaged",
+ exc_info=True,
+ )
+ return raw_id
+ return managed_id
+
+
+async def rewrite_response_ids(
+ provider: str,
+ method: str,
+ route: str,
+ body: dict,
+ user_api_key_dict: UserAPIKeyAuth,
+ prisma_client: Any,
+ managed_files_hook: Any,
+) -> dict:
+ """
+ Mint managed IDs for raw provider values listed in
+ ``BUILTIN_OUTPUT_ID_FIELD_MAP`` and swap them into *body*.
+
+ Returns the same *body* object (unchanged) when no map entry exists for
+ this ``(provider, method, route)`` combination.
+ Returns a shallow-copy of *body* with swapped values when any field is
+ rewritten.
+ """
+ from litellm.proxy.auth.auth_utils import normalize_request_route
+
+ # Strip passthrough prefix then normalize to get e.g. /v1/batches/{batch_id}
+ canonical = normalize_request_route(_canonical_path(route))
+ field_specs = BUILTIN_OUTPUT_ID_FIELD_MAP.get((provider, method, canonical))
+ if field_specs is None:
+ verbose_proxy_logger.debug(
+ "managed_id_rewriter: no output rewrite map for provider=%s method=%s route=%s",
+ provider,
+ method,
+ canonical,
+ )
+ return body
+
+ # Collection endpoints (POST /v1/batches, /v1/responses) carry no resource
+ # id in the path; everything else (retrieve / cancel / delete) does. Only
+ # creates may degrade to a raw id on a cross-owner collision.
+ is_create_route = "{" not in canonical
+
+ mutated = dict(body) # shallow copy; only return if something changed
+ changed = False
+
+ def _record(field_name: str, raw_value: str, managed_id: str) -> None:
+ nonlocal changed
+ if managed_id != raw_value:
+ mutated[field_name] = managed_id
+ changed = True
+ verbose_proxy_logger.debug(
+ "managed_id_rewriter: output field rewritten field=%s route=%s method=%s",
+ field_name,
+ canonical,
+ method,
+ )
+
+ # File fields are rewritten first so that nested references (e.g. a batch's
+ # input_file_id) are already managed IDs when the object snapshot is
+ # captured below — keeping the DB-served list in sync with a direct GET.
+ for field_name, expected_prefix in field_specs:
+ if expected_prefix not in _FILE_PREFIXES:
+ continue
+ raw_value = mutated.get(field_name)
+ if not isinstance(raw_value, str) or not raw_value.startswith(expected_prefix):
+ continue
+ managed_id = await _mint_or_reuse_file(
+ raw_value,
+ provider,
+ user_api_key_dict,
+ prisma_client,
+ managed_files_hook,
+ # The file's own ``id`` carries the full upstream metadata; nested
+ # references do not, so only the former is persisted as a snapshot.
+ file_object_snapshot=body if field_name == "id" else None,
+ is_create_route=is_create_route,
+ )
+ _record(field_name, raw_value, managed_id)
+
+ for field_name, expected_prefix in field_specs:
+ if expected_prefix in _FILE_PREFIXES:
+ continue
+ raw_value = mutated.get(field_name)
+ if not isinstance(raw_value, str) or not raw_value.startswith(expected_prefix):
+ continue
+ purpose = "batch" if raw_value.startswith("batch_") else "response"
+ managed_id = await _mint_or_reuse_object(
+ raw_value,
+ provider,
+ purpose,
+ mutated,
+ user_api_key_dict,
+ prisma_client,
+ is_create_route,
+ )
+ _record(field_name, raw_value, managed_id)
+
+ verbose_proxy_logger.debug(
+ "managed_id_rewriter: output rewrite completed changed=%s provider=%s method=%s route=%s",
+ changed,
+ provider,
+ method,
+ canonical,
+ )
+ return mutated if changed else body
+
+
+# ---------------------------------------------------------------------------
+# List-route interception — serve listing entirely from DB
+# ---------------------------------------------------------------------------
+
+
+def is_passthrough_list_route(provider: str, method: str, route: str) -> bool:
+ """Return True when this is a GET list route whose results should be served
+ from the DB (user-scoped) rather than forwarded upstream."""
+ if method != "GET":
+ return False
+ from litellm.proxy.auth.auth_utils import normalize_request_route
+
+ canonical = normalize_request_route(_canonical_path(route))
+ return (provider, canonical) in _LIST_ROUTE_TABLE
+
+
+def _parse_file_object(file_object: Any) -> Any:
+ """Prisma may return ``Json`` columns as either a parsed dict or the raw
+ JSON string (depending on driver / row source). Mirror the handling used
+ elsewhere (see ``openai_files_endpoints/common_utils.py``) so callers can
+ treat the result uniformly.
+ """
+ if isinstance(file_object, str):
+ try:
+ return json.loads(file_object)
+ except (TypeError, ValueError):
+ return None
+ return file_object
+
+
+def _empty_list_response() -> Dict[str, Any]:
+ return {
+ "object": "list",
+ "data": [],
+ "first_id": None,
+ "last_id": None,
+ "has_more": False,
+ }
+
+
+def _parse_list_limit(query_params: Optional[Dict[str, Any]]) -> Tuple[int, int]:
+ params = query_params or {}
+ try:
+ raw_limit = int(params.get("limit", 20))
+ except (TypeError, ValueError):
+ raw_limit = 20
+ # Fetch one extra to cheaply detect has_more.
+ return raw_limit, min(raw_limit, 100) + 1
+
+
+async def _build_list_where_with_cursor(
+ prisma_client: Any,
+ resource_kind: str,
+ provider: str,
+ owner_filter: Dict[str, Any],
+ query_params: Optional[Dict[str, Any]],
+) -> Tuple[Dict[str, Any], str]:
+ """Return a Prisma ``where`` clause and fetch order for a list query."""
+ params = query_params or {}
+ after_id: Optional[str] = params.get("after")
+ before_id: Optional[str] = params.get("before")
+ where: Dict[str, Any] = dict(owner_filter)
+ fetch_order = "desc"
+
+ cursor_id = after_id or before_id
+ # A cursor minted for a different provider would resolve to that provider's
+ # created_at boundary and silently skip/repeat this provider's rows, so
+ # ignore it and serve the unscoped first page instead.
+ if not cursor_id or not _managed_id_matches_provider(cursor_id, provider):
+ return where, fetch_order
+
+ cursor_table = (
+ prisma_client.db.litellm_managedfiletable
+ if resource_kind == "files"
+ else prisma_client.db.litellm_managedobjecttable
+ )
+ cursor_field = (
+ "unified_file_id" if resource_kind == "files" else "unified_object_id"
+ )
+ try:
+ cursor_row = await cursor_table.find_first(
+ where={**owner_filter, cursor_field: cursor_id}
+ )
+ if cursor_row is not None:
+ if after_id:
+ op = "lt"
+ else:
+ op = "gt"
+ fetch_order = "asc"
+ # created_at is not unique, so the boundary must also compare the
+ # unique id (the secondary sort key) to avoid skipping or repeating
+ # rows that share the cursor row's timestamp across a page boundary.
+ boundary = {
+ "OR": [
+ {"created_at": {op: cursor_row.created_at}},
+ {
+ "AND": [
+ {"created_at": cursor_row.created_at},
+ {cursor_field: {op: cursor_id}},
+ ]
+ },
+ ]
+ }
+ where = {"AND": [where, boundary]} if where else boundary
+ except Exception:
+ pass
+ return where, fetch_order
+
+
+async def _fetch_list_rows(
+ prisma_client: Any,
+ resource_kind: str,
+ where: Dict[str, Any],
+ fetch_order: str,
+ fetch_limit: int,
+) -> Optional[List[Any]]:
+ # created_at is not unique, so a second sort on the unique id column gives a
+ # total order, keeping the limit+1 page boundary and cursor deterministic
+ # across rows that share a created_at timestamp.
+ try:
+ if resource_kind == "files":
+ return await prisma_client.db.litellm_managedfiletable.find_many(
+ where=where,
+ order=[{"created_at": fetch_order}, {"unified_file_id": fetch_order}],
+ take=fetch_limit,
+ )
+ return await prisma_client.db.litellm_managedobjecttable.find_many(
+ where={**where, "file_purpose": "batch"},
+ order=[{"created_at": fetch_order}, {"unified_object_id": fetch_order}],
+ take=fetch_limit,
+ )
+ except Exception:
+ verbose_proxy_logger.warning(
+ "managed_id_rewriter: list DB query failed", exc_info=True
+ )
+ return None
+
+
+async def _fetch_provider_scoped_list_rows(
+ prisma_client: Any,
+ resource_kind: str,
+ provider: str,
+ where: Dict[str, Any],
+ fetch_order: str,
+ raw_limit: int,
+ fetch_limit: int,
+) -> Tuple[List[Any], bool]:
+ """Fetch one page of list rows scoped to *provider* at the DB level.
+
+ Both resource kinds carry a provider-distinguishing value that the query
+ filters on directly: object rows namespace ``model_object_id`` as
+ ``passthrough:{provider}:{raw}`` (see ``_mint_or_reuse_object``) and file
+ rows carry ``_passthrough_provider:{provider}`` in ``flat_model_file_ids``
+ (see ``_mint_or_reuse_file``), since the file table has no provider column.
+ Pushing the scope into the query means a single DB round-trip serves the
+ page, with no application-layer scanning that could truncate large pools.
+
+ A DB failure returns an empty page (fail closed) so the caller never falls
+ through to the upstream provider.
+ """
+ scoped_where = dict(where)
+ if resource_kind == "files":
+ scoped_where["flat_model_file_ids"] = {
+ "has": _passthrough_provider_marker(provider)
+ }
+ else:
+ scoped_where["model_object_id"] = {"startswith": f"passthrough:{provider}:"}
+
+ rows = await _fetch_list_rows(
+ prisma_client, resource_kind, scoped_where, fetch_order, fetch_limit
+ )
+ if rows is None:
+ return [], False
+
+ effective_limit = min(raw_limit, 100)
+ has_more = len(rows) > effective_limit
+ page = rows[:effective_limit]
+ if fetch_order == "asc":
+ page = list(reversed(page))
+ return page, has_more
+
+
+def _serialize_file_list_item(row: Any) -> Dict[str, Any]:
+ item: Dict[str, Any] = {
+ "id": row.unified_file_id,
+ "object": "file",
+ "created_at": int(row.created_at.timestamp()) if row.created_at else None,
+ }
+ file_object = _parse_file_object(row.file_object)
+ if isinstance(file_object, dict):
+ item.update(file_object)
+ item["id"] = row.unified_file_id # managed ID always wins over stored raw id
+ return item
+
+
+def _serialize_batch_list_item(row: Any) -> Dict[str, Any]:
+ item: Dict[str, Any] = {}
+ file_object = _parse_file_object(row.file_object)
+ if isinstance(file_object, dict):
+ item.update(file_object)
+ item["id"] = row.unified_object_id # managed ID always wins
+ item["object"] = "batch"
+ return item
+
+
+def _list_boundary_ids(
+ rows: List[Any], resource_kind: str
+) -> Tuple[Optional[str], Optional[str]]:
+ if not rows:
+ return None, None
+ id_attr = "unified_file_id" if resource_kind == "files" else "unified_object_id"
+ return getattr(rows[0], id_attr), getattr(rows[-1], id_attr)
+
+
+async def list_passthrough_ids_from_db(
+ provider: str,
+ route: str,
+ user_api_key_dict: UserAPIKeyAuth,
+ prisma_client: Any,
+ query_params: Optional[Dict[str, Any]] = None,
+) -> Optional[Dict[str, Any]]:
+ """Query the DB for managed IDs the caller owns and return an OpenAI-style
+ paginated list response.
+
+ Returns ``None`` when ``prisma_client`` is unavailable or the route is not
+ a recognised list route (caller should fall through to upstream).
+
+ Pagination params ``after``, ``before``, and ``limit`` are read from
+ ``query_params`` to match the OpenAI Batches / Files list API.
+
+ Ownership scoping:
+ - Proxy admins / master key: see **all** rows.
+ - Regular users: only rows matching their ``user_id`` / ``team_id``.
+ """
+ if prisma_client is None:
+ return None
+
+ from litellm.proxy.auth.auth_utils import normalize_request_route
+
+ canonical = normalize_request_route(_canonical_path(route))
+ resource_kind = _LIST_ROUTE_TABLE.get((provider, canonical))
+ if resource_kind is None:
+ return None
+
+ owner_filter = build_owner_filter(user_api_key_dict)
+ if owner_filter is None:
+ verbose_proxy_logger.warning(
+ "managed_id_rewriter: list denied — caller has no user_id or team_id"
+ )
+ return _empty_list_response()
+
+ raw_limit, fetch_limit = _parse_list_limit(query_params)
+ where, fetch_order = await _build_list_where_with_cursor(
+ prisma_client, resource_kind, provider, owner_filter, query_params
+ )
+ page, has_more = await _fetch_provider_scoped_list_rows(
+ prisma_client,
+ resource_kind,
+ provider,
+ where,
+ fetch_order,
+ raw_limit,
+ fetch_limit,
+ )
+ if resource_kind == "files":
+ data = [_serialize_file_list_item(row) for row in page]
+ else:
+ data = [_serialize_batch_list_item(row) for row in page]
+
+ first_id, last_id = _list_boundary_ids(page, resource_kind)
+ verbose_proxy_logger.debug(
+ "managed_id_rewriter: list served from DB provider=%s kind=%s count=%d admin=%s",
+ provider,
+ resource_kind,
+ len(data),
+ owner_filter == {},
+ )
+ return {
+ "object": "list",
+ "data": data,
+ "first_id": first_id,
+ "last_id": last_id,
+ "has_more": has_more,
+ }
+
+
+# ---------------------------------------------------------------------------
+# INPUT path extractors — all delegate to _resolve_one
+# ---------------------------------------------------------------------------
+
+
+async def rewrite_path_ids(
+ path: str,
+ provider: str,
+ user_api_key_dict: UserAPIKeyAuth,
+ prisma_client: Any,
+ managed_files_hook: Any,
+) -> str:
+ """
+ Walk URL path segments and resolve any passthrough managed IDs to raw
+ provider IDs. Returns *path* unchanged when no managed IDs are found.
+ """
+ budget = _RawIdGuardBudget()
+ segments = path.split("/")
+ new_segments: List[str] = []
+ changed = False
+ for seg in segments:
+ decoded_seg = unquote(seg)
+ if is_managed(decoded_seg):
+ raw = await _resolve_one(
+ decoded_seg,
+ provider,
+ user_api_key_dict,
+ prisma_client,
+ managed_files_hook,
+ )
+ new_segments.append(quote(raw, safe="-_.~"))
+ changed = True
+ else:
+ await _guard_raw_provider_id(
+ decoded_seg, provider, user_api_key_dict, prisma_client, budget
+ )
+ new_segments.append(seg)
+ if changed:
+ verbose_proxy_logger.debug(
+ "managed_id_rewriter: path ids rewritten provider=%s", provider
+ )
+ return "/".join(new_segments) if changed else path
+
+
+async def rewrite_query_ids(
+ params: Optional[Dict[str, Any]],
+ provider: str,
+ user_api_key_dict: UserAPIKeyAuth,
+ prisma_client: Any,
+ managed_files_hook: Any,
+) -> Optional[Dict[str, Any]]:
+ """
+ Walk query param values and resolve any passthrough managed IDs.
+ Returns *params* unchanged (same object) when nothing is resolved.
+ """
+ if not params:
+ return params
+ budget = _RawIdGuardBudget()
+ mutated = dict(params)
+ rewritten_keys: List[str] = []
+ for key, val in list(mutated.items()):
+ if isinstance(val, str):
+ if is_managed(val):
+ mutated[key] = await _resolve_one(
+ val, provider, user_api_key_dict, prisma_client, managed_files_hook
+ )
+ rewritten_keys.append(key)
+ else:
+ await _guard_raw_provider_id(
+ val, provider, user_api_key_dict, prisma_client, budget
+ )
+ if rewritten_keys:
+ verbose_proxy_logger.debug(
+ "managed_id_rewriter: query ids rewritten provider=%s keys=%s",
+ provider,
+ rewritten_keys,
+ )
+ return mutated if rewritten_keys else params
+
+
+async def rewrite_body_ids(
+ body: Optional[Dict[str, Any]],
+ provider: str,
+ user_api_key_dict: UserAPIKeyAuth,
+ prisma_client: Any,
+ managed_files_hook: Any,
+) -> Optional[Dict[str, Any]]:
+ """
+ Recursively walk a request body dict/list and resolve any passthrough
+ managed IDs. Skips litellm internal keys (``litellm_*``).
+ Returns *body* unchanged (same object) when nothing is resolved.
+ """
+ if not body:
+ return body
+
+ budget = _RawIdGuardBudget()
+
+ async def _walk(node: Any, depth: int) -> Any:
+ if depth >= _MAX_BODY_REWRITE_DEPTH:
+ return node
+ if isinstance(node, dict):
+ result: Dict[str, Any] = {}
+ changed_inner = False
+ for k, v in node.items():
+ # Skip litellm internal injection keys (e.g. litellm_logging_obj)
+ if isinstance(k, str) and k.startswith("litellm_"):
+ result[k] = v
+ continue
+ new_v = await _walk(v, depth + 1)
+ result[k] = new_v
+ if new_v is not v:
+ changed_inner = True
+ return result if changed_inner else node
+ elif isinstance(node, list):
+ new_list = [await _walk(item, depth + 1) for item in node]
+ if any(n is not o for n, o in zip(new_list, node)):
+ return new_list
+ return node
+ elif isinstance(node, str):
+ if is_managed(node):
+ return await _resolve_one(
+ node, provider, user_api_key_dict, prisma_client, managed_files_hook
+ )
+ await _guard_raw_provider_id(
+ node, provider, user_api_key_dict, prisma_client, budget
+ )
+ return node
+ return node
+
+ rewritten = await _walk(body, 0)
+ if rewritten is not body:
+ verbose_proxy_logger.debug(
+ "managed_id_rewriter: body ids rewritten provider=%s", provider
+ )
+ return rewritten
diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
index 5aa0f6cb184..f08e021630d 100644
--- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
+++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
@@ -1,3065 +1,3252 @@
-import ast
-import asyncio
-import copy
-import json
-import posixpath
-import traceback
-from base64 import b64encode
-from datetime import datetime
-from typing import Any, Dict, List, Mapping, Optional, Tuple, Union, cast
-from urllib.parse import urlencode, urlparse
-
-import httpx
-from fastapi import (
- APIRouter,
- Depends,
- FastAPI,
- HTTPException,
- Request,
- Response,
- UploadFile,
- WebSocket,
- status,
-)
-from fastapi.responses import StreamingResponse
-from starlette.datastructures import UploadFile as StarletteUploadFile
-from starlette.websockets import WebSocketState
-from websockets.asyncio.client import connect
-from websockets.exceptions import (
- ConnectionClosedError,
- ConnectionClosedOK,
- InvalidStatus,
-)
-
-import litellm
-from litellm._logging import verbose_proxy_logger
-from litellm._uuid import uuid
-from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG
-from litellm.integrations.custom_logger import CustomLogger
-from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
-from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
-from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
-from litellm.passthrough import BasePassthroughUtils
-from litellm.proxy._types import (
- ConfigFieldInfo,
- ConfigFieldUpdate,
- LiteLLMRoutes,
- PassThroughEndpointResponse,
- PassThroughGenericEndpoint,
- ProxyException,
- UserAPIKeyAuth,
-)
-from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
-from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
-from litellm.proxy.common_utils.http_parsing_utils import (
- _read_request_body,
- _safe_get_request_headers,
-)
-from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
-from litellm.proxy.utils import get_server_root_path, normalize_route_for_root_path
-from litellm.secret_managers.main import get_secret_str
-from litellm.types.llms.custom_http import httpxSpecialProvider
-from litellm.types.passthrough_endpoints.pass_through_endpoints import (
- EndpointType,
- LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY,
- LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY,
- PassthroughStandardLoggingPayload,
-)
-
-from .streaming_handler import PassThroughStreamingHandler
-from .success_handler import PassThroughEndpointLogging
-
-router = APIRouter()
-
-pass_through_endpoint_logging = PassThroughEndpointLogging()
-
-# Global registry to track registered pass-through routes and prevent memory leaks
-_registered_pass_through_routes: Dict[
- str, Dict[str, Union[str, List[str], Dict[str, Any]]]
-] = {}
-
-
-def get_response_body(response: httpx.Response) -> Optional[dict]:
- try:
- return response.json()
- except Exception:
- return None
-
-
-async def set_env_variables_in_header(custom_headers: Optional[dict]) -> Optional[dict]:
- """
- checks if any headers on config.yaml are defined as os.environ/COHERE_API_KEY etc
-
- only runs for headers defined on config.yaml
-
- example header can be
-
- {"Authorization": "Bearer os.environ/COHERE_API_KEY"}
- """
- if custom_headers is None:
- return None
- headers = {}
- for key, value in custom_headers.items():
- # langfuse Api requires base64 encoded headers - it's simpleer to just ask litellm users to set their langfuse public and secret keys
- # we can then get the b64 encoded keys here
- if key == "LANGFUSE_PUBLIC_KEY" or key == "LANGFUSE_SECRET_KEY":
- # langfuse requires b64 encoded headers - we construct that here
- _langfuse_public_key = custom_headers["LANGFUSE_PUBLIC_KEY"]
- _langfuse_secret_key = custom_headers["LANGFUSE_SECRET_KEY"]
- if isinstance(
- _langfuse_public_key, str
- ) and _langfuse_public_key.startswith("os.environ/"):
- _langfuse_public_key = get_secret_str(_langfuse_public_key)
- if isinstance(
- _langfuse_secret_key, str
- ) and _langfuse_secret_key.startswith("os.environ/"):
- _langfuse_secret_key = get_secret_str(_langfuse_secret_key)
- headers["Authorization"] = "Basic " + b64encode(
- f"{_langfuse_public_key}:{_langfuse_secret_key}".encode("utf-8")
- ).decode("ascii")
- else:
- # for all other headers
- headers[key] = value
- if isinstance(value, str) and "os.environ/" in value:
- verbose_proxy_logger.debug(
- "pass through endpoint - looking up 'os.environ/' variable"
- )
- # get string section that is os.environ/
- start_index = value.find("os.environ/")
- _variable_name = value[start_index:]
-
- verbose_proxy_logger.debug(
- "pass through endpoint - getting secret for variable name: %s",
- _variable_name,
- )
- _secret_value = get_secret_str(_variable_name)
- if _secret_value is not None:
- new_value = value.replace(_variable_name, _secret_value)
- headers[key] = new_value
- return headers
-
-
-async def chat_completion_pass_through_endpoint( # noqa: PLR0915
- fastapi_response: Response,
- request: Request,
- adapter_id: str,
- user_api_key_dict: UserAPIKeyAuth,
-):
- from litellm.proxy.proxy_server import (
- add_litellm_data_to_request,
- general_settings,
- llm_router,
- proxy_config,
- proxy_logging_obj,
- user_api_base,
- user_max_tokens,
- user_model,
- user_request_timeout,
- user_temperature,
- version,
- )
-
- data = {}
- try:
- body = await request.body()
- body_str = body.decode()
- try:
- data = ast.literal_eval(body_str)
- except Exception:
- data = json.loads(body_str)
-
- data["adapter_id"] = adapter_id
-
- verbose_proxy_logger.debug(
- "Request received by LiteLLM:\n{}".format(json.dumps(data, indent=4)),
- )
- data["model"] = (
- general_settings.get("completion_model", None) # server default
- or user_model # model name passed via cli args
- or data.get("model", None) # default passed in http request
- )
- if user_model:
- data["model"] = user_model
-
- data = await add_litellm_data_to_request(
- data=data, # type: ignore
- request=request,
- general_settings=general_settings,
- user_api_key_dict=user_api_key_dict,
- version=version,
- proxy_config=proxy_config,
- )
-
- # override with user settings, these are params passed via cli
- if user_temperature:
- data["temperature"] = user_temperature
- if user_request_timeout:
- data["request_timeout"] = user_request_timeout
- if user_max_tokens:
- data["max_tokens"] = user_max_tokens
- if user_api_base:
- data["api_base"] = user_api_base
-
- ### MODEL ALIAS MAPPING ###
- # check if model name in model alias map
- # get the actual model name
- if data["model"] in litellm.model_alias_map:
- data["model"] = litellm.model_alias_map[data["model"]]
-
- # Check key-specific aliases
- if (
- isinstance(data["model"], str)
- and user_api_key_dict.aliases
- and isinstance(user_api_key_dict.aliases, dict)
- and data["model"] in user_api_key_dict.aliases
- ):
- data["model"] = user_api_key_dict.aliases[data["model"]]
-
- ### CALL HOOKS ### - modify incoming data before calling the model
- data = await proxy_logging_obj.pre_call_hook( # type: ignore
- user_api_key_dict=user_api_key_dict, data=data, call_type="text_completion"
- )
-
- ### ROUTE THE REQUESTs ###
- router_model_names = llm_router.model_names if llm_router is not None else []
- # skip router if user passed their key
- if "api_key" in data:
- llm_response = asyncio.create_task(litellm.aadapter_completion(**data))
- elif (
- llm_router is not None and data["model"] in router_model_names
- ): # model in router model list
- llm_response = asyncio.create_task(llm_router.aadapter_completion(**data))
- elif (
- llm_router is not None
- and llm_router.model_group_alias is not None
- and data["model"] in llm_router.model_group_alias
- ): # model set in model_group_alias
- llm_response = asyncio.create_task(llm_router.aadapter_completion(**data))
- elif llm_router is not None and llm_router.has_model_id(
- data["model"]
- ): # model in router model list
- llm_response = asyncio.create_task(llm_router.aadapter_completion(**data))
- elif (
- llm_router is not None
- and data["model"] not in router_model_names
- and (
- llm_router.default_deployment is not None
- or len(llm_router.pattern_router.patterns) > 0
- )
- ): # check for wildcard routes or default deployment before checking deployment_names
- llm_response = asyncio.create_task(llm_router.aadapter_completion(**data))
- elif (
- llm_router is not None and data["model"] in llm_router.deployment_names
- ): # model in router deployments, calling a specific deployment on the router (lowest priority)
- llm_response = asyncio.create_task(
- llm_router.aadapter_completion(**data, specific_deployment=True)
- )
- elif user_model is not None: # `litellm --model `
- llm_response = asyncio.create_task(litellm.aadapter_completion(**data))
- else:
- raise HTTPException(
- status_code=status.HTTP_400_BAD_REQUEST,
- detail={
- "error": "completion: Invalid model name passed in model="
- + data.get("model", "")
- },
- )
-
- # Await the llm_response task
- response = await llm_response
-
- hidden_params = getattr(response, "_hidden_params", {}) or {}
- model_id = hidden_params.get("model_id", None) or ""
- cache_key = hidden_params.get("cache_key", None) or ""
- api_base = hidden_params.get("api_base", None) or ""
- response_cost = hidden_params.get("response_cost", None) or ""
-
- ### ALERTING ###
- asyncio.create_task(
- proxy_logging_obj.update_request_status(
- litellm_call_id=data.get("litellm_call_id", ""), status="success"
- )
- )
-
- verbose_proxy_logger.debug("final response: %s", response)
-
- fastapi_response.headers.update(
- ProxyBaseLLMRequestProcessing.get_custom_headers(
- user_api_key_dict=user_api_key_dict,
- model_id=model_id,
- cache_key=cache_key,
- api_base=api_base,
- version=version,
- response_cost=response_cost,
- )
- )
-
- verbose_proxy_logger.debug("\nResponse from Litellm:\n{}".format(response))
- return response
- except Exception as e:
- await proxy_logging_obj.post_call_failure_hook(
- user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data
- )
- verbose_proxy_logger.exception(
- "litellm.proxy.proxy_server.completion(): Exception occured - {}".format(
- str(e)
- )
- )
- error_msg = f"{str(e)}"
- raise ProxyException(
- message=getattr(e, "message", error_msg),
- type=getattr(e, "type", "None"),
- param=getattr(e, "param", "None"),
- code=getattr(e, "status_code", 500),
- )
-
-
-class HttpPassThroughEndpointHelpers(BasePassthroughUtils):
- @staticmethod
- def get_response_headers(
- headers: httpx.Headers,
- litellm_call_id: Optional[str] = None,
- custom_headers: Optional[dict] = None,
- ) -> dict:
- # Exclude headers that uvicorn writes itself (server, date) and
- # encoding/length headers that don't survive re-serialization.
- # If we forward the upstream's Server header, uvicorn adds its
- # own and strict HTTP parsers (e.g. aiohttp) reject the
- # response with "Duplicate 'Server' header found".
- excluded_headers = {
- "transfer-encoding",
- "content-encoding",
- "content-length",
- "server",
- "date",
- "connection",
- "keep-alive",
- }
-
- return_headers = {
- key: value
- for key, value in headers.items()
- if key.lower() not in excluded_headers
- }
- if litellm_call_id:
- return_headers["x-litellm-call-id"] = litellm_call_id
- if custom_headers:
- # Ensure custom headers don't override actual upstream response headers or let framework defaults (like content-length: 0) interfere.
- sanitized_custom_headers = {
- key: value
- for key, value in custom_headers.items()
- if key.lower() not in excluded_headers
- }
- return_headers.update(sanitized_custom_headers)
-
- return return_headers
-
- @staticmethod
- def get_endpoint_type(url: str) -> EndpointType:
- parsed_url = urlparse(url)
- if (
- ("generateContent") in url
- or ("streamGenerateContent") in url
- or ("rawPredict") in url
- or ("streamRawPredict") in url
- ):
- return EndpointType.VERTEX_AI
- elif parsed_url.hostname == "api.anthropic.com":
- return EndpointType.ANTHROPIC
- elif (
- parsed_url.hostname == "api.openai.com"
- or parsed_url.hostname == "openai.azure.com"
- or (parsed_url.hostname and "openai.com" in parsed_url.hostname)
- ):
- return EndpointType.OPENAI
- return EndpointType.GENERIC
-
- @staticmethod
- async def _make_non_streaming_http_request(
- request: Request,
- async_client: httpx.AsyncClient,
- url: str,
- headers: dict,
- requested_query_params: Optional[dict] = None,
- custom_body: Optional[dict] = None,
- ) -> httpx.Response:
- """
- Make a non-streaming HTTP request
-
- If request is GET, don't include a JSON body
- """
- if request.method == "GET":
- response = await async_client.request(
- method=request.method,
- url=url,
- headers=headers,
- params=requested_query_params,
- )
- else:
- response = await async_client.request(
- method=request.method,
- url=url,
- headers=headers,
- params=requested_query_params,
- json=custom_body,
- )
- return response
-
- @staticmethod
- async def non_streaming_http_request_handler(
- request: Request,
- async_client: httpx.AsyncClient,
- url: httpx.URL,
- headers: dict,
- requested_query_params: Optional[dict] = None,
- _parsed_body: Optional[dict] = None,
- forward_multipart: bool = False,
- ) -> httpx.Response:
- """
- Handle non-streaming HTTP requests
-
- Handles special cases when GET requests, multipart/form-data requests, and generic httpx requests
- """
- if request.method == "GET":
- response = await async_client.request(
- method=request.method,
- url=url,
- headers=headers,
- params=requested_query_params,
- )
- elif (
- HttpPassThroughEndpointHelpers.is_multipart(request) is True
- and forward_multipart
- ):
- # Forward multipart via make_multipart_http_request even when _parsed_body is
- # non-empty (pass_through_request always injects litellm_logging_obj, etc.).
- # forward_multipart is False when custom_body was supplied (JSON body despite
- # multipart content-type) — those requests use the generic json= path.
- return await HttpPassThroughEndpointHelpers.make_multipart_http_request(
- request=request,
- async_client=async_client,
- url=url,
- headers=headers,
- requested_query_params=requested_query_params,
- )
- else:
- # Generic httpx method
- response = await async_client.request(
- method=request.method,
- url=url,
- headers=headers,
- params=requested_query_params,
- json=_parsed_body,
- )
- return response
-
- @staticmethod
- def is_multipart(request: Request) -> bool:
- """Check if the request is a multipart/form-data request"""
- return "multipart/form-data" in request.headers.get("content-type", "")
-
- @staticmethod
- async def _build_request_files_from_upload_file(
- upload_file: Union[UploadFile, StarletteUploadFile],
- ) -> Tuple[Optional[str], bytes, Optional[str]]:
- """Build a request files dict from an UploadFile object"""
- file_content = await upload_file.read()
- return (upload_file.filename, file_content, upload_file.content_type)
-
- @staticmethod
- async def make_multipart_http_request(
- request: Request,
- async_client: httpx.AsyncClient,
- url: httpx.URL,
- headers: dict,
- requested_query_params: Optional[dict] = None,
- stream: bool = False,
- ) -> httpx.Response:
- """Process multipart/form-data requests, handling both files and form fields"""
- form_data = await request.form()
- files = {}
- form_data_dict = {}
-
- for field_name, field_value in form_data.items():
- if isinstance(field_value, (StarletteUploadFile, UploadFile)):
- files[field_name] = (
- await HttpPassThroughEndpointHelpers._build_request_files_from_upload_file(
- upload_file=field_value
- )
- )
- else:
- form_data_dict[field_name] = field_value
-
- # Remove content-type header - httpx will set it correctly with the new boundary
- # when it creates the multipart body from files/data parameters
- headers_copy = headers.copy()
- headers_copy.pop("content-type", None)
-
- # httpx.AsyncClient.request() does not accept stream=; use send() for streaming.
- if stream:
- req = async_client.build_request(
- request.method,
- url,
- headers=headers_copy,
- params=requested_query_params,
- files=files,
- data=form_data_dict,
- )
- return await async_client.send(req, stream=True)
-
- return await async_client.request(
- method=request.method,
- url=url,
- headers=headers_copy,
- params=requested_query_params,
- files=files,
- data=form_data_dict,
- )
-
- @staticmethod
- def _init_kwargs_for_pass_through_endpoint(
- request: Request,
- user_api_key_dict: UserAPIKeyAuth,
- passthrough_logging_payload: PassthroughStandardLoggingPayload,
- logging_obj: LiteLLMLoggingObj,
- _parsed_body: Optional[dict] = None,
- litellm_call_id: Optional[str] = None,
- ) -> dict:
- """
- Filter out litellm params from the request body
- """
- from litellm.types.utils import all_litellm_params
-
- _parsed_body = _parsed_body or {}
-
- litellm_params_in_body = {}
- for k in all_litellm_params:
- if k in _parsed_body:
- litellm_params_in_body[k] = _parsed_body.pop(k, None)
-
- _metadata = dict(
- LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(
- user_api_key_dict=user_api_key_dict
- )
- )
-
- litellm_metadata = litellm_params_in_body.pop("litellm_metadata", None)
- metadata = litellm_params_in_body.pop("metadata", None)
- if litellm_metadata:
- _metadata.update(litellm_metadata)
- if metadata:
- _metadata.update(metadata)
-
- _metadata = _update_metadata_with_tags_in_header(
- request=request,
- metadata=_metadata,
- )
-
- # Set internal keys after merging client-supplied metadata so a request
- # body that mirrors them cannot clobber the authenticated key or the
- # real parent span.
- _metadata["user_api_key"] = user_api_key_dict.api_key
- _metadata["litellm_parent_otel_span"] = user_api_key_dict.parent_otel_span
-
- kwargs = {
- "litellm_params": {
- **litellm_params_in_body, # type: ignore
- "metadata": _metadata,
- "proxy_server_request": {
- "url": str(request.url),
- "method": request.method,
- "body": copy.copy(_parsed_body), # use copy instead of deepcopy
- "headers": request.headers,
- },
- },
- "call_type": "pass_through_endpoint",
- "litellm_call_id": litellm_call_id,
- "passthrough_logging_payload": passthrough_logging_payload,
- }
-
- logging_obj.model_call_details["passthrough_logging_payload"] = (
- passthrough_logging_payload
- )
-
- return kwargs
-
- @staticmethod
- def construct_target_url_with_subpath(
- base_target: str, subpath: str, include_subpath: Optional[bool]
- ) -> str:
- """
- Helper function to construct the full target URL with subpath handling.
-
- Args:
- base_target: The base target URL
- subpath: The captured subpath from the request
- include_subpath: Whether to include the subpath in the target URL
-
- Returns:
- The constructed full target URL
- """
- if not include_subpath:
- return base_target
-
- if not subpath:
- return base_target
-
- # Ensure base_target ends with / and subpath doesn't start with /
- if not base_target.endswith("/"):
- base_target = base_target + "/"
- if subpath.startswith("/"):
- subpath = subpath[1:]
-
- # Resolve any '..' segments in the subpath so it cannot climb above
- # the base_target prefix that the operator configured. Preserve a
- # trailing slash on the original subpath since some upstreams treat
- # `/foo` and `/foo/` as different resources.
- trailing_slash = subpath.endswith("/")
- safe_subpath = posixpath.normpath("/" + subpath).lstrip("/")
- if safe_subpath == ".":
- safe_subpath = ""
- if trailing_slash and safe_subpath and not safe_subpath.endswith("/"):
- safe_subpath += "/"
-
- return base_target + safe_subpath
-
- @staticmethod
- def join_base_and_endpoint_path(base_url: httpx.URL, endpoint_path: str) -> str:
- """
- Combine the path component of ``base_url`` with ``endpoint_path``.
-
- Preserves any path prefix configured on the base URL and resolves
- ``..`` segments in the endpoint so the result stays within the base
- path. A trailing slash on ``endpoint_path`` is preserved.
- """
- trailing_slash = endpoint_path.endswith("/")
- base_path = base_url.path or ""
- if not base_path or base_path == "/":
- normalized_endpoint = posixpath.normpath("/" + endpoint_path.lstrip("/"))
- if trailing_slash and normalized_endpoint != "/":
- normalized_endpoint += "/"
- return normalized_endpoint
-
- base_path = base_path.rstrip("/")
- clean_endpoint = endpoint_path.lstrip("/")
- combined = posixpath.normpath(base_path + "/" + clean_endpoint)
- # If normalization climbs out of the base path, fall back to base.
- if combined != base_path and not combined.startswith(base_path + "/"):
- return base_path + "/"
- if trailing_slash and not combined.endswith("/"):
- combined += "/"
- return combined
-
- @staticmethod
- def _update_stream_param_based_on_request_body(
- parsed_body: dict,
- stream: Optional[bool] = None,
- ) -> Optional[bool]:
- """
- If stream is provided in the request body, use it.
- Otherwise, use the stream parameter passed to the `pass_through_request` function
- """
- if "stream" in parsed_body:
- return parsed_body.get("stream", stream)
- return stream
-
-
-async def pass_through_request( # noqa: PLR0915
- request: Request,
- target: str,
- custom_headers: dict,
- user_api_key_dict: UserAPIKeyAuth,
- custom_body: Optional[dict] = None,
- forward_headers: Optional[bool] = False,
- merge_query_params: Optional[bool] = False,
- query_params: Optional[dict] = None,
- default_query_params: Optional[dict] = None,
- stream: Optional[bool] = None,
- cost_per_request: Optional[float] = None,
- custom_llm_provider: Optional[str] = None,
- guardrails_config: Optional[dict] = None,
-):
- """
- Pass through endpoint handler, makes the httpx request for pass-through endpoints and ensures logging hooks are called
-
- Args:
- request: The incoming request
- target: The target URL
- custom_headers: The custom headers
- user_api_key_dict: The user API key dictionary
- custom_body: The custom body
- forward_headers: Whether to forward headers
- merge_query_params: Whether to merge query params
- query_params: The query params
- default_query_params: The default query params to be applied if not overridden by client
- stream: Whether to stream the response
- cost_per_request: Optional field - cost per request to the target endpoint
- custom_llm_provider: Optional field - custom LLM provider for the endpoint
- guardrails_config: Optional field - guardrails configuration for passthrough endpoint
- """
- from litellm.exceptions import ModifyResponseException
- from litellm.litellm_core_utils.litellm_logging import Logging
- from litellm.proxy.pass_through_endpoints.passthrough_guardrails import (
- PassthroughGuardrailHandler,
- )
- from litellm.proxy.proxy_server import proxy_logging_obj
-
- #########################################################
- # Initialize variables
- #########################################################
- litellm_call_id = str(uuid.uuid4())
- url: Optional[httpx.URL] = None
-
- # parsed request body
- _parsed_body: Optional[dict] = None
- # kwargs for pass through endpoint, contains metadata, litellm_params, call_type, litellm_call_id, passthrough_logging_payload
- kwargs: Optional[dict] = None
- logging_obj: Optional[Logging] = None
-
- #########################################################
- try:
- url = httpx.URL(target)
- headers = custom_headers
- headers = HttpPassThroughEndpointHelpers.forward_headers_from_request(
- request_headers=_safe_get_request_headers(request).copy(),
- headers=headers,
- forward_headers=forward_headers,
- )
-
- # Apply default query parameters if provided, regardless of merge_query_params setting
- if default_query_params or merge_query_params:
- # Determine what to merge based on settings
- request_params = dict(request.query_params) if merge_query_params else {}
-
- # Create a new URL with the merged query params
- url = url.copy_with(
- query=urlencode(
- HttpPassThroughEndpointHelpers.get_merged_query_parameters(
- existing_url=url,
- request_query_params=request_params,
- default_query_params=default_query_params,
- )
- ).encode("ascii")
- )
-
- endpoint_type: EndpointType = HttpPassThroughEndpointHelpers.get_endpoint_type(
- str(url)
- )
-
- # SigV4-signed callers (e.g. Bedrock) attach the exact bytes that were
- # signed via request.state; we must send those instead of re-encoding the
- # parsed dict (hooks mutate it, breaking the signature / Content-Length).
- # Tolerate request objects without `state` (test fixtures) and only honor
- # values httpx accepts for `content=`.
- _request_state = getattr(request, "state", None)
- state_raw_body: Optional[Union[str, bytes]] = (
- getattr(_request_state, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, None)
- if _request_state is not None
- else None
- )
- if state_raw_body is not None and not isinstance(
- state_raw_body, (str, bytes, bytearray)
- ):
- state_raw_body = None
-
- # Skip body parsing for multipart requests - make_multipart_http_request will handle it
- # But if custom_body is provided (e.g., JSON parsed despite multipart content-type), use it
- is_multipart = (
- HttpPassThroughEndpointHelpers.is_multipart(request) and not custom_body
- )
-
- if custom_body:
- _parsed_body = custom_body
- elif is_multipart:
- # Don't parse multipart body here - it will be handled by make_multipart_http_request
- _parsed_body = {}
- else:
- _parsed_body = await _read_request_body(request)
- verbose_proxy_logger.debug(
- "Pass through endpoint sending request to \nURL {}\nheaders: {}\nbody: {}\n".format(
- url, headers, _parsed_body
- )
- )
-
- ### COLLECT GUARDRAILS FOR PASSTHROUGH ENDPOINT ###
- # Passthrough endpoints are opt-in only for guardrails
- # When enabled, collect guardrails from org/team/key levels + passthrough-specific
- guardrails_to_run = PassthroughGuardrailHandler.collect_guardrails(
- user_api_key_dict=user_api_key_dict,
- passthrough_guardrails_config=guardrails_config,
- )
-
- # Add guardrails to metadata if any should run
- if guardrails_to_run and len(guardrails_to_run) > 0:
- if _parsed_body is None:
- _parsed_body = {}
- if "metadata" not in _parsed_body:
- _parsed_body["metadata"] = {}
- _parsed_body["metadata"]["guardrails"] = guardrails_to_run
- verbose_proxy_logger.debug(
- f"Added guardrails to passthrough request metadata: {guardrails_to_run}"
- )
-
- ## LOGGING OBJECT ## - initialize before pre_call_hook so guardrails can access it
- # Surface the requested model (when the body carries one) so logging/spans
- # read e.g. ``chat gpt-4o`` instead of ``chat unknown``.
- passthrough_model = (
- _parsed_body.get("model") if isinstance(_parsed_body, dict) else None
- ) or "unknown"
- start_time = datetime.now()
- logging_obj = Logging(
- model=passthrough_model,
- messages=[{"role": "user", "content": safe_dumps(_parsed_body)}],
- stream=False,
- call_type="pass_through_endpoint",
- start_time=start_time,
- litellm_call_id=litellm_call_id,
- function_id="1245",
- )
-
- # Store passthrough guardrails config on logging_obj for field targeting
- logging_obj.passthrough_guardrails_config = guardrails_config
-
- # Store logging_obj in data so guardrails can access it
- if _parsed_body is None:
- _parsed_body = {}
- _parsed_body["litellm_logging_obj"] = logging_obj
-
- ### CALL HOOKS ### - modify incoming data / reject request before calling the model
- _parsed_body = await proxy_logging_obj.pre_call_hook(
- user_api_key_dict=user_api_key_dict,
- data=_parsed_body,
- call_type="pass_through_endpoint",
- )
- async_client_obj = get_async_httpx_client(
- llm_provider=httpxSpecialProvider.PassThroughEndpoint,
- params={"timeout": 600},
- )
- async_client = async_client_obj.client
- passthrough_logging_payload = PassthroughStandardLoggingPayload(
- url=str(url),
- request_body=_parsed_body,
- request_method=getattr(request, "method", None),
- cost_per_request=cost_per_request,
- )
- kwargs = HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint(
- user_api_key_dict=user_api_key_dict,
- _parsed_body=_parsed_body,
- passthrough_logging_payload=passthrough_logging_payload,
- litellm_call_id=litellm_call_id,
- request=request,
- logging_obj=logging_obj,
- )
-
- # Store custom_llm_provider in kwargs and logging object if provided
- if custom_llm_provider:
- logging_obj.model_call_details["custom_llm_provider"] = custom_llm_provider
- logging_obj.model_call_details["litellm_params"] = kwargs.get(
- "litellm_params", {}
- )
-
- # done for supporting 'parallel_request_limiter.py' with pass-through endpoints
- logging_obj.update_environment_variables(
- model=passthrough_model,
- user="unknown",
- optional_params={},
- litellm_params=kwargs["litellm_params"],
- call_type="pass_through_endpoint",
- )
- logging_obj.model_call_details["litellm_call_id"] = litellm_call_id
-
- # combine url with query params for logging
- requested_query_params: Optional[dict] = query_params or dict(
- request.query_params
- )
-
- requested_query_params_str = None
- if requested_query_params:
- requested_query_params_str = "&".join(
- f"{k}={v}" for k, v in requested_query_params.items()
- )
-
- logging_url = str(url)
- if requested_query_params_str:
- if "?" in str(url):
- logging_url = str(url) + "&" + requested_query_params_str
- else:
- logging_url = str(url) + "?" + requested_query_params_str
-
- logging_obj.pre_call(
- input=[{"role": "user", "content": safe_dumps(_parsed_body)}],
- api_key="",
- additional_args={
- "complete_input_dict": _parsed_body,
- "api_base": str(logging_url),
- "headers": headers,
- },
- )
- stream = (
- HttpPassThroughEndpointHelpers._update_stream_param_based_on_request_body(
- parsed_body=_parsed_body,
- stream=stream,
- )
- )
-
- if stream:
- if is_multipart:
- response = (
- await HttpPassThroughEndpointHelpers.make_multipart_http_request(
- request=request,
- async_client=async_client,
- url=url,
- headers=headers,
- requested_query_params=requested_query_params,
- stream=True,
- )
- )
- else:
- # SigV4-signed callers (Bedrock) supply the exact pre-signed bytes;
- # otherwise httpx encodes the parsed JSON dict as before.
- body_kwargs: Dict[str, Any] = (
- {"content": state_raw_body}
- if state_raw_body is not None
- else {"json": _parsed_body}
- )
- req = async_client.build_request(
- request.method,
- url,
- params=requested_query_params,
- headers=headers,
- **body_kwargs,
- )
-
- response = await async_client.send(req, stream=stream)
-
- try:
- response.raise_for_status()
- except httpx.HTTPStatusError as e:
- raise HTTPException(
- status_code=e.response.status_code, detail=await e.response.aread()
- )
-
- return StreamingResponse(
- PassThroughStreamingHandler.chunk_processor(
- response=response,
- request_body=_parsed_body,
- litellm_logging_obj=logging_obj,
- endpoint_type=endpoint_type,
- start_time=start_time,
- passthrough_success_handler_obj=pass_through_endpoint_logging,
- url_route=str(url),
- ),
- headers=HttpPassThroughEndpointHelpers.get_response_headers(
- headers=response.headers,
- litellm_call_id=litellm_call_id,
- ),
- status_code=response.status_code,
- )
-
- if state_raw_body is not None:
- # SigV4-signed callers (Bedrock) require the exact pre-signed bytes
- # to be forwarded so the signature/Content-Length stay valid.
- response = await async_client.request(
- method=request.method,
- url=url,
- headers=headers,
- params=requested_query_params,
- content=state_raw_body,
- )
- else:
- response = (
- await HttpPassThroughEndpointHelpers.non_streaming_http_request_handler(
- request=request,
- async_client=async_client,
- url=url,
- headers=headers,
- requested_query_params=requested_query_params,
- _parsed_body=_parsed_body,
- forward_multipart=is_multipart,
- )
- )
- verbose_proxy_logger.debug("response.headers= %s", response.headers)
-
- if _is_streaming_response(response) is True:
- try:
- response.raise_for_status()
- except httpx.HTTPStatusError as e:
- raise HTTPException(
- status_code=e.response.status_code, detail=await e.response.aread()
- )
-
- return StreamingResponse(
- PassThroughStreamingHandler.chunk_processor(
- response=response,
- request_body=_parsed_body,
- litellm_logging_obj=logging_obj,
- endpoint_type=endpoint_type,
- start_time=start_time,
- passthrough_success_handler_obj=pass_through_endpoint_logging,
- url_route=str(url),
- ),
- headers=HttpPassThroughEndpointHelpers.get_response_headers(
- headers=response.headers,
- litellm_call_id=litellm_call_id,
- ),
- status_code=response.status_code,
- )
-
- try:
- response.raise_for_status()
- except httpx.HTTPStatusError as e:
- raise HTTPException(
- status_code=e.response.status_code, detail=e.response.text
- )
-
- if response.status_code >= 300:
- raise HTTPException(status_code=response.status_code, detail=response.text)
-
- content = await response.aread()
-
- ## POST-CALL GUARDRAILS ##
- _content_modified = False
- response_body: Optional[dict] = get_response_body(response)
- if response_body is not None and guardrails_to_run:
- # Build an enriched data dict: _parsed_body has been stripped of
- # `metadata` by both pre_call_hook and _init_kwargs_for_pass_through_endpoint,
- # so we re-attach the configured guardrails here so should_run_guardrail
- # sees them.
- hook_data = dict(_parsed_body or {})
- existing_metadata = hook_data.get("metadata")
- if not isinstance(existing_metadata, dict):
- existing_metadata = {}
- hook_data["metadata"] = {
- **existing_metadata,
- "guardrails": guardrails_to_run,
- }
- response_body = await proxy_logging_obj.post_call_success_hook(
- data=hook_data,
- user_api_key_dict=user_api_key_dict,
- response=response_body, # type: ignore[arg-type]
- )
- if isinstance(response_body, dict):
- content = json.dumps(response_body).encode("utf-8")
- _content_modified = True
- else:
- verbose_proxy_logger.debug(
- "pass_through_endpoint: post_call_success_hook returned %s, expected dict — using original response",
- type(response_body).__name__,
- )
- elif response_body is None:
- verbose_proxy_logger.debug(
- "pass_through_endpoint: response body not JSON-parseable, skipping post-call guardrails"
- )
-
- ## LOG SUCCESS
- passthrough_logging_payload["response_body"] = response_body
- end_time = datetime.now()
- asyncio.create_task(
- pass_through_endpoint_logging.pass_through_async_success_handler(
- httpx_response=response,
- response_body=response_body,
- url_route=str(url),
- result="",
- start_time=start_time,
- end_time=end_time,
- logging_obj=logging_obj,
- cache_hit=False,
- request_body=_parsed_body,
- custom_llm_provider=custom_llm_provider,
- **kwargs,
- )
- )
-
- ## CUSTOM HEADERS - `x-litellm-*`
- custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers(
- user_api_key_dict=user_api_key_dict,
- call_id=litellm_call_id,
- model_id=None,
- cache_key=None,
- api_base=str(url._uri_reference),
- )
-
- response_headers = HttpPassThroughEndpointHelpers.get_response_headers(
- headers=response.headers,
- custom_headers=custom_headers,
- )
- if _content_modified:
- response_headers.pop("content-length", None)
-
- return Response(
- content=content,
- status_code=response.status_code,
- headers=response_headers,
- )
- except ModifyResponseException as e:
- verbose_proxy_logger.info(
- "pass_through_endpoint: Guardrail %s modified response: %s",
- e.guardrail_name,
- str(e.message or "")[:200],
- )
- try:
- await proxy_logging_obj.post_call_failure_hook(
- user_api_key_dict=user_api_key_dict,
- original_exception=e,
- request_data=e.request_data,
- )
- except Exception:
- verbose_proxy_logger.warning(
- "pass_through_endpoint: post_call_failure_hook raised during guardrail block",
- exc_info=True,
- )
- error_body = {
- "error": {
- "message": e.message or "Response blocked by guardrail",
- "type": "content_filter",
- "guardrail_name": e.guardrail_name,
- "model": e.model,
- }
- }
- return Response(
- content=json.dumps(error_body),
- status_code=200,
- media_type="application/json",
- )
- except Exception as e:
- custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers(
- user_api_key_dict=user_api_key_dict,
- call_id=litellm_call_id,
- model_id=None,
- cache_key=None,
- api_base=str(url._uri_reference) if url else None,
- )
- verbose_proxy_logger.exception(
- "litellm.proxy.proxy_server.pass_through_endpoint(): Exception occured - {}".format(
- str(e)
- )
- )
-
- #########################################################
- # Monitoring: Trigger post_call_failure_hook
- # for pass through endpoint failure
- #########################################################
- request_payload: dict = _parsed_body or {}
- # add user_api_key_dict, litellm_call_id, passthrough_logging_payloa for logging
- if kwargs:
- for key, value in kwargs.items():
- request_payload[key] = value
- if logging_obj is not None:
- request_payload["litellm_logging_obj"] = logging_obj
-
- if (
- "model" not in request_payload
- and _parsed_body
- and isinstance(_parsed_body, dict)
- ):
- request_payload["model"] = _parsed_body.get("model", "")
- if "custom_llm_provider" not in request_payload and custom_llm_provider:
- request_payload["custom_llm_provider"] = custom_llm_provider
-
- await proxy_logging_obj.post_call_failure_hook(
- user_api_key_dict=user_api_key_dict,
- original_exception=e,
- request_data=request_payload,
- traceback_str=traceback.format_exc(
- limit=MAXIMUM_TRACEBACK_LINES_TO_LOG,
- ),
- )
-
- #########################################################
-
- if isinstance(e, HTTPException):
- raise ProxyException(
- message=getattr(e, "message", str(getattr(e, "detail", str(e)))),
- type=getattr(e, "type", "None"),
- param=getattr(e, "param", "None"),
- code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST),
- headers=custom_headers,
- )
- else:
- error_msg = f"{str(e)}"
- raise ProxyException(
- message=getattr(e, "message", error_msg),
- type=getattr(e, "type", "None"),
- param=getattr(e, "param", "None"),
- code=getattr(e, "status_code", 500),
- headers=custom_headers,
- )
-
-
-def _update_metadata_with_tags_in_header(request: Request, metadata: dict) -> dict:
- """
- If tags are in the request headers, add them to the metadata
-
- Used for google and vertex JS SDKs, and Azure passthrough
- Checks both 'tags' and 'x-litellm-tags' headers
- """
- tags_to_add = []
-
- # Check for 'tags' header first
- _tags = request.headers.get("tags")
- if _tags:
- tags_to_add.extend([tag.strip() for tag in _tags.split(",")])
-
- _tags = request.headers.get("x-litellm-tags")
- if _tags:
- tags_to_add.extend([tag.strip() for tag in _tags.split(",")])
-
- # Only add tags key if there are tags to add
- if tags_to_add:
- if "tags" not in metadata:
- metadata["tags"] = []
- metadata["tags"].extend(tags_to_add)
-
- return metadata
-
-
-async def _parse_request_data_by_content_type(
- request: Request,
-) -> Tuple[Optional[Any], Optional[Any], Optional[Any], Optional[Any]]:
- """
- Parse request data based on content type.
-
- Handles JSON, multipart/form-data, and URL-encoded form data.
-
- Returns:
- Tuple of (query_params_data, custom_body_data, file_data, stream)
- """
- content_type = request.headers.get("content-type", "")
-
- query_params_data = None
- custom_body_data = None
- file_data = None
- stream = None
-
- if "application/json" in content_type:
- # ✅ Handle JSON
- try:
- body = await request.json()
- query_params_data = body.get("query_params")
- custom_body_data = body.get("custom_body")
- stream = body.get("stream")
- except json.JSONDecodeError:
- # Handle requests with no body (e.g., DELETE requests)
- pass
- elif "multipart/form-data" in content_type:
- # ✅ Try to parse as JSON first (handles misconfigured clients sending JSON with multipart content-type)
- # If that fails, skip parsing - pass_through_request will handle actual multipart
- try:
- body = await request.json()
- # Successfully parsed as JSON - treat as JSON body
- query_params_data = body.get("query_params")
- custom_body_data = body.get("custom_body")
- stream = body.get("stream")
- # If custom_body is not set, use the entire body
- if custom_body_data is None and body:
- custom_body_data = body
- except (json.JSONDecodeError, Exception):
- # Not JSON - this is actual multipart data
- # Skip parsing here to avoid consuming the request body stream
- # make_multipart_http_request will handle it
- pass
-
- elif "application/x-www-form-urlencoded" in content_type:
- # ✅ Handle URL-encoded form data
- form = await request.form()
- query_params_data = form.get("query_params")
- custom_body_data = form.get("custom_body")
-
- else:
- # ✅ Fallback: maybe no body, just query params
- query_params_data = dict(request.query_params) or None
-
- return query_params_data, custom_body_data, file_data, stream
-
-
-def create_pass_through_route(
- endpoint,
- target: str,
- custom_headers: Optional[Mapping[str, Any]] = None,
- _forward_headers: Optional[bool] = False,
- _merge_query_params: Optional[bool] = False,
- dependencies: Optional[List] = None,
- include_subpath: Optional[bool] = False,
- cost_per_request: Optional[float] = None,
- custom_llm_provider: Optional[str] = None,
- is_streaming_request: Optional[bool] = False,
- query_params: Optional[dict] = None,
- default_query_params: Optional[dict] = None,
- guardrails: Optional[Dict[str, Any]] = None,
- config_file_path: Optional[str] = None,
-):
- # check if target is an adapter.py or a url
- from litellm._uuid import uuid
- from litellm.proxy.types_utils.utils import get_instance_fn
-
- try:
- if isinstance(target, CustomLogger):
- adapter = target
- else:
- adapter = get_instance_fn(value=target, config_file_path=config_file_path)
- adapter_id = str(uuid.uuid4())
- litellm.adapters = [{"id": adapter_id, "adapter": adapter}]
-
- async def endpoint_func( # type: ignore
- request: Request,
- fastapi_response: Response,
- user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
- subpath: str = "", # captures sub-paths when include_subpath=True
- ):
- return await chat_completion_pass_through_endpoint(
- fastapi_response=fastapi_response,
- request=request,
- adapter_id=adapter_id,
- user_api_key_dict=user_api_key_dict,
- )
-
- except Exception:
- verbose_proxy_logger.debug("Defaulting to target being a url.")
-
- async def endpoint_func( # type: ignore
- request: Request,
- fastapi_response: Response,
- user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
- subpath: str = "", # captures sub-paths when include_subpath=True
- ):
- from litellm.proxy.auth.auth_utils import ( # noqa: PLC0415
- get_request_route,
- )
- from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
- InitPassThroughEndpointHelpers,
- )
-
- path = get_request_route(request)
-
- # Parse request data based on content type
- (
- query_params_data,
- custom_body_data,
- file_data,
- stream,
- ) = await _parse_request_data_by_content_type(request)
-
- if not InitPassThroughEndpointHelpers.is_registered_pass_through_route(
- route=path
- ):
- raise HTTPException(
- status_code=404,
- detail=f"Pass-through endpoint {endpoint} not found. This could have been deleted or not yet added to the proxy.",
- )
-
- passthrough_params = (
- InitPassThroughEndpointHelpers.get_registered_pass_through_route(
- route=path, method=request.method
- )
- )
- target_params = {
- "target": target,
- "custom_headers": custom_headers,
- "forward_headers": _forward_headers,
- "merge_query_params": _merge_query_params,
- "cost_per_request": cost_per_request,
- "guardrails": None,
- }
-
- if passthrough_params is not None:
- target_params.update(passthrough_params.get("passthrough_params", {}))
-
- # Extract and cast parameters with proper types
- param_target = target_params.get("target") or target
- param_custom_headers = target_params.get("custom_headers", custom_headers)
- param_forward_headers = target_params.get(
- "forward_headers", _forward_headers
- )
- param_merge_query_params = target_params.get(
- "merge_query_params", _merge_query_params
- )
- param_cost_per_request = target_params.get(
- "cost_per_request", cost_per_request
- )
- param_guardrails = target_params.get("guardrails", None)
- param_default_query_params = target_params.get("default_query_params", None)
-
- # Construct the full target URL with subpath if needed
- full_target = (
- HttpPassThroughEndpointHelpers.construct_target_url_with_subpath(
- base_target=cast(str, param_target),
- subpath=subpath,
- include_subpath=include_subpath,
- )
- )
-
- # Ensure custom_headers is a dict. Botocore returns a HeadersDict
- # for SigV4-prepared requests, which is a Mapping but not a dict.
- headers_dict = (
- dict(param_custom_headers)
- if isinstance(param_custom_headers, Mapping)
- else {}
- )
-
- # Ensure query_params and custom_body are dicts or None
- final_query_params = (
- query_params_data if isinstance(query_params_data, dict) else {}
- )
- if query_params:
- final_query_params.update(query_params)
- # Programmatic callers set LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY on
- # request.state (see Bedrock proxy). Parsed JSON envelope otherwise.
- state_custom_body: Optional[dict] = getattr(
- request.state,
- LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY,
- None,
- )
- final_custom_body: Optional[dict] = None
- if isinstance(state_custom_body, dict):
- final_custom_body = state_custom_body
- elif isinstance(custom_body_data, dict):
- final_custom_body = custom_body_data
-
- try:
- return await pass_through_request( # type: ignore
- request=request,
- target=full_target,
- custom_headers=headers_dict,
- user_api_key_dict=user_api_key_dict,
- forward_headers=cast(Optional[bool], param_forward_headers),
- merge_query_params=cast(Optional[bool], param_merge_query_params),
- query_params=final_query_params,
- default_query_params=cast(
- Optional[dict], param_default_query_params
- ),
- stream=is_streaming_request or stream,
- custom_body=final_custom_body,
- cost_per_request=cast(Optional[float], param_cost_per_request),
- custom_llm_provider=custom_llm_provider,
- guardrails_config=cast(Optional[dict], param_guardrails),
- )
- finally:
- if hasattr(request.state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY):
- delattr(request.state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY)
- if hasattr(request.state, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY):
- delattr(request.state, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY)
-
- return endpoint_func
-
-
-def create_websocket_passthrough_route(
- endpoint: str,
- target: str,
- custom_headers: Optional[dict] = None,
- _forward_headers: Optional[bool] = False,
- dependencies: Optional[List] = None,
- cost_per_request: Optional[float] = None,
-):
- """
- Create a WebSocket passthrough route function.
-
- Args:
- endpoint: The endpoint path (for logging purposes)
- target: The target WebSocket URL (e.g., "wss://api.example.com/ws")
- custom_headers: Custom headers to include in the WebSocket connection
- _forward_headers: Whether to forward incoming headers
- dependencies: FastAPI dependencies to inject
-
- Returns:
- A WebSocket passthrough function that can be registered with app.websocket()
- """
- from litellm.proxy.auth.user_api_key_auth import user_api_key_auth_websocket
-
- async def websocket_endpoint_func(
- websocket: WebSocket,
- user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth_websocket),
- **kwargs, # For additional query parameters
- ):
- """
- WebSocket passthrough endpoint function.
-
- This function handles the WebSocket connection by:
- 1. Accepting the incoming WebSocket connection
- 2. Establishing a connection to the target WebSocket
- 3. Forwarding messages bidirectionally
- 4. Handling connection cleanup
- """
- return await websocket_passthrough_request(
- websocket=websocket,
- target=target,
- custom_headers=custom_headers or {},
- user_api_key_dict=user_api_key_dict,
- forward_headers=_forward_headers,
- endpoint=endpoint,
- cost_per_request=cost_per_request,
- accept_websocket=True, # Generic usage should accept the WebSocket
- )
-
- return websocket_endpoint_func
-
-
-async def websocket_passthrough_request( # noqa: PLR0915
- websocket: WebSocket,
- target: str,
- custom_headers: dict,
- user_api_key_dict: UserAPIKeyAuth,
- forward_headers: Optional[bool] = False,
- endpoint: Optional[str] = None,
- cost_per_request: Optional[float] = None,
- accept_websocket: bool = True,
-):
- """
- WebSocket passthrough request handler.
-
- Args:
- websocket: The incoming WebSocket connection
- target: The target WebSocket URL
- custom_headers: Custom headers to include in the connection
- user_api_key_dict: The user API key dictionary
- forward_headers: Whether to forward incoming headers
- endpoint: The endpoint path (for logging purposes)
- cost_per_request: Optional field - cost per request to the target endpoint
- """
- from litellm.litellm_core_utils.litellm_logging import Logging
- from litellm.proxy.proxy_server import proxy_logging_obj
- from litellm.types.passthrough_endpoints.pass_through_endpoints import (
- PassthroughStandardLoggingPayload,
- )
-
- # Initialize tracking variables
- start_time = datetime.now()
- websocket_messages: list[dict[str, Any]] = []
- litellm_call_id = str(uuid.uuid4())
-
- verbose_proxy_logger.info(
- f"WebSocket passthrough ({endpoint}): Starting WebSocket connection to {target}"
- )
-
- # Only accept the WebSocket if requested (for generic usage)
- if accept_websocket:
- await websocket.accept()
- verbose_proxy_logger.debug(
- f"WebSocket passthrough ({endpoint}): WebSocket connection accepted"
- )
-
- # Prepare headers for the upstream connection
- upstream_headers = custom_headers.copy()
-
- if forward_headers:
- # Forward relevant headers from the incoming request
- incoming_headers = dict(websocket.headers)
- for header_name, header_value in incoming_headers.items():
- # Only forward certain headers to avoid conflicts
- if header_name.lower() in [
- "authorization",
- "x-api-key",
- "x-goog-user-project",
- ]:
- upstream_headers[header_name] = header_value
-
- # Initialize logging object similar to HTTP passthrough
- logging_obj = Logging(
- model="unknown",
- messages=[{"role": "user", "content": "WebSocket connection"}],
- stream=True, # WebSockets are inherently streaming
- call_type="pass_through_endpoint",
- start_time=start_time,
- litellm_call_id=litellm_call_id,
- function_id="websocket_passthrough",
- )
-
- # Create passthrough logging payload
- passthrough_logging_payload = PassthroughStandardLoggingPayload(
- url=target,
- request_body={}, # WebSocket doesn't have a traditional request body
- request_method="WEBSOCKET",
- cost_per_request=cost_per_request,
- )
-
- # Create a dummy request object for WebSocket connections to maintain compatibility
- # with the existing _init_kwargs_for_pass_through_endpoint function
- class DummyRequest:
- def __init__(
- self, url: str, method: str = "WEBSOCKET", headers: Optional[dict] = None
- ):
- self.url = url
- self.method = method
- self.headers = headers or {}
-
- def __str__(self):
- return f"DummyRequest(url={self.url}, method={self.method})"
-
- dummy_request = DummyRequest(
- url=target,
- method="WEBSOCKET",
- headers=dict(websocket.headers) if hasattr(websocket, "headers") else {},
- )
-
- # Initialize kwargs for logging using the same pattern as HTTP passthrough
- kwargs = HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint(
- user_api_key_dict=user_api_key_dict,
- _parsed_body={}, # WebSocket doesn't have a traditional request body
- passthrough_logging_payload=passthrough_logging_payload,
- litellm_call_id=litellm_call_id,
- request=dummy_request, # type: ignore
- logging_obj=logging_obj,
- )
-
- # Update logging environment variables
- logging_obj.update_environment_variables(
- model="unknown",
- user="unknown",
- optional_params={},
- litellm_params=dict(kwargs.get("litellm_params", {})),
- call_type="pass_through_endpoint",
- )
- logging_obj.model_call_details["litellm_call_id"] = litellm_call_id
-
- # Pre-call logging
- logging_obj.pre_call(
- input=[{"role": "user", "content": "WebSocket connection"}],
- api_key="",
- additional_args={
- "complete_input_dict": {},
- "api_base": target,
- "headers": upstream_headers,
- },
- )
-
- ### CALL HOOKS ### - modify incoming data / reject request before calling the model
- websocket_data: dict[str, Any] = {}
- websocket_data = await proxy_logging_obj.pre_call_hook(
- user_api_key_dict=user_api_key_dict,
- data=websocket_data,
- call_type="pass_through_endpoint",
- )
-
- try:
- verbose_proxy_logger.debug(
- f"WebSocket passthrough ({endpoint}): Establishing upstream connection to {target}"
- )
- async with connect(
- target,
- additional_headers=upstream_headers,
- ) as upstream_ws:
- verbose_proxy_logger.info(
- f"WebSocket passthrough ({endpoint}): Upstream connection established successfully"
- )
-
- async def forward_client_to_upstream() -> None:
- """Forward messages from client to upstream WebSocket"""
- try:
- while True:
- message = await websocket.receive()
- message_type = message.get("type")
- if message_type == "websocket.disconnect":
- await upstream_ws.close()
- break
-
- text_data = message.get("text")
- bytes_data = message.get("bytes")
-
- if text_data is not None:
- # Try to extract model from client setup message for Vertex AI Live
- if endpoint and "/vertex_ai/live" in endpoint:
- verbose_proxy_logger.debug(
- f"WebSocket passthrough ({endpoint}): Processing client message for model extraction"
- )
- try:
- client_message = json.loads(text_data)
- if (
- isinstance(client_message, dict)
- and "setup" in client_message
- ):
- setup_data = client_message["setup"]
- verbose_proxy_logger.debug(
- f"WebSocket passthrough ({endpoint}): Found setup data in client message: {setup_data}"
- )
- if (
- isinstance(setup_data, dict)
- and "model" in setup_data
- ):
- extracted_model = (
- _extract_model_from_vertex_ai_setup(
- setup_data
- )
- )
- if extracted_model:
- kwargs["model"] = extracted_model
- kwargs["custom_llm_provider"] = (
- "vertex_ai-language-models"
- )
- # Update logging object with correct model
- logging_obj.model = extracted_model
- logging_obj.model_call_details[
- "model"
- ] = extracted_model
- logging_obj.model_call_details[
- "custom_llm_provider"
- ] = "vertex_ai"
- verbose_proxy_logger.info(
- f"WebSocket passthrough ({endpoint}): Successfully extracted model '{extracted_model}' and set provider to 'vertex_ai' from client setup message"
- )
- else:
- verbose_proxy_logger.warning(
- f"WebSocket passthrough ({endpoint}): Failed to extract model from client setup data: {setup_data}"
- )
- else:
- verbose_proxy_logger.debug(
- f"WebSocket passthrough ({endpoint}): Setup data does not contain model field: {setup_data}"
- )
- else:
- verbose_proxy_logger.debug(
- f"WebSocket passthrough ({endpoint}): Client message does not contain setup data"
- )
- except (json.JSONDecodeError, KeyError, TypeError) as e:
- verbose_proxy_logger.debug(
- f"WebSocket passthrough ({endpoint}): Client message is not a valid setup message: {e}"
- )
- pass # Not a JSON message or doesn't contain setup data
-
- await upstream_ws.send(text_data)
- elif bytes_data is not None:
- await upstream_ws.send(bytes_data)
- except asyncio.CancelledError:
- raise
- except Exception:
- verbose_proxy_logger.exception(
- f"WebSocket passthrough ({endpoint}): error forwarding client message"
- )
- await upstream_ws.close()
-
- async def forward_upstream_to_client() -> None:
- """Forward messages from upstream to client WebSocket"""
- try:
- # Wait for the first response from upstream
- raw_response = await upstream_ws.recv(decode=False)
- # Ensure raw_response is bytes before decoding
- if isinstance(raw_response, str):
- raw_response = raw_response.encode("ascii")
- setup_response = json.loads(raw_response.decode("ascii"))
- verbose_proxy_logger.debug(f"Setup response: {setup_response}")
-
- # Extract model and provider from setup response for Vertex AI Live
- if endpoint and "/vertex_ai/live" in endpoint:
- verbose_proxy_logger.debug(
- f"WebSocket passthrough ({endpoint}): Processing server setup response for model extraction"
- )
- extracted_model = _extract_model_from_vertex_ai_setup(
- setup_response
- )
- if extracted_model:
- kwargs["model"] = extracted_model
- kwargs["custom_llm_provider"] = "vertex_ai_language_models"
- # Update logging object with correct model
- logging_obj.model = extracted_model
- logging_obj.model_call_details["model"] = extracted_model
- logging_obj.model_call_details["custom_llm_provider"] = (
- "vertex_ai_language_models"
- )
- verbose_proxy_logger.debug(
- f"WebSocket passthrough ({endpoint}): Successfully extracted model '{extracted_model}' and set provider to 'vertex_ai' from server setup response"
- )
- else:
- verbose_proxy_logger.warning(
- f"WebSocket passthrough ({endpoint}): Failed to extract model from server setup response: {setup_response}"
- )
- else:
- verbose_proxy_logger.debug(
- f"WebSocket passthrough ({endpoint}): Not a Vertex AI Live endpoint, skipping model extraction"
- )
-
- # Send the setup response to the client
- await websocket.send_text(json.dumps(setup_response))
-
- # Now continuously forward messages from upstream to client
- async for upstream_message in upstream_ws:
- if isinstance(upstream_message, bytes):
- await websocket.send_bytes(upstream_message)
- # Parse and collect for cost tracking
- try:
- message_data = json.loads(upstream_message.decode())
- websocket_messages.append(message_data)
- except (json.JSONDecodeError, UnicodeDecodeError):
- pass
- else:
- await websocket.send_text(upstream_message)
- # Parse and collect for cost tracking
- try:
- message_data = json.loads(upstream_message)
- websocket_messages.append(message_data)
- except json.JSONDecodeError:
- pass
-
- except (ConnectionClosedOK, ConnectionClosedError) as e:
- verbose_proxy_logger.debug(
- f"Upstream WebSocket connection closed: {e}"
- )
- pass
- except asyncio.CancelledError:
- verbose_proxy_logger.debug(
- "asyncio.CancelledError in forward_upstream_to_client"
- )
- raise
- except Exception as e:
- verbose_proxy_logger.debug(
- f"Exception in forward_upstream_to_client: {e}"
- )
- verbose_proxy_logger.exception(
- f"WebSocket passthrough ({endpoint}): error forwarding upstream message"
- )
- raise
-
- # Create tasks for bidirectional message forwarding
- tasks = [
- asyncio.create_task(forward_client_to_upstream()),
- asyncio.create_task(forward_upstream_to_client()),
- ]
-
- done, pending = await asyncio.wait(
- tasks, return_when=asyncio.FIRST_COMPLETED
- )
-
- # Cancel remaining tasks
- for task in pending:
- task.cancel()
- try:
- await task
- except asyncio.CancelledError:
- pass
-
- # Check for exceptions in completed tasks
- for task in done:
- exception = task.exception()
- if exception is not None:
- raise exception
-
- end_time = datetime.now()
-
- # Update passthrough logging payload with response data
- passthrough_logging_payload["response_body"] = websocket_messages # type: ignore
- passthrough_logging_payload["end_time"] = end_time # type: ignore
-
- # Remove logging_obj from kwargs to avoid duplicate keyword argument
- success_kwargs = kwargs.copy()
- success_kwargs.pop("logging_obj", None)
-
- # # Add user authentication context for database logging
- # if user_api_key_dict:
- # success_kwargs.setdefault('litellm_params', {})
- # success_kwargs['litellm_params'].update({
- # 'proxy_server_request': {
- # 'body': {
- # 'user': user_api_key_dict.user_id,
- # 'team_id': user_api_key_dict.team_id,
- # 'end_user_id': user_api_key_dict.end_user_id,
- # }
- # }
- # })
- # # Also add the user_api_key for direct access
- # success_kwargs['user_api_key'] = user_api_key_dict.api_key
-
- # Create a dummy httpx.Response for WebSocket connections
- class MockWebSocketResponse:
- def __init__(self, target_url: str):
- self.status_code = 200
- self.text = "WebSocket connection successful"
- self.headers: dict[str, str] = {}
- self.request = MockWebSocketRequest(target_url)
-
- class MockWebSocketRequest:
- def __init__(self, target_url: str):
- self.method = "WEBSOCKET"
- self.url = target_url
-
- mock_response = MockWebSocketResponse(target)
-
- # Use the same success handler as HTTP passthrough endpoints
- asyncio.create_task(
- pass_through_endpoint_logging.pass_through_async_success_handler(
- httpx_response=mock_response, # type: ignore
- response_body=websocket_messages, # type: ignore
- url_route=endpoint or "",
- result="websocket_connection_successful",
- start_time=start_time,
- end_time=end_time,
- logging_obj=logging_obj,
- cache_hit=False,
- request_body={},
- **success_kwargs,
- )
- )
-
- # Call the proxy logging success hook
- if proxy_logging_obj:
- await proxy_logging_obj.post_call_success_hook(
- data={},
- user_api_key_dict=user_api_key_dict,
- response={"status": "websocket_connection_successful"}, # type: ignore
- )
-
- except InvalidStatus as exc:
- verbose_proxy_logger.exception(
- f"WebSocket passthrough ({endpoint}): upstream rejected WebSocket connection"
- )
-
- # Prepare request payload for logging
- request_payload = {}
- if kwargs:
- for key, value in kwargs.items():
- request_payload[key] = value
- if logging_obj is not None:
- request_payload["litellm_logging_obj"] = logging_obj
-
- # Log the connection failure using the same pattern as HTTP
- await proxy_logging_obj.post_call_failure_hook(
- user_api_key_dict=user_api_key_dict,
- original_exception=exc,
- request_data=request_payload,
- traceback_str=traceback.format_exc(
- limit=MAXIMUM_TRACEBACK_LINES_TO_LOG,
- ),
- )
-
- if websocket.client_state != WebSocketState.DISCONNECTED:
- await websocket.close(
- code=getattr(exc, "status_code", 1011),
- reason="Upstream connection rejected",
- )
- except Exception as e:
- verbose_proxy_logger.exception(
- f"WebSocket passthrough ({endpoint}): unexpected error while proxying WebSocket"
- )
-
- # Prepare request payload for logging
- request_payload = {}
- if kwargs:
- for key, value in kwargs.items():
- request_payload[key] = value
- if logging_obj is not None:
- request_payload["litellm_logging_obj"] = logging_obj
-
- # Log the unexpected error using the same pattern as HTTP
- await proxy_logging_obj.post_call_failure_hook(
- user_api_key_dict=user_api_key_dict,
- original_exception=e,
- request_data=request_payload,
- traceback_str=traceback.format_exc(
- limit=MAXIMUM_TRACEBACK_LINES_TO_LOG,
- ),
- )
-
- if websocket.client_state != WebSocketState.DISCONNECTED:
- await websocket.close(code=1011, reason="WebSocket passthrough error")
- finally:
- if websocket.client_state != WebSocketState.DISCONNECTED:
- await websocket.close()
-
-
-def _is_streaming_response(response: httpx.Response) -> bool:
- _content_type = response.headers.get("content-type")
- if _content_type is not None and "text/event-stream" in _content_type:
- return True
- return False
-
-
-def _extract_model_from_vertex_ai_setup(setup_response: dict) -> Optional[str]:
- """
- Extract the model name from Vertex AI Live setup response.
-
- The setup response can contain a model field in two formats:
- 1. Direct: {"model": "projects/.../models/gemini-2.0-flash-live-preview-04-09"}
- 2. Nested: {"setup": {"model": "projects/.../models/gemini-2.0-flash-live-preview-04-09"}}
-
- We extract just the model name: "gemini-2.0-flash-live-preview-04-09"
- """
- try:
- # Handle both direct model field and nested setup.model field
- model_path = None
- if isinstance(setup_response, dict):
- if "model" in setup_response:
- model_path = setup_response["model"]
- elif (
- "setup" in setup_response
- and isinstance(setup_response["setup"], dict)
- and "model" in setup_response["setup"]
- ):
- model_path = setup_response["setup"]["model"]
-
- if isinstance(model_path, str) and "/models/" in model_path:
- # Extract the model name after the last "/models/"
- model_name = model_path.split("/models/")[-1]
- return model_name
- except Exception as e:
- verbose_proxy_logger.debug(f"Error extracting model from setup response: {e}")
- return None
-
-
-class SafeRouteAdder:
- """
- Wrapper class for adding routes to FastAPI app.
- Only adds routes if they don't already exist on the app.
- """
-
- @staticmethod
- def _is_path_registered(app: FastAPI, path: str, methods: List[str]) -> bool:
- """
- Check if a path with any of the specified methods is already registered on the app.
-
- Args:
- app: The FastAPI application instance
- path: The path to check (e.g., "/v1/chat/completions")
- methods: List of HTTP methods to check (e.g., ["GET", "POST"])
-
- Returns:
- True if the path is already registered with any of the methods, False otherwise
- """
- for route in app.routes:
- # Use getattr to safely access route attributes
- route_path = getattr(route, "path", None)
- route_methods = getattr(route, "methods", None)
-
- if route_path == path and route_methods is not None:
- # Check if any of the methods overlap
- if any(method in route_methods for method in methods):
- return True
- return False
-
- @staticmethod
- def add_api_route_if_not_exists(
- app: FastAPI,
- path: str,
- endpoint: Any,
- methods: List[str],
- dependencies: Optional[List] = None,
- ) -> bool:
- """
- Add an API route to the app only if it doesn't already exist.
-
- Args:
- app: The FastAPI application instance
- path: The path for the route
- endpoint: The endpoint function/callable
- methods: List of HTTP methods
- dependencies: Optional list of dependencies
-
- Returns:
- True if route was added, False if it already existed
- """
- if SafeRouteAdder._is_path_registered(app=app, path=path, methods=methods):
- verbose_proxy_logger.debug(
- "Skipping route registration - path %s with methods %s already registered on app",
- path,
- methods,
- )
- return False
-
- app.add_api_route(
- path=path,
- endpoint=endpoint,
- methods=methods,
- dependencies=dependencies,
- )
- verbose_proxy_logger.debug(
- "Successfully added route: %s with methods %s",
- path,
- methods,
- )
- return True
-
-
-class InitPassThroughEndpointHelpers:
- @staticmethod
- def add_exact_path_route(
- app: FastAPI,
- path: str,
- target: str,
- custom_headers: Optional[dict],
- forward_headers: Optional[bool],
- merge_query_params: Optional[bool],
- dependencies: Optional[List],
- cost_per_request: Optional[float],
- endpoint_id: str,
- guardrails: Optional[dict] = None,
- methods: Optional[List[str]] = None,
- default_query_params: Optional[dict] = None,
- config_file_path: Optional[str] = None,
- ):
- """Add exact path route for pass-through endpoint"""
- # Default to all methods if none specified (backward compatibility)
- if methods is None or len(methods) == 0:
- methods = ["GET", "POST", "PUT", "DELETE", "PATCH"]
-
- # Create route key that includes methods for uniqueness
- methods_str = ",".join(sorted(methods))
- route_key = f"{endpoint_id}:exact:{path}:{methods_str}"
-
- # Check if this exact route is already registered
- if route_key in _registered_pass_through_routes:
- verbose_proxy_logger.debug(
- "Updating duplicate exact pass through endpoint: %s with methods %s (already registered)",
- path,
- methods,
- )
-
- verbose_proxy_logger.debug(
- "adding exact pass through endpoint: %s, methods: %s, dependencies: %s",
- path,
- methods,
- dependencies,
- )
-
- # Use SafeRouteAdder to only add route if it doesn't exist on the app
- SafeRouteAdder.add_api_route_if_not_exists(
- app=app,
- path=path,
- endpoint=create_pass_through_route( # type: ignore
- path,
- target,
- custom_headers,
- forward_headers,
- merge_query_params,
- dependencies,
- cost_per_request=cost_per_request,
- default_query_params=default_query_params,
- guardrails=guardrails,
- config_file_path=config_file_path,
- ),
- methods=methods,
- dependencies=dependencies,
- )
-
- # Always register/update the route metadata (headers, target) even if FastAPI route exists
- _registered_pass_through_routes[route_key] = {
- "endpoint_id": endpoint_id,
- "path": path,
- "type": "exact",
- "methods": methods,
- "passthrough_params": {
- "target": target,
- "custom_headers": custom_headers,
- "forward_headers": forward_headers,
- "merge_query_params": merge_query_params,
- "default_query_params": default_query_params,
- "dependencies": dependencies,
- "cost_per_request": cost_per_request,
- "guardrails": guardrails,
- },
- }
-
- @staticmethod
- def add_subpath_route(
- app: FastAPI,
- path: str,
- target: str,
- custom_headers: Optional[dict],
- forward_headers: Optional[bool],
- merge_query_params: Optional[bool],
- dependencies: Optional[List],
- cost_per_request: Optional[float],
- endpoint_id: str,
- guardrails: Optional[dict] = None,
- methods: Optional[List[str]] = None,
- default_query_params: Optional[dict] = None,
- config_file_path: Optional[str] = None,
- ):
- """Add wildcard route for sub-paths"""
- # Default to all methods if none specified (backward compatibility)
- if methods is None or len(methods) == 0:
- methods = ["GET", "POST", "PUT", "DELETE", "PATCH"]
-
- wildcard_path = f"{path}/{{subpath:path}}"
- methods_str = ",".join(sorted(methods))
- route_key = f"{endpoint_id}:subpath:{path}:{methods_str}"
-
- # Check if this subpath route is already registered
- if route_key in _registered_pass_through_routes:
- verbose_proxy_logger.debug(
- "Updating duplicate wildcard pass through endpoint: %s with methods %s (already registered)",
- wildcard_path,
- methods,
- )
-
- verbose_proxy_logger.debug(
- "adding wildcard pass through endpoint: %s, methods: %s, dependencies: %s",
- wildcard_path,
- methods,
- dependencies,
- )
-
- # Use SafeRouteAdder to only add route if it doesn't exist on the app
- SafeRouteAdder.add_api_route_if_not_exists(
- app=app,
- path=wildcard_path,
- endpoint=create_pass_through_route( # type: ignore
- path,
- target,
- custom_headers,
- forward_headers,
- merge_query_params,
- dependencies,
- include_subpath=True,
- cost_per_request=cost_per_request,
- default_query_params=default_query_params,
- guardrails=guardrails,
- config_file_path=config_file_path,
- ),
- methods=methods,
- dependencies=dependencies,
- )
-
- # Register the route to prevent duplicates only if it was added
- _registered_pass_through_routes[route_key] = {
- "endpoint_id": endpoint_id,
- "path": path,
- "type": "subpath",
- "methods": methods,
- "passthrough_params": {
- "target": target,
- "custom_headers": custom_headers,
- "forward_headers": forward_headers,
- "merge_query_params": merge_query_params,
- "default_query_params": default_query_params,
- "dependencies": dependencies,
- "cost_per_request": cost_per_request,
- "guardrails": guardrails,
- },
- }
-
- @staticmethod
- def remove_endpoint_routes(endpoint_id: str):
- """Remove all routes for a specific endpoint ID from the registry
- and clean up corresponding entries from LiteLLMRoutes.openai_routes."""
- keys_to_remove = [
- key
- for key, value in _registered_pass_through_routes.items()
- if value["endpoint_id"] == endpoint_id
- ]
- for key in keys_to_remove:
- route_info = _registered_pass_through_routes[key]
- path = route_info.get("path")
- if isinstance(path, str):
- openai_routes = LiteLLMRoutes.openai_routes.value
- if path in openai_routes:
- openai_routes.remove(path)
- if route_info.get("type") == "subpath":
- wildcard_path = path.rstrip("/") + "/*"
- if wildcard_path in openai_routes:
- openai_routes.remove(wildcard_path)
- del _registered_pass_through_routes[key]
- verbose_proxy_logger.debug(
- "Removed pass-through route from registry: %s", key
- )
-
- @staticmethod
- def clear_all_pass_through_routes():
- """Clear all pass-through routes from the registry"""
- _registered_pass_through_routes.clear()
-
- @staticmethod
- def get_all_registered_pass_through_routes() -> List[str]:
- """Get all registered pass-through endpoints from the registry"""
- return list(_registered_pass_through_routes.keys())
-
- @staticmethod
- def _build_full_path_with_root(path: str) -> str:
- """
- Build full path by prepending server root path if needed.
-
- Args:
- path: The relative path to build
-
- Returns:
- Full path with server root prepended (if root is not "/")
- """
- root_path = get_server_root_path()
- if root_path == "/":
- return path
- return f"{root_path}{path}"
-
- @staticmethod
- def is_registered_pass_through_route(route: str) -> bool:
- """
- Check if route is a registered pass-through endpoint from DB
-
- Uses the in-memory registry to avoid additional DB queries
- Optimized for minimal latency
-
- Args:
- route: The route to check
-
- Returns:
- bool: True if route is a registered pass-through endpoint, False otherwise
- """
- ## CHECK IF MAPPED PASS THROUGH ENDPOINT
- normalized_route = normalize_route_for_root_path(route)
- if normalized_route is not None:
- for mapped_route in LiteLLMRoutes.mapped_pass_through_routes.value:
- if normalized_route.startswith(mapped_route):
- return True
-
- # Fast path: check if any registered route key contains this path
- # Keys are in format: "{endpoint_id}:exact:{path}:{methods}" or "{endpoint_id}:subpath:{path}:{methods}"
- # For backward compatibility, also support old format: "{endpoint_id}:exact:{path}" or "{endpoint_id}:subpath:{path}"
- # Extract unique paths from keys for quick checking
- for key in _registered_pass_through_routes.keys():
- parts = key.split(":", 3) # Split into [endpoint_id, type, path, methods?]
- if len(parts) >= 3:
- route_type = parts[1]
- registered_path = (
- InitPassThroughEndpointHelpers._build_full_path_with_root(parts[2])
- )
- if route_type == "exact" and route == registered_path:
- return True
- elif route_type == "subpath":
- if route == registered_path or route.startswith(
- registered_path + "/"
- ):
- return True
-
- return False
-
- @staticmethod
- def get_registered_pass_through_route(
- route: str, method: Optional[str] = None
- ) -> Optional[Dict[str, Any]]:
- """Get passthrough params for a given route and optionally filter by HTTP method"""
- for key in _registered_pass_through_routes.keys():
- parts = key.split(":", 3) # Split into [endpoint_id, type, path, methods?]
- if len(parts) >= 3:
- route_type = parts[1]
- registered_path = (
- InitPassThroughEndpointHelpers._build_full_path_with_root(parts[2])
- )
-
- # Get the methods for this route
- route_methods = _registered_pass_through_routes[key].get("methods", [])
-
- # Check if path matches
- path_matches = False
- if route_type == "exact" and route == registered_path:
- path_matches = True
- elif route_type == "subpath":
- if route == registered_path or route.startswith(
- registered_path + "/"
- ):
- path_matches = True
-
- # If path matches and method filter is provided, check if method is allowed
- if path_matches:
- if method is None or not route_methods or method in route_methods:
- return _registered_pass_through_routes[key]
-
- return None
-
-
-def _get_combined_pass_through_endpoints(
- pass_through_endpoints: Union[List[Dict], List[PassThroughGenericEndpoint]],
- config_pass_through_endpoints: List[Dict],
-):
- """Get combined pass-through endpoints from db + config"""
- return pass_through_endpoints + config_pass_through_endpoints
-
-
-async def _register_pass_through_endpoint(
- endpoint: Union[Dict[str, Any], PassThroughGenericEndpoint],
- app: FastAPI,
- premium_user: bool,
- visited_endpoints: set[str],
- config_file_path: Optional[str] = None,
-) -> None:
- endpoint_data: Dict[str, Any]
- if isinstance(endpoint, PassThroughGenericEndpoint):
- endpoint_data = endpoint.model_dump()
- else:
- endpoint_data = endpoint
-
- if endpoint_data.get("id") is None:
- endpoint_data["id"] = str(uuid.uuid4())
- endpoint_id = cast(str, endpoint_data["id"])
-
- target = endpoint_data.get("target")
- path = endpoint_data.get("path")
- if path is None:
- raise ValueError("Path is required for pass-through endpoint")
-
- custom_headers = await set_env_variables_in_header(
- custom_headers=endpoint_data.get("headers")
- )
- forward_headers = endpoint_data.get("forward_headers")
- merge_query_params = endpoint_data.get("merge_query_params")
- default_query_params = endpoint_data.get("default_query_params")
- auth = endpoint_data.get("auth")
- dependencies = None
-
- if auth is not None and str(auth).lower() == "true":
- # Authentication on a pass-through endpoint used to be enterprise-only.
- # That left OSS with no safe configuration: auth=True raised at startup
- # unless the operator had a license. The safe option must always be free,
- # and unauthenticated forwarding should require explicit opt-in.
- dependencies = [Depends(user_api_key_auth)]
- if path not in LiteLLMRoutes.openai_routes.value:
- LiteLLMRoutes.openai_routes.value.append(path)
-
- if target is None:
- return
-
- guardrails = endpoint_data.get("guardrails")
- methods = endpoint_data.get("methods")
- cost_per_request = endpoint_data.get("cost_per_request")
-
- verbose_proxy_logger.debug(
- "Initializing pass through endpoint: %s (ID: %s)", path, endpoint_id
- )
- InitPassThroughEndpointHelpers.add_exact_path_route(
- app=app,
- path=path,
- target=target,
- custom_headers=custom_headers,
- forward_headers=forward_headers,
- merge_query_params=merge_query_params,
- dependencies=dependencies,
- cost_per_request=cost_per_request,
- endpoint_id=endpoint_id,
- guardrails=guardrails,
- methods=methods,
- default_query_params=default_query_params,
- config_file_path=config_file_path,
- )
-
- methods_for_key = methods if methods else ["GET", "POST", "PUT", "DELETE", "PATCH"]
- methods_str = ",".join(sorted(methods_for_key))
- visited_endpoints.add(f"{endpoint_id}:exact:{path}:{methods_str}")
-
- if endpoint_data.get("include_subpath", False) is True:
- if auth is not None and str(auth).lower() == "true":
- wildcard_path = path.rstrip("/") + "/*"
- if wildcard_path not in LiteLLMRoutes.openai_routes.value:
- LiteLLMRoutes.openai_routes.value.append(wildcard_path)
- InitPassThroughEndpointHelpers.add_subpath_route(
- app=app,
- path=path,
- target=target,
- custom_headers=custom_headers,
- forward_headers=forward_headers,
- merge_query_params=merge_query_params,
- dependencies=dependencies,
- cost_per_request=cost_per_request,
- endpoint_id=endpoint_id,
- guardrails=guardrails,
- methods=methods,
- default_query_params=default_query_params,
- config_file_path=config_file_path,
- )
- visited_endpoints.add(f"{endpoint_id}:subpath:{path}:{methods_str}")
-
- verbose_proxy_logger.debug(
- "Added new pass through endpoint: %s (ID: %s)", path, endpoint_id
- )
-
-
-async def initialize_pass_through_endpoints(
- pass_through_endpoints: Union[List[Dict], List[PassThroughGenericEndpoint]],
- config_file_path: Optional[str] = None,
-):
- """
- 1. Create a global list of pass-through endpoints (db + config)
- 2. Clear all existing pass-through endpoints from the FastAPI app routes
- 3. Add new endpoints to the in-memory registry
-
- Initialize a list of pass-through endpoints by adding them to the FastAPI app routes
-
- Args:
- pass_through_endpoints: List of pass-through endpoints to initialize
- config_file_path: Path to the operator's config.yaml when this call
- originates from a YAML-load. Threaded through to
- ``create_pass_through_route`` so an operator using
- ``s3://``/``gcs://`` ``custom_handler`` in their config still
- loads. Callers from the DB-overlay / runtime API path must leave
- this ``None`` so the runtime gate in ``get_instance_fn`` fires.
-
- Returns:
- None
- """
- verbose_proxy_logger.debug("initializing pass through endpoints")
- from litellm.proxy.proxy_server import (
- app,
- config_passthrough_endpoints,
- premium_user,
- )
-
- ## get combined pass-through endpoints from db + config
- combined_pass_through_endpoints: List[Union[Dict, PassThroughGenericEndpoint]]
-
- if config_passthrough_endpoints is not None:
- combined_pass_through_endpoints = _get_combined_pass_through_endpoints( # type: ignore
- pass_through_endpoints, config_passthrough_endpoints
- )
- else:
- combined_pass_through_endpoints = pass_through_endpoints # type: ignore
-
- ## clear all existing pass-through endpoints from the FastAPI app routes
- # InitPassThroughEndpointHelpers.clear_all_pass_through_routes()
-
- # get a list of all registered pass-through endpoints
- # mark the ones that are visited in the list
- # remove the ones that are not visited from the list
- registered_pass_through_endpoints = (
- InitPassThroughEndpointHelpers.get_all_registered_pass_through_routes()
- )
-
- visited_endpoints: set[str] = set()
-
- for endpoint in combined_pass_through_endpoints:
- await _register_pass_through_endpoint(
- endpoint=endpoint,
- app=app,
- premium_user=premium_user,
- visited_endpoints=visited_endpoints,
- config_file_path=config_file_path,
- )
-
- # remove the ones that are not visited from the list
- for endpoint_key in registered_pass_through_endpoints:
- if endpoint_key not in visited_endpoints:
- InitPassThroughEndpointHelpers.remove_endpoint_routes(endpoint_key)
-
-
-def _get_pass_through_endpoints_from_config() -> List[PassThroughGenericEndpoint]:
- """
- Get pass-through endpoints defined in the config file.
- These are read-only and cannot be edited via the UI.
- Malformed endpoints are logged and skipped; they do not crash the function.
- """
- from pydantic import ValidationError
-
- from litellm.proxy.proxy_server import config_passthrough_endpoints
-
- if config_passthrough_endpoints is None or len(config_passthrough_endpoints) == 0:
- return []
-
- returned_endpoints: List[PassThroughGenericEndpoint] = []
- for endpoint in config_passthrough_endpoints:
- try:
- if isinstance(endpoint, dict):
- endpoint_dict = dict(endpoint)
- endpoint_dict["is_from_config"] = True
- returned_endpoints.append(PassThroughGenericEndpoint(**endpoint_dict))
- elif isinstance(endpoint, PassThroughGenericEndpoint):
- # Create a copy with is_from_config=True
- endpoint_dict = endpoint.model_dump()
- endpoint_dict["is_from_config"] = True
- returned_endpoints.append(PassThroughGenericEndpoint(**endpoint_dict))
- except ValidationError as e:
- verbose_proxy_logger.warning(
- "Skipping malformed pass-through endpoint from config: %s",
- e,
- exc_info=False,
- )
-
- return returned_endpoints
-
-
-async def _get_pass_through_endpoints_from_db(
- endpoint_id: Optional[str] = None,
- user_api_key_dict: Optional[UserAPIKeyAuth] = None,
-) -> List[PassThroughGenericEndpoint]:
- from litellm.proxy._types import LitellmUserRoles
- from litellm.proxy.proxy_server import get_config_general_settings
-
- try:
- if user_api_key_dict is None:
- user_api_key_dict = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
- response: ConfigFieldInfo = await get_config_general_settings(
- field_name="pass_through_endpoints", user_api_key_dict=user_api_key_dict
- )
- except Exception:
- return []
-
- pass_through_endpoint_data: Optional[List] = response.field_value
- if pass_through_endpoint_data is None:
- return []
-
- returned_endpoints: List[PassThroughGenericEndpoint] = []
- if endpoint_id is None:
- # Return all endpoints from DB, mark as not from config
- for endpoint in pass_through_endpoint_data:
- if isinstance(endpoint, dict):
- endpoint_dict = dict(endpoint)
- endpoint_dict["is_from_config"] = False
- returned_endpoints.append(PassThroughGenericEndpoint(**endpoint_dict))
- elif isinstance(endpoint, PassThroughGenericEndpoint):
- endpoint_dict = endpoint.model_dump()
- endpoint_dict["is_from_config"] = False
- returned_endpoints.append(PassThroughGenericEndpoint(**endpoint_dict))
- else:
- # Find specific endpoint by ID
- found_endpoint = _find_endpoint_by_id(pass_through_endpoint_data, endpoint_id)
- if found_endpoint is not None:
- endpoint_dict = (
- found_endpoint.model_dump()
- if isinstance(found_endpoint, PassThroughGenericEndpoint)
- else dict(found_endpoint)
- )
- endpoint_dict["is_from_config"] = False
- returned_endpoints.append(PassThroughGenericEndpoint(**endpoint_dict))
-
- return returned_endpoints
-
-
-async def _filter_endpoints_by_team_allowed_routes(
- team_id: str,
- pass_through_endpoints: List[PassThroughGenericEndpoint],
- prisma_client,
-) -> List[PassThroughGenericEndpoint]:
- """
- Filter pass-through endpoints based on team's allowed_passthrough_routes metadata.
-
- Args:
- team_id: The team ID to check permissions for
- pass_through_endpoints: List of endpoints to filter
- prisma_client: Database client
-
- Returns:
- Filtered list of endpoints based on team permissions
-
- Raises:
- HTTPException: If team is not found
- """
- # retrieve team from db
- team = await prisma_client.db.litellm_teamtable.find_unique(
- where={"team_id": team_id},
- )
- if team is None:
- raise HTTPException(
- status_code=404,
- detail={"error": "Team not found"},
- )
-
- # retrieve team metadata
- team_metadata = team.metadata
- if (
- team_metadata is not None
- and team_metadata.get("allowed_passthrough_routes") is not None
- ):
- ## FILTER pass_through_endpoints by allowed_passthrough_routes
- pass_through_endpoints = [
- endpoint
- for endpoint in pass_through_endpoints
- if endpoint.path in team_metadata.get("allowed_passthrough_routes")
- ]
-
- return pass_through_endpoints
-
-
-@router.get(
- "/config/pass_through_endpoint",
- dependencies=[Depends(user_api_key_auth)],
- response_model=PassThroughEndpointResponse,
-)
-@router.get(
- "/config/pass_through_endpoint/team/{team_id}",
- dependencies=[Depends(user_api_key_auth)],
- response_model=PassThroughEndpointResponse,
-)
-async def get_pass_through_endpoints(
- endpoint_id: Optional[str] = None,
- user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
- team_id: Optional[str] = None,
-):
- """
- GET configured pass through endpoint.
-
- If no endpoint_id given, return all configured endpoints.
- """ ## Get existing pass-through endpoint field value
- from litellm.proxy._types import CommonProxyErrors
- from litellm.proxy.proxy_server import prisma_client
-
- if prisma_client is None:
- raise HTTPException(
- status_code=500,
- detail={"error": CommonProxyErrors.db_not_connected_error.value},
- )
-
- # Get endpoints from DB (editable via UI)
- db_endpoints = await _get_pass_through_endpoints_from_db(
- endpoint_id=endpoint_id, user_api_key_dict=user_api_key_dict
- )
-
- # Get endpoints from config file (read-only, not editable via UI)
- config_endpoints = _get_pass_through_endpoints_from_config()
-
- # Merge: config endpoints not in DB + all DB endpoints (DB overrides config for same path)
- db_paths = {ep.path for ep in db_endpoints}
- config_only_endpoints = [ep for ep in config_endpoints if ep.path not in db_paths]
- if endpoint_id is not None:
- # When filtering by endpoint_id, only return if found in DB (config endpoints use generated IDs)
- pass_through_endpoints = db_endpoints
- else:
- pass_through_endpoints = config_only_endpoints + db_endpoints
-
- if team_id is not None:
- pass_through_endpoints = await _filter_endpoints_by_team_allowed_routes(
- team_id=team_id,
- pass_through_endpoints=pass_through_endpoints,
- prisma_client=prisma_client,
- )
-
- return PassThroughEndpointResponse(endpoints=pass_through_endpoints)
-
-
-@router.post(
- "/config/pass_through_endpoint/{endpoint_id}",
- dependencies=[Depends(user_api_key_auth)],
-)
-async def update_pass_through_endpoints(
- endpoint_id: str,
- data: PassThroughGenericEndpoint,
- request: Request,
- user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
-):
- """
- Update a pass-through endpoint by ID.
- """
- from litellm.proxy.proxy_server import (
- get_config_general_settings,
- update_config_general_settings,
- )
-
- ## Get existing pass-through endpoint field value
- try:
- response: ConfigFieldInfo = await get_config_general_settings(
- field_name="pass_through_endpoints", user_api_key_dict=user_api_key_dict
- )
- except Exception:
- raise HTTPException(
- status_code=404,
- detail={"error": "No pass-through endpoints found"},
- )
-
- pass_through_endpoint_data: Optional[List] = response.field_value
- if pass_through_endpoint_data is None:
- raise HTTPException(
- status_code=404,
- detail={"error": "No pass-through endpoints found"},
- )
-
- # Find the endpoint to update
- found_endpoint = _find_endpoint_by_id(pass_through_endpoint_data, endpoint_id)
-
- if found_endpoint is None:
- raise HTTPException(
- status_code=404,
- detail={"error": f"Endpoint with ID '{endpoint_id}' not found"},
- )
-
- # Find the index for updating the list
- endpoint_index = None
- for idx, endpoint in enumerate(pass_through_endpoint_data):
- _endpoint = (
- PassThroughGenericEndpoint(**endpoint)
- if isinstance(endpoint, dict)
- else endpoint
- )
- if _endpoint.id == endpoint_id:
- endpoint_index = idx
- break
-
- if endpoint_index is None:
- raise HTTPException(
- status_code=404,
- detail={
- "error": f"Could not find index for endpoint with ID '{endpoint_id}'"
- },
- )
-
- # Get the update data as dict, excluding None values for partial updates
- # Exclude is_from_config as it's a response-only field (computed at read time)
- update_data = data.model_dump(exclude_none=True, exclude={"is_from_config"})
-
- # Start with existing endpoint data
- endpoint_dict = found_endpoint.model_dump()
-
- # Update with new data (only non-None values)
- endpoint_dict.update(update_data)
-
- # Preserve existing ID if not provided in update and endpoint has ID
- if "id" not in update_data and found_endpoint.id is not None:
- endpoint_dict["id"] = found_endpoint.id
-
- # Remove is_from_config before saving - it's a response-only field (computed at read time)
- endpoint_dict.pop("is_from_config", None)
-
- # Create updated endpoint object
- updated_endpoint = PassThroughGenericEndpoint(**endpoint_dict)
-
- # Update the list
- pass_through_endpoint_data[endpoint_index] = endpoint_dict
-
- # Remove old routes from registry before they get re-registered
- InitPassThroughEndpointHelpers.remove_endpoint_routes(endpoint_id)
-
- ## Update db
- updated_data = ConfigFieldUpdate(
- field_name="pass_through_endpoints",
- field_value=pass_through_endpoint_data,
- config_type="general_settings",
- )
-
- await update_config_general_settings(
- data=updated_data, user_api_key_dict=user_api_key_dict
- )
-
- # Re-register the route with updated headers
- _custom_headers: Optional[dict] = updated_endpoint.headers or {}
- _custom_headers = await set_env_variables_in_header(custom_headers=_custom_headers)
-
- if updated_endpoint.include_subpath:
- InitPassThroughEndpointHelpers.add_subpath_route(
- app=request.app,
- path=updated_endpoint.path,
- target=updated_endpoint.target,
- custom_headers=_custom_headers,
- forward_headers=None, # Defaults not available in model? assuming None logic handles it
- merge_query_params=None,
- dependencies=None,
- cost_per_request=updated_endpoint.cost_per_request,
- endpoint_id=updated_endpoint.id or endpoint_id or "",
- guardrails=getattr(updated_endpoint, "guardrails", None),
- methods=updated_endpoint.methods,
- default_query_params=updated_endpoint.default_query_params,
- )
- else:
- InitPassThroughEndpointHelpers.add_exact_path_route(
- app=request.app,
- path=updated_endpoint.path,
- target=updated_endpoint.target,
- custom_headers=_custom_headers,
- forward_headers=None,
- merge_query_params=None,
- dependencies=None,
- cost_per_request=updated_endpoint.cost_per_request,
- endpoint_id=updated_endpoint.id or endpoint_id or "",
- guardrails=getattr(updated_endpoint, "guardrails", None),
- methods=updated_endpoint.methods,
- default_query_params=updated_endpoint.default_query_params,
- )
-
- return PassThroughEndpointResponse(
- endpoints=[updated_endpoint] if updated_endpoint else []
- )
-
-
-@router.post(
- "/config/pass_through_endpoint",
- dependencies=[Depends(user_api_key_auth)],
-)
-async def create_pass_through_endpoints(
- data: PassThroughGenericEndpoint,
- request: Request,
- user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
-):
- """
- Create new pass-through endpoint
- """
- from litellm._uuid import uuid
- from litellm.proxy.proxy_server import (
- get_config_general_settings,
- update_config_general_settings,
- )
-
- ## Get existing pass-through endpoint field value
-
- try:
- response: ConfigFieldInfo = await get_config_general_settings(
- field_name="pass_through_endpoints", user_api_key_dict=user_api_key_dict
- )
- except Exception:
- response = ConfigFieldInfo(
- field_name="pass_through_endpoints", field_value=None
- )
-
- ## Auto-generate ID if not provided
- # Exclude is_from_config as it's a response-only field (computed at read time)
- data_dict = data.model_dump(exclude={"is_from_config"})
- if data_dict.get("id") is None:
- data_dict["id"] = str(uuid.uuid4())
-
- if response.field_value is None:
- response.field_value = [data_dict]
- elif isinstance(response.field_value, List):
- response.field_value.append(data_dict)
-
- ## Update db
- updated_data = ConfigFieldUpdate(
- field_name="pass_through_endpoints",
- field_value=response.field_value,
- config_type="general_settings",
- )
- await update_config_general_settings(
- data=updated_data, user_api_key_dict=user_api_key_dict
- )
-
- # Return the created endpoint with the generated ID
- created_endpoint = PassThroughGenericEndpoint(**data_dict)
-
- # Register the new route
- _custom_headers: Optional[dict] = created_endpoint.headers or {}
- _custom_headers = await set_env_variables_in_header(custom_headers=_custom_headers)
-
- if created_endpoint.include_subpath:
- InitPassThroughEndpointHelpers.add_subpath_route(
- app=request.app,
- path=created_endpoint.path,
- target=created_endpoint.target,
- custom_headers=_custom_headers,
- forward_headers=None,
- merge_query_params=None,
- dependencies=None,
- cost_per_request=created_endpoint.cost_per_request,
- endpoint_id=created_endpoint.id or "",
- guardrails=getattr(created_endpoint, "guardrails", None),
- methods=created_endpoint.methods,
- default_query_params=created_endpoint.default_query_params,
- )
- else:
- InitPassThroughEndpointHelpers.add_exact_path_route(
- app=request.app,
- path=created_endpoint.path,
- target=created_endpoint.target,
- custom_headers=_custom_headers,
- forward_headers=None,
- merge_query_params=None,
- dependencies=None,
- cost_per_request=created_endpoint.cost_per_request,
- endpoint_id=created_endpoint.id or "",
- guardrails=getattr(created_endpoint, "guardrails", None),
- methods=created_endpoint.methods,
- default_query_params=created_endpoint.default_query_params,
- )
-
- return PassThroughEndpointResponse(endpoints=[created_endpoint])
-
-
-@router.delete(
- "/config/pass_through_endpoint",
- dependencies=[Depends(user_api_key_auth)],
- response_model=PassThroughEndpointResponse,
-)
-async def delete_pass_through_endpoints(
- endpoint_id: str,
- user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
-):
- """
- Delete a pass-through endpoint by ID.
-
- Returns - the deleted endpoint
- """
- from litellm.proxy.proxy_server import (
- get_config_general_settings,
- update_config_general_settings,
- )
-
- ## Get existing pass-through endpoint field value
-
- try:
- response: ConfigFieldInfo = await get_config_general_settings(
- field_name="pass_through_endpoints", user_api_key_dict=user_api_key_dict
- )
- except Exception:
- response = ConfigFieldInfo(
- field_name="pass_through_endpoints", field_value=None
- )
-
- ## Update field by removing endpoint
- pass_through_endpoint_data: Optional[List] = response.field_value
- if response.field_value is None or pass_through_endpoint_data is None:
- raise HTTPException(
- status_code=400,
- detail={"error": "There are no pass-through endpoints setup."},
- )
-
- # Find the endpoint to delete
- found_endpoint = _find_endpoint_by_id(pass_through_endpoint_data, endpoint_id)
-
- if found_endpoint is None:
- raise HTTPException(
- status_code=400,
- detail={
- "error": "Endpoint with ID '{}' was not found in pass-through endpoint list.".format(
- endpoint_id
- )
- },
- )
-
- # Find the index for deleting from the list
- endpoint_index = None
- for idx, endpoint in enumerate(pass_through_endpoint_data):
- _endpoint = (
- PassThroughGenericEndpoint(**endpoint)
- if isinstance(endpoint, dict)
- else endpoint
- )
- if _endpoint.id == endpoint_id:
- endpoint_index = idx
- break
-
- if endpoint_index is None:
- raise HTTPException(
- status_code=400,
- detail={
- "error": f"Could not find index for endpoint with ID '{endpoint_id}'"
- },
- )
-
- # Remove the endpoint
- pass_through_endpoint_data.pop(endpoint_index)
- response_obj = found_endpoint
-
- # Remove routes from registry
- InitPassThroughEndpointHelpers.remove_endpoint_routes(endpoint_id)
-
- ## Update db
- updated_data = ConfigFieldUpdate(
- field_name="pass_through_endpoints",
- field_value=pass_through_endpoint_data,
- config_type="general_settings",
- )
- await update_config_general_settings(
- data=updated_data, user_api_key_dict=user_api_key_dict
- )
-
- return PassThroughEndpointResponse(endpoints=[response_obj])
-
-
-def _find_endpoint_by_id(
- endpoints_data: List,
- endpoint_id: str,
-) -> Optional[PassThroughGenericEndpoint]:
- """
- Find an endpoint by ID.
-
- Args:
- endpoints_data: List of endpoint data (dicts or PassThroughGenericEndpoint objects)
- endpoint_id: ID to search for
-
- Returns:
- Found endpoint or None if not found
- """
- for endpoint in endpoints_data:
- _endpoint: Optional[PassThroughGenericEndpoint] = None
- if isinstance(endpoint, dict):
- _endpoint = PassThroughGenericEndpoint(**endpoint)
- elif isinstance(endpoint, PassThroughGenericEndpoint):
- _endpoint = endpoint
-
- # Only compare IDs to IDs
- if _endpoint is not None and _endpoint.id == endpoint_id:
- return _endpoint
-
- return None
-
-
-async def initialize_pass_through_endpoints_in_db():
- """
- Gets all pass-through endpoints from db and initializes them in the proxy server.
- """
- pass_through_endpoints = await _get_pass_through_endpoints_from_db()
- await initialize_pass_through_endpoints(
- pass_through_endpoints=pass_through_endpoints
- )
+import ast
+import asyncio
+import copy
+import json
+import posixpath
+import traceback
+from base64 import b64encode
+from datetime import datetime
+from typing import Any, Dict, List, Mapping, Optional, Tuple, Union, cast
+from urllib.parse import urlencode, urlparse
+
+import httpx
+from fastapi import (
+ APIRouter,
+ Depends,
+ FastAPI,
+ HTTPException,
+ Request,
+ Response,
+ UploadFile,
+ WebSocket,
+ status,
+)
+from fastapi.responses import StreamingResponse
+from starlette.datastructures import UploadFile as StarletteUploadFile
+from starlette.websockets import WebSocketState
+from websockets.asyncio.client import connect
+from websockets.exceptions import (
+ ConnectionClosedError,
+ ConnectionClosedOK,
+ InvalidStatus,
+)
+
+import litellm
+from litellm._logging import verbose_proxy_logger
+from litellm._uuid import uuid
+from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG
+from litellm.integrations.custom_logger import CustomLogger
+from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
+from litellm.llms.base_llm.managed_resources.utils import (
+ resolve_passthrough_managed_id_provider,
+)
+from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
+from litellm.passthrough import BasePassthroughUtils
+from litellm.proxy._types import (
+ ConfigFieldInfo,
+ ConfigFieldUpdate,
+ LiteLLMRoutes,
+ PassThroughEndpointResponse,
+ PassThroughGenericEndpoint,
+ ProxyException,
+ UserAPIKeyAuth,
+)
+from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
+from litellm.proxy.common_utils.http_parsing_utils import (
+ _read_request_body,
+ _safe_get_request_headers,
+)
+from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
+from litellm.proxy.utils import get_server_root_path, normalize_route_for_root_path
+from litellm.secret_managers.main import get_secret_str
+from litellm.types.llms.custom_http import httpxSpecialProvider
+from litellm.types.passthrough_endpoints.pass_through_endpoints import (
+ EndpointType,
+ LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY,
+ LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY,
+ PassthroughStandardLoggingPayload,
+)
+
+from .streaming_handler import PassThroughStreamingHandler
+from .success_handler import PassThroughEndpointLogging
+
+router = APIRouter()
+
+pass_through_endpoint_logging = PassThroughEndpointLogging()
+
+# Global registry to track registered pass-through routes and prevent memory leaks
+_registered_pass_through_routes: Dict[
+ str, Dict[str, Union[str, List[str], Dict[str, Any]]]
+] = {}
+
+
+def get_response_body(response: httpx.Response) -> Optional[dict]:
+ try:
+ return response.json()
+ except Exception:
+ return None
+
+
+async def set_env_variables_in_header(custom_headers: Optional[dict]) -> Optional[dict]:
+ """
+ checks if any headers on config.yaml are defined as os.environ/COHERE_API_KEY etc
+
+ only runs for headers defined on config.yaml
+
+ example header can be
+
+ {"Authorization": "Bearer os.environ/COHERE_API_KEY"}
+ """
+ if custom_headers is None:
+ return None
+ headers = {}
+ for key, value in custom_headers.items():
+ # langfuse Api requires base64 encoded headers - it's simpleer to just ask litellm users to set their langfuse public and secret keys
+ # we can then get the b64 encoded keys here
+ if key == "LANGFUSE_PUBLIC_KEY" or key == "LANGFUSE_SECRET_KEY":
+ # langfuse requires b64 encoded headers - we construct that here
+ _langfuse_public_key = custom_headers["LANGFUSE_PUBLIC_KEY"]
+ _langfuse_secret_key = custom_headers["LANGFUSE_SECRET_KEY"]
+ if isinstance(
+ _langfuse_public_key, str
+ ) and _langfuse_public_key.startswith("os.environ/"):
+ _langfuse_public_key = get_secret_str(_langfuse_public_key)
+ if isinstance(
+ _langfuse_secret_key, str
+ ) and _langfuse_secret_key.startswith("os.environ/"):
+ _langfuse_secret_key = get_secret_str(_langfuse_secret_key)
+ headers["Authorization"] = "Basic " + b64encode(
+ f"{_langfuse_public_key}:{_langfuse_secret_key}".encode("utf-8")
+ ).decode("ascii")
+ else:
+ # for all other headers
+ headers[key] = value
+ if isinstance(value, str) and "os.environ/" in value:
+ verbose_proxy_logger.debug(
+ "pass through endpoint - looking up 'os.environ/' variable"
+ )
+ # get string section that is os.environ/
+ start_index = value.find("os.environ/")
+ _variable_name = value[start_index:]
+
+ verbose_proxy_logger.debug(
+ "pass through endpoint - getting secret for variable name: %s",
+ _variable_name,
+ )
+ _secret_value = get_secret_str(_variable_name)
+ if _secret_value is not None:
+ new_value = value.replace(_variable_name, _secret_value)
+ headers[key] = new_value
+ return headers
+
+
+async def chat_completion_pass_through_endpoint( # noqa: PLR0915
+ fastapi_response: Response,
+ request: Request,
+ adapter_id: str,
+ user_api_key_dict: UserAPIKeyAuth,
+):
+ from litellm.proxy.proxy_server import (
+ add_litellm_data_to_request,
+ general_settings,
+ llm_router,
+ proxy_config,
+ proxy_logging_obj,
+ user_api_base,
+ user_max_tokens,
+ user_model,
+ user_request_timeout,
+ user_temperature,
+ version,
+ )
+
+ data = {}
+ try:
+ body = await request.body()
+ body_str = body.decode()
+ try:
+ data = ast.literal_eval(body_str)
+ except Exception:
+ data = json.loads(body_str)
+
+ data["adapter_id"] = adapter_id
+
+ verbose_proxy_logger.debug(
+ "Request received by LiteLLM:\n{}".format(json.dumps(data, indent=4)),
+ )
+ data["model"] = (
+ general_settings.get("completion_model", None) # server default
+ or user_model # model name passed via cli args
+ or data.get("model", None) # default passed in http request
+ )
+ if user_model:
+ data["model"] = user_model
+
+ data = await add_litellm_data_to_request(
+ data=data, # type: ignore
+ request=request,
+ general_settings=general_settings,
+ user_api_key_dict=user_api_key_dict,
+ version=version,
+ proxy_config=proxy_config,
+ )
+
+ # override with user settings, these are params passed via cli
+ if user_temperature:
+ data["temperature"] = user_temperature
+ if user_request_timeout:
+ data["request_timeout"] = user_request_timeout
+ if user_max_tokens:
+ data["max_tokens"] = user_max_tokens
+ if user_api_base:
+ data["api_base"] = user_api_base
+
+ ### MODEL ALIAS MAPPING ###
+ # check if model name in model alias map
+ # get the actual model name
+ if data["model"] in litellm.model_alias_map:
+ data["model"] = litellm.model_alias_map[data["model"]]
+
+ # Check key-specific aliases
+ if (
+ isinstance(data["model"], str)
+ and user_api_key_dict.aliases
+ and isinstance(user_api_key_dict.aliases, dict)
+ and data["model"] in user_api_key_dict.aliases
+ ):
+ data["model"] = user_api_key_dict.aliases[data["model"]]
+
+ ### CALL HOOKS ### - modify incoming data before calling the model
+ data = await proxy_logging_obj.pre_call_hook( # type: ignore
+ user_api_key_dict=user_api_key_dict, data=data, call_type="text_completion"
+ )
+
+ ### ROUTE THE REQUESTs ###
+ router_model_names = llm_router.model_names if llm_router is not None else []
+ # skip router if user passed their key
+ if "api_key" in data:
+ llm_response = asyncio.create_task(litellm.aadapter_completion(**data))
+ elif (
+ llm_router is not None and data["model"] in router_model_names
+ ): # model in router model list
+ llm_response = asyncio.create_task(llm_router.aadapter_completion(**data))
+ elif (
+ llm_router is not None
+ and llm_router.model_group_alias is not None
+ and data["model"] in llm_router.model_group_alias
+ ): # model set in model_group_alias
+ llm_response = asyncio.create_task(llm_router.aadapter_completion(**data))
+ elif llm_router is not None and llm_router.has_model_id(
+ data["model"]
+ ): # model in router model list
+ llm_response = asyncio.create_task(llm_router.aadapter_completion(**data))
+ elif (
+ llm_router is not None
+ and data["model"] not in router_model_names
+ and (
+ llm_router.default_deployment is not None
+ or len(llm_router.pattern_router.patterns) > 0
+ )
+ ): # check for wildcard routes or default deployment before checking deployment_names
+ llm_response = asyncio.create_task(llm_router.aadapter_completion(**data))
+ elif (
+ llm_router is not None and data["model"] in llm_router.deployment_names
+ ): # model in router deployments, calling a specific deployment on the router (lowest priority)
+ llm_response = asyncio.create_task(
+ llm_router.aadapter_completion(**data, specific_deployment=True)
+ )
+ elif user_model is not None: # `litellm --model `
+ llm_response = asyncio.create_task(litellm.aadapter_completion(**data))
+ else:
+ raise HTTPException(
+ status_code=status.HTTP_400_BAD_REQUEST,
+ detail={
+ "error": "completion: Invalid model name passed in model="
+ + data.get("model", "")
+ },
+ )
+
+ # Await the llm_response task
+ response = await llm_response
+
+ hidden_params = getattr(response, "_hidden_params", {}) or {}
+ model_id = hidden_params.get("model_id", None) or ""
+ cache_key = hidden_params.get("cache_key", None) or ""
+ api_base = hidden_params.get("api_base", None) or ""
+ response_cost = hidden_params.get("response_cost", None) or ""
+
+ ### ALERTING ###
+ asyncio.create_task(
+ proxy_logging_obj.update_request_status(
+ litellm_call_id=data.get("litellm_call_id", ""), status="success"
+ )
+ )
+
+ verbose_proxy_logger.debug("final response: %s", response)
+
+ fastapi_response.headers.update(
+ ProxyBaseLLMRequestProcessing.get_custom_headers(
+ user_api_key_dict=user_api_key_dict,
+ model_id=model_id,
+ cache_key=cache_key,
+ api_base=api_base,
+ version=version,
+ response_cost=response_cost,
+ )
+ )
+
+ verbose_proxy_logger.debug("\nResponse from Litellm:\n{}".format(response))
+ return response
+ except Exception as e:
+ await proxy_logging_obj.post_call_failure_hook(
+ user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data
+ )
+ verbose_proxy_logger.exception(
+ "litellm.proxy.proxy_server.completion(): Exception occured - {}".format(
+ str(e)
+ )
+ )
+ error_msg = f"{str(e)}"
+ raise ProxyException(
+ message=getattr(e, "message", error_msg),
+ type=getattr(e, "type", "None"),
+ param=getattr(e, "param", "None"),
+ code=getattr(e, "status_code", 500),
+ )
+
+
+class HttpPassThroughEndpointHelpers(BasePassthroughUtils):
+ @staticmethod
+ def get_response_headers(
+ headers: httpx.Headers,
+ litellm_call_id: Optional[str] = None,
+ custom_headers: Optional[dict] = None,
+ ) -> dict:
+ # Exclude headers that uvicorn writes itself (server, date) and
+ # encoding/length headers that don't survive re-serialization.
+ # If we forward the upstream's Server header, uvicorn adds its
+ # own and strict HTTP parsers (e.g. aiohttp) reject the
+ # response with "Duplicate 'Server' header found".
+ excluded_headers = {
+ "transfer-encoding",
+ "content-encoding",
+ "content-length",
+ "server",
+ "date",
+ "connection",
+ "keep-alive",
+ }
+
+ return_headers = {
+ key: value
+ for key, value in headers.items()
+ if key.lower() not in excluded_headers
+ }
+ if litellm_call_id:
+ return_headers["x-litellm-call-id"] = litellm_call_id
+ if custom_headers:
+ # Ensure custom headers don't override actual upstream response headers or let framework defaults (like content-length: 0) interfere.
+ sanitized_custom_headers = {
+ key: value
+ for key, value in custom_headers.items()
+ if key.lower() not in excluded_headers
+ }
+ return_headers.update(sanitized_custom_headers)
+
+ return return_headers
+
+ @staticmethod
+ def get_endpoint_type(url: str) -> EndpointType:
+ parsed_url = urlparse(url)
+ if (
+ ("generateContent") in url
+ or ("streamGenerateContent") in url
+ or ("rawPredict") in url
+ or ("streamRawPredict") in url
+ ):
+ return EndpointType.VERTEX_AI
+ elif parsed_url.hostname == "api.anthropic.com":
+ return EndpointType.ANTHROPIC
+ elif (
+ parsed_url.hostname == "api.openai.com"
+ or parsed_url.hostname == "openai.azure.com"
+ or (parsed_url.hostname and "openai.com" in parsed_url.hostname)
+ ):
+ return EndpointType.OPENAI
+ return EndpointType.GENERIC
+
+ @staticmethod
+ async def _make_non_streaming_http_request(
+ request: Request,
+ async_client: httpx.AsyncClient,
+ url: str,
+ headers: dict,
+ requested_query_params: Optional[dict] = None,
+ custom_body: Optional[dict] = None,
+ ) -> httpx.Response:
+ """
+ Make a non-streaming HTTP request
+
+ If request is GET, don't include a JSON body
+ """
+ if request.method == "GET":
+ response = await async_client.request(
+ method=request.method,
+ url=url,
+ headers=headers,
+ params=requested_query_params,
+ )
+ else:
+ response = await async_client.request(
+ method=request.method,
+ url=url,
+ headers=headers,
+ params=requested_query_params,
+ json=custom_body,
+ )
+ return response
+
+ @staticmethod
+ async def non_streaming_http_request_handler(
+ request: Request,
+ async_client: httpx.AsyncClient,
+ url: httpx.URL,
+ headers: dict,
+ requested_query_params: Optional[dict] = None,
+ _parsed_body: Optional[dict] = None,
+ forward_multipart: bool = False,
+ ) -> httpx.Response:
+ """
+ Handle non-streaming HTTP requests
+
+ Handles special cases when GET requests, multipart/form-data requests, and generic httpx requests
+ """
+ if request.method == "GET":
+ response = await async_client.request(
+ method=request.method,
+ url=url,
+ headers=headers,
+ params=requested_query_params,
+ )
+ elif (
+ HttpPassThroughEndpointHelpers.is_multipart(request) is True
+ and forward_multipart
+ ):
+ # Forward multipart via make_multipart_http_request even when _parsed_body is
+ # non-empty (pass_through_request always injects litellm_logging_obj, etc.).
+ # forward_multipart is False when custom_body was supplied (JSON body despite
+ # multipart content-type) — those requests use the generic json= path.
+ return await HttpPassThroughEndpointHelpers.make_multipart_http_request(
+ request=request,
+ async_client=async_client,
+ url=url,
+ headers=headers,
+ requested_query_params=requested_query_params,
+ )
+ else:
+ # Generic httpx method
+ response = await async_client.request(
+ method=request.method,
+ url=url,
+ headers=headers,
+ params=requested_query_params,
+ json=_parsed_body,
+ )
+ return response
+
+ @staticmethod
+ def is_multipart(request: Request) -> bool:
+ """Check if the request is a multipart/form-data request"""
+ return "multipart/form-data" in request.headers.get("content-type", "")
+
+ @staticmethod
+ async def _build_request_files_from_upload_file(
+ upload_file: Union[UploadFile, StarletteUploadFile],
+ ) -> Tuple[Optional[str], bytes, Optional[str]]:
+ """Build a request files dict from an UploadFile object"""
+ file_content = await upload_file.read()
+ return (upload_file.filename, file_content, upload_file.content_type)
+
+ @staticmethod
+ async def make_multipart_http_request(
+ request: Request,
+ async_client: httpx.AsyncClient,
+ url: httpx.URL,
+ headers: dict,
+ requested_query_params: Optional[dict] = None,
+ stream: bool = False,
+ ) -> httpx.Response:
+ """Process multipart/form-data requests, handling both files and form fields"""
+ form_data = await request.form()
+ files = {}
+ form_data_dict = {}
+
+ for field_name, field_value in form_data.items():
+ if isinstance(field_value, (StarletteUploadFile, UploadFile)):
+ files[field_name] = (
+ await HttpPassThroughEndpointHelpers._build_request_files_from_upload_file(
+ upload_file=field_value
+ )
+ )
+ else:
+ form_data_dict[field_name] = field_value
+
+ # Remove content-type header - httpx will set it correctly with the new boundary
+ # when it creates the multipart body from files/data parameters
+ headers_copy = headers.copy()
+ headers_copy.pop("content-type", None)
+
+ # httpx.AsyncClient.request() does not accept stream=; use send() for streaming.
+ if stream:
+ req = async_client.build_request(
+ request.method,
+ url,
+ headers=headers_copy,
+ params=requested_query_params,
+ files=files,
+ data=form_data_dict,
+ )
+ return await async_client.send(req, stream=True)
+
+ return await async_client.request(
+ method=request.method,
+ url=url,
+ headers=headers_copy,
+ params=requested_query_params,
+ files=files,
+ data=form_data_dict,
+ )
+
+ @staticmethod
+ def _init_kwargs_for_pass_through_endpoint(
+ request: Request,
+ user_api_key_dict: UserAPIKeyAuth,
+ passthrough_logging_payload: PassthroughStandardLoggingPayload,
+ logging_obj: LiteLLMLoggingObj,
+ _parsed_body: Optional[dict] = None,
+ litellm_call_id: Optional[str] = None,
+ ) -> dict:
+ """
+ Filter out litellm params from the request body
+ """
+ from litellm.types.utils import all_litellm_params
+
+ _parsed_body = _parsed_body or {}
+
+ litellm_params_in_body = {}
+ for k in all_litellm_params:
+ if k in _parsed_body:
+ litellm_params_in_body[k] = _parsed_body.pop(k, None)
+
+ _metadata = dict(
+ LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(
+ user_api_key_dict=user_api_key_dict
+ )
+ )
+
+ litellm_metadata = litellm_params_in_body.pop("litellm_metadata", None)
+ metadata = litellm_params_in_body.pop("metadata", None)
+ if litellm_metadata:
+ _metadata.update(litellm_metadata)
+ if metadata:
+ _metadata.update(metadata)
+
+ _metadata = _update_metadata_with_tags_in_header(
+ request=request,
+ metadata=_metadata,
+ )
+
+ # Set internal keys after merging client-supplied metadata so a request
+ # body that mirrors them cannot clobber the authenticated key or the
+ # real parent span.
+ _metadata["user_api_key"] = user_api_key_dict.api_key
+ _metadata["litellm_parent_otel_span"] = user_api_key_dict.parent_otel_span
+
+ kwargs = {
+ "litellm_params": {
+ **litellm_params_in_body, # type: ignore
+ "metadata": _metadata,
+ "proxy_server_request": {
+ "url": str(request.url),
+ "method": request.method,
+ "body": copy.copy(_parsed_body), # use copy instead of deepcopy
+ "headers": request.headers,
+ },
+ },
+ "call_type": "pass_through_endpoint",
+ "litellm_call_id": litellm_call_id,
+ "passthrough_logging_payload": passthrough_logging_payload,
+ }
+
+ logging_obj.model_call_details["passthrough_logging_payload"] = (
+ passthrough_logging_payload
+ )
+
+ return kwargs
+
+ @staticmethod
+ def construct_target_url_with_subpath(
+ base_target: str, subpath: str, include_subpath: Optional[bool]
+ ) -> str:
+ """
+ Helper function to construct the full target URL with subpath handling.
+
+ Args:
+ base_target: The base target URL
+ subpath: The captured subpath from the request
+ include_subpath: Whether to include the subpath in the target URL
+
+ Returns:
+ The constructed full target URL
+ """
+ if not include_subpath:
+ return base_target
+
+ if not subpath:
+ return base_target
+
+ # Ensure base_target ends with / and subpath doesn't start with /
+ if not base_target.endswith("/"):
+ base_target = base_target + "/"
+ if subpath.startswith("/"):
+ subpath = subpath[1:]
+
+ # Resolve any '..' segments in the subpath so it cannot climb above
+ # the base_target prefix that the operator configured. Preserve a
+ # trailing slash on the original subpath since some upstreams treat
+ # `/foo` and `/foo/` as different resources.
+ trailing_slash = subpath.endswith("/")
+ safe_subpath = posixpath.normpath("/" + subpath).lstrip("/")
+ if safe_subpath == ".":
+ safe_subpath = ""
+ if trailing_slash and safe_subpath and not safe_subpath.endswith("/"):
+ safe_subpath += "/"
+
+ return base_target + safe_subpath
+
+ @staticmethod
+ def join_base_and_endpoint_path(base_url: httpx.URL, endpoint_path: str) -> str:
+ """
+ Combine the path component of ``base_url`` with ``endpoint_path``.
+
+ Preserves any path prefix configured on the base URL and resolves
+ ``..`` segments in the endpoint so the result stays within the base
+ path. A trailing slash on ``endpoint_path`` is preserved.
+ """
+ trailing_slash = endpoint_path.endswith("/")
+ base_path = base_url.path or ""
+ if not base_path or base_path == "/":
+ normalized_endpoint = posixpath.normpath("/" + endpoint_path.lstrip("/"))
+ if trailing_slash and normalized_endpoint != "/":
+ normalized_endpoint += "/"
+ return normalized_endpoint
+
+ base_path = base_path.rstrip("/")
+ clean_endpoint = endpoint_path.lstrip("/")
+ combined = posixpath.normpath(base_path + "/" + clean_endpoint)
+ # If normalization climbs out of the base path, fall back to base.
+ if combined != base_path and not combined.startswith(base_path + "/"):
+ return base_path + "/"
+ if trailing_slash and not combined.endswith("/"):
+ combined += "/"
+ return combined
+
+ @staticmethod
+ def _update_stream_param_based_on_request_body(
+ parsed_body: dict,
+ stream: Optional[bool] = None,
+ ) -> Optional[bool]:
+ """
+ If stream is provided in the request body, use it.
+ Otherwise, use the stream parameter passed to the `pass_through_request` function
+ """
+ if "stream" in parsed_body:
+ return parsed_body.get("stream", stream)
+ return stream
+
+
+async def pass_through_request( # noqa: PLR0915
+ request: Request,
+ target: str,
+ custom_headers: dict,
+ user_api_key_dict: UserAPIKeyAuth,
+ custom_body: Optional[dict] = None,
+ forward_headers: Optional[bool] = False,
+ merge_query_params: Optional[bool] = False,
+ query_params: Optional[dict] = None,
+ default_query_params: Optional[dict] = None,
+ stream: Optional[bool] = None,
+ cost_per_request: Optional[float] = None,
+ custom_llm_provider: Optional[str] = None,
+ guardrails_config: Optional[dict] = None,
+):
+ """
+ Pass through endpoint handler, makes the httpx request for pass-through endpoints and ensures logging hooks are called
+
+ Args:
+ request: The incoming request
+ target: The target URL
+ custom_headers: The custom headers
+ user_api_key_dict: The user API key dictionary
+ custom_body: The custom body
+ forward_headers: Whether to forward headers
+ merge_query_params: Whether to merge query params
+ query_params: The query params
+ default_query_params: The default query params to be applied if not overridden by client
+ stream: Whether to stream the response
+ cost_per_request: Optional field - cost per request to the target endpoint
+ custom_llm_provider: Optional field - custom LLM provider for the endpoint
+ guardrails_config: Optional field - guardrails configuration for passthrough endpoint
+ """
+ from litellm.exceptions import ModifyResponseException
+ from litellm.litellm_core_utils.litellm_logging import Logging
+ from litellm.proxy.pass_through_endpoints.passthrough_guardrails import (
+ PassthroughGuardrailHandler,
+ )
+ from litellm.proxy.proxy_server import proxy_logging_obj
+
+ #########################################################
+ # Initialize variables
+ #########################################################
+ litellm_call_id = str(uuid.uuid4())
+ url: Optional[httpx.URL] = None
+
+ # parsed request body
+ _parsed_body: Optional[dict] = None
+ # kwargs for pass through endpoint, contains metadata, litellm_params, call_type, litellm_call_id, passthrough_logging_payload
+ kwargs: Optional[dict] = None
+ logging_obj: Optional[Logging] = None
+
+ #########################################################
+ try:
+ url = httpx.URL(target)
+ headers = custom_headers
+ headers = HttpPassThroughEndpointHelpers.forward_headers_from_request(
+ request_headers=_safe_get_request_headers(request).copy(),
+ headers=headers,
+ forward_headers=forward_headers,
+ )
+
+ # Apply default query parameters if provided, regardless of merge_query_params setting
+ if default_query_params or merge_query_params:
+ # Determine what to merge based on settings
+ request_params = dict(request.query_params) if merge_query_params else {}
+
+ # Create a new URL with the merged query params
+ url = url.copy_with(
+ query=urlencode(
+ HttpPassThroughEndpointHelpers.get_merged_query_parameters(
+ existing_url=url,
+ request_query_params=request_params,
+ default_query_params=default_query_params,
+ )
+ ).encode("ascii")
+ )
+
+ endpoint_type: EndpointType = HttpPassThroughEndpointHelpers.get_endpoint_type(
+ str(url)
+ )
+
+ # SigV4-signed callers (e.g. Bedrock) attach the exact bytes that were
+ # signed via request.state; we must send those instead of re-encoding the
+ # parsed dict (hooks mutate it, breaking the signature / Content-Length).
+ # Tolerate request objects without `state` (test fixtures) and only honor
+ # values httpx accepts for `content=`.
+ _request_state = getattr(request, "state", None)
+ state_raw_body: Optional[Union[str, bytes]] = (
+ getattr(_request_state, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, None)
+ if _request_state is not None
+ else None
+ )
+ if state_raw_body is not None and not isinstance(
+ state_raw_body, (str, bytes, bytearray)
+ ):
+ state_raw_body = None
+
+ # Skip body parsing for multipart requests - make_multipart_http_request will handle it
+ # But if custom_body is provided (e.g., JSON parsed despite multipart content-type), use it
+ is_multipart = (
+ HttpPassThroughEndpointHelpers.is_multipart(request) and not custom_body
+ )
+
+ if custom_body:
+ _parsed_body = custom_body
+ elif is_multipart:
+ # Don't parse multipart body here - it will be handled by make_multipart_http_request
+ _parsed_body = {}
+ else:
+ _parsed_body = await _read_request_body(request)
+ verbose_proxy_logger.debug(
+ "Pass through endpoint sending request to \nURL {}\nheaders: {}\nbody: {}\n".format(
+ url, headers, _parsed_body
+ )
+ )
+
+ ### COLLECT GUARDRAILS FOR PASSTHROUGH ENDPOINT ###
+ # Passthrough endpoints are opt-in only for guardrails
+ # When enabled, collect guardrails from org/team/key levels + passthrough-specific
+ guardrails_to_run = PassthroughGuardrailHandler.collect_guardrails(
+ user_api_key_dict=user_api_key_dict,
+ passthrough_guardrails_config=guardrails_config,
+ )
+
+ # Add guardrails to metadata if any should run
+ if guardrails_to_run and len(guardrails_to_run) > 0:
+ if _parsed_body is None:
+ _parsed_body = {}
+ if "metadata" not in _parsed_body:
+ _parsed_body["metadata"] = {}
+ _parsed_body["metadata"]["guardrails"] = guardrails_to_run
+ verbose_proxy_logger.debug(
+ f"Added guardrails to passthrough request metadata: {guardrails_to_run}"
+ )
+
+ ## LOGGING OBJECT ## - initialize before pre_call_hook so guardrails can access it
+ # Surface the requested model (when the body carries one) so logging/spans
+ # read e.g. ``chat gpt-4o`` instead of ``chat unknown``.
+ passthrough_model = (
+ _parsed_body.get("model") if isinstance(_parsed_body, dict) else None
+ ) or "unknown"
+ start_time = datetime.now()
+ logging_obj = Logging(
+ model=passthrough_model,
+ messages=[{"role": "user", "content": safe_dumps(_parsed_body)}],
+ stream=False,
+ call_type="pass_through_endpoint",
+ start_time=start_time,
+ litellm_call_id=litellm_call_id,
+ function_id="1245",
+ )
+
+ # Store passthrough guardrails config on logging_obj for field targeting
+ logging_obj.passthrough_guardrails_config = guardrails_config
+
+ # Store logging_obj in data so guardrails can access it
+ if _parsed_body is None:
+ _parsed_body = {}
+ _parsed_body["litellm_logging_obj"] = logging_obj
+
+ ### CALL HOOKS ### - modify incoming data / reject request before calling the model
+ _parsed_body = await proxy_logging_obj.pre_call_hook(
+ user_api_key_dict=user_api_key_dict,
+ data=_parsed_body,
+ call_type="pass_through_endpoint",
+ )
+ async_client_obj = get_async_httpx_client(
+ llm_provider=httpxSpecialProvider.PassThroughEndpoint,
+ params={"timeout": 600},
+ )
+ async_client = async_client_obj.client
+ passthrough_logging_payload = PassthroughStandardLoggingPayload(
+ url=str(url),
+ request_body=_parsed_body,
+ request_method=getattr(request, "method", None),
+ cost_per_request=cost_per_request,
+ )
+ kwargs = HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint(
+ user_api_key_dict=user_api_key_dict,
+ _parsed_body=_parsed_body,
+ passthrough_logging_payload=passthrough_logging_payload,
+ litellm_call_id=litellm_call_id,
+ request=request,
+ logging_obj=logging_obj,
+ )
+
+ # Store custom_llm_provider in kwargs and logging object if provided
+ if custom_llm_provider:
+ logging_obj.model_call_details["custom_llm_provider"] = custom_llm_provider
+ logging_obj.model_call_details["litellm_params"] = kwargs.get(
+ "litellm_params", {}
+ )
+
+ # done for supporting 'parallel_request_limiter.py' with pass-through endpoints
+ logging_obj.update_environment_variables(
+ model=passthrough_model,
+ user="unknown",
+ optional_params={},
+ litellm_params=kwargs["litellm_params"],
+ call_type="pass_through_endpoint",
+ )
+ logging_obj.model_call_details["litellm_call_id"] = litellm_call_id
+
+ # combine url with query params for logging
+ requested_query_params: Optional[dict] = query_params or dict(
+ request.query_params
+ )
+
+ ## PASSTHROUGH MANAGED ID RESOLUTION (INPUT) ##
+ # Resolve managed IDs in path, query params, and body back to raw
+ # provider IDs before forwarding upstream. Gated by feature flag and
+ # enterprise managed-files hook. Runs after pre_call_hook so
+ # guardrails have already seen the managed IDs.
+ from litellm.proxy.proxy_server import (
+ general_settings as proxy_general_settings,
+ )
+
+ _managed_id_provider = resolve_passthrough_managed_id_provider(
+ custom_llm_provider
+ )
+
+ if (
+ proxy_general_settings.get("passthrough_managed_object_ids", False)
+ and _managed_id_provider is not None
+ ):
+ verbose_proxy_logger.debug(
+ "pass_through_endpoint: managed-id input rewrite enabled for route=%s method=%s",
+ request.url.path,
+ request.method,
+ )
+ _passthrough_managed_hook = proxy_logging_obj.get_proxy_hook(
+ "managed_files"
+ )
+ if _passthrough_managed_hook is not None:
+ from litellm.proxy.pass_through_endpoints.managed_id_rewriter import (
+ rewrite_body_ids,
+ rewrite_path_ids,
+ rewrite_query_ids,
+ )
+ from litellm.proxy.proxy_server import (
+ prisma_client as _passthrough_prisma,
+ )
+
+ _original_path = url.path
+ _original_query_params = requested_query_params
+ _original_body = _parsed_body
+ _new_path = await rewrite_path_ids(
+ url.path,
+ _managed_id_provider,
+ user_api_key_dict,
+ _passthrough_prisma,
+ _passthrough_managed_hook,
+ )
+ if _new_path != url.path:
+ url = url.copy_with(path=_new_path)
+ requested_query_params = await rewrite_query_ids(
+ requested_query_params,
+ _managed_id_provider,
+ user_api_key_dict,
+ _passthrough_prisma,
+ _passthrough_managed_hook,
+ )
+ _parsed_body = await rewrite_body_ids(
+ _parsed_body,
+ _managed_id_provider,
+ user_api_key_dict,
+ _passthrough_prisma,
+ _passthrough_managed_hook,
+ )
+ verbose_proxy_logger.debug(
+ "pass_through_endpoint: managed-id input rewrite results path_changed=%s query_changed=%s body_changed=%s route=%s method=%s",
+ _new_path != _original_path,
+ requested_query_params is not _original_query_params,
+ _parsed_body is not _original_body,
+ request.url.path,
+ request.method,
+ )
+ else:
+ verbose_proxy_logger.debug(
+ "pass_through_endpoint: managed-id input rewrite skipped (managed_files hook not available) route=%s method=%s",
+ request.url.path,
+ request.method,
+ )
+
+ ## PASSTHROUGH MANAGED LIST (DB-only response) ##
+ # For GET /v1/files and GET /v1/batches passthrough routes, serve the
+ # listing entirely from our DB so each caller only sees their own IDs.
+ # Admins / master-key callers see all rows. Gated on the same
+ # conditions as INPUT/OUTPUT rewrite: feature flag, provider, AND
+ # the managed_files hook must be present. Without the hook no managed
+ # IDs are ever minted or stored, so the DB is empty and intercepting
+ # the list would silently hide the caller's real upstream files/batches.
+ if (
+ proxy_general_settings.get("passthrough_managed_object_ids", False)
+ and _managed_id_provider is not None
+ and request.method == "GET"
+ and proxy_logging_obj.get_proxy_hook("managed_files") is not None
+ ):
+ from litellm.proxy.auth.auth_utils import get_request_route
+ from litellm.proxy.pass_through_endpoints.managed_id_rewriter import (
+ is_passthrough_list_route,
+ list_passthrough_ids_from_db,
+ )
+ from litellm.proxy.proxy_server import (
+ prisma_client as _list_prisma,
+ )
+
+ if (
+ is_passthrough_list_route(
+ _managed_id_provider, request.method, get_request_route(request)
+ )
+ and _list_prisma is not None
+ ):
+ _list_result = await list_passthrough_ids_from_db(
+ provider=_managed_id_provider,
+ route=get_request_route(request),
+ user_api_key_dict=user_api_key_dict,
+ prisma_client=_list_prisma,
+ query_params=dict(request.query_params),
+ )
+ if _list_result is not None:
+ verbose_proxy_logger.debug(
+ "pass_through_endpoint: list served from DB route=%s count=%d",
+ request.url.path,
+ len(_list_result.get("data", [])),
+ )
+ return Response(
+ content=json.dumps(_list_result),
+ status_code=200,
+ media_type="application/json",
+ )
+
+ requested_query_params_str = None
+ if requested_query_params:
+ requested_query_params_str = "&".join(
+ f"{k}={v}" for k, v in requested_query_params.items()
+ )
+
+ logging_url = str(url)
+ if requested_query_params_str:
+ if "?" in str(url):
+ logging_url = str(url) + "&" + requested_query_params_str
+ else:
+ logging_url = str(url) + "?" + requested_query_params_str
+
+ logging_obj.pre_call(
+ input=[{"role": "user", "content": safe_dumps(_parsed_body)}],
+ api_key="",
+ additional_args={
+ "complete_input_dict": _parsed_body,
+ "api_base": str(logging_url),
+ "headers": headers,
+ },
+ )
+ stream = (
+ HttpPassThroughEndpointHelpers._update_stream_param_based_on_request_body(
+ parsed_body=_parsed_body or {},
+ stream=stream,
+ )
+ )
+
+ if stream:
+ if is_multipart:
+ response = (
+ await HttpPassThroughEndpointHelpers.make_multipart_http_request(
+ request=request,
+ async_client=async_client,
+ url=url,
+ headers=headers,
+ requested_query_params=requested_query_params,
+ stream=True,
+ )
+ )
+ else:
+ # SigV4-signed callers (Bedrock) supply the exact pre-signed bytes;
+ # otherwise httpx encodes the parsed JSON dict as before.
+ body_kwargs: Dict[str, Any] = (
+ {"content": state_raw_body}
+ if state_raw_body is not None
+ else {"json": _parsed_body}
+ )
+ req = async_client.build_request(
+ request.method,
+ url,
+ params=requested_query_params,
+ headers=headers,
+ **body_kwargs,
+ )
+
+ response = await async_client.send(req, stream=stream)
+
+ try:
+ response.raise_for_status()
+ except httpx.HTTPStatusError as e:
+ raise HTTPException(
+ status_code=e.response.status_code, detail=await e.response.aread()
+ )
+
+ return StreamingResponse(
+ PassThroughStreamingHandler.chunk_processor(
+ response=response,
+ request_body=_parsed_body,
+ litellm_logging_obj=logging_obj,
+ endpoint_type=endpoint_type,
+ start_time=start_time,
+ passthrough_success_handler_obj=pass_through_endpoint_logging,
+ url_route=str(url),
+ ),
+ headers=HttpPassThroughEndpointHelpers.get_response_headers(
+ headers=response.headers,
+ litellm_call_id=litellm_call_id,
+ ),
+ status_code=response.status_code,
+ )
+
+ if state_raw_body is not None:
+ # SigV4-signed callers (Bedrock) require the exact pre-signed bytes
+ # to be forwarded so the signature/Content-Length stay valid.
+ response = await async_client.request(
+ method=request.method,
+ url=url,
+ headers=headers,
+ params=requested_query_params,
+ content=state_raw_body,
+ )
+ else:
+ response = (
+ await HttpPassThroughEndpointHelpers.non_streaming_http_request_handler(
+ request=request,
+ async_client=async_client,
+ url=url,
+ headers=headers,
+ requested_query_params=requested_query_params,
+ _parsed_body=_parsed_body,
+ forward_multipart=is_multipart,
+ )
+ )
+ verbose_proxy_logger.debug("response.headers= %s", response.headers)
+
+ if _is_streaming_response(response) is True:
+ try:
+ response.raise_for_status()
+ except httpx.HTTPStatusError as e:
+ raise HTTPException(
+ status_code=e.response.status_code, detail=await e.response.aread()
+ )
+
+ return StreamingResponse(
+ PassThroughStreamingHandler.chunk_processor(
+ response=response,
+ request_body=_parsed_body,
+ litellm_logging_obj=logging_obj,
+ endpoint_type=endpoint_type,
+ start_time=start_time,
+ passthrough_success_handler_obj=pass_through_endpoint_logging,
+ url_route=str(url),
+ ),
+ headers=HttpPassThroughEndpointHelpers.get_response_headers(
+ headers=response.headers,
+ litellm_call_id=litellm_call_id,
+ ),
+ status_code=response.status_code,
+ )
+
+ try:
+ response.raise_for_status()
+ except httpx.HTTPStatusError as e:
+ raise HTTPException(
+ status_code=e.response.status_code, detail=e.response.text
+ )
+
+ if response.status_code >= 300:
+ raise HTTPException(status_code=response.status_code, detail=response.text)
+
+ content = await response.aread()
+
+ ## POST-CALL GUARDRAILS ##
+ _content_modified = False
+ response_body: Optional[dict] = get_response_body(response)
+ if response_body is not None and guardrails_to_run:
+ # Build an enriched data dict: _parsed_body has been stripped of
+ # `metadata` by both pre_call_hook and _init_kwargs_for_pass_through_endpoint,
+ # so we re-attach the configured guardrails here so should_run_guardrail
+ # sees them.
+ hook_data = dict(_parsed_body or {})
+ existing_metadata = hook_data.get("metadata")
+ if not isinstance(existing_metadata, dict):
+ existing_metadata = {}
+ hook_data["metadata"] = {
+ **existing_metadata,
+ "guardrails": guardrails_to_run,
+ }
+ response_body = await proxy_logging_obj.post_call_success_hook(
+ data=hook_data,
+ user_api_key_dict=user_api_key_dict,
+ response=response_body, # type: ignore[arg-type]
+ )
+ if isinstance(response_body, dict):
+ content = json.dumps(response_body).encode("utf-8")
+ _content_modified = True
+ else:
+ verbose_proxy_logger.debug(
+ "pass_through_endpoint: post_call_success_hook returned %s, expected dict — using original response",
+ type(response_body).__name__,
+ )
+ elif response_body is None:
+ verbose_proxy_logger.debug(
+ "pass_through_endpoint: response body not JSON-parseable, skipping post-call guardrails"
+ )
+
+ ## PASSTHROUGH MANAGED ID MINTING (OUTPUT) ##
+ # Mint managed IDs for raw provider IDs in the response body and swap
+ # them before the response reaches the client. Runs after guardrails
+ # so guardrails see the raw IDs (cleaner) and the client receives the
+ # managed IDs. Gated by feature flag and enterprise managed-files hook.
+ if (
+ proxy_general_settings.get("passthrough_managed_object_ids", False)
+ and _managed_id_provider is not None
+ and isinstance(response_body, dict)
+ and response.status_code < 300
+ ):
+ verbose_proxy_logger.debug(
+ "pass_through_endpoint: managed-id output rewrite enabled for route=%s method=%s status=%s",
+ request.url.path,
+ request.method,
+ response.status_code,
+ )
+ _passthrough_managed_hook = proxy_logging_obj.get_proxy_hook(
+ "managed_files"
+ )
+ if _passthrough_managed_hook is not None:
+ from litellm.proxy.auth.auth_utils import get_request_route
+ from litellm.proxy.pass_through_endpoints.managed_id_rewriter import (
+ rewrite_response_ids,
+ )
+ from litellm.proxy.proxy_server import (
+ prisma_client as _passthrough_prisma,
+ )
+
+ _new_body = await rewrite_response_ids(
+ provider=_managed_id_provider,
+ method=request.method,
+ route=get_request_route(request),
+ body=response_body,
+ user_api_key_dict=user_api_key_dict,
+ prisma_client=_passthrough_prisma,
+ managed_files_hook=_passthrough_managed_hook,
+ )
+ if _new_body is not response_body:
+ response_body = _new_body
+ content = json.dumps(response_body).encode("utf-8")
+ _content_modified = True
+ verbose_proxy_logger.debug(
+ "pass_through_endpoint: managed-id output rewrite applied route=%s method=%s",
+ request.url.path,
+ request.method,
+ )
+ else:
+ verbose_proxy_logger.debug(
+ "pass_through_endpoint: managed-id output rewrite no-op route=%s method=%s",
+ request.url.path,
+ request.method,
+ )
+ else:
+ verbose_proxy_logger.debug(
+ "pass_through_endpoint: managed-id output rewrite skipped (managed_files hook not available) route=%s method=%s",
+ request.url.path,
+ request.method,
+ )
+
+ ## LOG SUCCESS
+ passthrough_logging_payload["response_body"] = response_body
+ end_time = datetime.now()
+ asyncio.create_task(
+ pass_through_endpoint_logging.pass_through_async_success_handler(
+ httpx_response=response,
+ response_body=response_body,
+ url_route=str(url),
+ result="",
+ start_time=start_time,
+ end_time=end_time,
+ logging_obj=logging_obj,
+ cache_hit=False,
+ request_body=_parsed_body or {},
+ custom_llm_provider=custom_llm_provider,
+ **kwargs,
+ )
+ )
+
+ ## CUSTOM HEADERS - `x-litellm-*`
+ custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers(
+ user_api_key_dict=user_api_key_dict,
+ call_id=litellm_call_id,
+ model_id=None,
+ cache_key=None,
+ api_base=str(url._uri_reference),
+ )
+
+ response_headers = HttpPassThroughEndpointHelpers.get_response_headers(
+ headers=response.headers,
+ custom_headers=custom_headers,
+ )
+ if _content_modified:
+ response_headers.pop("content-length", None)
+
+ return Response(
+ content=content,
+ status_code=response.status_code,
+ headers=response_headers,
+ )
+ except ModifyResponseException as e:
+ verbose_proxy_logger.info(
+ "pass_through_endpoint: Guardrail %s modified response: %s",
+ e.guardrail_name,
+ str(e.message or "")[:200],
+ )
+ try:
+ await proxy_logging_obj.post_call_failure_hook(
+ user_api_key_dict=user_api_key_dict,
+ original_exception=e,
+ request_data=e.request_data,
+ )
+ except Exception:
+ verbose_proxy_logger.warning(
+ "pass_through_endpoint: post_call_failure_hook raised during guardrail block",
+ exc_info=True,
+ )
+ error_body = {
+ "error": {
+ "message": e.message or "Response blocked by guardrail",
+ "type": "content_filter",
+ "guardrail_name": e.guardrail_name,
+ "model": e.model,
+ }
+ }
+ return Response(
+ content=json.dumps(error_body),
+ status_code=200,
+ media_type="application/json",
+ )
+ except Exception as e:
+ custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers(
+ user_api_key_dict=user_api_key_dict,
+ call_id=litellm_call_id,
+ model_id=None,
+ cache_key=None,
+ api_base=str(url._uri_reference) if url else None,
+ )
+ verbose_proxy_logger.exception(
+ "litellm.proxy.proxy_server.pass_through_endpoint(): Exception occured - {}".format(
+ str(e)
+ )
+ )
+
+ #########################################################
+ # Monitoring: Trigger post_call_failure_hook
+ # for pass through endpoint failure
+ #########################################################
+ request_payload: dict = _parsed_body or {}
+ # add user_api_key_dict, litellm_call_id, passthrough_logging_payloa for logging
+ if kwargs:
+ for key, value in kwargs.items():
+ request_payload[key] = value
+ if logging_obj is not None:
+ request_payload["litellm_logging_obj"] = logging_obj
+
+ if (
+ "model" not in request_payload
+ and _parsed_body
+ and isinstance(_parsed_body, dict)
+ ):
+ request_payload["model"] = _parsed_body.get("model", "")
+ if "custom_llm_provider" not in request_payload and custom_llm_provider:
+ request_payload["custom_llm_provider"] = custom_llm_provider
+
+ await proxy_logging_obj.post_call_failure_hook(
+ user_api_key_dict=user_api_key_dict,
+ original_exception=e,
+ request_data=request_payload,
+ traceback_str=traceback.format_exc(
+ limit=MAXIMUM_TRACEBACK_LINES_TO_LOG,
+ ),
+ )
+
+ #########################################################
+
+ if isinstance(e, HTTPException):
+ raise ProxyException(
+ message=getattr(e, "message", str(getattr(e, "detail", str(e)))),
+ type=getattr(e, "type", "None"),
+ param=getattr(e, "param", "None"),
+ code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST),
+ headers=custom_headers,
+ )
+ else:
+ error_msg = f"{str(e)}"
+ raise ProxyException(
+ message=getattr(e, "message", error_msg),
+ type=getattr(e, "type", "None"),
+ param=getattr(e, "param", "None"),
+ code=getattr(e, "status_code", 500),
+ headers=custom_headers,
+ )
+
+
+def _update_metadata_with_tags_in_header(request: Request, metadata: dict) -> dict:
+ """
+ If tags are in the request headers, add them to the metadata
+
+ Used for google and vertex JS SDKs, and Azure passthrough
+ Checks both 'tags' and 'x-litellm-tags' headers
+ """
+ tags_to_add = []
+
+ # Check for 'tags' header first
+ _tags = request.headers.get("tags")
+ if _tags:
+ tags_to_add.extend([tag.strip() for tag in _tags.split(",")])
+
+ _tags = request.headers.get("x-litellm-tags")
+ if _tags:
+ tags_to_add.extend([tag.strip() for tag in _tags.split(",")])
+
+ # Only add tags key if there are tags to add
+ if tags_to_add:
+ if "tags" not in metadata:
+ metadata["tags"] = []
+ metadata["tags"].extend(tags_to_add)
+
+ return metadata
+
+
+async def _parse_request_data_by_content_type(
+ request: Request,
+) -> Tuple[Optional[Any], Optional[Any], Optional[Any], Optional[Any]]:
+ """
+ Parse request data based on content type.
+
+ Handles JSON, multipart/form-data, and URL-encoded form data.
+
+ Returns:
+ Tuple of (query_params_data, custom_body_data, file_data, stream)
+ """
+ content_type = request.headers.get("content-type", "")
+
+ query_params_data = None
+ custom_body_data = None
+ file_data = None
+ stream = None
+
+ if "application/json" in content_type:
+ # ✅ Handle JSON
+ try:
+ body = await request.json()
+ query_params_data = body.get("query_params")
+ custom_body_data = body.get("custom_body")
+ stream = body.get("stream")
+ except json.JSONDecodeError:
+ # Handle requests with no body (e.g., DELETE requests)
+ pass
+ elif "multipart/form-data" in content_type:
+ # ✅ Try to parse as JSON first (handles misconfigured clients sending JSON with multipart content-type)
+ # If that fails, skip parsing - pass_through_request will handle actual multipart
+ try:
+ body = await request.json()
+ # Successfully parsed as JSON - treat as JSON body
+ query_params_data = body.get("query_params")
+ custom_body_data = body.get("custom_body")
+ stream = body.get("stream")
+ # If custom_body is not set, use the entire body
+ if custom_body_data is None and body:
+ custom_body_data = body
+ except (json.JSONDecodeError, Exception):
+ # Not JSON - this is actual multipart data
+ # Skip parsing here to avoid consuming the request body stream
+ # make_multipart_http_request will handle it
+ pass
+
+ elif "application/x-www-form-urlencoded" in content_type:
+ # ✅ Handle URL-encoded form data
+ form = await request.form()
+ query_params_data = form.get("query_params")
+ custom_body_data = form.get("custom_body")
+
+ else:
+ # ✅ Fallback: maybe no body, just query params
+ query_params_data = dict(request.query_params) or None
+
+ return query_params_data, custom_body_data, file_data, stream
+
+
+def create_pass_through_route(
+ endpoint,
+ target: str,
+ custom_headers: Optional[Mapping[str, Any]] = None,
+ _forward_headers: Optional[bool] = False,
+ _merge_query_params: Optional[bool] = False,
+ dependencies: Optional[List] = None,
+ include_subpath: Optional[bool] = False,
+ cost_per_request: Optional[float] = None,
+ custom_llm_provider: Optional[str] = None,
+ is_streaming_request: Optional[bool] = False,
+ query_params: Optional[dict] = None,
+ default_query_params: Optional[dict] = None,
+ guardrails: Optional[Dict[str, Any]] = None,
+ config_file_path: Optional[str] = None,
+):
+ # check if target is an adapter.py or a url
+ from litellm._uuid import uuid
+ from litellm.proxy.types_utils.utils import get_instance_fn
+
+ try:
+ if isinstance(target, CustomLogger):
+ adapter = target
+ else:
+ adapter = get_instance_fn(value=target, config_file_path=config_file_path)
+ adapter_id = str(uuid.uuid4())
+ litellm.adapters = [{"id": adapter_id, "adapter": adapter}]
+
+ async def endpoint_func( # type: ignore
+ request: Request,
+ fastapi_response: Response,
+ user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
+ subpath: str = "", # captures sub-paths when include_subpath=True
+ ):
+ return await chat_completion_pass_through_endpoint(
+ fastapi_response=fastapi_response,
+ request=request,
+ adapter_id=adapter_id,
+ user_api_key_dict=user_api_key_dict,
+ )
+
+ except Exception:
+ verbose_proxy_logger.debug("Defaulting to target being a url.")
+
+ async def endpoint_func( # type: ignore
+ request: Request,
+ fastapi_response: Response,
+ user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
+ subpath: str = "", # captures sub-paths when include_subpath=True
+ ):
+ from litellm.proxy.auth.auth_utils import ( # noqa: PLC0415
+ get_request_route,
+ )
+ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
+ InitPassThroughEndpointHelpers,
+ )
+
+ path = get_request_route(request)
+
+ # Parse request data based on content type
+ (
+ query_params_data,
+ custom_body_data,
+ file_data,
+ stream,
+ ) = await _parse_request_data_by_content_type(request)
+
+ if not InitPassThroughEndpointHelpers.is_registered_pass_through_route(
+ route=path
+ ):
+ raise HTTPException(
+ status_code=404,
+ detail=f"Pass-through endpoint {endpoint} not found. This could have been deleted or not yet added to the proxy.",
+ )
+
+ passthrough_params = (
+ InitPassThroughEndpointHelpers.get_registered_pass_through_route(
+ route=path, method=request.method
+ )
+ )
+ target_params = {
+ "target": target,
+ "custom_headers": custom_headers,
+ "forward_headers": _forward_headers,
+ "merge_query_params": _merge_query_params,
+ "cost_per_request": cost_per_request,
+ "guardrails": None,
+ }
+
+ if passthrough_params is not None:
+ target_params.update(passthrough_params.get("passthrough_params", {}))
+
+ # Extract and cast parameters with proper types
+ param_target = target_params.get("target") or target
+ param_custom_headers = target_params.get("custom_headers", custom_headers)
+ param_forward_headers = target_params.get(
+ "forward_headers", _forward_headers
+ )
+ param_merge_query_params = target_params.get(
+ "merge_query_params", _merge_query_params
+ )
+ param_cost_per_request = target_params.get(
+ "cost_per_request", cost_per_request
+ )
+ param_guardrails = target_params.get("guardrails", None)
+ param_default_query_params = target_params.get("default_query_params", None)
+
+ # Construct the full target URL with subpath if needed
+ full_target = (
+ HttpPassThroughEndpointHelpers.construct_target_url_with_subpath(
+ base_target=cast(str, param_target),
+ subpath=subpath,
+ include_subpath=include_subpath,
+ )
+ )
+
+ # Ensure custom_headers is a dict. Botocore returns a HeadersDict
+ # for SigV4-prepared requests, which is a Mapping but not a dict.
+ headers_dict = (
+ dict(param_custom_headers)
+ if isinstance(param_custom_headers, Mapping)
+ else {}
+ )
+
+ # Ensure query_params and custom_body are dicts or None
+ final_query_params = (
+ query_params_data if isinstance(query_params_data, dict) else {}
+ )
+ if query_params:
+ final_query_params.update(query_params)
+ # Programmatic callers set LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY on
+ # request.state (see Bedrock proxy). Parsed JSON envelope otherwise.
+ state_custom_body: Optional[dict] = getattr(
+ request.state,
+ LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY,
+ None,
+ )
+ final_custom_body: Optional[dict] = None
+ if isinstance(state_custom_body, dict):
+ final_custom_body = state_custom_body
+ elif isinstance(custom_body_data, dict):
+ final_custom_body = custom_body_data
+
+ try:
+ return await pass_through_request( # type: ignore
+ request=request,
+ target=full_target,
+ custom_headers=headers_dict,
+ user_api_key_dict=user_api_key_dict,
+ forward_headers=cast(Optional[bool], param_forward_headers),
+ merge_query_params=cast(Optional[bool], param_merge_query_params),
+ query_params=final_query_params,
+ default_query_params=cast(
+ Optional[dict], param_default_query_params
+ ),
+ stream=is_streaming_request or stream,
+ custom_body=final_custom_body,
+ cost_per_request=cast(Optional[float], param_cost_per_request),
+ custom_llm_provider=custom_llm_provider,
+ guardrails_config=cast(Optional[dict], param_guardrails),
+ )
+ finally:
+ if hasattr(request.state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY):
+ delattr(request.state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY)
+ if hasattr(request.state, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY):
+ delattr(request.state, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY)
+
+ return endpoint_func
+
+
+def create_websocket_passthrough_route(
+ endpoint: str,
+ target: str,
+ custom_headers: Optional[dict] = None,
+ _forward_headers: Optional[bool] = False,
+ dependencies: Optional[List] = None,
+ cost_per_request: Optional[float] = None,
+):
+ """
+ Create a WebSocket passthrough route function.
+
+ Args:
+ endpoint: The endpoint path (for logging purposes)
+ target: The target WebSocket URL (e.g., "wss://api.example.com/ws")
+ custom_headers: Custom headers to include in the WebSocket connection
+ _forward_headers: Whether to forward incoming headers
+ dependencies: FastAPI dependencies to inject
+
+ Returns:
+ A WebSocket passthrough function that can be registered with app.websocket()
+ """
+ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth_websocket
+
+ async def websocket_endpoint_func(
+ websocket: WebSocket,
+ user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth_websocket),
+ **kwargs, # For additional query parameters
+ ):
+ """
+ WebSocket passthrough endpoint function.
+
+ This function handles the WebSocket connection by:
+ 1. Accepting the incoming WebSocket connection
+ 2. Establishing a connection to the target WebSocket
+ 3. Forwarding messages bidirectionally
+ 4. Handling connection cleanup
+ """
+ return await websocket_passthrough_request(
+ websocket=websocket,
+ target=target,
+ custom_headers=custom_headers or {},
+ user_api_key_dict=user_api_key_dict,
+ forward_headers=_forward_headers,
+ endpoint=endpoint,
+ cost_per_request=cost_per_request,
+ accept_websocket=True, # Generic usage should accept the WebSocket
+ )
+
+ return websocket_endpoint_func
+
+
+async def websocket_passthrough_request( # noqa: PLR0915
+ websocket: WebSocket,
+ target: str,
+ custom_headers: dict,
+ user_api_key_dict: UserAPIKeyAuth,
+ forward_headers: Optional[bool] = False,
+ endpoint: Optional[str] = None,
+ cost_per_request: Optional[float] = None,
+ accept_websocket: bool = True,
+):
+ """
+ WebSocket passthrough request handler.
+
+ Args:
+ websocket: The incoming WebSocket connection
+ target: The target WebSocket URL
+ custom_headers: Custom headers to include in the connection
+ user_api_key_dict: The user API key dictionary
+ forward_headers: Whether to forward incoming headers
+ endpoint: The endpoint path (for logging purposes)
+ cost_per_request: Optional field - cost per request to the target endpoint
+ """
+ from litellm.litellm_core_utils.litellm_logging import Logging
+ from litellm.proxy.proxy_server import proxy_logging_obj
+ from litellm.types.passthrough_endpoints.pass_through_endpoints import (
+ PassthroughStandardLoggingPayload,
+ )
+
+ # Initialize tracking variables
+ start_time = datetime.now()
+ websocket_messages: list[dict[str, Any]] = []
+ litellm_call_id = str(uuid.uuid4())
+
+ verbose_proxy_logger.info(
+ f"WebSocket passthrough ({endpoint}): Starting WebSocket connection to {target}"
+ )
+
+ # Only accept the WebSocket if requested (for generic usage)
+ if accept_websocket:
+ await websocket.accept()
+ verbose_proxy_logger.debug(
+ f"WebSocket passthrough ({endpoint}): WebSocket connection accepted"
+ )
+
+ # Prepare headers for the upstream connection
+ upstream_headers = custom_headers.copy()
+
+ if forward_headers:
+ # Forward relevant headers from the incoming request
+ incoming_headers = dict(websocket.headers)
+ for header_name, header_value in incoming_headers.items():
+ # Only forward certain headers to avoid conflicts
+ if header_name.lower() in [
+ "authorization",
+ "x-api-key",
+ "x-goog-user-project",
+ ]:
+ upstream_headers[header_name] = header_value
+
+ # Initialize logging object similar to HTTP passthrough
+ logging_obj = Logging(
+ model="unknown",
+ messages=[{"role": "user", "content": "WebSocket connection"}],
+ stream=True, # WebSockets are inherently streaming
+ call_type="pass_through_endpoint",
+ start_time=start_time,
+ litellm_call_id=litellm_call_id,
+ function_id="websocket_passthrough",
+ )
+
+ # Create passthrough logging payload
+ passthrough_logging_payload = PassthroughStandardLoggingPayload(
+ url=target,
+ request_body={}, # WebSocket doesn't have a traditional request body
+ request_method="WEBSOCKET",
+ cost_per_request=cost_per_request,
+ )
+
+ # Create a dummy request object for WebSocket connections to maintain compatibility
+ # with the existing _init_kwargs_for_pass_through_endpoint function
+ class DummyRequest:
+ def __init__(
+ self, url: str, method: str = "WEBSOCKET", headers: Optional[dict] = None
+ ):
+ self.url = url
+ self.method = method
+ self.headers = headers or {}
+
+ def __str__(self):
+ return f"DummyRequest(url={self.url}, method={self.method})"
+
+ dummy_request = DummyRequest(
+ url=target,
+ method="WEBSOCKET",
+ headers=dict(websocket.headers) if hasattr(websocket, "headers") else {},
+ )
+
+ # Initialize kwargs for logging using the same pattern as HTTP passthrough
+ kwargs = HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint(
+ user_api_key_dict=user_api_key_dict,
+ _parsed_body={}, # WebSocket doesn't have a traditional request body
+ passthrough_logging_payload=passthrough_logging_payload,
+ litellm_call_id=litellm_call_id,
+ request=dummy_request, # type: ignore
+ logging_obj=logging_obj,
+ )
+
+ # Update logging environment variables
+ logging_obj.update_environment_variables(
+ model="unknown",
+ user="unknown",
+ optional_params={},
+ litellm_params=dict(kwargs.get("litellm_params", {})),
+ call_type="pass_through_endpoint",
+ )
+ logging_obj.model_call_details["litellm_call_id"] = litellm_call_id
+
+ # Pre-call logging
+ logging_obj.pre_call(
+ input=[{"role": "user", "content": "WebSocket connection"}],
+ api_key="",
+ additional_args={
+ "complete_input_dict": {},
+ "api_base": target,
+ "headers": upstream_headers,
+ },
+ )
+
+ ### CALL HOOKS ### - modify incoming data / reject request before calling the model
+ websocket_data: dict[str, Any] = {}
+ websocket_data = await proxy_logging_obj.pre_call_hook(
+ user_api_key_dict=user_api_key_dict,
+ data=websocket_data,
+ call_type="pass_through_endpoint",
+ )
+
+ try:
+ verbose_proxy_logger.debug(
+ f"WebSocket passthrough ({endpoint}): Establishing upstream connection to {target}"
+ )
+ async with connect(
+ target,
+ additional_headers=upstream_headers,
+ ) as upstream_ws:
+ verbose_proxy_logger.info(
+ f"WebSocket passthrough ({endpoint}): Upstream connection established successfully"
+ )
+
+ async def forward_client_to_upstream() -> None:
+ """Forward messages from client to upstream WebSocket"""
+ try:
+ while True:
+ message = await websocket.receive()
+ message_type = message.get("type")
+ if message_type == "websocket.disconnect":
+ await upstream_ws.close()
+ break
+
+ text_data = message.get("text")
+ bytes_data = message.get("bytes")
+
+ if text_data is not None:
+ # Try to extract model from client setup message for Vertex AI Live
+ if endpoint and "/vertex_ai/live" in endpoint:
+ verbose_proxy_logger.debug(
+ f"WebSocket passthrough ({endpoint}): Processing client message for model extraction"
+ )
+ try:
+ client_message = json.loads(text_data)
+ if (
+ isinstance(client_message, dict)
+ and "setup" in client_message
+ ):
+ setup_data = client_message["setup"]
+ verbose_proxy_logger.debug(
+ f"WebSocket passthrough ({endpoint}): Found setup data in client message: {setup_data}"
+ )
+ if (
+ isinstance(setup_data, dict)
+ and "model" in setup_data
+ ):
+ extracted_model = (
+ _extract_model_from_vertex_ai_setup(
+ setup_data
+ )
+ )
+ if extracted_model:
+ kwargs["model"] = extracted_model
+ kwargs["custom_llm_provider"] = (
+ "vertex_ai-language-models"
+ )
+ # Update logging object with correct model
+ logging_obj.model = extracted_model
+ logging_obj.model_call_details[
+ "model"
+ ] = extracted_model
+ logging_obj.model_call_details[
+ "custom_llm_provider"
+ ] = "vertex_ai"
+ verbose_proxy_logger.info(
+ f"WebSocket passthrough ({endpoint}): Successfully extracted model '{extracted_model}' and set provider to 'vertex_ai' from client setup message"
+ )
+ else:
+ verbose_proxy_logger.warning(
+ f"WebSocket passthrough ({endpoint}): Failed to extract model from client setup data: {setup_data}"
+ )
+ else:
+ verbose_proxy_logger.debug(
+ f"WebSocket passthrough ({endpoint}): Setup data does not contain model field: {setup_data}"
+ )
+ else:
+ verbose_proxy_logger.debug(
+ f"WebSocket passthrough ({endpoint}): Client message does not contain setup data"
+ )
+ except (json.JSONDecodeError, KeyError, TypeError) as e:
+ verbose_proxy_logger.debug(
+ f"WebSocket passthrough ({endpoint}): Client message is not a valid setup message: {e}"
+ )
+ pass # Not a JSON message or doesn't contain setup data
+
+ await upstream_ws.send(text_data)
+ elif bytes_data is not None:
+ await upstream_ws.send(bytes_data)
+ except asyncio.CancelledError:
+ raise
+ except Exception:
+ verbose_proxy_logger.exception(
+ f"WebSocket passthrough ({endpoint}): error forwarding client message"
+ )
+ await upstream_ws.close()
+
+ async def forward_upstream_to_client() -> None:
+ """Forward messages from upstream to client WebSocket"""
+ try:
+ # Wait for the first response from upstream
+ raw_response = await upstream_ws.recv(decode=False)
+ # Ensure raw_response is bytes before decoding
+ if isinstance(raw_response, str):
+ raw_response = raw_response.encode("ascii")
+ setup_response = json.loads(raw_response.decode("ascii"))
+ verbose_proxy_logger.debug(f"Setup response: {setup_response}")
+
+ # Extract model and provider from setup response for Vertex AI Live
+ if endpoint and "/vertex_ai/live" in endpoint:
+ verbose_proxy_logger.debug(
+ f"WebSocket passthrough ({endpoint}): Processing server setup response for model extraction"
+ )
+ extracted_model = _extract_model_from_vertex_ai_setup(
+ setup_response
+ )
+ if extracted_model:
+ kwargs["model"] = extracted_model
+ kwargs["custom_llm_provider"] = "vertex_ai_language_models"
+ # Update logging object with correct model
+ logging_obj.model = extracted_model
+ logging_obj.model_call_details["model"] = extracted_model
+ logging_obj.model_call_details["custom_llm_provider"] = (
+ "vertex_ai_language_models"
+ )
+ verbose_proxy_logger.debug(
+ f"WebSocket passthrough ({endpoint}): Successfully extracted model '{extracted_model}' and set provider to 'vertex_ai' from server setup response"
+ )
+ else:
+ verbose_proxy_logger.warning(
+ f"WebSocket passthrough ({endpoint}): Failed to extract model from server setup response: {setup_response}"
+ )
+ else:
+ verbose_proxy_logger.debug(
+ f"WebSocket passthrough ({endpoint}): Not a Vertex AI Live endpoint, skipping model extraction"
+ )
+
+ # Send the setup response to the client
+ await websocket.send_text(json.dumps(setup_response))
+
+ # Now continuously forward messages from upstream to client
+ async for upstream_message in upstream_ws:
+ if isinstance(upstream_message, bytes):
+ await websocket.send_bytes(upstream_message)
+ # Parse and collect for cost tracking
+ try:
+ message_data = json.loads(upstream_message.decode())
+ websocket_messages.append(message_data)
+ except (json.JSONDecodeError, UnicodeDecodeError):
+ pass
+ else:
+ await websocket.send_text(upstream_message)
+ # Parse and collect for cost tracking
+ try:
+ message_data = json.loads(upstream_message)
+ websocket_messages.append(message_data)
+ except json.JSONDecodeError:
+ pass
+
+ except (ConnectionClosedOK, ConnectionClosedError) as e:
+ verbose_proxy_logger.debug(
+ f"Upstream WebSocket connection closed: {e}"
+ )
+ pass
+ except asyncio.CancelledError:
+ verbose_proxy_logger.debug(
+ "asyncio.CancelledError in forward_upstream_to_client"
+ )
+ raise
+ except Exception as e:
+ verbose_proxy_logger.debug(
+ f"Exception in forward_upstream_to_client: {e}"
+ )
+ verbose_proxy_logger.exception(
+ f"WebSocket passthrough ({endpoint}): error forwarding upstream message"
+ )
+ raise
+
+ # Create tasks for bidirectional message forwarding
+ tasks = [
+ asyncio.create_task(forward_client_to_upstream()),
+ asyncio.create_task(forward_upstream_to_client()),
+ ]
+
+ done, pending = await asyncio.wait(
+ tasks, return_when=asyncio.FIRST_COMPLETED
+ )
+
+ # Cancel remaining tasks
+ for task in pending:
+ task.cancel()
+ try:
+ await task
+ except asyncio.CancelledError:
+ pass
+
+ # Check for exceptions in completed tasks
+ for task in done:
+ exception = task.exception()
+ if exception is not None:
+ raise exception
+
+ end_time = datetime.now()
+
+ # Update passthrough logging payload with response data
+ passthrough_logging_payload["response_body"] = websocket_messages # type: ignore
+ passthrough_logging_payload["end_time"] = end_time # type: ignore
+
+ # Remove logging_obj from kwargs to avoid duplicate keyword argument
+ success_kwargs = kwargs.copy()
+ success_kwargs.pop("logging_obj", None)
+
+ # # Add user authentication context for database logging
+ # if user_api_key_dict:
+ # success_kwargs.setdefault('litellm_params', {})
+ # success_kwargs['litellm_params'].update({
+ # 'proxy_server_request': {
+ # 'body': {
+ # 'user': user_api_key_dict.user_id,
+ # 'team_id': user_api_key_dict.team_id,
+ # 'end_user_id': user_api_key_dict.end_user_id,
+ # }
+ # }
+ # })
+ # # Also add the user_api_key for direct access
+ # success_kwargs['user_api_key'] = user_api_key_dict.api_key
+
+ # Create a dummy httpx.Response for WebSocket connections
+ class MockWebSocketResponse:
+ def __init__(self, target_url: str):
+ self.status_code = 200
+ self.text = "WebSocket connection successful"
+ self.headers: dict[str, str] = {}
+ self.request = MockWebSocketRequest(target_url)
+
+ class MockWebSocketRequest:
+ def __init__(self, target_url: str):
+ self.method = "WEBSOCKET"
+ self.url = target_url
+
+ mock_response = MockWebSocketResponse(target)
+
+ # Use the same success handler as HTTP passthrough endpoints
+ asyncio.create_task(
+ pass_through_endpoint_logging.pass_through_async_success_handler(
+ httpx_response=mock_response, # type: ignore
+ response_body=websocket_messages, # type: ignore
+ url_route=endpoint or "",
+ result="websocket_connection_successful",
+ start_time=start_time,
+ end_time=end_time,
+ logging_obj=logging_obj,
+ cache_hit=False,
+ request_body={},
+ **success_kwargs,
+ )
+ )
+
+ # Call the proxy logging success hook
+ if proxy_logging_obj:
+ await proxy_logging_obj.post_call_success_hook(
+ data={},
+ user_api_key_dict=user_api_key_dict,
+ response={"status": "websocket_connection_successful"}, # type: ignore
+ )
+
+ except InvalidStatus as exc:
+ verbose_proxy_logger.exception(
+ f"WebSocket passthrough ({endpoint}): upstream rejected WebSocket connection"
+ )
+
+ # Prepare request payload for logging
+ request_payload = {}
+ if kwargs:
+ for key, value in kwargs.items():
+ request_payload[key] = value
+ if logging_obj is not None:
+ request_payload["litellm_logging_obj"] = logging_obj
+
+ # Log the connection failure using the same pattern as HTTP
+ await proxy_logging_obj.post_call_failure_hook(
+ user_api_key_dict=user_api_key_dict,
+ original_exception=exc,
+ request_data=request_payload,
+ traceback_str=traceback.format_exc(
+ limit=MAXIMUM_TRACEBACK_LINES_TO_LOG,
+ ),
+ )
+
+ if websocket.client_state != WebSocketState.DISCONNECTED:
+ await websocket.close(
+ code=getattr(exc, "status_code", 1011),
+ reason="Upstream connection rejected",
+ )
+ except Exception as e:
+ verbose_proxy_logger.exception(
+ f"WebSocket passthrough ({endpoint}): unexpected error while proxying WebSocket"
+ )
+
+ # Prepare request payload for logging
+ request_payload = {}
+ if kwargs:
+ for key, value in kwargs.items():
+ request_payload[key] = value
+ if logging_obj is not None:
+ request_payload["litellm_logging_obj"] = logging_obj
+
+ # Log the unexpected error using the same pattern as HTTP
+ await proxy_logging_obj.post_call_failure_hook(
+ user_api_key_dict=user_api_key_dict,
+ original_exception=e,
+ request_data=request_payload,
+ traceback_str=traceback.format_exc(
+ limit=MAXIMUM_TRACEBACK_LINES_TO_LOG,
+ ),
+ )
+
+ if websocket.client_state != WebSocketState.DISCONNECTED:
+ await websocket.close(code=1011, reason="WebSocket passthrough error")
+ finally:
+ if websocket.client_state != WebSocketState.DISCONNECTED:
+ await websocket.close()
+
+
+def _is_streaming_response(response: httpx.Response) -> bool:
+ _content_type = response.headers.get("content-type")
+ if _content_type is not None and "text/event-stream" in _content_type:
+ return True
+ return False
+
+
+def _extract_model_from_vertex_ai_setup(setup_response: dict) -> Optional[str]:
+ """
+ Extract the model name from Vertex AI Live setup response.
+
+ The setup response can contain a model field in two formats:
+ 1. Direct: {"model": "projects/.../models/gemini-2.0-flash-live-preview-04-09"}
+ 2. Nested: {"setup": {"model": "projects/.../models/gemini-2.0-flash-live-preview-04-09"}}
+
+ We extract just the model name: "gemini-2.0-flash-live-preview-04-09"
+ """
+ try:
+ # Handle both direct model field and nested setup.model field
+ model_path = None
+ if isinstance(setup_response, dict):
+ if "model" in setup_response:
+ model_path = setup_response["model"]
+ elif (
+ "setup" in setup_response
+ and isinstance(setup_response["setup"], dict)
+ and "model" in setup_response["setup"]
+ ):
+ model_path = setup_response["setup"]["model"]
+
+ if isinstance(model_path, str) and "/models/" in model_path:
+ # Extract the model name after the last "/models/"
+ model_name = model_path.split("/models/")[-1]
+ return model_name
+ except Exception as e:
+ verbose_proxy_logger.debug(f"Error extracting model from setup response: {e}")
+ return None
+
+
+class SafeRouteAdder:
+ """
+ Wrapper class for adding routes to FastAPI app.
+ Only adds routes if they don't already exist on the app.
+ """
+
+ @staticmethod
+ def _is_path_registered(app: FastAPI, path: str, methods: List[str]) -> bool:
+ """
+ Check if a path with any of the specified methods is already registered on the app.
+
+ Args:
+ app: The FastAPI application instance
+ path: The path to check (e.g., "/v1/chat/completions")
+ methods: List of HTTP methods to check (e.g., ["GET", "POST"])
+
+ Returns:
+ True if the path is already registered with any of the methods, False otherwise
+ """
+ for route in app.routes:
+ # Use getattr to safely access route attributes
+ route_path = getattr(route, "path", None)
+ route_methods = getattr(route, "methods", None)
+
+ if route_path == path and route_methods is not None:
+ # Check if any of the methods overlap
+ if any(method in route_methods for method in methods):
+ return True
+ return False
+
+ @staticmethod
+ def add_api_route_if_not_exists(
+ app: FastAPI,
+ path: str,
+ endpoint: Any,
+ methods: List[str],
+ dependencies: Optional[List] = None,
+ ) -> bool:
+ """
+ Add an API route to the app only if it doesn't already exist.
+
+ Args:
+ app: The FastAPI application instance
+ path: The path for the route
+ endpoint: The endpoint function/callable
+ methods: List of HTTP methods
+ dependencies: Optional list of dependencies
+
+ Returns:
+ True if route was added, False if it already existed
+ """
+ if SafeRouteAdder._is_path_registered(app=app, path=path, methods=methods):
+ verbose_proxy_logger.debug(
+ "Skipping route registration - path %s with methods %s already registered on app",
+ path,
+ methods,
+ )
+ return False
+
+ app.add_api_route(
+ path=path,
+ endpoint=endpoint,
+ methods=methods,
+ dependencies=dependencies,
+ )
+ verbose_proxy_logger.debug(
+ "Successfully added route: %s with methods %s",
+ path,
+ methods,
+ )
+ return True
+
+
+class InitPassThroughEndpointHelpers:
+ @staticmethod
+ def add_exact_path_route(
+ app: FastAPI,
+ path: str,
+ target: str,
+ custom_headers: Optional[dict],
+ forward_headers: Optional[bool],
+ merge_query_params: Optional[bool],
+ dependencies: Optional[List],
+ cost_per_request: Optional[float],
+ endpoint_id: str,
+ guardrails: Optional[dict] = None,
+ methods: Optional[List[str]] = None,
+ default_query_params: Optional[dict] = None,
+ config_file_path: Optional[str] = None,
+ ):
+ """Add exact path route for pass-through endpoint"""
+ # Default to all methods if none specified (backward compatibility)
+ if methods is None or len(methods) == 0:
+ methods = ["GET", "POST", "PUT", "DELETE", "PATCH"]
+
+ # Create route key that includes methods for uniqueness
+ methods_str = ",".join(sorted(methods))
+ route_key = f"{endpoint_id}:exact:{path}:{methods_str}"
+
+ # Check if this exact route is already registered
+ if route_key in _registered_pass_through_routes:
+ verbose_proxy_logger.debug(
+ "Updating duplicate exact pass through endpoint: %s with methods %s (already registered)",
+ path,
+ methods,
+ )
+
+ verbose_proxy_logger.debug(
+ "adding exact pass through endpoint: %s, methods: %s, dependencies: %s",
+ path,
+ methods,
+ dependencies,
+ )
+
+ # Use SafeRouteAdder to only add route if it doesn't exist on the app
+ SafeRouteAdder.add_api_route_if_not_exists(
+ app=app,
+ path=path,
+ endpoint=create_pass_through_route( # type: ignore
+ path,
+ target,
+ custom_headers,
+ forward_headers,
+ merge_query_params,
+ dependencies,
+ cost_per_request=cost_per_request,
+ default_query_params=default_query_params,
+ guardrails=guardrails,
+ config_file_path=config_file_path,
+ ),
+ methods=methods,
+ dependencies=dependencies,
+ )
+
+ # Always register/update the route metadata (headers, target) even if FastAPI route exists
+ _registered_pass_through_routes[route_key] = {
+ "endpoint_id": endpoint_id,
+ "path": path,
+ "type": "exact",
+ "methods": methods,
+ "passthrough_params": {
+ "target": target,
+ "custom_headers": custom_headers,
+ "forward_headers": forward_headers,
+ "merge_query_params": merge_query_params,
+ "default_query_params": default_query_params,
+ "dependencies": dependencies,
+ "cost_per_request": cost_per_request,
+ "guardrails": guardrails,
+ },
+ }
+
+ @staticmethod
+ def add_subpath_route(
+ app: FastAPI,
+ path: str,
+ target: str,
+ custom_headers: Optional[dict],
+ forward_headers: Optional[bool],
+ merge_query_params: Optional[bool],
+ dependencies: Optional[List],
+ cost_per_request: Optional[float],
+ endpoint_id: str,
+ guardrails: Optional[dict] = None,
+ methods: Optional[List[str]] = None,
+ default_query_params: Optional[dict] = None,
+ config_file_path: Optional[str] = None,
+ ):
+ """Add wildcard route for sub-paths"""
+ # Default to all methods if none specified (backward compatibility)
+ if methods is None or len(methods) == 0:
+ methods = ["GET", "POST", "PUT", "DELETE", "PATCH"]
+
+ wildcard_path = f"{path}/{{subpath:path}}"
+ methods_str = ",".join(sorted(methods))
+ route_key = f"{endpoint_id}:subpath:{path}:{methods_str}"
+
+ # Check if this subpath route is already registered
+ if route_key in _registered_pass_through_routes:
+ verbose_proxy_logger.debug(
+ "Updating duplicate wildcard pass through endpoint: %s with methods %s (already registered)",
+ wildcard_path,
+ methods,
+ )
+
+ verbose_proxy_logger.debug(
+ "adding wildcard pass through endpoint: %s, methods: %s, dependencies: %s",
+ wildcard_path,
+ methods,
+ dependencies,
+ )
+
+ # Use SafeRouteAdder to only add route if it doesn't exist on the app
+ SafeRouteAdder.add_api_route_if_not_exists(
+ app=app,
+ path=wildcard_path,
+ endpoint=create_pass_through_route( # type: ignore
+ path,
+ target,
+ custom_headers,
+ forward_headers,
+ merge_query_params,
+ dependencies,
+ include_subpath=True,
+ cost_per_request=cost_per_request,
+ default_query_params=default_query_params,
+ guardrails=guardrails,
+ config_file_path=config_file_path,
+ ),
+ methods=methods,
+ dependencies=dependencies,
+ )
+
+ # Register the route to prevent duplicates only if it was added
+ _registered_pass_through_routes[route_key] = {
+ "endpoint_id": endpoint_id,
+ "path": path,
+ "type": "subpath",
+ "methods": methods,
+ "passthrough_params": {
+ "target": target,
+ "custom_headers": custom_headers,
+ "forward_headers": forward_headers,
+ "merge_query_params": merge_query_params,
+ "default_query_params": default_query_params,
+ "dependencies": dependencies,
+ "cost_per_request": cost_per_request,
+ "guardrails": guardrails,
+ },
+ }
+
+ @staticmethod
+ def remove_endpoint_routes(endpoint_id: str):
+ """Remove all routes for a specific endpoint ID from the registry
+ and clean up corresponding entries from LiteLLMRoutes.openai_routes."""
+ keys_to_remove = [
+ key
+ for key, value in _registered_pass_through_routes.items()
+ if value["endpoint_id"] == endpoint_id
+ ]
+ for key in keys_to_remove:
+ route_info = _registered_pass_through_routes[key]
+ path = route_info.get("path")
+ if isinstance(path, str):
+ openai_routes = LiteLLMRoutes.openai_routes.value
+ if path in openai_routes:
+ openai_routes.remove(path)
+ if route_info.get("type") == "subpath":
+ wildcard_path = path.rstrip("/") + "/*"
+ if wildcard_path in openai_routes:
+ openai_routes.remove(wildcard_path)
+ del _registered_pass_through_routes[key]
+ verbose_proxy_logger.debug(
+ "Removed pass-through route from registry: %s", key
+ )
+
+ @staticmethod
+ def clear_all_pass_through_routes():
+ """Clear all pass-through routes from the registry"""
+ _registered_pass_through_routes.clear()
+
+ @staticmethod
+ def get_all_registered_pass_through_routes() -> List[str]:
+ """Get all registered pass-through endpoints from the registry"""
+ return list(_registered_pass_through_routes.keys())
+
+ @staticmethod
+ def _build_full_path_with_root(path: str) -> str:
+ """
+ Build full path by prepending server root path if needed.
+
+ Args:
+ path: The relative path to build
+
+ Returns:
+ Full path with server root prepended (if root is not "/")
+ """
+ root_path = get_server_root_path()
+ if root_path == "/":
+ return path
+ return f"{root_path}{path}"
+
+ @staticmethod
+ def is_registered_pass_through_route(route: str) -> bool:
+ """
+ Check if route is a registered pass-through endpoint from DB
+
+ Uses the in-memory registry to avoid additional DB queries
+ Optimized for minimal latency
+
+ Args:
+ route: The route to check
+
+ Returns:
+ bool: True if route is a registered pass-through endpoint, False otherwise
+ """
+ ## CHECK IF MAPPED PASS THROUGH ENDPOINT
+ normalized_route = normalize_route_for_root_path(route)
+ if normalized_route is not None:
+ for mapped_route in LiteLLMRoutes.mapped_pass_through_routes.value:
+ if normalized_route.startswith(mapped_route):
+ return True
+
+ # Fast path: check if any registered route key contains this path
+ # Keys are in format: "{endpoint_id}:exact:{path}:{methods}" or "{endpoint_id}:subpath:{path}:{methods}"
+ # For backward compatibility, also support old format: "{endpoint_id}:exact:{path}" or "{endpoint_id}:subpath:{path}"
+ # Extract unique paths from keys for quick checking
+ for key in _registered_pass_through_routes.keys():
+ parts = key.split(":", 3) # Split into [endpoint_id, type, path, methods?]
+ if len(parts) >= 3:
+ route_type = parts[1]
+ registered_path = (
+ InitPassThroughEndpointHelpers._build_full_path_with_root(parts[2])
+ )
+ if route_type == "exact" and route == registered_path:
+ return True
+ elif route_type == "subpath":
+ if route == registered_path or route.startswith(
+ registered_path + "/"
+ ):
+ return True
+
+ return False
+
+ @staticmethod
+ def get_registered_pass_through_route(
+ route: str, method: Optional[str] = None
+ ) -> Optional[Dict[str, Any]]:
+ """Get passthrough params for a given route and optionally filter by HTTP method"""
+ for key in _registered_pass_through_routes.keys():
+ parts = key.split(":", 3) # Split into [endpoint_id, type, path, methods?]
+ if len(parts) >= 3:
+ route_type = parts[1]
+ registered_path = (
+ InitPassThroughEndpointHelpers._build_full_path_with_root(parts[2])
+ )
+
+ # Get the methods for this route
+ route_methods = _registered_pass_through_routes[key].get("methods", [])
+
+ # Check if path matches
+ path_matches = False
+ if route_type == "exact" and route == registered_path:
+ path_matches = True
+ elif route_type == "subpath":
+ if route == registered_path or route.startswith(
+ registered_path + "/"
+ ):
+ path_matches = True
+
+ # If path matches and method filter is provided, check if method is allowed
+ if path_matches:
+ if method is None or not route_methods or method in route_methods:
+ return _registered_pass_through_routes[key]
+
+ return None
+
+
+def _get_combined_pass_through_endpoints(
+ pass_through_endpoints: Union[List[Dict], List[PassThroughGenericEndpoint]],
+ config_pass_through_endpoints: List[Dict],
+):
+ """Get combined pass-through endpoints from db + config"""
+ return pass_through_endpoints + config_pass_through_endpoints
+
+
+async def _register_pass_through_endpoint(
+ endpoint: Union[Dict[str, Any], PassThroughGenericEndpoint],
+ app: FastAPI,
+ premium_user: bool,
+ visited_endpoints: set[str],
+ config_file_path: Optional[str] = None,
+) -> None:
+ endpoint_data: Dict[str, Any]
+ if isinstance(endpoint, PassThroughGenericEndpoint):
+ endpoint_data = endpoint.model_dump()
+ else:
+ endpoint_data = endpoint
+
+ if endpoint_data.get("id") is None:
+ endpoint_data["id"] = str(uuid.uuid4())
+ endpoint_id = cast(str, endpoint_data["id"])
+
+ target = endpoint_data.get("target")
+ path = endpoint_data.get("path")
+ if path is None:
+ raise ValueError("Path is required for pass-through endpoint")
+
+ custom_headers = await set_env_variables_in_header(
+ custom_headers=endpoint_data.get("headers")
+ )
+ forward_headers = endpoint_data.get("forward_headers")
+ merge_query_params = endpoint_data.get("merge_query_params")
+ default_query_params = endpoint_data.get("default_query_params")
+ auth = endpoint_data.get("auth")
+ dependencies = None
+
+ if auth is not None and str(auth).lower() == "true":
+ # Authentication on a pass-through endpoint used to be enterprise-only.
+ # That left OSS with no safe configuration: auth=True raised at startup
+ # unless the operator had a license. The safe option must always be free,
+ # and unauthenticated forwarding should require explicit opt-in.
+ dependencies = [Depends(user_api_key_auth)]
+ if path not in LiteLLMRoutes.openai_routes.value:
+ LiteLLMRoutes.openai_routes.value.append(path)
+
+ if target is None:
+ return
+
+ guardrails = endpoint_data.get("guardrails")
+ methods = endpoint_data.get("methods")
+ cost_per_request = endpoint_data.get("cost_per_request")
+
+ verbose_proxy_logger.debug(
+ "Initializing pass through endpoint: %s (ID: %s)", path, endpoint_id
+ )
+ InitPassThroughEndpointHelpers.add_exact_path_route(
+ app=app,
+ path=path,
+ target=target,
+ custom_headers=custom_headers,
+ forward_headers=forward_headers,
+ merge_query_params=merge_query_params,
+ dependencies=dependencies,
+ cost_per_request=cost_per_request,
+ endpoint_id=endpoint_id,
+ guardrails=guardrails,
+ methods=methods,
+ default_query_params=default_query_params,
+ config_file_path=config_file_path,
+ )
+
+ methods_for_key = methods if methods else ["GET", "POST", "PUT", "DELETE", "PATCH"]
+ methods_str = ",".join(sorted(methods_for_key))
+ visited_endpoints.add(f"{endpoint_id}:exact:{path}:{methods_str}")
+
+ if endpoint_data.get("include_subpath", False) is True:
+ if auth is not None and str(auth).lower() == "true":
+ wildcard_path = path.rstrip("/") + "/*"
+ if wildcard_path not in LiteLLMRoutes.openai_routes.value:
+ LiteLLMRoutes.openai_routes.value.append(wildcard_path)
+ InitPassThroughEndpointHelpers.add_subpath_route(
+ app=app,
+ path=path,
+ target=target,
+ custom_headers=custom_headers,
+ forward_headers=forward_headers,
+ merge_query_params=merge_query_params,
+ dependencies=dependencies,
+ cost_per_request=cost_per_request,
+ endpoint_id=endpoint_id,
+ guardrails=guardrails,
+ methods=methods,
+ default_query_params=default_query_params,
+ config_file_path=config_file_path,
+ )
+ visited_endpoints.add(f"{endpoint_id}:subpath:{path}:{methods_str}")
+
+ verbose_proxy_logger.debug(
+ "Added new pass through endpoint: %s (ID: %s)", path, endpoint_id
+ )
+
+
+async def initialize_pass_through_endpoints(
+ pass_through_endpoints: Union[List[Dict], List[PassThroughGenericEndpoint]],
+ config_file_path: Optional[str] = None,
+):
+ """
+ 1. Create a global list of pass-through endpoints (db + config)
+ 2. Clear all existing pass-through endpoints from the FastAPI app routes
+ 3. Add new endpoints to the in-memory registry
+
+ Initialize a list of pass-through endpoints by adding them to the FastAPI app routes
+
+ Args:
+ pass_through_endpoints: List of pass-through endpoints to initialize
+ config_file_path: Path to the operator's config.yaml when this call
+ originates from a YAML-load. Threaded through to
+ ``create_pass_through_route`` so an operator using
+ ``s3://``/``gcs://`` ``custom_handler`` in their config still
+ loads. Callers from the DB-overlay / runtime API path must leave
+ this ``None`` so the runtime gate in ``get_instance_fn`` fires.
+
+ Returns:
+ None
+ """
+ verbose_proxy_logger.debug("initializing pass through endpoints")
+ from litellm.proxy.proxy_server import (
+ app,
+ config_passthrough_endpoints,
+ premium_user,
+ )
+
+ ## get combined pass-through endpoints from db + config
+ combined_pass_through_endpoints: List[Union[Dict, PassThroughGenericEndpoint]]
+
+ if config_passthrough_endpoints is not None:
+ combined_pass_through_endpoints = _get_combined_pass_through_endpoints( # type: ignore
+ pass_through_endpoints, config_passthrough_endpoints
+ )
+ else:
+ combined_pass_through_endpoints = pass_through_endpoints # type: ignore
+
+ ## clear all existing pass-through endpoints from the FastAPI app routes
+ # InitPassThroughEndpointHelpers.clear_all_pass_through_routes()
+
+ # get a list of all registered pass-through endpoints
+ # mark the ones that are visited in the list
+ # remove the ones that are not visited from the list
+ registered_pass_through_endpoints = (
+ InitPassThroughEndpointHelpers.get_all_registered_pass_through_routes()
+ )
+
+ visited_endpoints: set[str] = set()
+
+ for endpoint in combined_pass_through_endpoints:
+ await _register_pass_through_endpoint(
+ endpoint=endpoint,
+ app=app,
+ premium_user=premium_user,
+ visited_endpoints=visited_endpoints,
+ config_file_path=config_file_path,
+ )
+
+ # remove the ones that are not visited from the list
+ for endpoint_key in registered_pass_through_endpoints:
+ if endpoint_key not in visited_endpoints:
+ InitPassThroughEndpointHelpers.remove_endpoint_routes(endpoint_key)
+
+
+def _get_pass_through_endpoints_from_config() -> List[PassThroughGenericEndpoint]:
+ """
+ Get pass-through endpoints defined in the config file.
+ These are read-only and cannot be edited via the UI.
+ Malformed endpoints are logged and skipped; they do not crash the function.
+ """
+ from pydantic import ValidationError
+
+ from litellm.proxy.proxy_server import config_passthrough_endpoints
+
+ if config_passthrough_endpoints is None or len(config_passthrough_endpoints) == 0:
+ return []
+
+ returned_endpoints: List[PassThroughGenericEndpoint] = []
+ for endpoint in config_passthrough_endpoints:
+ try:
+ if isinstance(endpoint, dict):
+ endpoint_dict = dict(endpoint)
+ endpoint_dict["is_from_config"] = True
+ returned_endpoints.append(PassThroughGenericEndpoint(**endpoint_dict))
+ elif isinstance(endpoint, PassThroughGenericEndpoint):
+ # Create a copy with is_from_config=True
+ endpoint_dict = endpoint.model_dump()
+ endpoint_dict["is_from_config"] = True
+ returned_endpoints.append(PassThroughGenericEndpoint(**endpoint_dict))
+ except ValidationError as e:
+ verbose_proxy_logger.warning(
+ "Skipping malformed pass-through endpoint from config: %s",
+ e,
+ exc_info=False,
+ )
+
+ return returned_endpoints
+
+
+async def _get_pass_through_endpoints_from_db(
+ endpoint_id: Optional[str] = None,
+ user_api_key_dict: Optional[UserAPIKeyAuth] = None,
+) -> List[PassThroughGenericEndpoint]:
+ from litellm.proxy._types import LitellmUserRoles
+ from litellm.proxy.proxy_server import get_config_general_settings
+
+ try:
+ if user_api_key_dict is None:
+ user_api_key_dict = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
+ response: ConfigFieldInfo = await get_config_general_settings(
+ field_name="pass_through_endpoints", user_api_key_dict=user_api_key_dict
+ )
+ except Exception:
+ return []
+
+ pass_through_endpoint_data: Optional[List] = response.field_value
+ if pass_through_endpoint_data is None:
+ return []
+
+ returned_endpoints: List[PassThroughGenericEndpoint] = []
+ if endpoint_id is None:
+ # Return all endpoints from DB, mark as not from config
+ for endpoint in pass_through_endpoint_data:
+ if isinstance(endpoint, dict):
+ endpoint_dict = dict(endpoint)
+ endpoint_dict["is_from_config"] = False
+ returned_endpoints.append(PassThroughGenericEndpoint(**endpoint_dict))
+ elif isinstance(endpoint, PassThroughGenericEndpoint):
+ endpoint_dict = endpoint.model_dump()
+ endpoint_dict["is_from_config"] = False
+ returned_endpoints.append(PassThroughGenericEndpoint(**endpoint_dict))
+ else:
+ # Find specific endpoint by ID
+ found_endpoint = _find_endpoint_by_id(pass_through_endpoint_data, endpoint_id)
+ if found_endpoint is not None:
+ endpoint_dict = (
+ found_endpoint.model_dump()
+ if isinstance(found_endpoint, PassThroughGenericEndpoint)
+ else dict(found_endpoint)
+ )
+ endpoint_dict["is_from_config"] = False
+ returned_endpoints.append(PassThroughGenericEndpoint(**endpoint_dict))
+
+ return returned_endpoints
+
+
+async def _filter_endpoints_by_team_allowed_routes(
+ team_id: str,
+ pass_through_endpoints: List[PassThroughGenericEndpoint],
+ prisma_client,
+) -> List[PassThroughGenericEndpoint]:
+ """
+ Filter pass-through endpoints based on team's allowed_passthrough_routes metadata.
+
+ Args:
+ team_id: The team ID to check permissions for
+ pass_through_endpoints: List of endpoints to filter
+ prisma_client: Database client
+
+ Returns:
+ Filtered list of endpoints based on team permissions
+
+ Raises:
+ HTTPException: If team is not found
+ """
+ # retrieve team from db
+ team = await prisma_client.db.litellm_teamtable.find_unique(
+ where={"team_id": team_id},
+ )
+ if team is None:
+ raise HTTPException(
+ status_code=404,
+ detail={"error": "Team not found"},
+ )
+
+ # retrieve team metadata
+ team_metadata = team.metadata
+ if (
+ team_metadata is not None
+ and team_metadata.get("allowed_passthrough_routes") is not None
+ ):
+ ## FILTER pass_through_endpoints by allowed_passthrough_routes
+ pass_through_endpoints = [
+ endpoint
+ for endpoint in pass_through_endpoints
+ if endpoint.path in team_metadata.get("allowed_passthrough_routes")
+ ]
+
+ return pass_through_endpoints
+
+
+@router.get(
+ "/config/pass_through_endpoint",
+ dependencies=[Depends(user_api_key_auth)],
+ response_model=PassThroughEndpointResponse,
+)
+@router.get(
+ "/config/pass_through_endpoint/team/{team_id}",
+ dependencies=[Depends(user_api_key_auth)],
+ response_model=PassThroughEndpointResponse,
+)
+async def get_pass_through_endpoints(
+ endpoint_id: Optional[str] = None,
+ user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
+ team_id: Optional[str] = None,
+):
+ """
+ GET configured pass through endpoint.
+
+ If no endpoint_id given, return all configured endpoints.
+ """ ## Get existing pass-through endpoint field value
+ from litellm.proxy._types import CommonProxyErrors
+ from litellm.proxy.proxy_server import prisma_client
+
+ if prisma_client is None:
+ raise HTTPException(
+ status_code=500,
+ detail={"error": CommonProxyErrors.db_not_connected_error.value},
+ )
+
+ # Get endpoints from DB (editable via UI)
+ db_endpoints = await _get_pass_through_endpoints_from_db(
+ endpoint_id=endpoint_id, user_api_key_dict=user_api_key_dict
+ )
+
+ # Get endpoints from config file (read-only, not editable via UI)
+ config_endpoints = _get_pass_through_endpoints_from_config()
+
+ # Merge: config endpoints not in DB + all DB endpoints (DB overrides config for same path)
+ db_paths = {ep.path for ep in db_endpoints}
+ config_only_endpoints = [ep for ep in config_endpoints if ep.path not in db_paths]
+ if endpoint_id is not None:
+ # When filtering by endpoint_id, only return if found in DB (config endpoints use generated IDs)
+ pass_through_endpoints = db_endpoints
+ else:
+ pass_through_endpoints = config_only_endpoints + db_endpoints
+
+ if team_id is not None:
+ pass_through_endpoints = await _filter_endpoints_by_team_allowed_routes(
+ team_id=team_id,
+ pass_through_endpoints=pass_through_endpoints,
+ prisma_client=prisma_client,
+ )
+
+ return PassThroughEndpointResponse(endpoints=pass_through_endpoints)
+
+
+@router.post(
+ "/config/pass_through_endpoint/{endpoint_id}",
+ dependencies=[Depends(user_api_key_auth)],
+)
+async def update_pass_through_endpoints(
+ endpoint_id: str,
+ data: PassThroughGenericEndpoint,
+ request: Request,
+ user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
+):
+ """
+ Update a pass-through endpoint by ID.
+ """
+ from litellm.proxy.proxy_server import (
+ get_config_general_settings,
+ update_config_general_settings,
+ )
+
+ ## Get existing pass-through endpoint field value
+ try:
+ response: ConfigFieldInfo = await get_config_general_settings(
+ field_name="pass_through_endpoints", user_api_key_dict=user_api_key_dict
+ )
+ except Exception:
+ raise HTTPException(
+ status_code=404,
+ detail={"error": "No pass-through endpoints found"},
+ )
+
+ pass_through_endpoint_data: Optional[List] = response.field_value
+ if pass_through_endpoint_data is None:
+ raise HTTPException(
+ status_code=404,
+ detail={"error": "No pass-through endpoints found"},
+ )
+
+ # Find the endpoint to update
+ found_endpoint = _find_endpoint_by_id(pass_through_endpoint_data, endpoint_id)
+
+ if found_endpoint is None:
+ raise HTTPException(
+ status_code=404,
+ detail={"error": f"Endpoint with ID '{endpoint_id}' not found"},
+ )
+
+ # Find the index for updating the list
+ endpoint_index = None
+ for idx, endpoint in enumerate(pass_through_endpoint_data):
+ _endpoint = (
+ PassThroughGenericEndpoint(**endpoint)
+ if isinstance(endpoint, dict)
+ else endpoint
+ )
+ if _endpoint.id == endpoint_id:
+ endpoint_index = idx
+ break
+
+ if endpoint_index is None:
+ raise HTTPException(
+ status_code=404,
+ detail={
+ "error": f"Could not find index for endpoint with ID '{endpoint_id}'"
+ },
+ )
+
+ # Get the update data as dict, excluding None values for partial updates
+ # Exclude is_from_config as it's a response-only field (computed at read time)
+ update_data = data.model_dump(exclude_none=True, exclude={"is_from_config"})
+
+ # Start with existing endpoint data
+ endpoint_dict = found_endpoint.model_dump()
+
+ # Update with new data (only non-None values)
+ endpoint_dict.update(update_data)
+
+ # Preserve existing ID if not provided in update and endpoint has ID
+ if "id" not in update_data and found_endpoint.id is not None:
+ endpoint_dict["id"] = found_endpoint.id
+
+ # Remove is_from_config before saving - it's a response-only field (computed at read time)
+ endpoint_dict.pop("is_from_config", None)
+
+ # Create updated endpoint object
+ updated_endpoint = PassThroughGenericEndpoint(**endpoint_dict)
+
+ # Update the list
+ pass_through_endpoint_data[endpoint_index] = endpoint_dict
+
+ # Remove old routes from registry before they get re-registered
+ InitPassThroughEndpointHelpers.remove_endpoint_routes(endpoint_id)
+
+ ## Update db
+ updated_data = ConfigFieldUpdate(
+ field_name="pass_through_endpoints",
+ field_value=pass_through_endpoint_data,
+ config_type="general_settings",
+ )
+
+ await update_config_general_settings(
+ data=updated_data, user_api_key_dict=user_api_key_dict
+ )
+
+ # Re-register the route with updated headers
+ _custom_headers: Optional[dict] = updated_endpoint.headers or {}
+ _custom_headers = await set_env_variables_in_header(custom_headers=_custom_headers)
+
+ if updated_endpoint.include_subpath:
+ InitPassThroughEndpointHelpers.add_subpath_route(
+ app=request.app,
+ path=updated_endpoint.path,
+ target=updated_endpoint.target,
+ custom_headers=_custom_headers,
+ forward_headers=None, # Defaults not available in model? assuming None logic handles it
+ merge_query_params=None,
+ dependencies=None,
+ cost_per_request=updated_endpoint.cost_per_request,
+ endpoint_id=updated_endpoint.id or endpoint_id or "",
+ guardrails=getattr(updated_endpoint, "guardrails", None),
+ methods=updated_endpoint.methods,
+ default_query_params=updated_endpoint.default_query_params,
+ )
+ else:
+ InitPassThroughEndpointHelpers.add_exact_path_route(
+ app=request.app,
+ path=updated_endpoint.path,
+ target=updated_endpoint.target,
+ custom_headers=_custom_headers,
+ forward_headers=None,
+ merge_query_params=None,
+ dependencies=None,
+ cost_per_request=updated_endpoint.cost_per_request,
+ endpoint_id=updated_endpoint.id or endpoint_id or "",
+ guardrails=getattr(updated_endpoint, "guardrails", None),
+ methods=updated_endpoint.methods,
+ default_query_params=updated_endpoint.default_query_params,
+ )
+
+ return PassThroughEndpointResponse(
+ endpoints=[updated_endpoint] if updated_endpoint else []
+ )
+
+
+@router.post(
+ "/config/pass_through_endpoint",
+ dependencies=[Depends(user_api_key_auth)],
+)
+async def create_pass_through_endpoints(
+ data: PassThroughGenericEndpoint,
+ request: Request,
+ user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
+):
+ """
+ Create new pass-through endpoint
+ """
+ from litellm._uuid import uuid
+ from litellm.proxy.proxy_server import (
+ get_config_general_settings,
+ update_config_general_settings,
+ )
+
+ ## Get existing pass-through endpoint field value
+
+ try:
+ response: ConfigFieldInfo = await get_config_general_settings(
+ field_name="pass_through_endpoints", user_api_key_dict=user_api_key_dict
+ )
+ except Exception:
+ response = ConfigFieldInfo(
+ field_name="pass_through_endpoints", field_value=None
+ )
+
+ ## Auto-generate ID if not provided
+ # Exclude is_from_config as it's a response-only field (computed at read time)
+ data_dict = data.model_dump(exclude={"is_from_config"})
+ if data_dict.get("id") is None:
+ data_dict["id"] = str(uuid.uuid4())
+
+ if response.field_value is None:
+ response.field_value = [data_dict]
+ elif isinstance(response.field_value, List):
+ response.field_value.append(data_dict)
+
+ ## Update db
+ updated_data = ConfigFieldUpdate(
+ field_name="pass_through_endpoints",
+ field_value=response.field_value,
+ config_type="general_settings",
+ )
+ await update_config_general_settings(
+ data=updated_data, user_api_key_dict=user_api_key_dict
+ )
+
+ # Return the created endpoint with the generated ID
+ created_endpoint = PassThroughGenericEndpoint(**data_dict)
+
+ # Register the new route
+ _custom_headers: Optional[dict] = created_endpoint.headers or {}
+ _custom_headers = await set_env_variables_in_header(custom_headers=_custom_headers)
+
+ if created_endpoint.include_subpath:
+ InitPassThroughEndpointHelpers.add_subpath_route(
+ app=request.app,
+ path=created_endpoint.path,
+ target=created_endpoint.target,
+ custom_headers=_custom_headers,
+ forward_headers=None,
+ merge_query_params=None,
+ dependencies=None,
+ cost_per_request=created_endpoint.cost_per_request,
+ endpoint_id=created_endpoint.id or "",
+ guardrails=getattr(created_endpoint, "guardrails", None),
+ methods=created_endpoint.methods,
+ default_query_params=created_endpoint.default_query_params,
+ )
+ else:
+ InitPassThroughEndpointHelpers.add_exact_path_route(
+ app=request.app,
+ path=created_endpoint.path,
+ target=created_endpoint.target,
+ custom_headers=_custom_headers,
+ forward_headers=None,
+ merge_query_params=None,
+ dependencies=None,
+ cost_per_request=created_endpoint.cost_per_request,
+ endpoint_id=created_endpoint.id or "",
+ guardrails=getattr(created_endpoint, "guardrails", None),
+ methods=created_endpoint.methods,
+ default_query_params=created_endpoint.default_query_params,
+ )
+
+ return PassThroughEndpointResponse(endpoints=[created_endpoint])
+
+
+@router.delete(
+ "/config/pass_through_endpoint",
+ dependencies=[Depends(user_api_key_auth)],
+ response_model=PassThroughEndpointResponse,
+)
+async def delete_pass_through_endpoints(
+ endpoint_id: str,
+ user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
+):
+ """
+ Delete a pass-through endpoint by ID.
+
+ Returns - the deleted endpoint
+ """
+ from litellm.proxy.proxy_server import (
+ get_config_general_settings,
+ update_config_general_settings,
+ )
+
+ ## Get existing pass-through endpoint field value
+
+ try:
+ response: ConfigFieldInfo = await get_config_general_settings(
+ field_name="pass_through_endpoints", user_api_key_dict=user_api_key_dict
+ )
+ except Exception:
+ response = ConfigFieldInfo(
+ field_name="pass_through_endpoints", field_value=None
+ )
+
+ ## Update field by removing endpoint
+ pass_through_endpoint_data: Optional[List] = response.field_value
+ if response.field_value is None or pass_through_endpoint_data is None:
+ raise HTTPException(
+ status_code=400,
+ detail={"error": "There are no pass-through endpoints setup."},
+ )
+
+ # Find the endpoint to delete
+ found_endpoint = _find_endpoint_by_id(pass_through_endpoint_data, endpoint_id)
+
+ if found_endpoint is None:
+ raise HTTPException(
+ status_code=400,
+ detail={
+ "error": "Endpoint with ID '{}' was not found in pass-through endpoint list.".format(
+ endpoint_id
+ )
+ },
+ )
+
+ # Find the index for deleting from the list
+ endpoint_index = None
+ for idx, endpoint in enumerate(pass_through_endpoint_data):
+ _endpoint = (
+ PassThroughGenericEndpoint(**endpoint)
+ if isinstance(endpoint, dict)
+ else endpoint
+ )
+ if _endpoint.id == endpoint_id:
+ endpoint_index = idx
+ break
+
+ if endpoint_index is None:
+ raise HTTPException(
+ status_code=400,
+ detail={
+ "error": f"Could not find index for endpoint with ID '{endpoint_id}'"
+ },
+ )
+
+ # Remove the endpoint
+ pass_through_endpoint_data.pop(endpoint_index)
+ response_obj = found_endpoint
+
+ # Remove routes from registry
+ InitPassThroughEndpointHelpers.remove_endpoint_routes(endpoint_id)
+
+ ## Update db
+ updated_data = ConfigFieldUpdate(
+ field_name="pass_through_endpoints",
+ field_value=pass_through_endpoint_data,
+ config_type="general_settings",
+ )
+ await update_config_general_settings(
+ data=updated_data, user_api_key_dict=user_api_key_dict
+ )
+
+ return PassThroughEndpointResponse(endpoints=[response_obj])
+
+
+def _find_endpoint_by_id(
+ endpoints_data: List,
+ endpoint_id: str,
+) -> Optional[PassThroughGenericEndpoint]:
+ """
+ Find an endpoint by ID.
+
+ Args:
+ endpoints_data: List of endpoint data (dicts or PassThroughGenericEndpoint objects)
+ endpoint_id: ID to search for
+
+ Returns:
+ Found endpoint or None if not found
+ """
+ for endpoint in endpoints_data:
+ _endpoint: Optional[PassThroughGenericEndpoint] = None
+ if isinstance(endpoint, dict):
+ _endpoint = PassThroughGenericEndpoint(**endpoint)
+ elif isinstance(endpoint, PassThroughGenericEndpoint):
+ _endpoint = endpoint
+
+ # Only compare IDs to IDs
+ if _endpoint is not None and _endpoint.id == endpoint_id:
+ return _endpoint
+
+ return None
+
+
+async def initialize_pass_through_endpoints_in_db():
+ """
+ Gets all pass-through endpoints from db and initializes them in the proxy server.
+ """
+ pass_through_endpoints = await _get_pass_through_endpoints_from_db()
+ await initialize_pass_through_endpoints(
+ pass_through_endpoints=pass_through_endpoints
+ )
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 2a085077434..b296792cd09 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -15644,7 +15644,6 @@ async def get_routes():
app.include_router(router)
app.include_router(response_router)
-app.include_router(batches_router)
app.include_router(public_endpoints_router)
app.include_router(rerank_router)
app.include_router(ocr_router)
@@ -15657,6 +15656,7 @@ app.include_router(fine_tuning_router)
app.include_router(credential_router)
app.include_router(llm_passthrough_router)
app.include_router(pass_through_router)
+app.include_router(batches_router)
app.include_router(health_router)
app.include_router(key_management_router)
app.include_router(internal_user_router)
diff --git a/litellm/types/utils.py b/litellm/types/utils.py
index 8f471b62b5e..5574d616fac 100644
--- a/litellm/types/utils.py
+++ b/litellm/types/utils.py
@@ -3602,6 +3602,10 @@ class SpecialEnums(Enum):
"litellm:custom_llm_provider:{};model_id:{};video_id:{}"
)
+ LITELLM_PASSTHROUGH_MANAGED_ID_COMPLETE_STR = (
+ "litellm_proxy:passthrough;provider:{};unified_id,{};raw_id,{}"
+ )
+
class ServiceTier(Enum):
"""Enum for service tier types used in cost calculations."""
diff --git a/tests/pass_through_unit_tests/test_passthrough_managed_ids.py b/tests/pass_through_unit_tests/test_passthrough_managed_ids.py
new file mode 100644
index 00000000000..8cf07da3ce4
--- /dev/null
+++ b/tests/pass_through_unit_tests/test_passthrough_managed_ids.py
@@ -0,0 +1,2087 @@
+"""
+Unit tests for passthrough managed IDs (Scope A).
+
+Tests cover:
+ - managed_id_codec: encode / decode / is_managed round-trip and rejection cases.
+ - managed_id_rewriter._resolve_one: cross-route 404, access-check 403, unknown ID 404,
+ raw pass-through.
+ - managed_id_rewriter.rewrite_response_ids: file create swap, batch create swap,
+ dedup reuse (no duplicate row), null field skip.
+ - managed_id_rewriter.rewrite_path_ids / rewrite_query_ids / rewrite_body_ids:
+ INPUT swap and raw pass-through.
+ - Flag-off: feature flag disabled → no swap at all.
+ - Cross-route: managed ID minted for 'openai' rejected on a different provider.
+ - Forged: unknown base64 → 404.
+"""
+
+from __future__ import annotations
+
+import base64
+import json
+import sys
+import os
+from typing import Any
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+sys.path.insert(0, os.path.abspath("../.."))
+
+import litellm
+from litellm.llms.base_llm.managed_resources.utils import (
+ resolve_passthrough_managed_id_provider,
+)
+from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy.pass_through_endpoints.managed_id_codec import (
+ decode,
+ encode,
+ is_managed,
+ new_managed_id,
+)
+from litellm.proxy.pass_through_endpoints.managed_id_rewriter import (
+ _MAX_RAW_ID_GUARD_LOOKUPS,
+ _canonical_path,
+ _passthrough_provider_marker,
+ _resolve_one,
+ is_passthrough_list_route,
+ list_passthrough_ids_from_db,
+ rewrite_body_ids,
+ rewrite_path_ids,
+ rewrite_query_ids,
+ rewrite_response_ids,
+)
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _user(user_id: str = "user-1", team_id: str = "team-1") -> UserAPIKeyAuth:
+ return UserAPIKeyAuth(user_id=user_id, team_id=team_id)
+
+
+def _admin_user() -> UserAPIKeyAuth:
+ u = UserAPIKeyAuth(user_id="admin", user_role="proxy_admin")
+ return u
+
+
+def _prisma_client() -> MagicMock:
+ """Return a MagicMock prisma_client with async db methods."""
+ pc = MagicMock()
+ pc.db = MagicMock()
+ pc.db.litellm_managedfiletable = MagicMock()
+ pc.db.litellm_managedfiletable.find_first = AsyncMock(return_value=None)
+ pc.db.litellm_managedfiletable.find_many = AsyncMock(return_value=[])
+ pc.db.litellm_managedfiletable.create = AsyncMock(return_value=None)
+ pc.db.litellm_managedobjecttable = MagicMock()
+ pc.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=None)
+ pc.db.litellm_managedobjecttable.upsert = AsyncMock(return_value=None)
+ pc.db.litellm_managedobjecttable.update = AsyncMock(return_value=None)
+ return pc
+
+
+def _managed_files_hook(store_side_effect: Any = None) -> MagicMock:
+ hook = MagicMock()
+ hook.get_unified_file_id = AsyncMock(return_value=None)
+ hook.store_unified_file_id = AsyncMock(side_effect=store_side_effect)
+ return hook
+
+
+def _owner_scoped_file_find_many(row: Any):
+ """Return a ``find_many`` that mimics Prisma owner-scoping for the managed
+ file table: an owner-scoped query (one carrying ``created_by`` / ``team_id``
+ / ``OR``) returns ``[]`` because the caller does not own *row*, while an
+ unscoped (global) query returns ``[row]``. This reproduces the cross-tenant
+ bypass that a caller-scoped dedup lookup allowed (the scoped query misses the
+ other tenant's row, so a fresh managed ID gets minted for the attacker)."""
+
+ async def _impl(*args: Any, where: Any = None, **kwargs: Any) -> Any:
+ where = where or {}
+ if "created_by" in where or "team_id" in where or "OR" in where:
+ return []
+ return [row]
+
+ return _impl
+
+
+# ---------------------------------------------------------------------------
+# managed_id_codec — unit tests
+# ---------------------------------------------------------------------------
+
+
+class TestCodec:
+ def test_encode_decode_roundtrip(self):
+ managed_id = encode("openai", "uuid-abc", "file-xyz")
+ payload = decode(managed_id)
+ assert payload is not None
+ assert payload.provider == "openai"
+ assert payload.unified_uuid == "uuid-abc"
+ assert payload.raw_provider_id == "file-xyz"
+
+ def test_is_managed_true(self):
+ assert is_managed(encode("openai", "u1", "file-abc")) is True
+
+ def test_is_managed_false_for_raw_ids(self):
+ assert is_managed("file-abc123") is False
+ assert is_managed("batch_xyz") is False
+ assert is_managed("resp_abc") is False
+
+ def test_decode_returns_none_for_garbage(self):
+ assert decode("not-base64!!!") is None
+ assert decode("") is None
+ assert decode("abc") is None
+
+ def test_decode_returns_none_for_wrong_type(self):
+ assert decode(None) is None # type: ignore[arg-type]
+ assert decode(42) is None # type: ignore[arg-type]
+
+ def test_decode_returns_none_for_unified_endpoint_id(self):
+ # A unified-endpoint ID: starts with litellm_proxy: but lacks passthrough;
+ plaintext = "litellm_proxy:application/octet-stream;unified_id,123;target_model_names,gpt-4"
+ unified_id = base64.urlsafe_b64encode(plaintext.encode()).decode().rstrip("=")
+ assert decode(unified_id) is None
+
+ def test_new_managed_id_produces_valid_id(self):
+ mid = new_managed_id("openai", "batch_abc")
+ payload = decode(mid)
+ assert payload is not None
+ assert payload.provider == "openai"
+ assert payload.raw_provider_id == "batch_abc"
+
+ def test_encode_padding_insensitive(self):
+ """Encoded IDs with varying lengths all decode correctly."""
+ for raw in ("file-x", "file-ab", "file-abc", "file-abcd"):
+ mid = encode("openai", "u", raw)
+ p = decode(mid)
+ assert p is not None and p.raw_provider_id == raw
+
+
+# ---------------------------------------------------------------------------
+# resolve_passthrough_managed_id_provider — provider scope mapping
+# ---------------------------------------------------------------------------
+
+
+class TestManagedIdProviderScope:
+ """Managed-ID scoping is keyed on the explicit forwarded provider, and both
+ azure and azure_ai must collapse to a single 'azure' scope so an ID minted
+ while routing as one resolves while routing as the other."""
+
+ def test_openai_scope(self):
+ assert resolve_passthrough_managed_id_provider("openai") == "openai"
+ assert (
+ resolve_passthrough_managed_id_provider(litellm.LlmProviders.OPENAI)
+ == "openai"
+ )
+
+ def test_azure_scope(self):
+ assert resolve_passthrough_managed_id_provider("azure") == "azure"
+ assert (
+ resolve_passthrough_managed_id_provider(litellm.LlmProviders.AZURE)
+ == "azure"
+ )
+
+ def test_azure_ai_collapses_to_azure(self):
+ assert resolve_passthrough_managed_id_provider("azure_ai") == "azure"
+ assert (
+ resolve_passthrough_managed_id_provider(litellm.LlmProviders.AZURE_AI)
+ == "azure"
+ )
+
+ def test_azure_ai_id_resolves_on_azure_route(self):
+ """End-to-end consequence of the collapse: an ID whose scope was
+ resolved from azure_ai shares the 'azure' namespace, so decoding +
+ cross-route checks line up with an azure-scoped ID."""
+ azure_ai_scope = resolve_passthrough_managed_id_provider("azure_ai")
+ azure_scope = resolve_passthrough_managed_id_provider("azure")
+ managed = new_managed_id(azure_ai_scope, "file-shared")
+ assert decode(managed).provider == azure_scope
+
+ def test_case_insensitive(self):
+ assert resolve_passthrough_managed_id_provider("AZURE") == "azure"
+ assert resolve_passthrough_managed_id_provider("OpenAI") == "openai"
+
+ def test_namespaced_provider_suffix(self):
+ assert resolve_passthrough_managed_id_provider("foo.azure") == "azure"
+ assert resolve_passthrough_managed_id_provider("foo.azure_ai") == "azure"
+ assert resolve_passthrough_managed_id_provider("foo.openai") == "openai"
+
+ def test_non_openai_azure_providers_not_scoped(self):
+ """Managed IDs only apply to explicit openai/azure pass-through; any
+ other provider (or a missing one) must return None so a third-party
+ OpenAI-compatible endpoint never triggers managed-ID minting."""
+ for provider in (None, "", "cohere", "vllm", "anthropic", "gemini", "bedrock"):
+ assert resolve_passthrough_managed_id_provider(provider) is None
+
+
+# ---------------------------------------------------------------------------
+# _canonical_path
+# ---------------------------------------------------------------------------
+
+
+class TestCanonicalPath:
+ def test_strips_openai_prefix(self):
+ assert _canonical_path("/openai/v1/batches/batch_x") == "/v1/batches/batch_x"
+
+ def test_strips_openai_passthrough_prefix(self):
+ assert _canonical_path("/openai_passthrough/v1/files") == "/v1/files"
+
+ def test_leaves_bare_path_unchanged(self):
+ assert _canonical_path("/v1/responses") == "/v1/responses"
+
+ def test_strips_azure_openai_prefix(self):
+ assert _canonical_path("/azure/openai/files") == "/v1/files"
+
+ def test_strips_azure_openai_batch_with_id(self):
+ assert (
+ _canonical_path("/azure/openai/batches/batch_abc123")
+ == "/v1/batches/batch_abc123"
+ )
+
+ def test_strips_azure_openai_responses(self):
+ assert _canonical_path("/azure/openai/responses") == "/v1/responses"
+
+ def test_strips_azure_ai_openai_prefix(self):
+ assert _canonical_path("/azure_ai/openai/files") == "/v1/files"
+
+ def test_strips_azure_ai_openai_batch_cancel(self):
+ assert (
+ _canonical_path("/azure_ai/openai/batches/batch_abc/cancel")
+ == "/v1/batches/batch_abc/cancel"
+ )
+
+ def test_azure_path_already_carrying_v1_is_not_doubled(self):
+ assert _canonical_path("/azure/openai/v1/files") == "/v1/files"
+ assert (
+ _canonical_path("/azure/openai/v1/batches/batch_abc")
+ == "/v1/batches/batch_abc"
+ )
+
+ def test_strips_azure_openai_file_with_id(self):
+ assert _canonical_path("/azure/openai/files/file-abc") == "/v1/files/file-abc"
+
+
+# ---------------------------------------------------------------------------
+# _resolve_one
+# ---------------------------------------------------------------------------
+
+
+class TestResolveOne:
+ @pytest.mark.asyncio
+ async def test_raw_id_passes_through(self):
+ result = await _resolve_one("file-abc", "openai", _user(), None, None)
+ assert result == "file-abc"
+
+ @pytest.mark.asyncio
+ async def test_cross_route_raises_404(self):
+ mid = encode("anthropic", "u", "file-abc")
+ from fastapi import HTTPException
+
+ with pytest.raises(HTTPException) as exc_info:
+ await _resolve_one(mid, "openai", _user(), None, None)
+ assert exc_info.value.status_code == 404
+
+ @pytest.mark.asyncio
+ async def test_unknown_managed_id_raises_404(self):
+ mid = encode("openai", "u", "file-abc")
+ pc = _prisma_client()
+ hook = _managed_files_hook()
+ # Both lookups return None → 404
+ from fastapi import HTTPException
+
+ with pytest.raises(HTTPException) as exc_info:
+ await _resolve_one(mid, "openai", _user(), pc, hook)
+ assert exc_info.value.status_code == 404
+
+ @pytest.mark.asyncio
+ async def test_access_denied_raises_403(self):
+ mid = encode("openai", "u", "file-abc")
+ hook = _managed_files_hook()
+ file_row = MagicMock()
+ file_row.created_by = "other-user"
+ file_row.team_id = "other-team"
+ hook.get_unified_file_id = AsyncMock(return_value=file_row)
+ from fastapi import HTTPException
+
+ with pytest.raises(HTTPException) as exc_info:
+ await _resolve_one(mid, "openai", _user("user-1", "team-1"), None, hook)
+ assert exc_info.value.status_code == 403
+
+ @pytest.mark.asyncio
+ async def test_valid_file_id_resolves(self):
+ mid = encode("openai", "u", "file-xyz")
+ hook = _managed_files_hook()
+ file_row = MagicMock()
+ file_row.created_by = "user-1"
+ file_row.team_id = "team-1"
+ hook.get_unified_file_id = AsyncMock(return_value=file_row)
+ result = await _resolve_one(mid, "openai", _user(), None, hook)
+ assert result == "file-xyz"
+
+ @pytest.mark.asyncio
+ async def test_valid_batch_id_resolves_via_object_table(self):
+ mid = encode("openai", "u", "batch_abc")
+ pc = _prisma_client()
+ obj_row = MagicMock()
+ obj_row.created_by = "user-1"
+ obj_row.team_id = "team-1"
+ pc.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=obj_row)
+ result = await _resolve_one(mid, "openai", _user(), pc, None)
+ assert result == "batch_abc"
+
+ @pytest.mark.asyncio
+ async def test_admin_can_access_any_resource(self):
+ mid = encode("openai", "u", "file-xyz")
+ hook = _managed_files_hook()
+ file_row = MagicMock()
+ file_row.created_by = "other-user"
+ file_row.team_id = "other-team"
+ hook.get_unified_file_id = AsyncMock(return_value=file_row)
+ result = await _resolve_one(mid, "openai", _admin_user(), None, hook)
+ assert result == "file-xyz"
+
+
+# ---------------------------------------------------------------------------
+# rewrite_response_ids — OUTPUT
+# ---------------------------------------------------------------------------
+
+
+class TestRewriteResponseIds:
+ @pytest.mark.asyncio
+ async def test_file_create_mints_managed_id(self):
+ pc = _prisma_client()
+ hook = _managed_files_hook()
+ body = {"id": "file-abc123", "object": "file"}
+ result = await rewrite_response_ids(
+ provider="openai",
+ method="POST",
+ route="/openai/v1/files",
+ body=body,
+ user_api_key_dict=_user(),
+ prisma_client=pc,
+ managed_files_hook=hook,
+ )
+ assert result is not body # mutated copy
+ assert result["id"] != "file-abc123"
+ payload = decode(result["id"])
+ assert payload is not None
+ assert payload.raw_provider_id == "file-abc123"
+ hook.store_unified_file_id.assert_awaited_once()
+
+ @pytest.mark.asyncio
+ async def test_file_create_persist_failure_leaves_raw_id(self):
+ """If the DB write fails, the response must keep the raw provider ID
+ (which still resolves upstream) rather than swap in a managed ID that no
+ DB row backs and that would 404 on every later resolve."""
+ pc = _prisma_client()
+ hook = _managed_files_hook(store_side_effect=Exception("db down"))
+ body = {"id": "file-abc123", "object": "file"}
+ result = await rewrite_response_ids(
+ provider="openai",
+ method="POST",
+ route="/openai/v1/files",
+ body=body,
+ user_api_key_dict=_user(),
+ prisma_client=pc,
+ managed_files_hook=hook,
+ )
+ hook.store_unified_file_id.assert_awaited_once()
+ assert result["id"] == "file-abc123"
+ assert decode(result["id"]) is None
+
+ @pytest.mark.asyncio
+ async def test_batch_create_mints_id_and_input_file_id(self):
+ pc = _prisma_client()
+ hook = _managed_files_hook()
+ body = {
+ "id": "batch_xyz",
+ "input_file_id": "file-abc",
+ "output_file_id": None,
+ "error_file_id": None,
+ }
+ result = await rewrite_response_ids(
+ provider="openai",
+ method="POST",
+ route="/openai/v1/batches",
+ body=body,
+ user_api_key_dict=_user(),
+ prisma_client=pc,
+ managed_files_hook=hook,
+ )
+ assert decode(result["id"]).raw_provider_id == "batch_xyz" # type: ignore[union-attr]
+ assert decode(result["input_file_id"]).raw_provider_id == "file-abc" # type: ignore[union-attr]
+ # Null fields skipped
+ assert result["output_file_id"] is None
+ assert result["error_file_id"] is None
+
+ @pytest.mark.asyncio
+ async def test_response_create_mints_id(self):
+ pc = _prisma_client()
+ hook = _managed_files_hook()
+ body = {"id": "resp_abc", "object": "response"}
+ result = await rewrite_response_ids(
+ provider="openai",
+ method="POST",
+ route="/openai/v1/responses",
+ body=body,
+ user_api_key_dict=_user(),
+ prisma_client=pc,
+ managed_files_hook=hook,
+ )
+ assert decode(result["id"]).raw_provider_id == "resp_abc" # type: ignore[union-attr]
+
+ @pytest.mark.asyncio
+ async def test_azure_response_create_mints_id(self):
+ pc = _prisma_client()
+ hook = _managed_files_hook()
+ body = {
+ "id": "resp_0dce2668af072bdc006a195db1f96c8194b6217f8e0d0b3ccd",
+ "object": "response",
+ "status": "completed",
+ }
+ result = await rewrite_response_ids(
+ provider="azure",
+ method="POST",
+ route="/azure/openai/responses",
+ body=body,
+ user_api_key_dict=_user(),
+ prisma_client=pc,
+ managed_files_hook=hook,
+ )
+ assert (
+ decode(result["id"]).raw_provider_id # type: ignore[union-attr]
+ == "resp_0dce2668af072bdc006a195db1f96c8194b6217f8e0d0b3ccd"
+ )
+
+ @pytest.mark.asyncio
+ async def test_no_map_entry_returns_body_unchanged(self):
+ pc = _prisma_client()
+ hook = _managed_files_hook()
+ body = {"id": "msg_xyz", "object": "message"}
+ result = await rewrite_response_ids(
+ provider="openai",
+ method="POST",
+ route="/openai/v1/chat/completions",
+ body=body,
+ user_api_key_dict=_user(),
+ prisma_client=pc,
+ managed_files_hook=hook,
+ )
+ assert result is body # same object, unchanged
+
+ @pytest.mark.asyncio
+ async def test_dedup_reuses_existing_file_row(self):
+ """File uploaded via passthrough, then referenced in a batch — no new row."""
+ existing_managed_id = new_managed_id("openai", "file-abc")
+ existing_row = MagicMock()
+ existing_row.unified_file_id = existing_managed_id
+ existing_row.created_by = "user-1"
+ existing_row.team_id = "team-1"
+
+ pc = _prisma_client()
+ # Dedup lookup finds existing row
+ pc.db.litellm_managedfiletable.find_many = AsyncMock(
+ return_value=[existing_row]
+ )
+ hook = _managed_files_hook()
+ body = {
+ "id": "batch_xyz",
+ "input_file_id": "file-abc",
+ "output_file_id": None,
+ "error_file_id": None,
+ }
+ result = await rewrite_response_ids(
+ provider="openai",
+ method="POST",
+ route="/openai/v1/batches",
+ body=body,
+ user_api_key_dict=_user(),
+ prisma_client=pc,
+ managed_files_hook=hook,
+ )
+ # input_file_id should be the SAME managed ID already in DB
+ assert result["input_file_id"] == existing_managed_id
+ # store_unified_file_id should NOT have been called (reused existing)
+ hook.store_unified_file_id.assert_not_awaited()
+
+ @pytest.mark.asyncio
+ async def test_dedup_skips_cross_provider_file_row(self):
+ """Same raw file ID for a different provider must mint a new managed ID."""
+ azure_managed_id = new_managed_id("azure", "file-abc")
+ existing_row = MagicMock()
+ existing_row.unified_file_id = azure_managed_id
+
+ pc = _prisma_client()
+ pc.db.litellm_managedfiletable.find_many = AsyncMock(
+ return_value=[existing_row]
+ )
+ hook = _managed_files_hook()
+ body = {"id": "file-abc", "object": "file"}
+ result = await rewrite_response_ids(
+ provider="openai",
+ method="POST",
+ route="/openai/v1/files",
+ body=body,
+ user_api_key_dict=_user(),
+ prisma_client=pc,
+ managed_files_hook=hook,
+ )
+ assert decode(result["id"]).provider == "openai"
+ assert decode(result["id"]).raw_provider_id == "file-abc"
+ assert result["id"] != azure_managed_id
+ hook.store_unified_file_id.assert_awaited_once()
+
+ @pytest.mark.asyncio
+ async def test_dedup_reuses_same_provider_row_amid_collision(self):
+ """When OpenAI and Azure both issued the same raw file ID, an Azure call
+ must reuse the existing Azure managed row deterministically rather than
+ mint a duplicate, even when the cross-provider OpenAI row is returned
+ first by the DB."""
+ raw_id = "file-collision"
+ openai_row = MagicMock()
+ openai_row.unified_file_id = new_managed_id("openai", raw_id)
+ openai_row.created_by = "user-1"
+ openai_row.team_id = "team-1"
+ azure_managed_id = new_managed_id("azure", raw_id)
+ azure_row = MagicMock()
+ azure_row.unified_file_id = azure_managed_id
+ azure_row.created_by = "user-1"
+ azure_row.team_id = "team-1"
+
+ pc = _prisma_client()
+ # Cross-provider row listed first to expose any non-deterministic pick.
+ pc.db.litellm_managedfiletable.find_many = AsyncMock(
+ return_value=[openai_row, azure_row]
+ )
+ hook = _managed_files_hook()
+ body = {"id": raw_id, "object": "file"}
+ result = await rewrite_response_ids(
+ provider="azure",
+ method="GET",
+ route=f"/azure/openai/files/{raw_id}",
+ body=body,
+ user_api_key_dict=_user(),
+ prisma_client=pc,
+ managed_files_hook=hook,
+ )
+ assert result["id"] == azure_managed_id
+ hook.store_unified_file_id.assert_not_awaited()
+
+ @pytest.mark.asyncio
+ async def test_cross_owner_file_retrieve_raises_404(self):
+ """
+ A caller who fetches another tenant's raw ``file-...`` ID through
+ GET /openai/v1/files/{file_id} (which bypasses the managed-ID input gate)
+ must be denied with a 404 — the response path must NOT mint a fresh
+ managed ID for that file under the attacker.
+ """
+ from fastapi import HTTPException
+
+ pc = _prisma_client()
+ other_owner_row = MagicMock()
+ other_owner_row.created_by = "victim"
+ other_owner_row.team_id = "victim-team"
+ other_owner_row.unified_file_id = encode("openai", "victim", "file-victim")
+ pc.db.litellm_managedfiletable.find_many = _owner_scoped_file_find_many(
+ other_owner_row
+ )
+ hook = _managed_files_hook()
+
+ body = {"id": "file-victim", "object": "file"}
+ with pytest.raises(HTTPException) as exc_info:
+ await rewrite_response_ids(
+ provider="openai",
+ method="GET",
+ route="/openai/v1/files/file-victim",
+ body=body,
+ user_api_key_dict=_user("attacker", "attacker-team"),
+ prisma_client=pc,
+ managed_files_hook=hook,
+ )
+ assert exc_info.value.status_code == 404
+ # Must not mint / persist a managed ID for the attacker.
+ hook.store_unified_file_id.assert_not_awaited()
+
+ @pytest.mark.asyncio
+ async def test_cross_owner_file_delete_raises_404(self):
+ """DELETE is also a non-create route: cross-owner raw file IDs are denied."""
+ from fastapi import HTTPException
+
+ pc = _prisma_client()
+ other_owner_row = MagicMock()
+ other_owner_row.created_by = "victim"
+ other_owner_row.team_id = "victim-team"
+ other_owner_row.unified_file_id = encode("openai", "victim", "file-victim")
+ pc.db.litellm_managedfiletable.find_many = _owner_scoped_file_find_many(
+ other_owner_row
+ )
+ hook = _managed_files_hook()
+
+ body = {"id": "file-victim", "object": "file", "deleted": True}
+ with pytest.raises(HTTPException) as exc_info:
+ await rewrite_response_ids(
+ provider="openai",
+ method="DELETE",
+ route="/openai/v1/files/file-victim",
+ body=body,
+ user_api_key_dict=_user("attacker", "attacker-team"),
+ prisma_client=pc,
+ managed_files_hook=hook,
+ )
+ assert exc_info.value.status_code == 404
+ hook.store_unified_file_id.assert_not_awaited()
+
+ @pytest.mark.asyncio
+ async def test_cross_owner_file_create_leaves_raw_id(self):
+ """
+ On the create (POST /v1/files) path a cross-owner dedup hit must NOT 404
+ the caller's own successful upload; leave the raw ID unmanaged instead
+ (mirrors the batch/response create behaviour).
+ """
+ pc = _prisma_client()
+ other_owner_row = MagicMock()
+ other_owner_row.created_by = "victim"
+ other_owner_row.team_id = "victim-team"
+ other_owner_row.unified_file_id = encode("openai", "victim", "file-shared")
+ pc.db.litellm_managedfiletable.find_many = _owner_scoped_file_find_many(
+ other_owner_row
+ )
+ hook = _managed_files_hook()
+
+ body = {"id": "file-shared", "object": "file"}
+ result = await rewrite_response_ids(
+ provider="openai",
+ method="POST",
+ route="/openai/v1/files",
+ body=body,
+ user_api_key_dict=_user("uploader", "uploader-team"),
+ prisma_client=pc,
+ managed_files_hook=hook,
+ )
+ assert result["id"] == "file-shared"
+ hook.store_unified_file_id.assert_not_awaited()
+
+ @pytest.mark.asyncio
+ async def test_team_member_reuses_shared_file_row(self):
+ """A teammate of the file owner can reuse the existing managed file row
+ (the cross-tenant guard scopes by team, not just the creating user)."""
+ existing_managed_id = new_managed_id("openai", "file-team")
+ existing_row = MagicMock()
+ existing_row.unified_file_id = existing_managed_id
+ existing_row.created_by = "owner-user"
+ existing_row.team_id = "shared-team"
+
+ pc = _prisma_client()
+ pc.db.litellm_managedfiletable.find_many = AsyncMock(
+ return_value=[existing_row]
+ )
+ hook = _managed_files_hook()
+
+ body = {"id": "file-team", "object": "file"}
+ result = await rewrite_response_ids(
+ provider="openai",
+ method="GET",
+ route="/openai/v1/files/file-team",
+ body=body,
+ user_api_key_dict=_user("teammate", "shared-team"),
+ prisma_client=pc,
+ managed_files_hook=hook,
+ )
+ assert result["id"] == existing_managed_id
+ hook.store_unified_file_id.assert_not_awaited()
+
+ @pytest.mark.asyncio
+ async def test_openai_passthrough_prefix_normalised(self):
+ """Routes under /openai_passthrough/ work the same as /openai/."""
+ pc = _prisma_client()
+ hook = _managed_files_hook()
+ body = {"id": "file-abc", "object": "file"}
+ result = await rewrite_response_ids(
+ provider="openai",
+ method="POST",
+ route="/openai_passthrough/v1/files",
+ body=body,
+ user_api_key_dict=_user(),
+ prisma_client=pc,
+ managed_files_hook=hook,
+ )
+ assert decode(result["id"]).raw_provider_id == "file-abc" # type: ignore[union-attr]
+
+ @pytest.mark.asyncio
+ async def test_batch_reuse_refreshes_stored_snapshot(self):
+ """Retrieving a completed batch must refresh the stored snapshot so the
+ DB-served list reflects fields (e.g. output_file_id) that were null at
+ creation time. The dedup-reuse path must update file_object, not just
+ return the existing id with a stale snapshot."""
+ existing_managed_id = new_managed_id("openai", "batch_done")
+ existing_row = MagicMock()
+ existing_row.unified_object_id = existing_managed_id
+ existing_row.created_by = "user-1"
+ existing_row.team_id = "team-1"
+
+ pc = _prisma_client()
+ pc.db.litellm_managedobjecttable.find_first = AsyncMock(
+ return_value=existing_row
+ )
+
+ completed_body = {
+ "id": "batch_done",
+ "object": "batch",
+ "status": "completed",
+ "output_file_id": "file-out",
+ "error_file_id": None,
+ }
+ result = await rewrite_response_ids(
+ provider="openai",
+ method="GET",
+ route="/openai/v1/batches/batch_done",
+ body=completed_body,
+ user_api_key_dict=_user("user-1", "team-1"),
+ prisma_client=pc,
+ managed_files_hook=None,
+ )
+
+ # Reuses the existing managed id (no new row minted)
+ assert result["id"] == existing_managed_id
+ pc.db.litellm_managedobjecttable.upsert.assert_not_awaited()
+ # The stored snapshot is refreshed with the completed batch body
+ pc.db.litellm_managedobjecttable.update.assert_awaited_once()
+ update_kwargs = pc.db.litellm_managedobjecttable.update.call_args.kwargs
+ assert update_kwargs["where"] == {"unified_object_id": existing_managed_id}
+ stored = json.loads(update_kwargs["data"]["file_object"])
+ assert stored["status"] == "completed"
+ # output_file_id is itself rewritten to a managed id wrapping the raw id
+ assert decode(stored["output_file_id"]).raw_provider_id == "file-out"
+
+ @pytest.mark.asyncio
+ async def test_cross_provider_batch_collision_mints_new_id(self):
+ """
+ If OpenAI and Azure independently issue the same raw batch ID, the
+ Azure call must mint its own row keyed by 'passthrough:azure:batch_shared'
+ and must NOT raise 404. The namespaced model_object_id prevents a
+ UniqueConstraintViolation on the @unique column.
+ """
+ pc = _prisma_client()
+ # Both providers return no existing row (different namespaced keys)
+ pc.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=None)
+ pc.db.litellm_managedobjecttable.upsert = AsyncMock(return_value=None)
+
+ body = {"id": "batch_shared", "object": "batch", "input_file_id": None}
+ result = await rewrite_response_ids(
+ provider="azure",
+ method="POST",
+ route="/azure/openai/batches",
+ body=body,
+ user_api_key_dict=_user("user-azure", "team-azure"),
+ prisma_client=pc,
+ managed_files_hook=None,
+ )
+ # Must mint a fresh azure-scoped managed ID
+ assert decode(result["id"]) is not None
+ assert decode(result["id"]).provider == "azure"
+ assert decode(result["id"]).raw_provider_id == "batch_shared"
+
+ # Verify the upsert stored the namespaced model_object_id
+ call_data = pc.db.litellm_managedobjecttable.upsert.call_args.kwargs["data"]
+ assert (
+ call_data["create"]["model_object_id"] == "passthrough:azure:batch_shared"
+ )
+
+ @pytest.mark.asyncio
+ async def test_batch_create_persist_failure_leaves_raw_id(self):
+ """If the object upsert fails, the batch response must keep the raw
+ provider ID rather than return a managed ID with no backing DB row that
+ would 404 on every subsequent resolve."""
+ pc = _prisma_client()
+ pc.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=None)
+ pc.db.litellm_managedobjecttable.upsert = AsyncMock(
+ side_effect=Exception("db down")
+ )
+ body = {"id": "batch_xyz", "object": "batch", "input_file_id": None}
+ result = await rewrite_response_ids(
+ provider="openai",
+ method="POST",
+ route="/openai/v1/batches",
+ body=body,
+ user_api_key_dict=_user(),
+ prisma_client=pc,
+ managed_files_hook=None,
+ )
+ pc.db.litellm_managedobjecttable.upsert.assert_awaited_once()
+ assert result["id"] == "batch_xyz"
+ assert decode(result["id"]) is None
+
+ @pytest.mark.asyncio
+ async def test_concurrent_create_converges_on_winner_managed_id(self):
+ """
+ Two callers minting the same namespaced object row race: the dedup lookup
+ finds nothing for both, but the @unique model_object_id lets only one
+ insert win. The loser's upsert raises, and it must re-read the winner's
+ row and return that managed ID rather than silently keeping the raw ID
+ (which would leave the two callers divergent for the same upstream batch).
+ """
+ pc = _prisma_client()
+ winner_managed_id = encode("openai", "winner-uuid", "batch_race")
+ winner_row = MagicMock()
+ winner_row.created_by = "user-1"
+ winner_row.team_id = "team-1"
+ winner_row.unified_object_id = winner_managed_id
+ # First (dedup) lookup misses; post-collision re-read finds the winner.
+ pc.db.litellm_managedobjecttable.find_first = AsyncMock(
+ side_effect=[None, winner_row]
+ )
+ pc.db.litellm_managedobjecttable.upsert = AsyncMock(
+ side_effect=Exception("UniqueConstraintViolation: model_object_id")
+ )
+
+ body = {"id": "batch_race", "object": "batch", "input_file_id": None}
+ result = await rewrite_response_ids(
+ provider="openai",
+ method="POST",
+ route="/openai/v1/batches",
+ body=body,
+ user_api_key_dict=_user(),
+ prisma_client=pc,
+ managed_files_hook=None,
+ )
+ # The loser converges on the winner's managed ID, not the raw batch ID.
+ assert result["id"] == winner_managed_id
+ assert decode(result["id"]).raw_provider_id == "batch_race"
+ assert pc.db.litellm_managedobjecttable.find_first.await_count == 2
+
+ @pytest.mark.asyncio
+ async def test_concurrent_create_race_with_cross_owner_winner_retrieve_404(self):
+ """
+ If the row that wins the insert race on a non-create (retrieve) route is
+ owned by a different tenant, the loser must be denied with 404 rather
+ than handed the raw ID — the post-collision re-read runs the same access
+ check as the initial dedup hit.
+ """
+ from fastapi import HTTPException
+
+ pc = _prisma_client()
+ winner_row = MagicMock()
+ winner_row.created_by = "other-user"
+ winner_row.team_id = "other-team"
+ winner_row.unified_object_id = encode("openai", "other-uuid", "batch_race")
+ pc.db.litellm_managedobjecttable.find_first = AsyncMock(
+ side_effect=[None, winner_row]
+ )
+ pc.db.litellm_managedobjecttable.upsert = AsyncMock(
+ side_effect=Exception("UniqueConstraintViolation: model_object_id")
+ )
+
+ body = {"id": "batch_race", "object": "batch", "input_file_id": None}
+ with pytest.raises(HTTPException) as exc_info:
+ await rewrite_response_ids(
+ provider="openai",
+ method="GET",
+ route="/openai/v1/batches/batch_race",
+ body=body,
+ user_api_key_dict=_user("attacker", "attacker-team"),
+ prisma_client=pc,
+ managed_files_hook=None,
+ )
+ assert exc_info.value.status_code == 404
+
+ @pytest.mark.asyncio
+ async def test_cross_provider_batch_collision_dedup_uses_namespaced_key(self):
+ """
+ When OpenAI already has a row for batch_shared, an Azure request must
+ look up 'passthrough:azure:batch_shared' (not 'batch_shared'), find
+ nothing, and mint a new row — not raise 404 or reuse the OpenAI row.
+ """
+ pc = _prisma_client()
+ # Simulate: OpenAI row exists under 'passthrough:openai:batch_shared',
+ # but Azure lookup for 'passthrough:azure:batch_shared' returns None.
+ pc.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=None)
+ pc.db.litellm_managedobjecttable.upsert = AsyncMock(return_value=None)
+
+ body = {"id": "batch_shared", "object": "batch", "input_file_id": None}
+ result = await rewrite_response_ids(
+ provider="azure",
+ method="POST",
+ route="/azure/openai/batches",
+ body=body,
+ user_api_key_dict=_user("user-azure", "team-azure"),
+ prisma_client=pc,
+ managed_files_hook=None,
+ )
+ # The dedup lookup must use the namespaced key
+ lookup_where = pc.db.litellm_managedobjecttable.find_first.call_args.kwargs[
+ "where"
+ ]
+ assert lookup_where["model_object_id"] == "passthrough:azure:batch_shared"
+ # Result is a valid azure-scoped managed ID
+ assert decode(result["id"]).provider == "azure"
+
+ @pytest.mark.asyncio
+ async def test_cross_owner_object_collision_returns_raw_id_not_404(self):
+ """
+ On the OUTPUT (mint) path, if the namespaced key is already owned by a
+ different caller (e.g. two upstream accounts under one provider name
+ issued the same raw batch ID), the caller's successful upstream create
+ must NOT be turned into a 404. Leave their raw ID unmanaged instead.
+ """
+ pc = _prisma_client()
+ other_owner_row = MagicMock()
+ other_owner_row.created_by = "other-user"
+ other_owner_row.team_id = "other-team"
+ other_owner_row.unified_object_id = encode(
+ "azure", "other-user", "batch_shared"
+ )
+ pc.db.litellm_managedobjecttable.find_first = AsyncMock(
+ return_value=other_owner_row
+ )
+ pc.db.litellm_managedobjecttable.upsert = AsyncMock(return_value=None)
+
+ body = {"id": "batch_shared", "object": "batch", "input_file_id": None}
+ result = await rewrite_response_ids(
+ provider="azure",
+ method="POST",
+ route="/azure/openai/batches",
+ body=body,
+ user_api_key_dict=_user("user-azure", "team-azure"),
+ prisma_client=pc,
+ managed_files_hook=None,
+ )
+ # Caller gets their raw batch ID back, unmanaged; not a 404, and not
+ # the other owner's managed ID.
+ assert result["id"] == "batch_shared"
+ # No new row is minted (would violate the @unique model_object_id).
+ pc.db.litellm_managedobjecttable.upsert.assert_not_awaited()
+
+ @pytest.mark.asyncio
+ async def test_cross_owner_object_retrieve_raises_404(self):
+ """
+ On a retrieve route, a caller who supplies another owner's raw batch ID
+ (which bypasses the managed-ID input gate) must be denied with a 404 —
+ the upstream object must NOT be echoed back with its raw ID.
+ """
+ from fastapi import HTTPException
+
+ pc = _prisma_client()
+ other_owner_row = MagicMock()
+ other_owner_row.created_by = "other-user"
+ other_owner_row.team_id = "other-team"
+ other_owner_row.unified_object_id = encode("openai", "other-user", "batch_xyz")
+ pc.db.litellm_managedobjecttable.find_first = AsyncMock(
+ return_value=other_owner_row
+ )
+
+ body = {"id": "batch_xyz", "object": "batch", "input_file_id": None}
+ with pytest.raises(HTTPException) as exc_info:
+ await rewrite_response_ids(
+ provider="openai",
+ method="GET",
+ route="/openai/v1/batches/batch_xyz",
+ body=body,
+ user_api_key_dict=_user("attacker", "attacker-team"),
+ prisma_client=pc,
+ managed_files_hook=None,
+ )
+ assert exc_info.value.status_code == 404
+ # Must not silently mint a row for the attacker either.
+ pc.db.litellm_managedobjecttable.upsert.assert_not_awaited()
+
+ @pytest.mark.asyncio
+ async def test_cross_owner_response_delete_raises_404(self):
+ """A delete route is also a non-create route: cross-owner access is denied."""
+ from fastapi import HTTPException
+
+ pc = _prisma_client()
+ other_owner_row = MagicMock()
+ other_owner_row.created_by = "other-user"
+ other_owner_row.team_id = "other-team"
+ other_owner_row.unified_object_id = encode("openai", "other-user", "resp_abc")
+ pc.db.litellm_managedobjecttable.find_first = AsyncMock(
+ return_value=other_owner_row
+ )
+
+ body = {"id": "resp_abc", "object": "response"}
+ with pytest.raises(HTTPException) as exc_info:
+ await rewrite_response_ids(
+ provider="openai",
+ method="DELETE",
+ route="/openai/v1/responses/resp_abc",
+ body=body,
+ user_api_key_dict=_user("attacker", "attacker-team"),
+ prisma_client=pc,
+ managed_files_hook=None,
+ )
+ assert exc_info.value.status_code == 404
+
+ @pytest.mark.asyncio
+ async def test_batch_retrieve_swaps_output_file_id(self):
+ pc = _prisma_client()
+ hook = _managed_files_hook()
+ body = {
+ "id": "batch_xyz",
+ "input_file_id": "file-in",
+ "output_file_id": "file-out",
+ "error_file_id": "file-err",
+ }
+ result = await rewrite_response_ids(
+ provider="openai",
+ method="GET",
+ route="/openai/v1/batches/batch_xyz",
+ body=body,
+ user_api_key_dict=_user(),
+ prisma_client=pc,
+ managed_files_hook=hook,
+ )
+ assert decode(result["output_file_id"]).raw_provider_id == "file-out" # type: ignore[union-attr]
+ assert decode(result["error_file_id"]).raw_provider_id == "file-err" # type: ignore[union-attr]
+
+ @pytest.mark.asyncio
+ async def test_file_create_persists_metadata_for_list(self):
+ """The file's upstream metadata is stored so the DB-served list returns
+ the same fields as a direct file GET (managed ID swapped in)."""
+ pc = _prisma_client()
+ hook = _managed_files_hook()
+ body = {
+ "id": "file-abc123",
+ "object": "file",
+ "bytes": 120,
+ "created_at": 1234567890,
+ "filename": "train.jsonl",
+ "purpose": "batch",
+ "status": "processed",
+ }
+ result = await rewrite_response_ids(
+ provider="openai",
+ method="POST",
+ route="/openai/v1/files",
+ body=body,
+ user_api_key_dict=_user(),
+ prisma_client=pc,
+ managed_files_hook=hook,
+ )
+ stored = hook.store_unified_file_id.call_args.kwargs["file_object"]
+ assert stored is not None
+ assert stored.filename == "train.jsonl"
+ assert stored.bytes == 120
+ assert stored.purpose == "batch"
+ # Managed ID is swapped into the persisted metadata (never the raw one).
+ assert stored.id == result["id"]
+ assert decode(stored.id).raw_provider_id == "file-abc123" # type: ignore[union-attr]
+
+ @pytest.mark.asyncio
+ async def test_file_create_without_metadata_stores_no_file_object(self):
+ """A minimal file response (no bytes/filename) falls back to storing the
+ row without metadata rather than raising."""
+ pc = _prisma_client()
+ hook = _managed_files_hook()
+ body = {"id": "file-abc123", "object": "file"}
+ await rewrite_response_ids(
+ provider="openai",
+ method="POST",
+ route="/openai/v1/files",
+ body=body,
+ user_api_key_dict=_user(),
+ prisma_client=pc,
+ managed_files_hook=hook,
+ )
+ hook.store_unified_file_id.assert_awaited_once()
+ assert hook.store_unified_file_id.call_args.kwargs["file_object"] is None
+
+ @pytest.mark.asyncio
+ async def test_file_create_persists_provider_marker_for_list_scope(self):
+ """The minted file row must carry the provider marker (it flows into
+ flat_model_file_ids), or the DB-pushed provider scope in
+ list_passthrough_ids_from_db would never match it."""
+ pc = _prisma_client()
+ hook = _managed_files_hook()
+ await rewrite_response_ids(
+ provider="azure",
+ method="POST",
+ route="/azure/openai/files",
+ body={"id": "file-abc123", "object": "file"},
+ user_api_key_dict=_user(),
+ prisma_client=pc,
+ managed_files_hook=hook,
+ )
+ mappings = hook.store_unified_file_id.call_args.kwargs["model_mappings"]
+ assert _passthrough_provider_marker("azure") in mappings.values()
+ assert _passthrough_provider_marker("openai") not in mappings.values()
+
+ @pytest.mark.asyncio
+ async def test_batch_snapshot_stores_managed_nested_file_ids(self):
+ """The persisted batch snapshot must carry the managed nested file ID so
+ the list response matches the rewritten direct GET response."""
+ import json as _json
+
+ pc = _prisma_client()
+ hook = _managed_files_hook()
+ body = {
+ "id": "batch_xyz",
+ "object": "batch",
+ "input_file_id": "file-in",
+ "output_file_id": None,
+ "error_file_id": None,
+ }
+ result = await rewrite_response_ids(
+ provider="openai",
+ method="POST",
+ route="/openai/v1/batches",
+ body=body,
+ user_api_key_dict=_user(),
+ prisma_client=pc,
+ managed_files_hook=hook,
+ )
+ stored = pc.db.litellm_managedobjecttable.upsert.call_args.kwargs["data"][
+ "create"
+ ]["file_object"]
+ snapshot = _json.loads(stored)
+ assert snapshot["input_file_id"] == result["input_file_id"]
+ assert decode(snapshot["input_file_id"]).raw_provider_id == "file-in" # type: ignore[union-attr]
+
+
+# ---------------------------------------------------------------------------
+# rewrite_path_ids — INPUT
+# ---------------------------------------------------------------------------
+
+
+class TestRewritePathIds:
+ @pytest.mark.asyncio
+ async def test_raw_segment_passes_through(self):
+ result = await rewrite_path_ids(
+ "/v1/batches/batch_abc", "openai", _user(), None, None
+ )
+ assert result == "/v1/batches/batch_abc"
+
+ @pytest.mark.asyncio
+ async def test_managed_segment_is_resolved(self):
+ mid = encode("openai", "u", "batch_abc")
+ hook = _managed_files_hook()
+ pc = _prisma_client()
+ obj_row = MagicMock()
+ obj_row.created_by = "user-1"
+ obj_row.team_id = "team-1"
+ pc.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=obj_row)
+ result = await rewrite_path_ids(
+ f"/v1/batches/{mid}", "openai", _user(), pc, hook
+ )
+ assert result == "/v1/batches/batch_abc"
+
+ @pytest.mark.asyncio
+ async def test_cross_route_in_path_raises_404(self):
+ mid = encode("anthropic", "u", "batch_abc")
+ from fastapi import HTTPException
+
+ with pytest.raises(HTTPException) as exc_info:
+ await rewrite_path_ids(f"/v1/batches/{mid}", "openai", _user(), None, None)
+ assert exc_info.value.status_code == 404
+
+
+# ---------------------------------------------------------------------------
+# rewrite_query_ids — INPUT
+# ---------------------------------------------------------------------------
+
+
+class TestRewriteQueryIds:
+ @pytest.mark.asyncio
+ async def test_raw_params_pass_through(self):
+ params = {"limit": "10", "after": "batch_xyz"}
+ result = await rewrite_query_ids(params, "openai", _user(), None, None)
+ assert result is params # unchanged same object
+
+ @pytest.mark.asyncio
+ async def test_none_returns_none(self):
+ result = await rewrite_query_ids(None, "openai", _user(), None, None)
+ assert result is None
+
+ @pytest.mark.asyncio
+ async def test_managed_param_is_resolved(self):
+ mid = encode("openai", "u", "file-abc")
+ hook = _managed_files_hook()
+ file_row = MagicMock()
+ file_row.created_by = "user-1"
+ file_row.team_id = "team-1"
+ hook.get_unified_file_id = AsyncMock(return_value=file_row)
+ params = {"file_id": mid}
+ result = await rewrite_query_ids(params, "openai", _user(), None, hook)
+ assert result is not params
+ assert result["file_id"] == "file-abc" # type: ignore[index]
+
+
+# ---------------------------------------------------------------------------
+# rewrite_body_ids — INPUT
+# ---------------------------------------------------------------------------
+
+
+class TestRewriteBodyIds:
+ @pytest.mark.asyncio
+ async def test_raw_body_passes_through(self):
+ body = {"input_file_id": "file-abc", "model": "gpt-4o"}
+ result = await rewrite_body_ids(body, "openai", _user(), None, None)
+ assert result is body
+
+ @pytest.mark.asyncio
+ async def test_none_returns_none(self):
+ result = await rewrite_body_ids(None, "openai", _user(), None, None)
+ assert result is None
+
+ @pytest.mark.asyncio
+ async def test_managed_id_in_body_resolved(self):
+ mid = encode("openai", "u", "file-xyz")
+ hook = _managed_files_hook()
+ file_row = MagicMock()
+ file_row.created_by = "user-1"
+ file_row.team_id = "team-1"
+ hook.get_unified_file_id = AsyncMock(return_value=file_row)
+ body = {"input_file_id": mid}
+ result = await rewrite_body_ids(body, "openai", _user(), None, hook)
+ assert result is not body
+ assert result["input_file_id"] == "file-xyz" # type: ignore[index]
+
+ @pytest.mark.asyncio
+ async def test_litellm_internal_key_preserved(self):
+ """litellm_logging_obj and similar keys are never walked."""
+ logging_obj = object()
+ body = {"litellm_logging_obj": logging_obj, "model": "gpt-4o"}
+ result = await rewrite_body_ids(body, "openai", _user(), None, None)
+ # Internal key preserved by reference
+ assert result["litellm_logging_obj"] is logging_obj # type: ignore[index]
+
+ @pytest.mark.asyncio
+ async def test_nested_list_resolved(self):
+ """Managed IDs inside nested lists are resolved."""
+ mid = encode("openai", "u", "file-nested")
+ hook = _managed_files_hook()
+ file_row = MagicMock()
+ file_row.created_by = "user-1"
+ file_row.team_id = "team-1"
+ hook.get_unified_file_id = AsyncMock(return_value=file_row)
+ body = {"files": [mid, "raw-string"]}
+ result = await rewrite_body_ids(body, "openai", _user(), None, hook)
+ assert result["files"][0] == "file-nested" # type: ignore[index]
+ assert result["files"][1] == "raw-string" # type: ignore[index]
+
+ @pytest.mark.asyncio
+ async def test_forged_managed_id_raises_404(self):
+ """An unknown managed ID in the body raises 404 (not passed to upstream)."""
+ mid = encode("openai", "u", "file-forged")
+ hook = _managed_files_hook()
+ hook.get_unified_file_id = AsyncMock(return_value=None)
+ pc = _prisma_client()
+ body = {"input_file_id": mid}
+ from fastapi import HTTPException
+
+ with pytest.raises(HTTPException) as exc_info:
+ await rewrite_body_ids(body, "openai", _user(), pc, hook)
+ assert exc_info.value.status_code == 404
+
+ @pytest.mark.asyncio
+ async def test_cross_user_access_denied_in_body(self):
+ """A managed ID owned by a different user raises 403."""
+ mid = encode("openai", "u", "file-other")
+ hook = _managed_files_hook()
+ file_row = MagicMock()
+ file_row.created_by = "other-user"
+ file_row.team_id = "other-team"
+ hook.get_unified_file_id = AsyncMock(return_value=file_row)
+ body = {"input_file_id": mid}
+ from fastapi import HTTPException
+
+ with pytest.raises(HTTPException) as exc_info:
+ await rewrite_body_ids(
+ body, "openai", _user("user-1", "team-1"), None, hook
+ )
+ assert exc_info.value.status_code == 403
+
+ @pytest.mark.asyncio
+ async def test_deeply_nested_body_does_not_overflow_stack(self):
+ """A pathologically deep body must not blow the Python stack: rewriting
+ stops at the depth cap and returns the body unchanged instead of raising
+ RecursionError."""
+ node: Any = {"leaf": "raw-value"}
+ for _ in range(5000):
+ node = {"nested": node}
+
+ result = await rewrite_body_ids(node, "openai", _user(), None, None)
+ assert result is node
+
+ @pytest.mark.asyncio
+ async def test_managed_id_resolved_within_depth_cap(self):
+ """A managed ID nested well within the depth cap is still resolved, so
+ the cap never truncates legitimately-shaped bodies."""
+ mid = encode("openai", "u", "file-deep")
+ hook = _managed_files_hook()
+ file_row = MagicMock()
+ file_row.created_by = "user-1"
+ file_row.team_id = "team-1"
+ hook.get_unified_file_id = AsyncMock(return_value=file_row)
+
+ leaf = {"input_file_id": mid}
+ node: Any = leaf
+ for _ in range(20):
+ node = {"nested": node}
+
+ result = await rewrite_body_ids(node, "openai", _user(), None, hook)
+
+ cursor = result
+ for _ in range(20):
+ cursor = cursor["nested"] # type: ignore[index]
+ assert cursor["input_file_id"] == "file-deep" # type: ignore[index]
+
+
+# ---------------------------------------------------------------------------
+# Raw-provider-ID input guard — a raw ID recovered by decoding another tenant's
+# managed ID must NOT be forwarded upstream when it maps to a managed resource
+# the caller does not own (otherwise a DELETE / cancel runs upstream before the
+# response-side ownership check).
+# ---------------------------------------------------------------------------
+
+
+class TestRawProviderIdInputGuard:
+ @staticmethod
+ def _victim_file_row() -> MagicMock:
+ row = MagicMock()
+ row.created_by = "victim"
+ row.team_id = "victim-team"
+ row.unified_file_id = encode("openai", "victim", "file-victim")
+ return row
+
+ @staticmethod
+ def _victim_object_row() -> MagicMock:
+ row = MagicMock()
+ row.created_by = "victim"
+ row.team_id = "victim-team"
+ row.unified_object_id = encode("openai", "victim", "batch_victim")
+ return row
+
+ @pytest.mark.asyncio
+ async def test_raw_file_path_for_other_owner_denied(self):
+ """DELETE /openai/v1/files/file-victim with a raw ID that belongs to
+ another tenant's managed file is rejected (404) before forwarding."""
+ from fastapi import HTTPException
+
+ pc = _prisma_client()
+ pc.db.litellm_managedfiletable.find_many = AsyncMock(
+ return_value=[self._victim_file_row()]
+ )
+ with pytest.raises(HTTPException) as exc_info:
+ await rewrite_path_ids(
+ "/openai/v1/files/file-victim",
+ "openai",
+ _user("attacker", "attacker-team"),
+ pc,
+ _managed_files_hook(),
+ )
+ assert exc_info.value.status_code == 404
+
+ @pytest.mark.asyncio
+ async def test_raw_batch_cancel_path_for_other_owner_denied(self):
+ """POST /openai/v1/batches/batch_victim/cancel with another tenant's raw
+ batch ID is rejected (404) before the upstream cancel runs."""
+ from fastapi import HTTPException
+
+ pc = _prisma_client()
+ pc.db.litellm_managedobjecttable.find_first = AsyncMock(
+ return_value=self._victim_object_row()
+ )
+ with pytest.raises(HTTPException) as exc_info:
+ await rewrite_path_ids(
+ "/openai/v1/batches/batch_victim/cancel",
+ "openai",
+ _user("attacker", "attacker-team"),
+ pc,
+ _managed_files_hook(),
+ )
+ assert exc_info.value.status_code == 404
+
+ @pytest.mark.asyncio
+ async def test_raw_file_query_for_other_owner_denied(self):
+ from fastapi import HTTPException
+
+ pc = _prisma_client()
+ pc.db.litellm_managedfiletable.find_many = AsyncMock(
+ return_value=[self._victim_file_row()]
+ )
+ with pytest.raises(HTTPException) as exc_info:
+ await rewrite_query_ids(
+ {"file_id": "file-victim"},
+ "openai",
+ _user("attacker", "attacker-team"),
+ pc,
+ _managed_files_hook(),
+ )
+ assert exc_info.value.status_code == 404
+
+ @pytest.mark.asyncio
+ async def test_raw_file_body_for_other_owner_denied(self):
+ from fastapi import HTTPException
+
+ pc = _prisma_client()
+ pc.db.litellm_managedfiletable.find_many = AsyncMock(
+ return_value=[self._victim_file_row()]
+ )
+ with pytest.raises(HTTPException) as exc_info:
+ await rewrite_body_ids(
+ {"input_file_id": "file-victim"},
+ "openai",
+ _user("attacker", "attacker-team"),
+ pc,
+ _managed_files_hook(),
+ )
+ assert exc_info.value.status_code == 404
+
+ @pytest.mark.asyncio
+ async def test_raw_file_owned_by_caller_passes_through(self):
+ """A raw ID the caller does own is left untouched and forwarded — the
+ guard must not block legitimate raw-ID usage."""
+ pc = _prisma_client()
+ own_row = MagicMock()
+ own_row.created_by = "user-1"
+ own_row.team_id = "team-1"
+ own_row.unified_file_id = encode("openai", "u", "file-mine")
+ pc.db.litellm_managedfiletable.find_many = AsyncMock(return_value=[own_row])
+ result = await rewrite_path_ids(
+ "/openai/v1/files/file-mine",
+ "openai",
+ _user("user-1", "team-1"),
+ pc,
+ _managed_files_hook(),
+ )
+ assert result == "/openai/v1/files/file-mine"
+
+ @pytest.mark.asyncio
+ async def test_unmanaged_raw_id_passes_through(self):
+ """A raw ID with no managed row at all is a genuine opt-out and is
+ forwarded unchanged."""
+ pc = _prisma_client()
+ result = await rewrite_path_ids(
+ "/openai/v1/files/file-never-managed",
+ "openai",
+ _user("attacker", "attacker-team"),
+ pc,
+ _managed_files_hook(),
+ )
+ assert result == "/openai/v1/files/file-never-managed"
+
+ @pytest.mark.asyncio
+ async def test_cross_provider_raw_file_not_blocked(self):
+ """A raw ID whose only managed row belongs to a different provider is not
+ this provider's resource, so the guard does not deny it."""
+ pc = _prisma_client()
+ azure_row = MagicMock()
+ azure_row.created_by = "victim"
+ azure_row.team_id = "victim-team"
+ azure_row.unified_file_id = encode("azure", "victim", "file-victim")
+ pc.db.litellm_managedfiletable.find_many = AsyncMock(return_value=[azure_row])
+ result = await rewrite_path_ids(
+ "/openai/v1/files/file-victim",
+ "openai",
+ _user("attacker", "attacker-team"),
+ pc,
+ _managed_files_hook(),
+ )
+ assert result == "/openai/v1/files/file-victim"
+
+
+# ---------------------------------------------------------------------------
+# Raw-provider-ID guard amplification — a body packed with id-shaped strings
+# must not fan out into one (unindexed) DB scan per string. The guard de-dupes
+# repeats and caps the distinct lookups per request, failing closed instead of
+# skipping the guard.
+# ---------------------------------------------------------------------------
+
+
+class TestRawProviderIdGuardBudget:
+ @pytest.mark.asyncio
+ async def test_many_distinct_raw_ids_capped(self):
+ """A body with more distinct raw file IDs than the per-request budget is
+ rejected with 400, and the number of (unindexed) DB scans never exceeds
+ the cap."""
+ from fastapi import HTTPException
+
+ pc = _prisma_client()
+ body = {"ids": [f"file-{i}" for i in range(_MAX_RAW_ID_GUARD_LOOKUPS + 25)]}
+ with pytest.raises(HTTPException) as exc_info:
+ await rewrite_body_ids(
+ body, "openai", _user("attacker", "attacker-team"), pc, None
+ )
+ assert exc_info.value.status_code == 400
+ assert (
+ pc.db.litellm_managedfiletable.find_many.call_count
+ == _MAX_RAW_ID_GUARD_LOOKUPS
+ )
+
+ @pytest.mark.asyncio
+ async def test_repeated_raw_id_deduped(self):
+ """The same raw ID repeated many times issues exactly one DB lookup."""
+ pc = _prisma_client()
+ body = {"ids": ["file-dup"] * (_MAX_RAW_ID_GUARD_LOOKUPS * 5)}
+ result = await rewrite_body_ids(
+ body, "openai", _user("attacker", "attacker-team"), pc, None
+ )
+ assert result is body
+ assert pc.db.litellm_managedfiletable.find_many.call_count == 1
+
+ @pytest.mark.asyncio
+ async def test_distinct_ids_under_cap_not_rejected(self):
+ """A realistically-sized body (few distinct raw IDs) is never rejected and
+ each distinct ID is guarded once."""
+ pc = _prisma_client()
+ body = {"ids": [f"file-{i}" for i in range(5)]}
+ result = await rewrite_body_ids(
+ body, "openai", _user("user-1", "team-1"), pc, None
+ )
+ assert result is body
+ assert pc.db.litellm_managedfiletable.find_many.call_count == 5
+
+ @pytest.mark.asyncio
+ async def test_budget_is_per_input_surface(self):
+ """Each input surface (path / query / body) gets its own budget, so a
+ request distributing IDs across them is still bounded per surface."""
+ from fastapi import HTTPException
+
+ pc = _prisma_client()
+ params = {f"k{i}": f"file-{i}" for i in range(_MAX_RAW_ID_GUARD_LOOKUPS + 5)}
+ with pytest.raises(HTTPException) as exc_info:
+ await rewrite_query_ids(
+ params, "openai", _user("attacker", "attacker-team"), pc, None
+ )
+ assert exc_info.value.status_code == 400
+ assert (
+ pc.db.litellm_managedfiletable.find_many.call_count
+ == _MAX_RAW_ID_GUARD_LOOKUPS
+ )
+
+
+# ---------------------------------------------------------------------------
+# Flag-off: behaviour unchanged when passthrough_managed_object_ids is False
+# ---------------------------------------------------------------------------
+
+
+class TestFlagOff:
+ """
+ When the feature flag is off the pass_through_request code paths skip both
+ hooks entirely. Here we verify the rewriter modules themselves are pure
+ no-ops when called with no DB / hook: raw IDs pass through.
+ """
+
+ @pytest.mark.asyncio
+ async def test_raw_file_in_response_not_swapped_without_hook(self):
+ body = {"id": "file-abc", "object": "file"}
+ result = await rewrite_response_ids(
+ provider="openai",
+ method="POST",
+ route="/openai/v1/files",
+ body=body,
+ user_api_key_dict=_user(),
+ prisma_client=None,
+ managed_files_hook=None,
+ )
+ # Without DB/hook, _mint_or_reuse_file returns raw_id unchanged
+ assert result is body or result["id"] == "file-abc"
+
+ @pytest.mark.asyncio
+ async def test_decode_failure_body_untouched(self):
+ body = {"id": "file-abc123"}
+ result = await rewrite_body_ids(body, "openai", _user(), None, None)
+ assert result is body
+
+
+# ---------------------------------------------------------------------------
+# list_passthrough_ids_from_db — unit tests
+# ---------------------------------------------------------------------------
+
+
+def _prisma_with_list(file_rows=None, batch_rows=None) -> MagicMock:
+ """Return a prisma_client whose find_many honors the provider scope pushed
+ into the ``where`` clause, mirroring how Postgres would filter rows.
+
+ File rows are scoped via ``flat_model_file_ids: {has: }`` and object
+ rows via ``model_object_id: {startswith: passthrough::}``; the mock
+ applies the same predicate so a test feeding mixed-provider rows exercises
+ the real DB-pushdown contract instead of an unscoped passthrough."""
+ pc = _prisma_client()
+
+ def _file_filter(*args, where=None, take=None, **kwargs):
+ rows = list(file_rows or [])
+ marker = (where or {}).get("flat_model_file_ids", {}) or {}
+ marker = marker.get("has")
+ if marker is not None:
+ rows = [
+ r
+ for r in rows
+ if marker in (getattr(r, "flat_model_file_ids", None) or [])
+ ]
+ return rows if take is None else rows[:take]
+
+ def _batch_filter(*args, where=None, take=None, **kwargs):
+ rows = list(batch_rows or [])
+ prefix = (where or {}).get("model_object_id", {}) or {}
+ prefix = prefix.get("startswith")
+ if prefix is not None:
+ rows = [
+ r
+ for r in rows
+ if str(getattr(r, "model_object_id", "") or "").startswith(prefix)
+ ]
+ return rows if take is None else rows[:take]
+
+ if file_rows is not None:
+ pc.db.litellm_managedfiletable.find_many = AsyncMock(side_effect=_file_filter)
+ if batch_rows is not None:
+ pc.db.litellm_managedobjecttable.find_many = AsyncMock(
+ side_effect=_batch_filter
+ )
+ return pc
+
+
+def _fake_file_row(
+ unified_id: str, created_by: str = "user-1", team_id: str = "team-1"
+):
+ row = MagicMock()
+ row.unified_file_id = unified_id
+ row.created_by = created_by
+ row.team_id = team_id
+ row.file_object = {"filename": "test.jsonl", "bytes": 42, "purpose": "batch"}
+ payload = decode(unified_id)
+ row.flat_model_file_ids = (
+ [payload.raw_provider_id, _passthrough_provider_marker(payload.provider)]
+ if payload is not None
+ else []
+ )
+
+ import datetime
+
+ row.created_at = datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc)
+ return row
+
+
+def _fake_batch_row(
+ unified_id: str, created_by: str = "user-1", team_id: str = "team-1"
+):
+ row = MagicMock()
+ row.unified_object_id = unified_id
+ row.created_by = created_by
+ row.team_id = team_id
+ row.file_object = {"status": "completed", "input_file_id": "file-managed-1"}
+ row.file_purpose = "batch"
+ payload = decode(unified_id)
+ row.model_object_id = (
+ f"passthrough:{payload.provider}:{payload.raw_provider_id}"
+ if payload is not None
+ else None
+ )
+
+ import datetime
+
+ row.created_at = datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc)
+ return row
+
+
+class TestListPassthroughIdsFromDb:
+ """Tests for list_passthrough_ids_from_db and is_passthrough_list_route."""
+
+ def test_is_passthrough_list_route_files(self):
+ assert is_passthrough_list_route("openai", "GET", "/openai/v1/files") is True
+
+ def test_is_passthrough_list_route_batches(self):
+ assert (
+ is_passthrough_list_route("azure", "GET", "/azure/openai/batches") is True
+ )
+
+ def test_is_passthrough_list_route_not_for_post(self):
+ assert is_passthrough_list_route("openai", "POST", "/openai/v1/files") is False
+
+ def test_is_passthrough_list_route_not_for_single_resource(self):
+ # GET /v1/files/{file_id} is not a list route
+ assert (
+ is_passthrough_list_route("openai", "GET", "/openai/v1/files/file-abc")
+ is False
+ )
+
+ def test_is_passthrough_list_route_azure_ai_prefix(self):
+ assert (
+ is_passthrough_list_route("azure", "GET", "/azure_ai/openai/files") is True
+ )
+
+ def test_is_passthrough_list_route_azure_path_already_carrying_v1(self):
+ assert (
+ is_passthrough_list_route("azure", "GET", "/azure/openai/v1/files") is True
+ )
+ assert (
+ is_passthrough_list_route("azure", "GET", "/azure/openai/v1/batches")
+ is True
+ )
+
+ def test_is_passthrough_list_route_not_for_azure_single_resource(self):
+ assert (
+ is_passthrough_list_route("azure", "GET", "/azure/openai/files/file-abc")
+ is False
+ )
+
+ @pytest.mark.asyncio
+ async def test_list_files_returns_owned_rows(self):
+ managed_id = new_managed_id("openai", "file-abc")
+ fake_row = _fake_file_row(managed_id)
+ pc = _prisma_with_list(file_rows=[fake_row])
+
+ result = await list_passthrough_ids_from_db(
+ provider="openai",
+ route="/openai/v1/files",
+ user_api_key_dict=_user("user-1", "team-1"),
+ prisma_client=pc,
+ )
+
+ assert result is not None
+ assert result["object"] == "list"
+ assert len(result["data"]) == 1
+ assert result["data"][0]["id"] == managed_id
+ assert result["data"][0]["object"] == "file"
+ assert result["first_id"] == managed_id
+
+ @pytest.mark.asyncio
+ async def test_list_batches_returns_owned_rows(self):
+ managed_id = new_managed_id("openai", "batch_abc")
+ fake_row = _fake_batch_row(managed_id)
+ pc = _prisma_with_list(batch_rows=[fake_row])
+
+ result = await list_passthrough_ids_from_db(
+ provider="openai",
+ route="/openai/v1/batches",
+ user_api_key_dict=_user("user-1", "team-1"),
+ prisma_client=pc,
+ )
+
+ assert result is not None
+ assert result["object"] == "list"
+ assert len(result["data"]) == 1
+ assert result["data"][0]["id"] == managed_id
+ assert result["data"][0]["object"] == "batch"
+
+ @pytest.mark.asyncio
+ async def test_list_files_admin_gets_all_rows(self):
+ """Admin should receive all rows; the where filter passed to DB is {}."""
+ rows = [
+ _fake_file_row(new_managed_id("openai", "file-1")),
+ _fake_file_row(new_managed_id("openai", "file-2")),
+ ]
+ pc = _prisma_with_list(file_rows=rows)
+
+ result = await list_passthrough_ids_from_db(
+ provider="openai",
+ route="/openai/v1/files",
+ user_api_key_dict=_admin_user(),
+ prisma_client=pc,
+ )
+
+ assert result is not None
+ assert len(result["data"]) == 2
+ # Admin adds no owner scoping, but the provider scope is always pushed
+ # to the DB; the only where clause is the provider marker filter.
+ call_kwargs = pc.db.litellm_managedfiletable.find_many.call_args.kwargs
+ assert call_kwargs["where"] == {
+ "flat_model_file_ids": {"has": _passthrough_provider_marker("openai")}
+ }
+
+ @pytest.mark.asyncio
+ async def test_list_files_user_scoped_where(self):
+ """Regular user should get a where clause scoped to their user_id / team_id."""
+ pc = _prisma_with_list(file_rows=[])
+
+ await list_passthrough_ids_from_db(
+ provider="openai",
+ route="/openai/v1/files",
+ user_api_key_dict=_user("user-2", "team-2"),
+ prisma_client=pc,
+ )
+
+ call_kwargs = pc.db.litellm_managedfiletable.find_many.call_args.kwargs
+ where = call_kwargs["where"]
+ # The OR clause should scope to user-2 or team-2
+ assert "OR" in where
+ entries = where["OR"]
+ assert {"created_by": "user-2"} in entries
+ assert {"team_id": "team-2"} in entries
+
+ @pytest.mark.asyncio
+ async def test_list_has_more_flag(self):
+ """has_more is True when DB returns limit+1 rows."""
+ rows = [
+ _fake_file_row(new_managed_id("openai", f"file-{i}")) for i in range(21)
+ ] # limit=20, fetch 21
+ pc = _prisma_with_list(file_rows=rows)
+
+ result = await list_passthrough_ids_from_db(
+ provider="openai",
+ route="/openai/v1/files",
+ user_api_key_dict=_admin_user(),
+ prisma_client=pc,
+ query_params={"limit": "20"},
+ )
+
+ assert result is not None
+ assert result["has_more"] is True
+ assert len(result["data"]) == 20 # extra row trimmed
+
+ @pytest.mark.asyncio
+ async def test_list_returns_none_for_non_list_route(self):
+ pc = _prisma_with_list()
+
+ result = await list_passthrough_ids_from_db(
+ provider="openai",
+ route="/openai/v1/files/file-abc", # single-resource, not a list
+ user_api_key_dict=_user(),
+ prisma_client=pc,
+ )
+
+ assert result is None
+
+ @pytest.mark.asyncio
+ async def test_list_db_error_returns_empty_not_none(self):
+ """DB failure must return an empty list, not None (which would fall through
+ to the upstream provider and leak the provider-wide listing)."""
+ pc = _prisma_with_list()
+ pc.db.litellm_managedfiletable.find_many = AsyncMock(
+ side_effect=Exception("db down")
+ )
+
+ result = await list_passthrough_ids_from_db(
+ provider="openai",
+ route="/openai/v1/files",
+ user_api_key_dict=_admin_user(),
+ prisma_client=pc,
+ )
+
+ # Must not return None (which would fall through to upstream)
+ assert result is not None
+ assert result["data"] == []
+ assert result["has_more"] is False
+
+ @pytest.mark.asyncio
+ async def test_list_returns_empty_for_caller_without_identity(self):
+ """Caller with neither user_id nor team_id should get an empty list."""
+ pc = _prisma_with_list(
+ file_rows=[_fake_file_row(new_managed_id("openai", "file-1"))]
+ )
+ anon = UserAPIKeyAuth() # no user_id, no team_id, not admin
+
+ result = await list_passthrough_ids_from_db(
+ provider="openai",
+ route="/openai/v1/files",
+ user_api_key_dict=anon,
+ prisma_client=pc,
+ )
+
+ assert result is not None
+ assert result["data"] == []
+
+ @pytest.mark.asyncio
+ async def test_list_files_pushes_provider_scope_to_db(self):
+ """File listing scopes by provider at the DB level via the provider
+ marker in flat_model_file_ids, so a single query serves the page and a
+ mixed-provider pool can never truncate or leak the other provider.
+
+ A large azure-only pool must return an empty openai page with
+ has_more=False in exactly one DB round-trip.
+ """
+ azure_rows = [
+ _fake_file_row(new_managed_id("azure", f"file-{i}")) for i in range(50)
+ ]
+ pc = _prisma_with_list(file_rows=azure_rows)
+
+ result = await list_passthrough_ids_from_db(
+ provider="openai", # asking for openai but DB only has azure rows
+ route="/openai/v1/files",
+ user_api_key_dict=_admin_user(),
+ prisma_client=pc,
+ query_params={"limit": "20"},
+ )
+
+ assert result is not None
+ assert result["data"] == []
+ assert result["has_more"] is False
+ where = pc.db.litellm_managedfiletable.find_many.call_args.kwargs["where"]
+ assert where["flat_model_file_ids"] == {
+ "has": _passthrough_provider_marker("openai")
+ }
+ assert pc.db.litellm_managedfiletable.find_many.await_count == 1
+
+ @pytest.mark.asyncio
+ async def test_list_ignores_cross_provider_cursor(self):
+ """An ``after`` cursor minted for a different provider must not shift the
+ created_at boundary: it would skip/repeat this provider's rows. The
+ cursor is ignored and the unscoped first page is served."""
+ import datetime
+
+ azure_row = _fake_file_row(new_managed_id("azure", "file-azure"))
+ pc = _prisma_with_list(file_rows=[azure_row])
+
+ cursor_row = MagicMock()
+ cursor_row.created_at = datetime.datetime(
+ 2025, 6, 1, tzinfo=datetime.timezone.utc
+ )
+ pc.db.litellm_managedfiletable.find_first = AsyncMock(return_value=cursor_row)
+
+ result = await list_passthrough_ids_from_db(
+ provider="azure",
+ route="/azure/openai/files",
+ user_api_key_dict=_admin_user(),
+ prisma_client=pc,
+ query_params={"after": new_managed_id("openai", "file-openai")},
+ )
+
+ assert result is not None
+ where = pc.db.litellm_managedfiletable.find_many.call_args.kwargs["where"]
+ assert "created_at" not in where
+ assert "OR" not in where and "AND" not in where
+
+ @pytest.mark.asyncio
+ async def test_list_applies_same_provider_cursor(self):
+ """An ``after`` cursor minted for the same provider advances pagination
+ past the cursor row using a compound (created_at, id) boundary so rows
+ sharing the cursor row's timestamp are not skipped."""
+ import datetime
+
+ azure_row = _fake_file_row(new_managed_id("azure", "file-azure"))
+ pc = _prisma_with_list(file_rows=[azure_row])
+
+ cursor_row = MagicMock()
+ cursor_row.created_at = datetime.datetime(
+ 2025, 6, 1, tzinfo=datetime.timezone.utc
+ )
+ pc.db.litellm_managedfiletable.find_first = AsyncMock(return_value=cursor_row)
+
+ cursor_id = new_managed_id("azure", "file-cursor")
+ result = await list_passthrough_ids_from_db(
+ provider="azure",
+ route="/azure/openai/files",
+ user_api_key_dict=_admin_user(),
+ prisma_client=pc,
+ query_params={"after": cursor_id},
+ )
+
+ assert result is not None
+ where = pc.db.litellm_managedfiletable.find_many.call_args.kwargs["where"]
+ assert "created_at" not in where
+ assert where["OR"] == [
+ {"created_at": {"lt": cursor_row.created_at}},
+ {
+ "AND": [
+ {"created_at": cursor_row.created_at},
+ {"unified_file_id": {"lt": cursor_id}},
+ ]
+ },
+ ]
+
+ @pytest.mark.asyncio
+ async def test_list_cursor_does_not_drop_created_at_ties(self):
+ """Regression: paginating a pool whose rows all share one created_at must
+ return every row exactly once. A timestamp-only ``lt`` cursor boundary
+ would skip every tied row after the first page; the compound
+ (created_at, id) boundary keeps the walk complete."""
+ import datetime
+
+ shared_ts = datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc)
+ rows = [_fake_file_row(new_managed_id("azure", f"file-{i}")) for i in range(5)]
+ for row in rows:
+ row.created_at = shared_ts
+ all_ids = {row.unified_file_id for row in rows}
+
+ def _matches(row, where):
+ for key, cond in where.items():
+ if key == "AND":
+ if not all(_matches(row, c) for c in cond):
+ return False
+ elif key == "OR":
+ if not any(_matches(row, c) for c in cond):
+ return False
+ elif key == "flat_model_file_ids":
+ marker = (cond or {}).get("has")
+ if marker not in (getattr(row, "flat_model_file_ids", None) or []):
+ return False
+ else:
+ actual = getattr(row, key, None)
+ if isinstance(cond, dict):
+ for op, val in cond.items():
+ if op == "lt" and not (actual is not None and actual < val):
+ return False
+ if op == "gt" and not (actual is not None and actual > val):
+ return False
+ if op == "startswith" and not str(actual or "").startswith(
+ val
+ ):
+ return False
+ elif actual != cond:
+ return False
+ return True
+
+ def _find_many(*_a, where=None, order=None, take=None, **_k):
+ matched = [r for r in rows if _matches(r, where or {})]
+ for spec in reversed(order or []):
+ ((field, direction),) = spec.items()
+ matched.sort(
+ key=lambda r: getattr(r, field), reverse=(direction == "desc")
+ )
+ return matched if take is None else matched[:take]
+
+ def _find_first(*_a, where=None, **_k):
+ return next((r for r in rows if _matches(r, where or {})), None)
+
+ pc = _prisma_client()
+ pc.db.litellm_managedfiletable.find_many = AsyncMock(side_effect=_find_many)
+ pc.db.litellm_managedfiletable.find_first = AsyncMock(side_effect=_find_first)
+
+ collected: list = []
+ after = None
+ for _ in range(len(rows) + 2):
+ params = {"limit": "2"}
+ if after is not None:
+ params["after"] = after
+ result = await list_passthrough_ids_from_db(
+ provider="azure",
+ route="/azure/openai/files",
+ user_api_key_dict=_admin_user(),
+ prisma_client=pc,
+ query_params=params,
+ )
+ assert result is not None
+ collected.extend(item["id"] for item in result["data"])
+ if not result["has_more"]:
+ break
+ after = result["last_id"]
+
+ assert sorted(collected) == sorted(all_ids)
+ assert len(collected) == len(set(collected))
+
+ @pytest.mark.asyncio
+ async def test_list_files_filters_by_provider(self):
+ openai_row = _fake_file_row(new_managed_id("openai", "file-openai"))
+ azure_row = _fake_file_row(new_managed_id("azure", "file-azure"))
+ pc = _prisma_with_list(file_rows=[azure_row, openai_row])
+
+ result = await list_passthrough_ids_from_db(
+ provider="openai",
+ route="/openai/v1/files",
+ user_api_key_dict=_admin_user(),
+ prisma_client=pc,
+ )
+
+ assert result is not None
+ assert len(result["data"]) == 1
+ assert decode(result["data"][0]["id"]).provider == "openai"
+
+ @pytest.mark.asyncio
+ async def test_list_batches_pushes_provider_scope_to_db(self):
+ """Batch listing scopes by provider at the DB level via the namespaced
+ model_object_id, so a single query serves the page instead of scanning."""
+ batch_row = _fake_batch_row(new_managed_id("azure", "batch_abc"))
+ pc = _prisma_with_list(batch_rows=[batch_row])
+
+ result = await list_passthrough_ids_from_db(
+ provider="azure",
+ route="/azure/openai/batches",
+ user_api_key_dict=_admin_user(),
+ prisma_client=pc,
+ )
+
+ assert result is not None
+ assert len(result["data"]) == 1
+ where = pc.db.litellm_managedobjecttable.find_many.call_args.kwargs["where"]
+ assert where["model_object_id"] == {"startswith": "passthrough:azure:"}
+ assert pc.db.litellm_managedobjecttable.find_many.await_count == 1
From dc4f5b12efed15b9d4c97cdd2f08892db38a2f9d Mon Sep 17 00:00:00 2001
From: Shivam Rawat
Date: Sat, 30 May 2026 17:07:24 -0700
Subject: [PATCH 35/44] =?UTF-8?q?fix(proxy):=20enforce=20allowed=5Fpassthr?=
=?UTF-8?q?ough=5Froutes=20for=20auth=3Dtrue=20pass-thr=E2=80=A6=20(#29256?=
=?UTF-8?q?)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
* fix(proxy): enforce allowed_passthrough_routes for auth=true pass-through
Pass-through endpoints with auth=true were injected into openai_routes,
so teams with openai_routes access bypassed per-team allowed_passthrough_routes.
Gate auth-enforced pass-through at JWT, virtual-key, and non-admin route checks.
Co-authored-by: Cursor
* fix(proxy): clarify JWT passthrough denial
Co-authored-by: Cursor
* fix(proxy): make pass-through auth checks method-aware
Prevent allowlist bypass when the same path is registered with different auth settings per HTTP method.
Co-authored-by: Cursor
* Fix passthrough route auth checks
* fix(proxy): reject unregistered pass-through HTTP methods
Enforce method-aware JWT checks and return 405 when stale FastAPI routes accept requests outside the current pass-through registry.
Co-authored-by: Cursor
* fix(proxy): remove duplicate request_method in JWT team lookup
Fixes SyntaxError on proxy startup caused by passing request_method twice to find_team_with_model_access.
Co-authored-by: Cursor
* Fix passthrough route auth enforcement
* fix(proxy): raise passthrough-specific 403 directly in virtual-key path
* fix(proxy): load team for RBAC role-claim JWT passthrough gating
* Revert "chore(tests): migrate Bedrock CI to AWS account 941277531214 (#28728)" (#29326)
This reverts the Bedrock CI account migration (#28728). The original account
(888602223428) was put under an AWS security restriction after a leaked key
and has since been reactivated, while the replacement account (941277531214)
lacks access to several models the suites exercise (legacy Bedrock Claude 3
models, Cohere, Nova Canvas image gen, Bedrock batch inference, and flagship
Opus). Pointing CI back at the reactivated account restores that coverage.
This is the exact inverse of #28728: all hardcoded 941277531214 references go
back to 888602223428 (provisioned/imported-model ARNs, AgentCore runtime ARNs
and their suffixes, batch execution role ARN, and the example proxy config),
the S3 buckets revert to litellm-proxy and load-testing-oct, the guardrail IDs
revert to wf0hkdb5x07f and ff6ujrregl1q, the SageMaker endpoint and Knowledge
Base revert to their original ids, and the live-call tests go back to the
legacy model strings. The grid_spec fail_reason workaround for the unentitled
Opus cells is dropped while keeping the unrelated bedrock_effort_ceiling field
added after the migration.
The CircleCI AWS_ACCESS_KEY_ID / AWS_SECRET_ACCESS_KEY env vars still point at
941277531214 and must be set to the reactivated account's fresh credentials
separately via the CircleCI API; AWS_REGION_NAME stays us-west-2.
(cherry picked from commit f11c12d1574fb51acd958bc7173a74e994aacf19)
* fix(proxy): scope pass-through 405 to registry routes; grant rerank passthrough in rpm tests
The auth=true pass-through 405 guard fired for mapped provider routes
(e.g. /assemblyai/*) that are not in the in-memory registry, since
get_registered_pass_through_route returns None for them while
is_registered_pass_through_route matches via mapped_pass_through_routes.
Only raise 405 when the path is registered but the request method is not
allowed, so mapped provider pass-throughs fall through to the default
target params as before.
The rpm-limit pass-through tests register /v1/rerank with auth=true but
gave their keys no allowed_passthrough_routes, so the new default-deny
returned 403 before the rate limiter ran (non-deterministically,
depending on registry insertion order). Grant the keys explicit
passthrough access so the tests exercise rate limiting under the new
auth model.
* fix(proxy): guard request method lookup against scopes without a method
Starlette's Request.method property reads scope["method"] and raises
KeyError when the scope omits it (e.g. minimally-constructed test
requests). getattr only swallows AttributeError, so the new
_get_request_method helper propagated the KeyError up through
user_api_key_auth and surfaced as a ProxyException. Catch KeyError
(and AttributeError) and fall back to None.
* test(passthrough): pin SERVER_ROOT_PATH in unregistered-method test
test_custom_proxy.py sets os.environ['SERVER_ROOT_PATH'] = '/my-custom-path'
at module import with no cleanup. When that module is collected into the same
xdist worker as this test, the leaked root path is prepended to registered
pass-through paths, so is_registered_pass_through_route misses '/test/path'
and the handler returns 404 instead of the expected 405 (order-dependent).
Pin SERVER_ROOT_PATH to '' so the test is deterministic.
* test(passthrough): restore regression coverage for non-auth-enforced pass-through via llm_api_routes
* fix(proxy): record auth flag in pass-through registry for allowlist enforcement
Auth-enforced pass-through detection inferred enforcement from the FastAPI
dependency stored at registration time. The management create and update
endpoints register routes with dependencies=None even though auth defaults to
true, so is_auth_enforced_pass_through_route treated those DB-created routes as
unenforced. A key allowed for llm_api_routes could then call a management-created
auth-enabled pass-through route without matching allowed_passthrough_routes.
Store the auth setting on each registry entry and read it directly when deciding
whether the allowlist applies, instead of deriving it from dependency metadata.
* fix(proxy): include bool in pass-through registry value type for auth flag
The auth flag stored in _registered_pass_through_routes is a bool, which
was not part of the registry value Union, so mypy rejected the dict literal.
Add bool to the Union and narrow route_methods to a list before the
membership check so the in-operator stays valid.
* fix(proxy): preserve stored auth flag on pass-through endpoint update
model_dump(exclude_none=True) re-included the auth=True default whenever
a partial update omitted auth, silently flipping an existing auth=false
pass-through to auth-enforced and 403ing every team/key without
allowed_passthrough_routes. Merge only explicitly set fields via
exclude_unset so omitted fields keep their stored value.
---------
Co-authored-by: Cursor
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
---
litellm/proxy/auth/handle_jwt.py | 84 +++-
litellm/proxy/auth/route_checks.py | 105 ++++-
litellm/proxy/auth/user_api_key_auth.py | 3 +
.../pass_through_endpoints.py | 47 +-
.../reasoning_effort_grid/grid_spec.py | 1 +
.../test_reasoning_effort_grid.py | 1 +
.../test_pass_through_endpoints.py | 8 +-
.../proxy/auth/test_handle_jwt.py | 425 ++++++++++++++++++
.../proxy/auth/test_route_checks.py | 306 +++++++++++--
.../proxy/auth/test_user_api_key_auth.py | 2 +
.../test_pass_through_endpoints.py | 308 +++++++++++++
11 files changed, 1252 insertions(+), 38 deletions(-)
diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py
index 5b116c142b2..9838b4ba49b 100644
--- a/litellm/proxy/auth/handle_jwt.py
+++ b/litellm/proxy/auth/handle_jwt.py
@@ -1110,6 +1110,40 @@ class JWTAuthManager:
return all_team_ids
+ @staticmethod
+ def _team_has_passthrough_route_access(
+ team_object: Optional[LiteLLM_TeamTable],
+ route: str,
+ request_method: Optional[str] = None,
+ ) -> bool:
+ normalized_request_method = (
+ request_method.upper() if isinstance(request_method, str) else None
+ )
+ if not RouteChecks.is_auth_enforced_pass_through_route(
+ route=route,
+ method=normalized_request_method,
+ ):
+ return True
+
+ # JWT team selection is team-scoped; key metadata is not available here,
+ # so passthrough access is granted only by the selected team's metadata.
+ return RouteChecks.check_passthrough_route_access(
+ route=route,
+ user_api_key_dict=UserAPIKeyAuth(
+ team_metadata=(team_object.metadata or {}) if team_object else {}
+ ),
+ )
+
+ @staticmethod
+ def _raise_team_passthrough_route_denial(route: str) -> None:
+ raise HTTPException(
+ status_code=403,
+ detail=(
+ f"Team not allowed to access passthrough route {route}. "
+ "Configure `allowed_passthrough_routes` on the team."
+ ),
+ )
+
@staticmethod
async def find_team_with_model_access(
team_ids: Set[str],
@@ -1120,10 +1154,13 @@ class JWTAuthManager:
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Optional[Span],
proxy_logging_obj: ProxyLogging,
+ request_method: Optional[str] = None,
) -> Tuple[Optional[str], Optional[LiteLLM_TeamTable]]:
"""Find first team with access to the requested model"""
from litellm.proxy.proxy_server import llm_router
+ denied_auth_enforced_pass_through_route = False
+
if not team_ids:
if jwt_handler.litellm_jwtauth.enforce_team_based_model_access:
raise HTTPException(
@@ -1158,6 +1195,16 @@ class JWTAuthManager:
user_route=route,
litellm_proxy_roles=jwt_handler.litellm_jwtauth,
)
+ if (
+ is_allowed
+ and not JWTAuthManager._team_has_passthrough_route_access(
+ team_object=team_object,
+ route=route,
+ request_method=request_method,
+ )
+ ):
+ is_allowed = False
+ denied_auth_enforced_pass_through_route = True
verbose_proxy_logger.debug(
f"JWT team route check: team_id={team_id}, route={route}, is_allowed={is_allowed}"
)
@@ -1166,6 +1213,9 @@ class JWTAuthManager:
except Exception:
continue
+ if denied_auth_enforced_pass_through_route:
+ JWTAuthManager._raise_team_passthrough_route_denial(route=route)
+
if requested_model:
raise HTTPException(
status_code=403,
@@ -1581,7 +1631,7 @@ class JWTAuthManager:
return None, None, None
@staticmethod
- async def auth_builder(
+ async def auth_builder( # noqa: PLR0915
api_key: str,
jwt_handler: JWTHandler,
request_data: dict,
@@ -1592,6 +1642,7 @@ class JWTAuthManager:
parent_otel_span: Optional[Span],
proxy_logging_obj: ProxyLogging,
request_headers: Optional[dict] = None,
+ request_method: Optional[str] = None,
) -> JWTAuthBuilderResult:
"""Main authentication and authorization builder"""
# Check if OIDC UserInfo endpoint is enabled, but fall back to standard
@@ -1723,12 +1774,43 @@ class JWTAuthManager:
team_ids=all_team_ids,
requested_model=request_data.get("model"),
route=route,
+ request_method=request_method,
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
+
+ # The RBAC role-claim path (rbac_role == TEAM) sets team_id without
+ # loading team_object, so fetch it here before gating an auth-enforced
+ # passthrough route on the team's allowed_passthrough_routes.
+ if (
+ team_id
+ and team_object is None
+ and RouteChecks.is_auth_enforced_pass_through_route(
+ route=route,
+ method=(
+ request_method.upper() if isinstance(request_method, str) else None
+ ),
+ )
+ ):
+ team_object = await get_team_object(
+ team_id=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,
+ team_id_upsert=jwt_handler.litellm_jwtauth.team_id_upsert,
+ )
+
+ if team_id and not JWTAuthManager._team_has_passthrough_route_access(
+ team_object=team_object,
+ route=route,
+ request_method=request_method,
+ ):
+ JWTAuthManager._raise_team_passthrough_route_denial(route=route)
+
# Extract alias fields for resolution (if configured)
org_alias = jwt_handler.get_org_alias(token=jwt_valid_token, default_value=None)
diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py
index a9519aa6cc5..4be2a185ce4 100644
--- a/litellm/proxy/auth/route_checks.py
+++ b/litellm/proxy/auth/route_checks.py
@@ -59,6 +59,10 @@ _PROXY_ADMIN_VIEW_ONLY_BLOCKED_ROUTES = frozenset(
# paths directly because the request route carries the resolved key id.
_PROXY_ADMIN_VIEW_ONLY_BLOCKED_KEY_SUFFIXES = ("/regenerate", "/reset_spend")
+_AUTH_ENFORCED_PASS_THROUGH_ROUTE_GROUPS = frozenset(
+ ("openai_routes", "llm_api_routes")
+)
+
class RouteChecks:
@staticmethod
@@ -103,6 +107,8 @@ class RouteChecks:
if len(valid_token.allowed_routes) == 0:
return True
+ denied_auth_enforced_pass_through_route = False
+
# explicit check for allowed routes (exact match or prefix match)
for allowed_route in valid_token.allowed_routes:
if RouteChecks._route_matches_allowed_route(
@@ -121,7 +127,20 @@ class RouteChecks:
route=route,
allowed_routes=LiteLLMRoutes._member_map_[allowed_route].value,
):
- return True
+ if (
+ allowed_route in _AUTH_ENFORCED_PASS_THROUGH_ROUTE_GROUPS
+ and RouteChecks.is_auth_enforced_pass_through_route(
+ route=route,
+ method=RouteChecks._get_request_method(request=request),
+ )
+ ):
+ if RouteChecks.check_passthrough_route_access(
+ route=route, user_api_key_dict=valid_token
+ ):
+ return True
+ denied_auth_enforced_pass_through_route = True
+ else:
+ return True
################################################
# For llm_api_routes, also check registered pass-through endpoints
@@ -134,7 +153,17 @@ class RouteChecks:
if InitPassThroughEndpointHelpers.is_registered_pass_through_route(
route=route
):
- return True
+ if RouteChecks.is_auth_enforced_pass_through_route(
+ route=route,
+ method=RouteChecks._get_request_method(request=request),
+ ):
+ if RouteChecks.check_passthrough_route_access(
+ route=route, user_api_key_dict=valid_token
+ ):
+ return True
+ denied_auth_enforced_pass_through_route = True
+ else:
+ return True
# Method-aware carve-out: allow GET on the two
# read-only MCP-server discovery endpoints
@@ -158,6 +187,9 @@ class RouteChecks:
):
return True
+ if denied_auth_enforced_pass_through_route:
+ raise RouteChecks._auth_pass_through_denied_exception(route=route)
+
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail=f"Virtual key is not allowed to call this route. Only allowed to call routes: {valid_token.allowed_routes}. Tried to call route: {route}",
@@ -228,7 +260,14 @@ class RouteChecks:
route=route,
)
- if RouteChecks.is_llm_api_route(route=route):
+ if RouteChecks.is_auth_enforced_pass_through_route(
+ route=route,
+ method=RouteChecks._get_request_method(request=request),
+ ):
+ RouteChecks._require_auth_pass_through_access(
+ route=route, valid_token=valid_token
+ )
+ elif RouteChecks.is_llm_api_route(route=route):
pass
elif RouteChecks.is_info_route(route=route):
# check if user allowed to call an info route
@@ -624,6 +663,66 @@ class RouteChecks:
return False
+ @staticmethod
+ def _get_request_method(request: Optional[Request]) -> Optional[str]:
+ if request is None:
+ return None
+
+ try:
+ method = request.method
+ except (AttributeError, KeyError):
+ return None
+ if not isinstance(method, str):
+ return None
+
+ return method.upper()
+
+ @staticmethod
+ def is_auth_enforced_pass_through_route(
+ route: str, method: Optional[str] = None
+ ) -> bool:
+ """
+ True for config/DB pass-through endpoints registered with auth=true.
+
+ These routes are injected into ``openai_routes`` for spend/budget hooks but
+ must not inherit blanket ``openai_routes`` RBAC; access is gated by
+ ``allowed_passthrough_routes`` on the key or team.
+ """
+ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
+ InitPassThroughEndpointHelpers,
+ )
+
+ route_info = InitPassThroughEndpointHelpers.get_registered_pass_through_route(
+ route=route, method=method
+ )
+ if route_info is None:
+ return False
+ return route_info.get("auth") is True
+
+ @staticmethod
+ def _auth_pass_through_denied_exception(route: str) -> HTTPException:
+ return HTTPException(
+ status_code=status.HTTP_403_FORBIDDEN,
+ detail=(
+ f"Key/team not allowed to access passthrough route {route}. "
+ "Configure `allowed_passthrough_routes` on the team or key."
+ ),
+ )
+
+ @staticmethod
+ def _require_auth_pass_through_access(
+ route: str,
+ valid_token: UserAPIKeyAuth,
+ ) -> None:
+ """
+ Require an explicit ``allowed_passthrough_routes`` match for auth=true pass-through.
+ """
+ if RouteChecks.check_passthrough_route_access(
+ route=route, user_api_key_dict=valid_token
+ ):
+ return
+ raise RouteChecks._auth_pass_through_denied_exception(route=route)
+
@staticmethod
def check_passthrough_route_access(
route: str, user_api_key_dict: UserAPIKeyAuth
diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py
index 828b719e299..2f9a1411400 100644
--- a/litellm/proxy/auth/user_api_key_auth.py
+++ b/litellm/proxy/auth/user_api_key_auth.py
@@ -925,6 +925,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
proxy_logging_obj=proxy_logging_obj,
parent_otel_span=parent_otel_span,
request_headers=_safe_get_request_headers(request),
+ request_method=RouteChecks._get_request_method(
+ request=request
+ ),
)
is_proxy_admin = result["is_proxy_admin"]
diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
index f08e021630d..9d68132b37d 100644
--- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
+++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
@@ -78,7 +78,7 @@ pass_through_endpoint_logging = PassThroughEndpointLogging()
# Global registry to track registered pass-through routes and prevent memory leaks
_registered_pass_through_routes: Dict[
- str, Dict[str, Union[str, List[str], Dict[str, Any]]]
+ str, Dict[str, Union[str, bool, List[str], Dict[str, Any]]]
] = {}
@@ -1539,6 +1539,17 @@ def create_pass_through_route(
route=path, method=request.method
)
)
+ if (
+ passthrough_params is None
+ and InitPassThroughEndpointHelpers.get_registered_pass_through_route(
+ route=path
+ )
+ is not None
+ ):
+ raise HTTPException(
+ status_code=status.HTTP_405_METHOD_NOT_ALLOWED,
+ detail=f"Method {request.method} is not allowed for pass-through endpoint {path}.",
+ )
target_params = {
"target": target,
"custom_headers": custom_headers,
@@ -2266,6 +2277,7 @@ class InitPassThroughEndpointHelpers:
methods: Optional[List[str]] = None,
default_query_params: Optional[dict] = None,
config_file_path: Optional[str] = None,
+ auth: bool = False,
):
"""Add exact path route for pass-through endpoint"""
# Default to all methods if none specified (backward compatibility)
@@ -2317,6 +2329,7 @@ class InitPassThroughEndpointHelpers:
"path": path,
"type": "exact",
"methods": methods,
+ "auth": auth,
"passthrough_params": {
"target": target,
"custom_headers": custom_headers,
@@ -2344,6 +2357,7 @@ class InitPassThroughEndpointHelpers:
methods: Optional[List[str]] = None,
default_query_params: Optional[dict] = None,
config_file_path: Optional[str] = None,
+ auth: bool = False,
):
"""Add wildcard route for sub-paths"""
# Default to all methods if none specified (backward compatibility)
@@ -2396,6 +2410,7 @@ class InitPassThroughEndpointHelpers:
"path": path,
"type": "subpath",
"methods": methods,
+ "auth": auth,
"passthrough_params": {
"target": target,
"custom_headers": custom_headers,
@@ -2514,8 +2529,15 @@ class InitPassThroughEndpointHelpers:
InitPassThroughEndpointHelpers._build_full_path_with_root(parts[2])
)
- # Get the methods for this route
- route_methods = _registered_pass_through_routes[key].get("methods", [])
+ # Get the methods for this route. Prefer the registered metadata,
+ # but keep supporting test fixtures / older registry entries that
+ # only encoded methods in the route key.
+ methods_entry = _registered_pass_through_routes[key].get("methods", [])
+ route_methods: List[str] = (
+ methods_entry if isinstance(methods_entry, list) else []
+ )
+ if not route_methods and len(parts) == 4:
+ route_methods = parts[3].split(",")
# Check if path matches
path_matches = False
@@ -2573,8 +2595,9 @@ async def _register_pass_through_endpoint(
default_query_params = endpoint_data.get("default_query_params")
auth = endpoint_data.get("auth")
dependencies = None
+ auth_enforced = auth is not None and str(auth).lower() == "true"
- if auth is not None and str(auth).lower() == "true":
+ if auth_enforced:
# Authentication on a pass-through endpoint used to be enterprise-only.
# That left OSS with no safe configuration: auth=True raised at startup
# unless the operator had a license. The safe option must always be free,
@@ -2607,6 +2630,7 @@ async def _register_pass_through_endpoint(
methods=methods,
default_query_params=default_query_params,
config_file_path=config_file_path,
+ auth=auth_enforced,
)
methods_for_key = methods if methods else ["GET", "POST", "PUT", "DELETE", "PATCH"]
@@ -2632,6 +2656,7 @@ async def _register_pass_through_endpoint(
methods=methods,
default_query_params=default_query_params,
config_file_path=config_file_path,
+ auth=auth_enforced,
)
visited_endpoints.add(f"{endpoint_id}:subpath:{path}:{methods_str}")
@@ -2955,14 +2980,18 @@ async def update_pass_through_endpoints(
},
)
- # Get the update data as dict, excluding None values for partial updates
+ # Only merge fields the caller explicitly sent so omitted fields keep their
+ # stored value. Without exclude_unset, defaults like auth=True would overwrite
+ # an existing auth=false entry on any unrelated edit.
# Exclude is_from_config as it's a response-only field (computed at read time)
- update_data = data.model_dump(exclude_none=True, exclude={"is_from_config"})
+ update_data = data.model_dump(
+ exclude_unset=True, exclude_none=True, exclude={"is_from_config"}
+ )
# Start with existing endpoint data
endpoint_dict = found_endpoint.model_dump()
- # Update with new data (only non-None values)
+ # Update with new data (only explicitly provided values)
endpoint_dict.update(update_data)
# Preserve existing ID if not provided in update and endpoint has ID
@@ -3010,6 +3039,7 @@ async def update_pass_through_endpoints(
guardrails=getattr(updated_endpoint, "guardrails", None),
methods=updated_endpoint.methods,
default_query_params=updated_endpoint.default_query_params,
+ auth=updated_endpoint.auth,
)
else:
InitPassThroughEndpointHelpers.add_exact_path_route(
@@ -3025,6 +3055,7 @@ async def update_pass_through_endpoints(
guardrails=getattr(updated_endpoint, "guardrails", None),
methods=updated_endpoint.methods,
default_query_params=updated_endpoint.default_query_params,
+ auth=updated_endpoint.auth,
)
return PassThroughEndpointResponse(
@@ -3103,6 +3134,7 @@ async def create_pass_through_endpoints(
guardrails=getattr(created_endpoint, "guardrails", None),
methods=created_endpoint.methods,
default_query_params=created_endpoint.default_query_params,
+ auth=created_endpoint.auth,
)
else:
InitPassThroughEndpointHelpers.add_exact_path_route(
@@ -3118,6 +3150,7 @@ async def create_pass_through_endpoints(
guardrails=getattr(created_endpoint, "guardrails", None),
methods=created_endpoint.methods,
default_query_params=created_endpoint.default_query_params,
+ auth=created_endpoint.auth,
)
return PassThroughEndpointResponse(endpoints=[created_endpoint])
diff --git a/tests/llm_translation/reasoning_effort_grid/grid_spec.py b/tests/llm_translation/reasoning_effort_grid/grid_spec.py
index 51a027ec101..a08013cd439 100644
--- a/tests/llm_translation/reasoning_effort_grid/grid_spec.py
+++ b/tests/llm_translation/reasoning_effort_grid/grid_spec.py
@@ -1,6 +1,7 @@
from dataclasses import dataclass, field
from typing import Dict, FrozenSet, List, Optional, Tuple
+
OMIT = object()
diff --git a/tests/llm_translation/reasoning_effort_grid/test_reasoning_effort_grid.py b/tests/llm_translation/reasoning_effort_grid/test_reasoning_effort_grid.py
index a5ef1b7ea3c..551ab8459d1 100644
--- a/tests/llm_translation/reasoning_effort_grid/test_reasoning_effort_grid.py
+++ b/tests/llm_translation/reasoning_effort_grid/test_reasoning_effort_grid.py
@@ -15,6 +15,7 @@ from .grid_spec import (
all_cells,
)
+
_PROMPT_MESSAGES: List[Dict[str, str]] = [
{"role": "user", "content": "Step by step, calculate 47 * 53. Show your work."}
]
diff --git a/tests/local_testing/test_pass_through_endpoints.py b/tests/local_testing/test_pass_through_endpoints.py
index bd96ff04f7e..68ba62bcbab 100644
--- a/tests/local_testing/test_pass_through_endpoints.py
+++ b/tests/local_testing/test_pass_through_endpoints.py
@@ -218,7 +218,9 @@ async def test_pass_through_endpoint_rpm_limit(
for mock_api_key in mock_api_keys:
cache_value = UserAPIKeyAuth(
- token=hash_token(mock_api_key), rpm_limit=rpm_limit
+ token=hash_token(mock_api_key),
+ rpm_limit=rpm_limit,
+ metadata={"allowed_passthrough_routes": ["/v1/rerank"]},
)
user_api_key_cache.set_cache(key=hash_token(mock_api_key), value=cache_value)
@@ -320,7 +322,9 @@ async def test_pass_through_endpoint_sequential_rpm_limit(
for mock_api_key in mock_api_keys:
cache_value = UserAPIKeyAuth(
- token=hash_token(mock_api_key), rpm_limit=rpm_limit
+ token=hash_token(mock_api_key),
+ rpm_limit=rpm_limit,
+ metadata={"allowed_passthrough_routes": ["/v1/rerank"]},
)
user_api_key_cache.set_cache(key=hash_token(mock_api_key), value=cache_value)
diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py
index 90c5d4f4fc5..14119f7ad4e 100644
--- a/tests/test_litellm/proxy/auth/test_handle_jwt.py
+++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py
@@ -1,6 +1,7 @@
from typing import Optional
from unittest.mock import AsyncMock, MagicMock, patch
+from fastapi import HTTPException
import pytest
from litellm.proxy._types import (
@@ -132,6 +133,141 @@ async def test_map_user_to_teams_null_inputs():
await JWTAuthManager.map_user_to_teams(user_object=None, team_object=None)
+@pytest.mark.asyncio
+async def test_find_team_with_model_access_reports_passthrough_allowlist_denial():
+ jwt_handler = JWTHandler()
+ jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth()
+ team = LiteLLM_TeamTable(
+ team_id="team-a",
+ models=["gpt-4"],
+ metadata={},
+ )
+
+ with (
+ patch(
+ "litellm.proxy.auth.handle_jwt.get_team_object",
+ new_callable=AsyncMock,
+ return_value=team,
+ ),
+ patch(
+ "litellm.proxy.auth.handle_jwt.can_team_access_model",
+ new_callable=AsyncMock,
+ return_value=True,
+ ),
+ patch(
+ "litellm.proxy.auth.handle_jwt.allowed_routes_check",
+ return_value=True,
+ ),
+ patch(
+ "litellm.proxy.auth.handle_jwt.RouteChecks.is_auth_enforced_pass_through_route",
+ return_value=True,
+ ) as mock_is_auth_enforced_pass_through_route,
+ patch(
+ "litellm.proxy.auth.handle_jwt.RouteChecks.check_passthrough_route_access",
+ return_value=False,
+ ) as mock_passthrough_check,
+ ):
+ with pytest.raises(HTTPException) as exc_info:
+ await JWTAuthManager.find_team_with_model_access(
+ team_ids={"team-a"},
+ requested_model="gpt-4",
+ route="/my-pass-through",
+ request_method="POST",
+ jwt_handler=jwt_handler,
+ prisma_client=None,
+ user_api_key_cache=MagicMock(),
+ parent_otel_span=None,
+ proxy_logging_obj=MagicMock(),
+ )
+
+ assert exc_info.value.status_code == 403
+ assert "allowed_passthrough_routes" in exc_info.value.detail
+ assert "requested model" not in exc_info.value.detail
+ mock_is_auth_enforced_pass_through_route.assert_called_once_with(
+ route="/my-pass-through", method="POST"
+ )
+
+ user_api_key_dict = mock_passthrough_check.call_args.kwargs["user_api_key_dict"]
+ assert user_api_key_dict.metadata == {}
+ assert user_api_key_dict.team_metadata == {}
+
+
+@pytest.mark.asyncio
+async def test_find_team_with_model_access_uses_request_method_for_passthrough_auth():
+ jwt_handler = JWTHandler()
+ jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth()
+ team = LiteLLM_TeamTable(
+ team_id="team-a",
+ models=["gpt-4"],
+ metadata={},
+ )
+ mock_registered_routes = {
+ "test-uuid-1:exact:/custom:GET": {
+ "endpoint_id": "test-uuid-1",
+ "path": "/custom",
+ "type": "exact",
+ "methods": ["GET"],
+ "auth": False,
+ },
+ "test-uuid-2:exact:/custom:POST": {
+ "endpoint_id": "test-uuid-2",
+ "path": "/custom",
+ "type": "exact",
+ "methods": ["POST"],
+ "auth": True,
+ },
+ }
+
+ with (
+ patch(
+ "litellm.proxy.auth.handle_jwt.get_team_object",
+ new_callable=AsyncMock,
+ return_value=team,
+ ),
+ patch(
+ "litellm.proxy.auth.handle_jwt.allowed_routes_check",
+ return_value=True,
+ ),
+ patch(
+ "litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes",
+ mock_registered_routes,
+ ),
+ patch(
+ "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_server_root_path",
+ return_value="/",
+ ),
+ ):
+ team_id, team_obj = await JWTAuthManager.find_team_with_model_access(
+ team_ids={"team-a"},
+ requested_model=None,
+ route="/custom",
+ jwt_handler=jwt_handler,
+ prisma_client=None,
+ user_api_key_cache=MagicMock(),
+ parent_otel_span=None,
+ proxy_logging_obj=MagicMock(),
+ request_method="GET",
+ )
+ assert team_id == "team-a"
+ assert team_obj == team
+
+ with pytest.raises(HTTPException) as exc_info:
+ await JWTAuthManager.find_team_with_model_access(
+ team_ids={"team-a"},
+ requested_model=None,
+ route="/custom",
+ jwt_handler=jwt_handler,
+ prisma_client=None,
+ user_api_key_cache=MagicMock(),
+ parent_otel_span=None,
+ proxy_logging_obj=MagicMock(),
+ request_method="POST",
+ )
+
+ assert exc_info.value.status_code == 403
+ assert "allowed_passthrough_routes" in exc_info.value.detail
+
+
@pytest.mark.asyncio
async def test_auth_builder_proxy_admin_user_role():
"""Test that is_proxy_admin is True when user_object.user_role is PROXY_ADMIN"""
@@ -1576,6 +1712,295 @@ async def test_auth_builder_uses_team_from_header_e2e():
assert result["team_object"] == team_object
+@pytest.mark.asyncio
+async def test_auth_builder_header_team_denies_auth_passthrough_without_allowlist():
+ """Header-selected JWT teams must enforce team allowed_passthrough_routes."""
+ from litellm.caching import DualCache
+ from litellm.proxy.utils import ProxyLogging
+
+ jwt_handler = JWTHandler()
+ user_api_key_cache = DualCache()
+ jwt_handler.update_environment(
+ prisma_client=None,
+ user_api_key_cache=user_api_key_cache,
+ litellm_jwtauth=LiteLLM_JWTAuth(
+ team_ids_jwt_field="groups",
+ user_id_jwt_field="sub",
+ ),
+ )
+
+ team_object = LiteLLM_TeamTable(team_id="team-2", metadata={})
+
+ with (
+ patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt,
+ patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock),
+ patch.object(
+ JWTAuthManager,
+ "check_admin_access",
+ new_callable=AsyncMock,
+ return_value=None,
+ ),
+ patch(
+ "litellm.proxy.auth.handle_jwt.get_team_object",
+ new_callable=AsyncMock,
+ return_value=team_object,
+ ),
+ patch.object(
+ JWTAuthManager,
+ "get_objects",
+ new_callable=AsyncMock,
+ ) as mock_get_objects,
+ patch(
+ "litellm.proxy.auth.handle_jwt.RouteChecks.is_auth_enforced_pass_through_route",
+ return_value=True,
+ ),
+ patch(
+ "litellm.proxy.auth.handle_jwt.RouteChecks.check_passthrough_route_access",
+ return_value=False,
+ ) as mock_passthrough_check,
+ ):
+ mock_auth_jwt.return_value = {
+ "sub": "user-1",
+ "scope": "",
+ "groups": ["team-1", "team-2"],
+ }
+
+ with pytest.raises(HTTPException) as exc_info:
+ await JWTAuthManager.auth_builder(
+ api_key="jwt-token",
+ jwt_handler=jwt_handler,
+ request_data={"model": "gpt-4"},
+ general_settings={},
+ route="/my-pass-through",
+ prisma_client=None,
+ user_api_key_cache=user_api_key_cache,
+ parent_otel_span=None,
+ proxy_logging_obj=ProxyLogging(user_api_key_cache=user_api_key_cache),
+ request_headers={"x-litellm-team-id": "team-2"},
+ request_method="POST",
+ )
+
+ assert exc_info.value.status_code == 403
+ assert "allowed_passthrough_routes" in exc_info.value.detail
+ mock_get_objects.assert_not_called()
+ user_api_key_dict = mock_passthrough_check.call_args.kwargs["user_api_key_dict"]
+ assert user_api_key_dict.team_metadata == {}
+
+
+@pytest.mark.asyncio
+async def test_auth_builder_specific_team_denies_auth_passthrough_without_allowlist():
+ """JWT-field-selected teams must enforce team allowed_passthrough_routes."""
+ from litellm.caching import DualCache
+ from litellm.proxy.utils import ProxyLogging
+
+ jwt_handler = JWTHandler()
+ user_api_key_cache = DualCache()
+ jwt_handler.update_environment(
+ prisma_client=None,
+ user_api_key_cache=user_api_key_cache,
+ litellm_jwtauth=LiteLLM_JWTAuth(
+ team_id_jwt_field="team_id",
+ user_id_jwt_field="sub",
+ ),
+ )
+
+ team_object = LiteLLM_TeamTable(team_id="team-1", metadata={})
+
+ with (
+ patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt,
+ patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock),
+ patch.object(
+ JWTAuthManager,
+ "check_admin_access",
+ new_callable=AsyncMock,
+ return_value=None,
+ ),
+ patch(
+ "litellm.proxy.auth.handle_jwt.get_team_object",
+ new_callable=AsyncMock,
+ return_value=team_object,
+ ),
+ patch.object(
+ JWTAuthManager,
+ "get_objects",
+ new_callable=AsyncMock,
+ ) as mock_get_objects,
+ patch(
+ "litellm.proxy.auth.handle_jwt.RouteChecks.is_auth_enforced_pass_through_route",
+ return_value=True,
+ ),
+ patch(
+ "litellm.proxy.auth.handle_jwt.RouteChecks.check_passthrough_route_access",
+ return_value=False,
+ ) as mock_passthrough_check,
+ ):
+ mock_auth_jwt.return_value = {
+ "sub": "user-1",
+ "scope": "",
+ "team_id": "team-1",
+ }
+
+ with pytest.raises(HTTPException) as exc_info:
+ await JWTAuthManager.auth_builder(
+ api_key="jwt-token",
+ jwt_handler=jwt_handler,
+ request_data={"model": "gpt-4"},
+ general_settings={},
+ route="/my-pass-through",
+ prisma_client=None,
+ user_api_key_cache=user_api_key_cache,
+ parent_otel_span=None,
+ proxy_logging_obj=ProxyLogging(user_api_key_cache=user_api_key_cache),
+ request_method="POST",
+ )
+
+ assert exc_info.value.status_code == 403
+ assert "allowed_passthrough_routes" in exc_info.value.detail
+ mock_get_objects.assert_not_called()
+ user_api_key_dict = mock_passthrough_check.call_args.kwargs["user_api_key_dict"]
+ assert user_api_key_dict.team_metadata == {}
+
+
+@pytest.mark.asyncio
+async def test_auth_builder_rbac_team_loads_team_for_passthrough_allowlist():
+ """RBAC role-claim teams (team_object unset) must load team metadata before gating."""
+ from litellm.caching import DualCache
+ from litellm.proxy.utils import ProxyLogging
+
+ jwt_handler = JWTHandler()
+ user_api_key_cache = DualCache()
+ jwt_handler.update_environment(
+ prisma_client=None,
+ user_api_key_cache=user_api_key_cache,
+ litellm_jwtauth=LiteLLM_JWTAuth(),
+ )
+
+ team_object = LiteLLM_TeamTable(
+ team_id="team-rbac",
+ metadata={"allowed_passthrough_routes": ["/my-pass-through"]},
+ )
+
+ with (
+ patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt,
+ patch.object(jwt_handler, "get_rbac_role", return_value=LitellmUserRoles.TEAM),
+ patch.object(jwt_handler, "get_object_id", return_value="team-rbac"),
+ patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock),
+ patch.object(
+ JWTAuthManager,
+ "check_admin_access",
+ new_callable=AsyncMock,
+ return_value=None,
+ ),
+ patch(
+ "litellm.proxy.auth.handle_jwt.get_team_object",
+ new_callable=AsyncMock,
+ return_value=team_object,
+ ) as mock_get_team,
+ patch.object(
+ JWTAuthManager,
+ "get_objects",
+ new_callable=AsyncMock,
+ return_value=(None, None, None, None),
+ ),
+ patch.object(JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock),
+ patch.object(
+ JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock
+ ),
+ patch(
+ "litellm.proxy.auth.handle_jwt.RouteChecks.is_auth_enforced_pass_through_route",
+ return_value=True,
+ ),
+ patch(
+ "litellm.proxy.auth.handle_jwt.RouteChecks.check_passthrough_route_access",
+ return_value=True,
+ ) as mock_passthrough_check,
+ ):
+ mock_auth_jwt.return_value = {"scope": ""}
+
+ result = await JWTAuthManager.auth_builder(
+ api_key="jwt-token",
+ jwt_handler=jwt_handler,
+ request_data={"model": "gpt-4"},
+ general_settings={},
+ route="/my-pass-through",
+ prisma_client=None,
+ user_api_key_cache=user_api_key_cache,
+ parent_otel_span=None,
+ proxy_logging_obj=ProxyLogging(user_api_key_cache=user_api_key_cache),
+ request_method="POST",
+ )
+
+ assert result["team_id"] == "team-rbac"
+ mock_get_team.assert_awaited_once()
+ assert mock_get_team.await_args.kwargs["team_id"] == "team-rbac"
+ user_api_key_dict = mock_passthrough_check.call_args.kwargs["user_api_key_dict"]
+ assert user_api_key_dict.team_metadata == {
+ "allowed_passthrough_routes": ["/my-pass-through"]
+ }
+
+
+@pytest.mark.asyncio
+async def test_auth_builder_rbac_team_denies_passthrough_without_allowlist():
+ """RBAC role-claim teams without an allowlist are still denied for passthrough."""
+ from litellm.caching import DualCache
+ from litellm.proxy.utils import ProxyLogging
+
+ jwt_handler = JWTHandler()
+ user_api_key_cache = DualCache()
+ jwt_handler.update_environment(
+ prisma_client=None,
+ user_api_key_cache=user_api_key_cache,
+ litellm_jwtauth=LiteLLM_JWTAuth(),
+ )
+
+ team_object = LiteLLM_TeamTable(team_id="team-rbac", metadata={})
+
+ with (
+ patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt,
+ patch.object(jwt_handler, "get_rbac_role", return_value=LitellmUserRoles.TEAM),
+ patch.object(jwt_handler, "get_object_id", return_value="team-rbac"),
+ patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock),
+ patch.object(
+ JWTAuthManager,
+ "check_admin_access",
+ new_callable=AsyncMock,
+ return_value=None,
+ ),
+ patch(
+ "litellm.proxy.auth.handle_jwt.get_team_object",
+ new_callable=AsyncMock,
+ return_value=team_object,
+ ) as mock_get_team,
+ patch(
+ "litellm.proxy.auth.handle_jwt.RouteChecks.is_auth_enforced_pass_through_route",
+ return_value=True,
+ ),
+ patch(
+ "litellm.proxy.auth.handle_jwt.RouteChecks.check_passthrough_route_access",
+ return_value=False,
+ ),
+ ):
+ mock_auth_jwt.return_value = {"scope": ""}
+
+ with pytest.raises(HTTPException) as exc_info:
+ await JWTAuthManager.auth_builder(
+ api_key="jwt-token",
+ jwt_handler=jwt_handler,
+ request_data={"model": "gpt-4"},
+ general_settings={},
+ route="/my-pass-through",
+ prisma_client=None,
+ user_api_key_cache=user_api_key_cache,
+ parent_otel_span=None,
+ proxy_logging_obj=ProxyLogging(user_api_key_cache=user_api_key_cache),
+ request_method="POST",
+ )
+
+ assert exc_info.value.status_code == 403
+ assert "allowed_passthrough_routes" in exc_info.value.detail
+ mock_get_team.assert_awaited_once()
+
+
@pytest.mark.asyncio
async def test_auth_builder_admin_on_llm_route_honors_team_header():
"""JWT proxy_admin + x-litellm-team-id on an LLM API route -> team context is
diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py
index b308a665062..ad9295d6b19 100644
--- a/tests/test_litellm/proxy/auth/test_route_checks.py
+++ b/tests/test_litellm/proxy/auth/test_route_checks.py
@@ -224,6 +224,28 @@ def test_virtual_key_mcp_routes_allows_v1_mcp_server():
assert result is True
+def test_auth_enforced_passthrough_check_does_not_apply_to_info_routes():
+ """Auth-enforced passthrough gating only applies to OpenAI/LLM route groups."""
+
+ valid_token = UserAPIKeyAuth(
+ user_id="test_user",
+ allowed_routes=["info_routes"],
+ )
+
+ with patch.object(
+ RouteChecks,
+ "is_auth_enforced_pass_through_route",
+ return_value=True,
+ ) as mock_is_auth_enforced_pass_through_route:
+ result = RouteChecks.is_virtual_key_allowed_to_call_route(
+ route="/team/info",
+ valid_token=valid_token,
+ )
+
+ assert result is True
+ mock_is_auth_enforced_pass_through_route.assert_not_called()
+
+
@pytest.mark.parametrize(
"route",
[
@@ -686,24 +708,88 @@ def test_anthropic_count_tokens_route_accessible_to_internal_users():
def test_virtual_key_llm_api_routes_allows_registered_pass_through_endpoints():
"""
- Test that virtual keys with llm_api_routes permission can access registered pass-through endpoints.
-
- This tests the scenario where a pass-through endpoint is registered from the DB
- (e.g., /azure-assistant) and a virtual key with llm_api_routes permission should be able to access
- both the exact path and subpaths (e.g., /azure-assistant/openai/assistants).
+ Virtual keys with llm_api_routes can access auth=true pass-through endpoints only when
+ allowed_passthrough_routes is configured on the key or team.
"""
- # Mock the registered pass-through routes
mock_registered_routes = {
- "test-uuid-1:exact:/azure-assistant": {
+ "test-uuid-1:exact:/azure-assistant:DELETE,GET,PATCH,POST,PUT": {
"endpoint_id": "test-uuid-1",
"path": "/azure-assistant",
"type": "exact",
+ "auth": True,
},
- "test-uuid-2:subpath:/custom-endpoint": {
+ "test-uuid-2:subpath:/custom-endpoint:DELETE,GET,PATCH,POST,PUT": {
"endpoint_id": "test-uuid-2",
"path": "/custom-endpoint",
"type": "subpath",
+ "auth": True,
+ },
+ }
+
+ with (
+ patch(
+ "litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes",
+ mock_registered_routes,
+ ),
+ patch(
+ "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_server_root_path",
+ return_value="/",
+ ),
+ ):
+ valid_token = UserAPIKeyAuth(
+ user_id="test_user",
+ allowed_routes=["llm_api_routes"],
+ metadata={
+ "allowed_passthrough_routes": [
+ "/azure-assistant",
+ "/custom-endpoint",
+ ]
+ },
+ )
+
+ assert (
+ RouteChecks.is_virtual_key_allowed_to_call_route(
+ route="/azure-assistant",
+ valid_token=valid_token,
+ )
+ is True
+ )
+ assert (
+ RouteChecks.is_virtual_key_allowed_to_call_route(
+ route="/custom-endpoint/openai/assistants",
+ valid_token=valid_token,
+ )
+ is True
+ )
+ assert (
+ RouteChecks.is_virtual_key_allowed_to_call_route(
+ route="/custom-endpoint",
+ valid_token=valid_token,
+ )
+ is True
+ )
+
+
+def test_virtual_key_llm_api_routes_allows_non_auth_enforced_pass_through_endpoints():
+ """
+ Virtual keys with llm_api_routes can access registered pass-through endpoints that
+ are NOT auth-enforced (auth=false) without configuring allowed_passthrough_routes.
+ This is the original behaviour and must not regress.
+ """
+
+ mock_registered_routes = {
+ "test-uuid-1:exact:/azure-assistant:DELETE,GET,PATCH,POST,PUT": {
+ "endpoint_id": "test-uuid-1",
+ "path": "/azure-assistant",
+ "type": "exact",
+ "auth": False,
+ },
+ "test-uuid-2:subpath:/custom-endpoint:DELETE,GET,PATCH,POST,PUT": {
+ "endpoint_id": "test-uuid-2",
+ "path": "/custom-endpoint",
+ "type": "subpath",
+ "auth": False,
},
}
@@ -717,32 +803,202 @@ def test_virtual_key_llm_api_routes_allows_registered_pass_through_endpoints():
return_value="/",
),
):
- # Create a virtual key with llm_api_routes permission
valid_token = UserAPIKeyAuth(
user_id="test_user",
allowed_routes=["llm_api_routes"],
)
- # Test exact match for registered pass-through endpoint
- result1 = RouteChecks.is_virtual_key_allowed_to_call_route(
- route="/azure-assistant",
- valid_token=valid_token,
+ assert (
+ RouteChecks.is_virtual_key_allowed_to_call_route(
+ route="/azure-assistant",
+ valid_token=valid_token,
+ )
+ is True
+ )
+ assert (
+ RouteChecks.is_virtual_key_allowed_to_call_route(
+ route="/custom-endpoint/openai/assistants",
+ valid_token=valid_token,
+ )
+ is True
+ )
+ assert (
+ RouteChecks.is_virtual_key_allowed_to_call_route(
+ route="/custom-endpoint",
+ valid_token=valid_token,
+ )
+ is True
)
- assert result1 is True
- # Test subpath for registered pass-through endpoint with subpath type
- result2 = RouteChecks.is_virtual_key_allowed_to_call_route(
- route="/custom-endpoint/openai/assistants",
- valid_token=valid_token,
- )
- assert result2 is True
- # Test exact match for subpath type
- result3 = RouteChecks.is_virtual_key_allowed_to_call_route(
- route="/custom-endpoint",
- valid_token=valid_token,
+def test_virtual_key_llm_api_routes_denies_auth_pass_through_without_allowlist():
+ """auth=true pass-through must not be reachable via llm_api_routes alone."""
+
+ mock_registered_routes = {
+ "test-uuid-1:exact:/azure-assistant:GET,POST": {
+ "endpoint_id": "test-uuid-1",
+ "path": "/azure-assistant",
+ "type": "exact",
+ "auth": True,
+ },
+ }
+
+ with (
+ patch(
+ "litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes",
+ mock_registered_routes,
+ ),
+ patch(
+ "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_server_root_path",
+ return_value="/",
+ ),
+ ):
+ valid_token = UserAPIKeyAuth(
+ user_id="test_user",
+ allowed_routes=["llm_api_routes"],
+ )
+
+ with pytest.raises(HTTPException) as exc_info:
+ RouteChecks.is_virtual_key_allowed_to_call_route(
+ route="/azure-assistant",
+ valid_token=valid_token,
+ )
+ assert exc_info.value.status_code == 403
+ assert "allowed_passthrough_routes" in exc_info.value.detail
+
+
+def test_virtual_key_llm_api_routes_uses_method_specific_auth_setting():
+ """Same-path pass-through routes must be checked against the request method."""
+
+ mock_registered_routes = {
+ "test-uuid-1:exact:/custom:GET": {
+ "endpoint_id": "test-uuid-1",
+ "path": "/custom",
+ "type": "exact",
+ "methods": ["GET"],
+ "auth": False,
+ },
+ "test-uuid-2:exact:/custom:POST": {
+ "endpoint_id": "test-uuid-2",
+ "path": "/custom",
+ "type": "exact",
+ "methods": ["POST"],
+ "auth": True,
+ },
+ }
+
+ with (
+ patch(
+ "litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes",
+ mock_registered_routes,
+ ),
+ patch(
+ "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_server_root_path",
+ return_value="/",
+ ),
+ ):
+ valid_token = UserAPIKeyAuth(
+ user_id="test_user",
+ allowed_routes=["llm_api_routes"],
+ )
+
+ get_request = MagicMock(spec=Request)
+ get_request.method = "GET"
+ assert (
+ RouteChecks.is_virtual_key_allowed_to_call_route(
+ route="/custom",
+ valid_token=valid_token,
+ request=get_request,
+ )
+ is True
+ )
+
+ post_request = MagicMock(spec=Request)
+ post_request.method = "POST"
+ with pytest.raises(HTTPException) as exc_info:
+ RouteChecks.is_virtual_key_allowed_to_call_route(
+ route="/custom",
+ valid_token=valid_token,
+ request=post_request,
+ )
+
+ assert exc_info.value.status_code == 403
+
+
+def test_non_proxy_admin_denies_auth_pass_through_without_allowlist():
+ """Internal users must not bypass allowed_passthrough_routes via openai_routes."""
+
+ mock_registered_routes = {
+ "test-uuid-1:exact:/my-pass-through:GET,POST": {
+ "endpoint_id": "test-uuid-1",
+ "path": "/my-pass-through",
+ "type": "exact",
+ "auth": True,
+ },
+ }
+
+ valid_token = UserAPIKeyAuth(
+ user_id="test_user",
+ user_role=LitellmUserRoles.INTERNAL_USER.value,
+ )
+
+ with (
+ patch(
+ "litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes",
+ mock_registered_routes,
+ ),
+ patch(
+ "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_server_root_path",
+ return_value="/",
+ ),
+ ):
+ with pytest.raises(HTTPException) as exc_info:
+ RouteChecks.non_proxy_admin_allowed_routes_check(
+ user_obj=None,
+ _user_role=LitellmUserRoles.INTERNAL_USER.value,
+ route="/my-pass-through",
+ request=MagicMock(spec=Request),
+ valid_token=valid_token,
+ request_data={},
+ )
+ assert exc_info.value.status_code == 403
+ assert "allowed_passthrough_routes" in exc_info.value.detail
+
+
+def test_non_proxy_admin_allows_auth_pass_through_with_team_allowlist():
+ mock_registered_routes = {
+ "test-uuid-1:exact:/my-pass-through:GET,POST": {
+ "endpoint_id": "test-uuid-1",
+ "path": "/my-pass-through",
+ "type": "exact",
+ "auth": True,
+ },
+ }
+
+ valid_token = UserAPIKeyAuth(
+ user_id="test_user",
+ user_role=LitellmUserRoles.INTERNAL_USER.value,
+ team_metadata={"allowed_passthrough_routes": ["/my-pass-through"]},
+ )
+
+ with (
+ patch(
+ "litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes",
+ mock_registered_routes,
+ ),
+ patch(
+ "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_server_root_path",
+ return_value="/",
+ ),
+ ):
+ RouteChecks.non_proxy_admin_allowed_routes_check(
+ user_obj=None,
+ _user_role=LitellmUserRoles.INTERNAL_USER.value,
+ route="/my-pass-through",
+ request=MagicMock(spec=Request),
+ valid_token=valid_token,
+ request_data={},
)
- assert result3 is True
def test_virtual_key_without_llm_api_routes_cannot_access_pass_through():
diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py
index defd3bbcdcd..a3452ac8024 100644
--- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py
+++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py
@@ -1513,6 +1513,7 @@ class TestJWTOAuth2Coexistence:
mock_request = MagicMock()
mock_request.url.path = "/v1/chat/completions"
+ mock_request.method = "POST"
mock_request.headers = {"authorization": f"Bearer {jwt_token}"}
mock_request.query_params = {}
@@ -1546,6 +1547,7 @@ class TestJWTOAuth2Coexistence:
mock_oauth2.assert_not_called()
# JWT auth SHOULD be called
mock_jwt_auth.assert_called_once()
+ assert mock_jwt_auth.call_args.kwargs["request_method"] == "POST"
assert result.user_id == "jwt-human-user"
@pytest.mark.asyncio
diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py
index 344742ffe89..61299e2662a 100644
--- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py
+++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py
@@ -519,6 +519,64 @@ def test_add_subpath_route():
assert callable(call_args["endpoint"])
+@pytest.mark.asyncio
+async def test_pass_through_handler_rejects_unregistered_method():
+ """
+ Stale FastAPI routes can remain after an endpoint is updated from all methods
+ to a restricted method list. The handler must enforce the current registry.
+ """
+ from fastapi import HTTPException
+
+ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
+ create_pass_through_route,
+ )
+
+ endpoint_func = create_pass_through_route(
+ endpoint="/test/path",
+ target="http://example.com",
+ )
+ request = MagicMock(spec=Request)
+ request.method = "GET"
+
+ with (
+ patch.dict(os.environ, {"SERVER_ROOT_PATH": ""}),
+ patch(
+ "litellm.proxy.auth.auth_utils.get_request_route",
+ return_value="/test/path",
+ ),
+ patch(
+ "litellm.proxy.pass_through_endpoints.pass_through_endpoints._parse_request_data_by_content_type",
+ new_callable=AsyncMock,
+ return_value=({}, {}, None, False),
+ ),
+ patch(
+ "litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes",
+ {
+ "test-endpoint-id:exact:/test/path:POST": {
+ "endpoint_id": "test-endpoint-id",
+ "path": "/test/path",
+ "type": "exact",
+ "methods": ["POST"],
+ "passthrough_params": {
+ "target": "http://example.com",
+ "custom_headers": {},
+ "forward_headers": False,
+ "merge_query_params": False,
+ },
+ }
+ },
+ ),
+ ):
+ with pytest.raises(HTTPException) as exc_info:
+ await endpoint_func(
+ request=request,
+ fastapi_response=MagicMock(),
+ user_api_key_dict=MagicMock(),
+ )
+
+ assert exc_info.value.status_code == 405
+
+
@pytest.mark.asyncio
async def test_initialize_pass_through_endpoints_with_include_subpath():
"""
@@ -1171,6 +1229,256 @@ async def test_update_pass_through_endpoint():
assert updated_data["cost_per_request"] == 0.75
+@pytest.mark.asyncio
+async def test_create_pass_through_endpoint_auth_true_enforces_allowlist():
+ """
+ Regression: a pass-through endpoint created through the management API with
+ auth=true (the model default) must be treated as allowlist-enforced. The
+ create path registers FastAPI routes with dependencies=None, so deriving
+ enforcement from dependency metadata let a key with broad llm_api_routes
+ access call the route without an allowed_passthrough_routes match.
+ """
+ from fastapi import HTTPException
+
+ from litellm.proxy._types import (
+ ConfigFieldInfo,
+ PassThroughGenericEndpoint,
+ UserAPIKeyAuth,
+ )
+ from litellm.proxy.auth.route_checks import RouteChecks
+ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
+ create_pass_through_endpoints,
+ )
+
+ registry: dict = {}
+
+ with (
+ patch(
+ "litellm.proxy.proxy_server.get_config_general_settings"
+ ) as mock_get_config,
+ patch("litellm.proxy.proxy_server.update_config_general_settings"),
+ patch(
+ "litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes",
+ registry,
+ ),
+ patch(
+ "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_server_root_path",
+ return_value="/",
+ ),
+ ):
+ mock_get_config.return_value = ConfigFieldInfo(
+ field_name="pass_through_endpoints", field_value=[]
+ )
+
+ # auth is not passed -> defaults to True on PassThroughGenericEndpoint
+ endpoint = PassThroughGenericEndpoint(
+ path="/secure-passthrough",
+ target="http://example.com/api",
+ methods=["POST"],
+ )
+ await create_pass_through_endpoints(
+ data=endpoint,
+ request=MagicMock(spec=Request),
+ user_api_key_dict=MagicMock(spec=UserAPIKeyAuth),
+ )
+
+ assert any(value.get("auth") is True for value in registry.values())
+ assert (
+ RouteChecks.is_auth_enforced_pass_through_route(
+ route="/secure-passthrough", method="POST"
+ )
+ is True
+ )
+
+ post_request = MagicMock(spec=Request)
+ post_request.method = "POST"
+
+ without_allowlist = UserAPIKeyAuth(
+ user_id="u", allowed_routes=["llm_api_routes"]
+ )
+ with pytest.raises(HTTPException) as exc_info:
+ RouteChecks.is_virtual_key_allowed_to_call_route(
+ route="/secure-passthrough",
+ valid_token=without_allowlist,
+ request=post_request,
+ )
+ assert exc_info.value.status_code == 403
+ assert "allowed_passthrough_routes" in exc_info.value.detail
+
+ with_allowlist = UserAPIKeyAuth(
+ user_id="u",
+ allowed_routes=["llm_api_routes"],
+ metadata={"allowed_passthrough_routes": ["/secure-passthrough"]},
+ )
+ assert (
+ RouteChecks.is_virtual_key_allowed_to_call_route(
+ route="/secure-passthrough",
+ valid_token=with_allowlist,
+ request=post_request,
+ )
+ is True
+ )
+
+
+@pytest.mark.asyncio
+async def test_update_pass_through_endpoint_auth_true_enforces_allowlist():
+ """
+ Regression: editing a pass-through endpoint through the management API must
+ keep an auth=true route allowlist-enforced. remove_endpoint_routes drops the
+ old registry entry, so the re-registration has to record the auth flag.
+ """
+ from fastapi import HTTPException
+
+ from litellm.proxy._types import (
+ ConfigFieldInfo,
+ PassThroughGenericEndpoint,
+ UserAPIKeyAuth,
+ )
+ from litellm.proxy.auth.route_checks import RouteChecks
+ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
+ update_pass_through_endpoints,
+ )
+
+ registry: dict = {}
+ existing_endpoint_id = "edit-me-123"
+ existing_endpoints = [
+ {
+ "id": existing_endpoint_id,
+ "path": "/edited-passthrough",
+ "target": "http://example.com/api",
+ "auth": True,
+ "methods": ["POST"],
+ }
+ ]
+
+ with (
+ patch(
+ "litellm.proxy.proxy_server.get_config_general_settings"
+ ) as mock_get_config,
+ patch("litellm.proxy.proxy_server.update_config_general_settings"),
+ patch(
+ "litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes",
+ registry,
+ ),
+ patch(
+ "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_server_root_path",
+ return_value="/",
+ ),
+ ):
+ mock_get_config.return_value = ConfigFieldInfo(
+ field_name="pass_through_endpoints", field_value=existing_endpoints
+ )
+
+ update_data = PassThroughGenericEndpoint(
+ path="/edited-passthrough",
+ target="http://newapi.com/v2",
+ methods=["POST"],
+ )
+ await update_pass_through_endpoints(
+ endpoint_id=existing_endpoint_id,
+ data=update_data,
+ request=MagicMock(spec=Request),
+ user_api_key_dict=MagicMock(spec=UserAPIKeyAuth),
+ )
+
+ assert (
+ RouteChecks.is_auth_enforced_pass_through_route(
+ route="/edited-passthrough", method="POST"
+ )
+ is True
+ )
+
+ post_request = MagicMock(spec=Request)
+ post_request.method = "POST"
+
+ without_allowlist = UserAPIKeyAuth(
+ user_id="u", allowed_routes=["llm_api_routes"]
+ )
+ with pytest.raises(HTTPException) as exc_info:
+ RouteChecks.is_virtual_key_allowed_to_call_route(
+ route="/edited-passthrough",
+ valid_token=without_allowlist,
+ request=post_request,
+ )
+ assert exc_info.value.status_code == 403
+ assert "allowed_passthrough_routes" in exc_info.value.detail
+
+
+@pytest.mark.asyncio
+async def test_update_pass_through_endpoint_preserves_auth_false():
+ """
+ Regression: editing an unrelated field on an auth=false pass-through must not
+ silently flip it to auth=true. auth defaults to True on the request model, so a
+ naive exclude_none merge would overwrite the stored auth=false and start
+ rejecting every team/key that lacks allowed_passthrough_routes.
+ """
+ from litellm.proxy._types import (
+ ConfigFieldInfo,
+ PassThroughGenericEndpoint,
+ UserAPIKeyAuth,
+ )
+ from litellm.proxy.auth.route_checks import RouteChecks
+ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
+ update_pass_through_endpoints,
+ )
+
+ registry: dict = {}
+ existing_endpoint_id = "public-forwarder-123"
+ existing_endpoints = [
+ {
+ "id": existing_endpoint_id,
+ "path": "/public-passthrough",
+ "target": "http://example.com/api",
+ "auth": False,
+ "methods": ["POST"],
+ }
+ ]
+
+ with (
+ patch(
+ "litellm.proxy.proxy_server.get_config_general_settings"
+ ) as mock_get_config,
+ patch(
+ "litellm.proxy.proxy_server.update_config_general_settings"
+ ) as mock_update_config,
+ patch(
+ "litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes",
+ registry,
+ ),
+ patch(
+ "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_server_root_path",
+ return_value="/",
+ ),
+ ):
+ mock_get_config.return_value = ConfigFieldInfo(
+ field_name="pass_through_endpoints", field_value=existing_endpoints
+ )
+
+ update_data = PassThroughGenericEndpoint(
+ path="/public-passthrough",
+ target="http://newapi.com/v2",
+ methods=["POST"],
+ )
+ result = await update_pass_through_endpoints(
+ endpoint_id=existing_endpoint_id,
+ data=update_data,
+ request=MagicMock(spec=Request),
+ user_api_key_dict=MagicMock(spec=UserAPIKeyAuth),
+ )
+
+ assert result.endpoints[0].auth is False
+
+ persisted = mock_update_config.call_args[1]["data"].field_value[0]
+ assert persisted["auth"] is False
+
+ assert (
+ RouteChecks.is_auth_enforced_pass_through_route(
+ route="/public-passthrough", method="POST"
+ )
+ is False
+ )
+
+
@pytest.mark.asyncio
async def test_update_pass_through_endpoint_not_found():
"""
From 90b510447540a69fef469a95e024e56fd4684617 Mon Sep 17 00:00:00 2001
From: ryan-crabbe-berri
Date: Sat, 30 May 2026 17:35:31 -0700
Subject: [PATCH 36/44] feat(mcp/auth): additive key access-group grants +
opt-in member assignment (#29313)
* fix(mcp): make key.access_group_ids grants additive over team ceiling
A key whose unified access_group_ids grant a private MCP server was
having that grant intersected against its team's MCP ceiling, so a key
in a team scoped to other servers (or with no own scope) lost the
granted server entirely. Resolve access_group_ids once as ungated
additive grants and union them on top of the key/team ceiling instead
of folding them into the key scope that gets intersected.
* test(mcp): align key access-group tests with additive-grant model
The previous commit moved key.access_group_ids resolution out of the
intersected key ceiling (_get_allowed_mcp_servers_for_key) and into the
ungated additive grant path (_get_key_access_group_mcp_server_extras),
unioned on top of the team ceiling. Five tests from #28890/#29195 still
asserted the old gated / in-key-scope contract and failed:
- _get_allowed_mcp_servers_for_key now returns the object_permission
ceiling only and never resolves access_group_ids; two tests now assert
the group resolver is not called from that path (with and without an
object_permission present).
- The extras path is ungated, so a group whose assigned_team_ids /
assigned_key_ids exclude the caller still contributes its servers.
- The end-to-end test asserts the grant surfaces via the extras path
rather than the base key path.
- Dropped test_key_access_group_ids_empty_returns_no_extras; the empty
case is already covered by the extras family's no-groups test.
* feat(auth): gate member access-group assignment on keys behind opt-in
Non-admin team members could attach access_group_ids to keys they create
or update, letting them self-grant resources (MCP servers/models) the team
admin never intended. Add an opt-in KEY_ACCESS_GROUP_ASSIGNMENT team-member
permission (default-deny) enforced at /key/generate and /key/update; proxy
and team admins bypass. Surfaces automatically as a checkbox in the team
Member Permissions UI.
* fix(auth): gate access-group assignment on /key/regenerate too
RegenerateKeyRequest inherits access_group_ids and prepare_key_update_data
persists it, so a non-admin key owner could self-grant access groups by
regenerating. Apply the same opt-in member gate using the existing key's
team.
* test(auth): cover member access-group gate and additive MCP grants
Add unit tests for enforce_member_can_assign_access_groups (deny without
opt-in, allow with opt-in, and proxy-admin / team-admin / non-team-key
bypasses) and for _get_key_access_group_mcp_server_extras (no-auth and
no-resolved-servers return empty, resolved ids are expanded, errors
degrade to no grants).
---
.../mcp_server/auth/user_api_key_auth_mcp.py | 130 +++++------
litellm/proxy/_types.py | 6 +
.../key_management_endpoints.py | 33 +++
.../team_member_permission_checks.py | 64 ++++++
.../auth/test_user_api_key_auth_mcp.py | 209 ++++++++++++------
.../test_team_member_permission_checks.py | 109 +++++++++
.../team/permission_definitions.tsx | 2 +
7 files changed, 412 insertions(+), 141 deletions(-)
diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
index 6fc0a97838b..2aacab80f57 100644
--- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
+++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
@@ -566,7 +566,7 @@ class MCPRequestHandler:
)
)
- key_access_group_extras = (
+ key_access_group_grants = (
await MCPRequestHandler._get_key_access_group_mcp_server_extras(
user_api_key_auth
)
@@ -577,11 +577,11 @@ class MCPRequestHandler:
#########################################################
key_set = set(allowed_mcp_servers_for_key)
team_set = set(allowed_mcp_servers_for_team)
- extras_set = set(key_access_group_extras)
+ grants_set = set(key_access_group_grants)
- has_lower_level_mcp_restrictions = bool(key_set or team_set or extras_set)
+ has_lower_level_mcp_restrictions = bool(key_set or team_set or grants_set)
- # 1. Team-gated base scope.
+ # 1. Key/team ceiling. An empty set means "this level does not restrict".
if not team_set:
base = key_set # no team restriction
elif not key_set:
@@ -589,9 +589,10 @@ class MCPRequestHandler:
else:
base = key_set & team_set # both restrict → intersect
- # 2. Extend with access-group extras (LIT-3189 — bypasses team
- # ceiling, gated by group's assigned_team_ids / assigned_key_ids).
- allowed_mcp_servers: List[str] = list(base | extras_set)
+ # 2. Add the key's access-group grants on top. These are additive:
+ # attaching a group to the key grants its servers regardless of the
+ # team ceiling.
+ allowed_mcp_servers: List[str] = list(base | grants_set)
#########################################################
# Check end_user permissions if end_user_id is set
@@ -890,52 +891,12 @@ class MCPRequestHandler:
) -> List[str]:
"""
Resolve the key's unified `access_group_ids` (LiteLLM_AccessGroupTable) to
- MCP server IDs, gated by the access group's `assigned_team_ids` /
- `assigned_key_ids`. These servers extend the team's MCP scope rather
- than being capped by it. Tag-style `mcp_access_groups` (per-server tags)
- are intentionally not handled here — they have no assignment fields and
- remain subject to the team ceiling.
- """
- if user_api_key_auth is None:
- return []
- try:
- from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
- global_mcp_server_manager,
- )
- from litellm.proxy.auth.auth_checks import (
- get_authorized_resources_from_key_access_groups,
- )
-
- raw_server_ids = await get_authorized_resources_from_key_access_groups(
- valid_token=user_api_key_auth,
- team_object=None,
- resource_field="access_mcp_server_ids",
- )
- if not raw_server_ids:
- return []
- # Permission entries may be server_ids OR names/aliases — expand to ids.
- return global_mcp_server_manager.expand_permission_list(raw_server_ids)
- except Exception as e:
- verbose_logger.warning(
- f"Failed to get key access group MCP server extras: {str(e)}"
- )
- return []
-
- @staticmethod
- async def _get_allowed_mcp_servers_for_key(
- user_api_key_auth: Optional[UserAPIKeyAuth] = None,
- ) -> List[str]:
- """
- Get allowed MCP servers for a key (the key's own scope).
-
- Unions two sources:
- - Legacy key.object_permission (mcp_servers, mcp_access_groups,
- mcp_tool_permissions).
- - Unified key.access_group_ids → access_group.access_mcp_server_ids.
- Mirrors the ungated fallback in can_key_call_model — the group is
- attached to the key itself, so it grants the key's own scope (no
- assigned_key_ids re-check). The gated, team-ceiling-busting override
- lives in _get_key_access_group_mcp_server_extras.
+ MCP server IDs as additive grants: a group attached to the key extends the
+ key's allowed servers on top of the key/team ceiling rather than being
+ capped by the team. Attaching the group to the key is itself the grant —
+ no `assigned_key_ids` / `assigned_team_ids` re-check. Tag-style
+ `mcp_access_groups` (per-server tags) live in the key's object_permission
+ scope, not here.
"""
if user_api_key_auth is None:
return []
@@ -945,7 +906,6 @@ class MCPRequestHandler:
)
from litellm.proxy.auth.auth_checks import (
_get_mcp_server_ids_from_access_groups,
- get_object_permission,
)
from litellm.proxy.proxy_server import (
prisma_client,
@@ -953,17 +913,48 @@ class MCPRequestHandler:
user_api_key_cache,
)
- # Unified key.access_group_ids → MCP servers (ungated: the group is
- # attached to the key, so it grants the key's own scope). Entries in
- # access_mcp_server_ids may be server_ids OR names/aliases, so expand
- # to ids here — matching the legacy object_permission path below.
- key_access_group_servers = global_mcp_server_manager.expand_permission_list(
- await _get_mcp_server_ids_from_access_groups(
- access_group_ids=user_api_key_auth.access_group_ids or [],
- prisma_client=prisma_client,
- user_api_key_cache=user_api_key_cache,
- proxy_logging_obj=proxy_logging_obj,
- )
+ raw_server_ids = await _get_mcp_server_ids_from_access_groups(
+ access_group_ids=user_api_key_auth.access_group_ids or [],
+ prisma_client=prisma_client,
+ user_api_key_cache=user_api_key_cache,
+ proxy_logging_obj=proxy_logging_obj,
+ )
+ if not raw_server_ids:
+ return []
+ # Permission entries may be server_ids OR names/aliases — expand to ids.
+ return global_mcp_server_manager.expand_permission_list(raw_server_ids)
+ except Exception as e:
+ verbose_logger.warning(
+ f"Failed to get key access group MCP server grants: {str(e)}"
+ )
+ return []
+
+ @staticmethod
+ async def _get_allowed_mcp_servers_for_key(
+ user_api_key_auth: Optional[UserAPIKeyAuth] = None,
+ ) -> List[str]:
+ """
+ Get the key's own MCP ceiling from its object_permission
+ (mcp_servers, tag-style mcp_access_groups, mcp_tool_permissions).
+
+ Unified key.access_group_ids are NOT resolved here — they are additive
+ grants handled by _get_key_access_group_mcp_server_extras and unioned on
+ top of the key/team ceiling, so they must not enter this scope (which is
+ intersected against the team).
+ """
+ if user_api_key_auth is None:
+ return []
+ try:
+ from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
+ global_mcp_server_manager,
+ )
+ from litellm.proxy.auth.auth_checks import (
+ get_object_permission,
+ )
+ from litellm.proxy.proxy_server import (
+ prisma_client,
+ proxy_logging_obj,
+ user_api_key_cache,
)
# Get key object permission (already loaded in main auth flow, or fetch from DB)
@@ -983,7 +974,7 @@ class MCPRequestHandler:
proxy_logging_obj=proxy_logging_obj,
)
if key_object_permission is None:
- return list(set(key_access_group_servers))
+ return []
# Permission entries may be server_ids OR names/aliases — expand to ids.
direct_mcp_servers = global_mcp_server_manager.expand_permission_list(
@@ -1005,12 +996,7 @@ class MCPRequestHandler:
)
# Combine all lists
- all_servers = (
- direct_mcp_servers
- + access_group_servers
- + tool_perm_servers
- + key_access_group_servers
- )
+ all_servers = direct_mcp_servers + access_group_servers + tool_perm_servers
return list(set(all_servers))
except Exception as e:
verbose_logger.warning(
diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
index 95294f48386..751f855ea34 100644
--- a/litellm/proxy/_types.py
+++ b/litellm/proxy/_types.py
@@ -242,6 +242,11 @@ class KeyManagementRoutes(str, enum.Enum):
TEAM_KEY_BULK_UPDATE = "/team/key/bulk_update"
KEY_RESET_SPEND = "/key/{key_id}/reset_spend"
+ # Field-level opt-in permission (not a real HTTP route). When present in a
+ # team's `team_member_permissions`, non-admin members of that team may set
+ # `access_group_ids` on keys they create/update. Default-deny.
+ KEY_ACCESS_GROUP_ASSIGNMENT = "/key/access_group_assignment"
+
# info and health routes
KEY_INFO = "/key/info"
KEY_HEALTH = "/key/health"
@@ -552,6 +557,7 @@ class LiteLLMRoutes(enum.Enum):
KeyManagementRoutes.SPEND_LOGS_V2.value,
KeyManagementRoutes.KEY_RESET_SPEND.value,
KeyManagementRoutes.KEY_ALIASES.value,
+ KeyManagementRoutes.KEY_ACCESS_GROUP_ASSIGNMENT.value,
]
management_routes = (
diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py
index f49fb1f37b5..0e645013b92 100644
--- a/litellm/proxy/management_endpoints/key_management_endpoints.py
+++ b/litellm/proxy/management_endpoints/key_management_endpoints.py
@@ -325,6 +325,14 @@ def _team_key_generation_check(
_team_key_generation.get("required_params"),
)
+ # Field-level opt-in: non-admin members may only assign access groups when
+ # the team has enabled KEY_ACCESS_GROUP_ASSIGNMENT.
+ TeamMemberPermissionChecks.enforce_member_can_assign_access_groups(
+ user_api_key_dict=user_api_key_dict,
+ team_table=team_table,
+ access_group_ids=data.access_group_ids,
+ )
+
return True
@@ -2267,6 +2275,14 @@ async def _validate_update_key_data(
detail=f"Team not found for team_id={data.team_id}. Non-admin users cannot set keys to non-existent teams.",
)
+ # Field-level opt-in: non-admin members may only assign access groups when
+ # the team has enabled KEY_ACCESS_GROUP_ASSIGNMENT.
+ TeamMemberPermissionChecks.enforce_member_can_assign_access_groups(
+ user_api_key_dict=user_api_key_dict,
+ team_table=team_obj,
+ access_group_ids=data.access_group_ids,
+ )
+
if team_obj is not None:
await _check_team_key_limits(
team_table=team_obj,
@@ -4511,6 +4527,23 @@ async def regenerate_key_fn( # noqa: PLR0915
detail={"error": "You are not authorized to regenerate this key"},
)
+ # Gate access_group_ids on regenerate, same as /key/generate and
+ # /key/update. Use the existing key's team since the body may omit it.
+ if data is not None and data.access_group_ids:
+ regenerate_team_table: Optional[LiteLLM_TeamTableCachedObj] = None
+ if _key_in_db.team_id is not None:
+ regenerate_team_table = await get_team_object(
+ team_id=_key_in_db.team_id,
+ prisma_client=prisma_client,
+ user_api_key_cache=user_api_key_cache,
+ check_db_only=True,
+ )
+ TeamMemberPermissionChecks.enforce_member_can_assign_access_groups(
+ user_api_key_dict=user_api_key_dict,
+ team_table=regenerate_team_table,
+ access_group_ids=data.access_group_ids,
+ )
+
verbose_proxy_logger.info(
"Key regeneration requested: key_alias=%s",
getattr(_key_in_db, "key_alias", None),
diff --git a/litellm/proxy/management_helpers/team_member_permission_checks.py b/litellm/proxy/management_helpers/team_member_permission_checks.py
index 50339210a6e..2272a37488f 100644
--- a/litellm/proxy/management_helpers/team_member_permission_checks.py
+++ b/litellm/proxy/management_helpers/team_member_permission_checks.py
@@ -154,6 +154,70 @@ class TeamMemberPermissionChecks:
return True
+ @staticmethod
+ def enforce_member_can_assign_access_groups(
+ user_api_key_dict: UserAPIKeyAuth,
+ team_table: Optional[LiteLLM_TeamTableCachedObj],
+ access_group_ids: Optional[List[str]],
+ ) -> None:
+ """
+ Field-level opt-in gate: a non-admin team member may only set
+ `access_group_ids` on a (team) key if their team has opted in by adding
+ `KEY_ACCESS_GROUP_ASSIGNMENT` to `team_member_permissions`.
+
+ Bypassed for proxy admins, team admins, and personal (non-team) keys.
+ Default-deny: members cannot self-assign access groups until enabled.
+
+ Raises HTTPException(403) when a gated member attempts the assignment.
+ """
+ from fastapi import HTTPException
+
+ from litellm.proxy.management_endpoints.key_management_endpoints import (
+ _get_user_in_team,
+ )
+
+ # No-op when the request does not assign any access groups.
+ if not access_group_ids:
+ return
+
+ # Proxy admins always bypass.
+ if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
+ return
+
+ # Personal (non-team) keys are out of scope for team-member gating.
+ if team_table is None:
+ return
+
+ team_member_object = _get_user_in_team(
+ team_table=team_table, user_id=user_api_key_dict.user_id
+ )
+
+ # Team admins always bypass (consistent with other member-permission checks).
+ if team_member_object is not None and team_member_object.role == "admin":
+ return
+
+ permissions = (
+ TeamMemberPermissionChecks._get_list_of_route_enum_as_str(
+ TeamMemberPermissionChecks.get_permissions_for_team_member(
+ team_member_object=team_member_object,
+ team_table=team_table,
+ )
+ )
+ if team_member_object is not None
+ else []
+ )
+
+ if KeyManagementRoutes.KEY_ACCESS_GROUP_ASSIGNMENT.value not in permissions:
+ raise HTTPException(
+ status_code=403,
+ detail=(
+ "Team members cannot assign access groups to keys for team "
+ f"{team_table.team_id}. Ask a team or proxy admin to enable the "
+ f"'{KeyManagementRoutes.KEY_ACCESS_GROUP_ASSIGNMENT.value}' team "
+ "member permission to allow this."
+ ),
+ )
+
@staticmethod
async def user_belongs_to_keys_team(
user_api_key_dict: UserAPIKeyAuth,
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py
index b3508e13127..95c826daa8e 100644
--- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py
@@ -200,6 +200,117 @@ class TestMCPRequestHandler:
result = await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth)
assert result == [] # Should handle exception gracefully
+ @pytest.mark.parametrize(
+ "key_servers,team_servers,grant_servers,expected,scenario",
+ [
+ # Key has no own scope, restrictive team ceiling {test}, server
+ # granted only via key.access_group_ids → caller sees team's server
+ # AND the grant (grant is added on top of the ceiling).
+ (
+ [],
+ ["test"],
+ ["context7"],
+ ["context7", "test"],
+ "grant_over_team_ceiling",
+ ),
+ # key {a} ∩ team {b} = {} ; the grant still surfaces, proving grants
+ # are unioned with the ceiling, not intersected against it.
+ (
+ ["a"],
+ ["b"],
+ ["context7"],
+ ["context7"],
+ "grant_survives_empty_intersection",
+ ),
+ # No grant → ceiling behavior is unchanged (no additive leakage).
+ (["x", "y"], ["x"], [], ["x"], "no_grant_keeps_intersection"),
+ ],
+ )
+ async def test_access_group_grants_are_additive_over_ceiling(
+ self, key_servers, team_servers, grant_servers, expected, scenario
+ ):
+ """Regression: key.access_group_ids grants are unioned on top of the
+ key/team MCP ceiling, so a grant reaches the caller even when the team
+ ceiling does not include it (and even when key ∩ team is empty)."""
+ mock_user_auth = UserAPIKeyAuth(
+ api_key="test-key",
+ user_id="test-user",
+ team_id="test-team",
+ access_group_ids=["grp-mcp"],
+ )
+ with (
+ patch.object(
+ MCPRequestHandler, "_get_allowed_mcp_servers_for_key"
+ ) as mock_key,
+ patch.object(
+ MCPRequestHandler, "_get_allowed_mcp_servers_for_team"
+ ) as mock_team,
+ patch.object(
+ MCPRequestHandler, "_get_key_access_group_mcp_server_extras"
+ ) as mock_grants,
+ ):
+ mock_key.return_value = key_servers
+ mock_team.return_value = team_servers
+ mock_grants.return_value = grant_servers
+ result = await MCPRequestHandler.get_allowed_mcp_servers(mock_user_auth)
+ assert sorted(result) == sorted(expected)
+
+ async def test_access_group_extras_returns_empty_when_no_auth(self):
+ """No auth object → no additive grants."""
+ result = await MCPRequestHandler._get_key_access_group_mcp_server_extras(None)
+ assert result == []
+
+ async def test_access_group_extras_returns_empty_without_access_group_ids(self):
+ """A key with no resolvable access groups yields no additive grants
+ (the `if not raw_server_ids: return []` branch)."""
+ auth = UserAPIKeyAuth(api_key="k", access_group_ids=[])
+ with (
+ patch(
+ "litellm.proxy.auth.auth_checks._get_mcp_server_ids_from_access_groups",
+ new=AsyncMock(return_value=[]),
+ ),
+ patch(
+ "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager"
+ ) as mock_mgr,
+ ):
+ result = await MCPRequestHandler._get_key_access_group_mcp_server_extras(
+ auth
+ )
+ assert result == []
+ # expand_permission_list must not be reached when there are no raw ids.
+ mock_mgr.expand_permission_list.assert_not_called()
+
+ async def test_access_group_extras_expands_resolved_server_ids(self):
+ """Resolved access-group server ids/names are expanded to server ids."""
+ auth = UserAPIKeyAuth(api_key="k", access_group_ids=["grp-mcp"])
+ with (
+ patch(
+ "litellm.proxy.auth.auth_checks._get_mcp_server_ids_from_access_groups",
+ new=AsyncMock(return_value=["alias-a", "srv-b"]),
+ ),
+ patch(
+ "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager"
+ ) as mock_mgr,
+ ):
+ mock_mgr.expand_permission_list.return_value = ["srv-a", "srv-b"]
+ result = await MCPRequestHandler._get_key_access_group_mcp_server_extras(
+ auth
+ )
+ assert sorted(result) == ["srv-a", "srv-b"]
+ mock_mgr.expand_permission_list.assert_called_once_with(["alias-a", "srv-b"])
+
+ async def test_access_group_extras_swallows_errors(self):
+ """Resolution failures degrade to no grants rather than raising."""
+ auth = UserAPIKeyAuth(api_key="k", access_group_ids=["grp-mcp"])
+ with patch(
+ "litellm.proxy.auth.auth_checks._get_mcp_server_ids_from_access_groups",
+ new=AsyncMock(side_effect=Exception("db down")),
+ ):
+ result = await MCPRequestHandler._get_key_access_group_mcp_server_extras(
+ auth
+ )
+ assert result == []
+
@pytest.mark.parametrize(
"headers,expected_api_key,expected_mcp_auth_header,expected_server_auth_headers",
[
@@ -3335,12 +3446,12 @@ async def test_mcp_key_access_group_extras_when_group_has_no_servers():
@pytest.mark.asyncio
-async def test_mcp_key_access_group_extras_when_group_authorizes_neither():
- """
- Escalation regression: team member attaches a foreign access group to their key.
- Group grants servers BUT assigned_team_ids/assigned_key_ids exclude this caller.
- No extras contributed.
- """
+async def test_mcp_key_access_group_extras_granted_even_when_group_authorizes_neither():
+ """Grants are ungated: attaching the group to the key is itself the grant, so its
+ servers are contributed even when assigned_team_ids/assigned_key_ids exclude this
+ caller. (A team member self-assigning a foreign group to reach past the team
+ ceiling is a known, accepted-for-now tradeoff; restricting who may set
+ key.access_group_ids is a separate concern.)"""
valid_token = UserAPIKeyAuth(
token="team-a-token",
access_group_ids=["team-b-mcp-group"],
@@ -3365,7 +3476,7 @@ async def test_mcp_key_access_group_extras_when_group_authorizes_neither():
result = await MCPRequestHandler._get_key_access_group_mcp_server_extras(
valid_token
)
- assert result == []
+ assert result == ["srv-finance-only"]
finally:
_stop_patches(patches)
@@ -3650,11 +3761,12 @@ async def test_get_allowed_mcp_servers_includes_team_access_group_extras_end_to_
@pytest.mark.asyncio
-async def test_key_access_group_ids_resolves_mcp_servers_ungated():
- """A teamless key whose unified access_group_ids grant an MCP server sees it
- even though the group lists the key in NEITHER assigned_key_ids NOR
- assigned_team_ids — the group is attached to the key, so it grants the key's
- own scope (ungated, mirroring can_key_call_model's fallback)."""
+async def test_allowed_mcp_servers_for_key_excludes_access_group_ids():
+ """The key's own ceiling (which is intersected against the team) must NOT resolve
+ access_group_ids — those are additive grants handled separately, so folding them
+ in here is exactly the bug this fix removes. A key with only access_group_ids and
+ no object_permission yields an empty ceiling, and the group resolver is never
+ called from this path."""
auth = UserAPIKeyAuth(
token="test-token-hash",
api_key="sk-test",
@@ -3671,15 +3783,16 @@ async def test_key_access_group_ids_resolves_mcp_servers_ungated():
):
result = await MCPRequestHandler._get_allowed_mcp_servers_for_key(auth)
- assert result == ["srv-stripe"]
- mock_resolver.assert_called_once()
- assert mock_resolver.call_args.kwargs["access_group_ids"] == ["mcp-premium"]
+ assert result == []
+ mock_resolver.assert_not_called()
@pytest.mark.asyncio
-async def test_key_access_group_ids_union_with_object_permission():
- """When both legacy key.object_permission and unified key.access_group_ids
- grant MCP servers, the final list is their union."""
+async def test_allowed_mcp_servers_for_key_uses_object_permission_not_access_groups():
+ """The key's own ceiling is built from object_permission alone. Even when the key
+ also carries access_group_ids that would resolve to other servers, those grants do
+ NOT enter this (intersected) scope — only the object_permission server comes back.
+ """
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
@@ -3714,82 +3827,40 @@ async def test_key_access_group_ids_union_with_object_permission():
"litellm.proxy.auth.auth_checks._get_mcp_server_ids_from_access_groups",
new_callable=AsyncMock,
return_value=["srv-stripe"],
- ),
+ ) as mock_resolver,
):
result = await MCPRequestHandler._get_allowed_mcp_servers_for_key(auth)
- assert set(result) == {"srv-direct", "srv-stripe"}
+ assert set(result) == {"srv-direct"}
+ mock_resolver.assert_not_called()
finally:
global_mcp_server_manager.registry.pop("srv-direct", None)
@pytest.mark.asyncio
-async def test_key_access_group_ids_empty_returns_no_extras():
- """Empty key.access_group_ids and no object_permission → resolver called with
- [], short-circuits without DB access, returns []."""
- auth = UserAPIKeyAuth(
- token="test-token-hash",
- api_key="sk-test",
- access_group_ids=[],
- )
-
- with (
- patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
- patch(
- "litellm.proxy.auth.auth_checks._get_mcp_server_ids_from_access_groups",
- new_callable=AsyncMock,
- return_value=[],
- ) as mock_resolver,
- ):
- result = await MCPRequestHandler._get_allowed_mcp_servers_for_key(auth)
-
- assert result == []
- mock_resolver.assert_called_once()
- assert mock_resolver.call_args.kwargs["access_group_ids"] == []
-
-
-@pytest.mark.asyncio
-async def test_get_allowed_mcp_servers_key_access_group_base_end_to_end():
- """End-to-end bug repro: a teamless key has an MCP-granting access group on
- its access_group_ids, but the group lists the key in NEITHER assigned_key_ids
- NOR assigned_team_ids. The gated extras path returns [] (no override), yet the
- ungated base key path grants the server → the key sees it through
- get_allowed_mcp_servers."""
+async def test_get_allowed_mcp_servers_surfaces_ungated_key_access_group_grant_end_to_end():
+ """End-to-end: a teamless key has an MCP-granting access group on its
+ access_group_ids. The grant is resolved ungated by the additive extras path and
+ surfaces through get_allowed_mcp_servers, even though the key's own ceiling
+ (object_permission) is empty."""
auth = UserAPIKeyAuth(
token="test-token",
api_key="sk-test",
access_group_ids=["mcp-group"],
)
- # Group grants the server but admits neither this key nor its (absent) team.
- fake_ag = _fake_mcp_access_group(
- access_group_id="mcp-group",
- access_mcp_server_ids=["srv-deepwiki"],
- assigned_team_ids=[],
- assigned_key_ids=[],
- )
patches = _patch_proxy_server_globals_for_mcp() + [
- # Ungated base resolver used by _get_allowed_mcp_servers_for_key.
patch(
"litellm.proxy.auth.auth_checks._get_mcp_server_ids_from_access_groups",
new_callable=AsyncMock,
return_value=["srv-deepwiki"],
),
- # Gated path (_get_key_access_group_mcp_server_extras) resolves the group
- # via get_access_object; empty assigned_* → it contributes nothing.
- patch(
- "litellm.proxy.auth.auth_checks.get_access_object",
- new_callable=AsyncMock,
- return_value=fake_ag,
- ),
]
_start_patches(patches)
try:
- # Sanity: the gated extras path alone denies (the old behavior).
extras = await MCPRequestHandler._get_key_access_group_mcp_server_extras(auth)
- assert extras == []
+ assert extras == ["srv-deepwiki"]
- # But the key now sees the server via the ungated base path.
result = await MCPRequestHandler.get_allowed_mcp_servers(auth)
assert result == ["srv-deepwiki"]
finally:
diff --git a/tests/test_litellm/proxy/management_helpers/test_team_member_permission_checks.py b/tests/test_litellm/proxy/management_helpers/test_team_member_permission_checks.py
index 6eb05aaf9d5..29aa75a0f0a 100644
--- a/tests/test_litellm/proxy/management_helpers/test_team_member_permission_checks.py
+++ b/tests/test_litellm/proxy/management_helpers/test_team_member_permission_checks.py
@@ -265,3 +265,112 @@ class TestCanTeamMemberExecuteKeyManagementEndpoint:
user_api_key_cache=MagicMock(),
existing_key_row=existing_key_row,
)
+
+
+class TestEnforceMemberCanAssignAccessGroups:
+ """Opt-in gate controlling whether a non-admin team member may set
+ `access_group_ids` on a key (generate/update/regenerate)."""
+
+ AG_PERMISSION = KeyManagementRoutes.KEY_ACCESS_GROUP_ASSIGNMENT.value
+
+ def _user(self, role="internal_user", user_id="user-a"):
+ u = MagicMock()
+ u.user_role = role
+ u.user_id = user_id
+ return u
+
+ def _team(self, team_member_permissions, team_id="team-a"):
+ team = MagicMock()
+ team.team_id = team_id
+ team.team_member_permissions = team_member_permissions
+ return team
+
+ def test_no_access_group_ids_is_noop(self, monkeypatch):
+ """When no access groups are requested the gate never raises, even
+ for a gated member with no opt-in permission."""
+ from litellm.proxy.management_endpoints import key_management_endpoints
+
+ monkeypatch.setattr(
+ key_management_endpoints,
+ "_get_user_in_team",
+ lambda **kwargs: Member(role="user", user_id="user-a"),
+ )
+
+ # Both None and empty list are no-ops.
+ for access_group_ids in (None, []):
+ TeamMemberPermissionChecks.enforce_member_can_assign_access_groups(
+ user_api_key_dict=self._user(),
+ team_table=self._team([]),
+ access_group_ids=access_group_ids,
+ )
+
+ def test_proxy_admin_bypasses(self, monkeypatch):
+ """Proxy admins may assign access groups regardless of team opt-in."""
+ from litellm.proxy._types import LitellmUserRoles
+
+ TeamMemberPermissionChecks.enforce_member_can_assign_access_groups(
+ user_api_key_dict=self._user(role=LitellmUserRoles.PROXY_ADMIN.value),
+ team_table=self._team([]),
+ access_group_ids=["ag-1"],
+ )
+
+ def test_personal_key_out_of_scope(self):
+ """Personal (non-team) keys are not gated by team-member permissions."""
+ TeamMemberPermissionChecks.enforce_member_can_assign_access_groups(
+ user_api_key_dict=self._user(),
+ team_table=None,
+ access_group_ids=["ag-1"],
+ )
+
+ def test_team_admin_bypasses(self, monkeypatch):
+ """Team admins may assign access groups even without the opt-in perm."""
+ from litellm.proxy.management_endpoints import key_management_endpoints
+
+ monkeypatch.setattr(
+ key_management_endpoints,
+ "_get_user_in_team",
+ lambda **kwargs: Member(role="admin", user_id="user-a"),
+ )
+
+ TeamMemberPermissionChecks.enforce_member_can_assign_access_groups(
+ user_api_key_dict=self._user(),
+ team_table=self._team([]),
+ access_group_ids=["ag-1"],
+ )
+
+ def test_member_denied_without_opt_in(self, monkeypatch):
+ """A non-admin member without the opt-in permission gets a 403."""
+ from fastapi import HTTPException
+
+ from litellm.proxy.management_endpoints import key_management_endpoints
+
+ monkeypatch.setattr(
+ key_management_endpoints,
+ "_get_user_in_team",
+ lambda **kwargs: Member(role="user", user_id="user-a"),
+ )
+
+ with pytest.raises(HTTPException) as exc:
+ TeamMemberPermissionChecks.enforce_member_can_assign_access_groups(
+ user_api_key_dict=self._user(),
+ team_table=self._team(["/key/generate", "/key/update"]),
+ access_group_ids=["ag-1"],
+ )
+ assert exc.value.status_code == 403
+ assert self.AG_PERMISSION in str(exc.value.detail)
+
+ def test_member_allowed_with_opt_in(self, monkeypatch):
+ """A non-admin member is allowed once the team opts in via the perm."""
+ from litellm.proxy.management_endpoints import key_management_endpoints
+
+ monkeypatch.setattr(
+ key_management_endpoints,
+ "_get_user_in_team",
+ lambda **kwargs: Member(role="user", user_id="user-a"),
+ )
+
+ TeamMemberPermissionChecks.enforce_member_can_assign_access_groups(
+ user_api_key_dict=self._user(),
+ team_table=self._team(["/key/generate", self.AG_PERMISSION]),
+ access_group_ids=["ag-1"],
+ )
diff --git a/ui/litellm-dashboard/src/components/team/permission_definitions.tsx b/ui/litellm-dashboard/src/components/team/permission_definitions.tsx
index 4a48baec32d..bcd910e3775 100644
--- a/ui/litellm-dashboard/src/components/team/permission_definitions.tsx
+++ b/ui/litellm-dashboard/src/components/team/permission_definitions.tsx
@@ -20,6 +20,8 @@ export const PERMISSION_DESCRIPTIONS: Record = {
"/key/list": "Member can list virtual keys belonging to this team",
"/key/block": "Member can block a virtual key belonging to this team",
"/key/unblock": "Member can unblock a virtual key belonging to this team",
+ "/key/access_group_assignment":
+ "Member can assign access groups to virtual keys for this team",
"/team/daily/activity":
"Member can view all team usage data (not just their own)",
"/spend/logs":
From 7d1bd9d9f4c313dc1e2cd4aeec5e70a199ee6529 Mon Sep 17 00:00:00 2001
From: yuneng-jiang
Date: Sat, 30 May 2026 17:48:16 -0700
Subject: [PATCH 37/44] fix(reset_budget): write only {spend, budget_reset_at}
and stop pre-zeroing counter (#29358)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
* fix(reset_budget): write only {spend, budget_reset_at} and stop pre-zeroing counter
ResetBudgetJob's batched update_data path shipped the full key/user/team
model on each reset. Prisma rejects object_permission_id and budget_limits
on the update input type, so any row carrying those fields detonated the
entire batch -- spend never reset, budget_reset_at never advanced. After
v1.84.0 started populating object_permission_id on UI-created keys, this
fires routinely.
_reset_budget_common also zeroed the cross-pod spend counter before the
DB write, so failed resets left enforcement reading 0 from the counter
while the DB still held the over-budget spend, admitting requests past
the cap until the counter naturally re-saturated from new reservations.
Switch the write to per-row narrow updates ({spend, budget_reset_at})
via db.batch_, and move the counter invalidation out of
_reset_budget_common so it only fires after the DB write commits. On
DB-write failure the counter is left untouched, enforcement continues
to block, and the next scheduler tick can retry without leaving a
bypass window.
Fixes #27730.
* fix(reset_budget): address Greptile review on #29358
- Strengthen the bypass-half regression test: replace the for-loop over
call_args_list (vacuously true when empty) with assert_not_called(),
so the test would actually flag a re-introduction of counter-zeroing
via any code path.
- Add the same explanatory docstring on _write_user_reset_updates and
_write_team_reset_updates that _write_key_reset_updates already has,
so all three helpers point future maintainers at #27730.
* test(reset_budget): update test_proxy_budget_reset for new batch-write path
Same shape as the previous test_reset_budget_job.py update: keys/users/teams
now write through prisma.db.batch_()..update, not update_data, so the
tests need a batcher mock and updated assertions. Adds:
- _wire_batcher_for_test helper that returns a list which accumulates per-row
batch updates captured from prisma_client.db.batch_().
- _attrify helper that wraps dict fixtures so getattr(item, "token") works
alongside the dict item-access the fake_reset_* mocks rely on. The new
narrow-write helpers use getattr to pull out the row's id, and would
silently skip plain dicts otherwise.
- Updates 3 partial_failure tests to assert against the batch-call list
(rows by id, payload contains only {spend, budget_reset_at}) instead of
update_data.assert_awaited_once + data_list inspection.
- Updates test_reset_budget_continues_other_categories_on_failure: only
budget + enduser still flow through update_data; key/user/team go through
the batch path now.
- Wires the batcher mock into 3 service_logger_*_success tests so commit()
is actually awaitable and the success hook fires.
These tests were silently passing locally only because the editable install
in .venv pointed at the main repo, not the worktree — running pytest with
PYTHONPATH overridden to the worktree (matching CI) reproduces the failures.
---
.../proxy/common_utils/reset_budget_job.py | 129 ++++++-----
.../test_proxy_budget_reset.py | 201 ++++++++++++-----
.../common_utils/test_reset_budget_job.py | 211 ++++++++++++++++--
3 files changed, 409 insertions(+), 132 deletions(-)
diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py
index 52bbeaf2ad3..40c8caa49e5 100644
--- a/litellm/proxy/common_utils/reset_budget_job.py
+++ b/litellm/proxy/common_utils/reset_budget_job.py
@@ -414,6 +414,72 @@ class ResetBudgetJob:
)
return [LiteLLM_EndUserTable(**row.dict()) for row in rows]
+ async def _write_key_reset_updates(
+ self, updated_keys: List[LiteLLM_VerificationToken]
+ ) -> None:
+ """
+ Write per-row {spend, budget_reset_at} updates for keys.
+
+ Avoids the batched full-model update path, which trips
+ prisma.errors.DataError on any row carrying object_permission_id or
+ budget_limits (see #27730). Both fields are rejected by Prisma's
+ update input type for LiteLLM_VerificationToken, and the failure
+ aborts the entire batch — silently leaving spend over the cap and
+ budget_reset_at unchanged forever.
+ """
+ batcher = self.prisma_client.db.batch_()
+ for k in updated_keys:
+ token = getattr(k, "token", None)
+ if token is None:
+ continue
+ batcher.litellm_verificationtoken.update(
+ where={"token": token},
+ data={"spend": 0, "budget_reset_at": k.budget_reset_at},
+ )
+ await batcher.commit()
+
+ async def _write_user_reset_updates(
+ self, updated_users: List[LiteLLM_UserTable]
+ ) -> None:
+ """
+ Write per-row {spend, budget_reset_at} updates for users.
+
+ Mirrors _write_key_reset_updates — avoids the full-model update path
+ that trips Prisma's DataError on rows carrying unrecognised fields
+ (see #27730).
+ """
+ batcher = self.prisma_client.db.batch_()
+ for u in updated_users:
+ user_id = getattr(u, "user_id", None)
+ if user_id is None:
+ continue
+ batcher.litellm_usertable.update(
+ where={"user_id": user_id},
+ data={"spend": 0, "budget_reset_at": u.budget_reset_at},
+ )
+ await batcher.commit()
+
+ async def _write_team_reset_updates(
+ self, updated_teams: List[LiteLLM_TeamTable]
+ ) -> None:
+ """
+ Write per-row {spend, budget_reset_at} updates for teams.
+
+ Mirrors _write_key_reset_updates — avoids the full-model update path
+ that trips Prisma's DataError on rows carrying unrecognised fields
+ (see #27730).
+ """
+ batcher = self.prisma_client.db.batch_()
+ for t in updated_teams:
+ team_id = getattr(t, "team_id", None)
+ if team_id is None:
+ continue
+ batcher.litellm_teamtable.update(
+ where={"team_id": team_id},
+ data={"spend": 0, "budget_reset_at": t.budget_reset_at},
+ )
+ await batcher.commit()
+
async def reset_budget_for_litellm_keys(self):
"""
Resets the budget for all the litellm keys
@@ -455,11 +521,7 @@ class ResetBudgetJob:
)
if updated_keys:
- await self.prisma_client.update_data(
- query_type="update_many",
- data_list=updated_keys,
- table_name="key",
- )
+ await self._write_key_reset_updates(updated_keys=updated_keys)
for k in updated_keys:
token = getattr(k, "token", None)
if token:
@@ -544,11 +606,7 @@ class ResetBudgetJob:
"Updated users %s", json.dumps(updated_users, indent=4, default=str)
)
if updated_users:
- await self.prisma_client.update_data(
- query_type="update_many",
- data_list=updated_users,
- table_name="user",
- )
+ await self._write_user_reset_updates(updated_users=updated_users)
for u in updated_users:
user_id = getattr(u, "user_id", None)
if user_id:
@@ -641,11 +699,7 @@ class ResetBudgetJob:
"Updated teams %s", json.dumps(updated_teams, indent=4, default=str)
)
if updated_teams:
- await self.prisma_client.update_data(
- query_type="update_many",
- data_list=updated_teams,
- table_name="team",
- )
+ await self._write_team_reset_updates(updated_teams=updated_teams)
for t in updated_teams:
team_id = getattr(t, "team_id", None)
if team_id:
@@ -816,49 +870,16 @@ class ResetBudgetJob:
"""
In-place, updates spend=0, and sets budget_reset_at to current_time + budget_duration
- Common logic for resetting budget for a team, user, or key
+ Common logic for resetting budget for a team, user, or key.
+
+ Spend-counter invalidation happens in the caller, AFTER the DB write
+ commits. Zeroing the counter here would open a bypass window when the
+ DB write fails: get_current_spend reads 0 from Redis while the DB
+ still holds the pre-reset value, admitting requests past the cap.
"""
try:
item.spend = 0.0
-
- # Reset the cross-pod spend counter.
- # Reset Redis directly (not via DualCache) so a Redis failure
- # doesn't silently leave a stale counter that get_current_spend
- # would read as authoritative, permanently blocking the user.
- from litellm.proxy.proxy_server import spend_counter_cache
-
- counter_key = None
- if item_type == "key" and hasattr(item, "token") and item.token is not None: # type: ignore[union-attr]
- counter_key = f"spend:key:{item.token}" # type: ignore[union-attr]
- elif (
- item_type == "team"
- and hasattr(item, "team_id")
- and item.team_id is not None # type: ignore[union-attr]
- ):
- counter_key = f"spend:team:{item.team_id}" # type: ignore[union-attr]
-
- if counter_key is not None:
- # Always reset in-memory (local fallback)
- spend_counter_cache.in_memory_cache.set_cache(
- key=counter_key, value=0.0
- )
- # Explicitly reset Redis with warning on failure
- if spend_counter_cache.redis_cache is not None:
- try:
- await spend_counter_cache.redis_cache.async_set_cache(
- key=counter_key, value=0.0
- )
- except Exception as redis_err:
- verbose_proxy_logger.warning(
- "Failed to reset spend counter in Redis for %s key=%s: %s. "
- "Budget may be over-enforced until counter expires.",
- item_type,
- counter_key,
- redis_err,
- )
-
if hasattr(item, "budget_duration") and item.budget_duration is not None:
- # Get standardized reset time based on budget duration
from litellm.proxy.common_utils.timezone_utils import (
get_budget_reset_time,
)
diff --git a/tests/litellm_utils_tests/test_proxy_budget_reset.py b/tests/litellm_utils_tests/test_proxy_budget_reset.py
index 6240bedd3e6..5c96eb619bf 100644
--- a/tests/litellm_utils_tests/test_proxy_budget_reset.py
+++ b/tests/litellm_utils_tests/test_proxy_budget_reset.py
@@ -22,6 +22,60 @@ from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob
# In a real-world scenario, these would be instances of LiteLLM_VerificationToken, LiteLLM_UserTable, etc.
+def _attrify(d: dict):
+ """
+ Wrap a dict so that attribute access (`.token`, `.user_id`, `.team_id`,
+ etc.) works alongside the existing item-access the fake_reset_* helpers
+ rely on. The reset job's narrow-write helpers use `getattr(item, "token",
+ None)` (et al), which returns None for plain dicts — that would silently
+ skip the row.
+ """
+ class _AttrDict(dict):
+ def __getattr__(self, k):
+ try:
+ return self[k]
+ except KeyError:
+ raise AttributeError(k)
+
+ def __setattr__(self, k, v):
+ self[k] = v
+
+ return _AttrDict(d)
+
+
+def _wire_batcher_for_test(prisma_client):
+ """
+ Wire prisma_client.db.batch_() to return a mock batcher whose .commit() is
+ awaitable and whose per-table .update() calls get captured. The reset job
+ writes key/user/team resets via prisma.db.batch_()..update — not via
+ prisma_client.update_data — so tests must let that batch path complete.
+
+ Returns the list that will accumulate {table, where, data} dicts from
+ each captured update call.
+ """
+ batch_calls = []
+
+ def make_batcher():
+ class _Table:
+ def __init__(self, table_name):
+ self._table_name = table_name
+
+ def update(self, where=None, data=None):
+ batch_calls.append(
+ {"table": self._table_name, "where": where, "data": data}
+ )
+
+ batcher = MagicMock()
+ batcher.litellm_verificationtoken = _Table("key")
+ batcher.litellm_usertable = _Table("user")
+ batcher.litellm_teamtable = _Table("team")
+ batcher.commit = AsyncMock(return_value=None)
+ return batcher
+
+ prisma_client.db.batch_ = MagicMock(side_effect=make_batcher)
+ return batch_calls
+
+
@pytest.mark.asyncio
async def test_reset_budget_keys_partial_failure():
"""
@@ -45,6 +99,9 @@ async def test_reset_budget_keys_partial_failure():
return_value=[key1, key2, key3, key4, key5, key6]
)
prisma_client.update_data = AsyncMock()
+ # Reset job writes key resets via prisma.db.batch_()..update — not
+ # via update_data — so wire that path.
+ batch_calls = _wire_batcher_for_test(prisma_client)
# Using a dummy logging object with async hooks mocked out.
proxy_logging_obj = MagicMock()
@@ -56,6 +113,15 @@ async def test_reset_budget_keys_partial_failure():
now = datetime.utcnow()
+ # token is needed because the new write path uses where={"token": ...}
+ # and _AttrDict makes getattr work alongside item access used by fake_reset_key.
+ for k in [key1, key2, key3, key4, key5, key6]:
+ k.setdefault("token", k["id"])
+ key1, key2, key3, key4, key5, key6 = (
+ _attrify(k) for k in [key1, key2, key3, key4, key5, key6]
+ )
+ prisma_client.get_data = AsyncMock(return_value=[key1, key2, key3, key4, key5, key6])
+
async def fake_reset_key(key, current_time):
if key["id"] == "key1":
# Simulate a failure on key1 (for example, this might be due to an invariant check)
@@ -80,17 +146,17 @@ async def test_reset_budget_keys_partial_failure():
# Assert that the helper was called for 6 keys
assert mock_reset_key.call_count == 6
- # Assert that update_data was called once with a list containing all 6 keys
- prisma_client.update_data.assert_awaited_once()
- update_call = prisma_client.update_data.call_args
- assert update_call.kwargs.get("table_name") == "key"
- updated_keys = update_call.kwargs.get("data_list", [])
- assert len(updated_keys) == 5
- assert updated_keys[0]["id"] == "key2"
- assert updated_keys[1]["id"] == "key3"
- assert updated_keys[2]["id"] == "key4"
- assert updated_keys[3]["id"] == "key5"
- assert updated_keys[4]["id"] == "key6"
+ # Assert that the new narrow write path got 5 batched updates (key1 failed).
+ # update_data must NOT have been called for keys.
+ prisma_client.update_data.assert_not_awaited()
+ key_writes = [c for c in batch_calls if c["table"] == "key"]
+ assert len(key_writes) == 5
+ written_ids = [c["where"]["token"] for c in key_writes]
+ assert written_ids == ["key2", "key3", "key4", "key5", "key6"]
+ # And every write must carry only {spend, budget_reset_at} — never the full row.
+ for c in key_writes:
+ assert set(c["data"].keys()) == {"spend", "budget_reset_at"}
+ assert c["data"]["spend"] == 0
# Verify that the failure logging hook was scheduled (due to the failure for key1)
failure_hook_calls = (
@@ -125,6 +191,7 @@ async def test_reset_budget_users_partial_failure():
return_value=[user1, user2, user3, user4, user5, user6]
)
prisma_client.update_data = AsyncMock()
+ batch_calls = _wire_batcher_for_test(prisma_client)
proxy_logging_obj = MagicMock()
proxy_logging_obj.service_logging_obj = MagicMock()
@@ -133,6 +200,15 @@ async def test_reset_budget_users_partial_failure():
job = ResetBudgetJob(proxy_logging_obj, prisma_client)
+ # user_id required for the new write path's where clause; _AttrDict so
+ # getattr(u, 'user_id') works alongside the dict access fake_reset_user uses.
+ for u in [user1, user2, user3, user4, user5, user6]:
+ u.setdefault("user_id", u["id"])
+ user1, user2, user3, user4, user5, user6 = (
+ _attrify(u) for u in [user1, user2, user3, user4, user5, user6]
+ )
+ prisma_client.get_data = AsyncMock(return_value=[user1, user2, user3, user4, user5, user6])
+
async def fake_reset_user(user, current_time):
if user["id"] == "user1":
raise Exception("Simulated failure for user1")
@@ -150,16 +226,14 @@ async def test_reset_budget_users_partial_failure():
await asyncio.sleep(0.1)
assert mock_reset_user.call_count == 6
- prisma_client.update_data.assert_awaited_once()
- update_call = prisma_client.update_data.call_args
- assert update_call.kwargs.get("table_name") == "user"
- updated_users = update_call.kwargs.get("data_list", [])
- assert len(updated_users) == 5
- assert updated_users[0]["id"] == "user2"
- assert updated_users[1]["id"] == "user3"
- assert updated_users[2]["id"] == "user4"
- assert updated_users[3]["id"] == "user5"
- assert updated_users[4]["id"] == "user6"
+ prisma_client.update_data.assert_not_awaited()
+ user_writes = [c for c in batch_calls if c["table"] == "user"]
+ assert len(user_writes) == 5
+ written_ids = [c["where"]["user_id"] for c in user_writes]
+ assert written_ids == ["user2", "user3", "user4", "user5", "user6"]
+ for c in user_writes:
+ assert set(c["data"].keys()) == {"spend", "budget_reset_at"}
+ assert c["data"]["spend"] == 0
failure_hook_calls = (
proxy_logging_obj.service_logging_obj.async_service_failure_hook.call_args_list
@@ -308,6 +382,7 @@ async def test_reset_budget_teams_partial_failure():
prisma_client = MagicMock()
prisma_client.get_data = AsyncMock(return_value=[team1, team2])
prisma_client.update_data = AsyncMock()
+ batch_calls = _wire_batcher_for_test(prisma_client)
proxy_logging_obj = MagicMock()
proxy_logging_obj.service_logging_obj = MagicMock()
@@ -316,6 +391,12 @@ async def test_reset_budget_teams_partial_failure():
job = ResetBudgetJob(proxy_logging_obj, prisma_client)
+ # team_id required for the new write path's where clause; _AttrDict for getattr.
+ for t in [team1, team2]:
+ t.setdefault("team_id", t["id"])
+ team1, team2 = _attrify(team1), _attrify(team2)
+ prisma_client.get_data = AsyncMock(return_value=[team1, team2])
+
async def fake_reset_team(team, current_time):
if team["id"] == "team1":
raise Exception("Simulated failure for team1")
@@ -333,12 +414,12 @@ async def test_reset_budget_teams_partial_failure():
await asyncio.sleep(0.1)
assert mock_reset_team.call_count == 2
- prisma_client.update_data.assert_awaited_once()
- update_call = prisma_client.update_data.call_args
- assert update_call.kwargs.get("table_name") == "team"
- updated_teams = update_call.kwargs.get("data_list", [])
- assert len(updated_teams) == 1
- assert updated_teams[0]["id"] == "team2"
+ prisma_client.update_data.assert_not_awaited()
+ team_writes = [c for c in batch_calls if c["table"] == "team"]
+ assert len(team_writes) == 1
+ assert team_writes[0]["where"] == {"team_id": "team2"}
+ assert set(team_writes[0]["data"].keys()) == {"spend", "budget_reset_at"}
+ assert team_writes[0]["data"]["spend"] == 0
failure_hook_calls = (
proxy_logging_obj.service_logging_obj.async_service_failure_hook.call_args_list
@@ -402,6 +483,18 @@ async def test_reset_budget_continues_other_categories_on_failure():
prisma_client.get_data = AsyncMock(side_effect=fake_get_data)
prisma_client.update_data = AsyncMock()
+ batch_calls = _wire_batcher_for_test(prisma_client)
+ # ID fields required by the new write path's where clauses; _AttrDict
+ # lets getattr() see them alongside the item-access fake_reset_* helpers use.
+ for k in [key1, key2]:
+ k.setdefault("token", k["id"])
+ for u in [user1, user2]:
+ u.setdefault("user_id", u["id"])
+ for t in [team1, team2]:
+ t.setdefault("team_id", t["id"])
+ key1, key2 = _attrify(key1), _attrify(key2)
+ user1, user2 = _attrify(user1), _attrify(user2)
+ team1, team2 = _attrify(team1), _attrify(team2)
# Mock db.litellm_verificationtoken.update_many (used by reset_budget_for_keys_linked_to_budgets)
prisma_client.db.litellm_verificationtoken.update_many = AsyncMock(
return_value={"count": 0}
@@ -488,32 +581,29 @@ async def test_reset_budget_continues_other_categories_on_failure():
"team_membership",
}
- # Verify that update_data was called three times (one per category, enduser update includes two)
- assert prisma_client.update_data.await_count == 5
+ # After the fix, keys/users/teams write via prisma.db.batch_()..update,
+ # so only budget + enduser still go through update_data.
calls = prisma_client.update_data.await_args_list
-
- # Check keys update: both keys succeed.
- keys_call = calls[0]
- assert keys_call.kwargs.get("table_name") == "key"
- assert len(keys_call.kwargs.get("data_list", [])) == 2
-
- # Check users update: only user2 succeeded.
- users_call = calls[1]
- assert users_call.kwargs.get("table_name") == "user"
- users_updated = users_call.kwargs.get("data_list", [])
- assert len(users_updated) == 1
- assert users_updated[0]["id"] == "user2"
-
- # Check teams update: both teams succeed.
- teams_call = calls[2]
- assert teams_call.kwargs.get("table_name") == "team"
- assert len(teams_call.kwargs.get("data_list", [])) == 2
+ update_data_tables = [c.kwargs.get("table_name") for c in calls]
+ assert sorted(update_data_tables) == ["budget", "enduser"]
# Check enduser update: enduser succeed.
- enduser_call = calls[4]
- assert enduser_call.kwargs.get("table_name") == "enduser"
+ enduser_call = next(c for c in calls if c.kwargs.get("table_name") == "enduser")
assert len(enduser_call.kwargs.get("data_list", [])) == 1
+ # Check the new batch write path: 2 keys + 1 user (user1 failed) + 2 teams.
+ key_writes = [c for c in batch_calls if c["table"] == "key"]
+ user_writes = [c for c in batch_calls if c["table"] == "user"]
+ team_writes = [c for c in batch_calls if c["table"] == "team"]
+ assert len(key_writes) == 2
+ assert len(user_writes) == 1
+ assert user_writes[0]["where"] == {"user_id": "user2"}
+ assert len(team_writes) == 2
+ # Every batched write must carry only the two reset fields, never the full row.
+ for c in key_writes + user_writes + team_writes:
+ assert set(c["data"].keys()) == {"spend", "budget_reset_at"}
+ assert c["data"]["spend"] == 0
+
# ---------------------------------------------------------------------------
# Additional tests for service logger behavior (keys, users, teams, endusers)
@@ -527,12 +617,13 @@ async def test_service_logger_keys_success():
logger success hook is called with the correct event metadata and no exception is logged.
"""
keys = [
- {"id": "key1", "spend": 10.0, "budget_duration": 60},
- {"id": "key2", "spend": 15.0, "budget_duration": 60},
+ {"id": "key1", "spend": 10.0, "budget_duration": 60, "token": "key1"},
+ {"id": "key2", "spend": 15.0, "budget_duration": 60, "token": "key2"},
]
prisma_client = MagicMock()
prisma_client.get_data = AsyncMock(return_value=keys)
prisma_client.update_data = AsyncMock()
+ _wire_batcher_for_test(prisma_client)
proxy_logging_obj = MagicMock()
proxy_logging_obj.service_logging_obj = MagicMock()
@@ -644,12 +735,13 @@ async def test_service_logger_users_success():
the correct metadata and no exception is logged.
"""
users = [
- {"id": "user1", "spend": 20.0, "budget_duration": 120},
- {"id": "user2", "spend": 25.0, "budget_duration": 120},
+ {"id": "user1", "spend": 20.0, "budget_duration": 120, "user_id": "user1"},
+ {"id": "user2", "spend": 25.0, "budget_duration": 120, "user_id": "user2"},
]
prisma_client = MagicMock()
prisma_client.get_data = AsyncMock(return_value=users)
prisma_client.update_data = AsyncMock()
+ _wire_batcher_for_test(prisma_client)
proxy_logging_obj = MagicMock()
proxy_logging_obj.service_logging_obj = MagicMock()
@@ -756,12 +848,13 @@ async def test_service_logger_teams_success():
the proper metadata and nothing is logged as an exception.
"""
teams = [
- {"id": "team1", "spend": 30.0, "budget_duration": 180},
- {"id": "team2", "spend": 35.0, "budget_duration": 180},
+ {"id": "team1", "spend": 30.0, "budget_duration": 180, "team_id": "team1"},
+ {"id": "team2", "spend": 35.0, "budget_duration": 180, "team_id": "team2"},
]
prisma_client = MagicMock()
prisma_client.get_data = AsyncMock(return_value=teams)
prisma_client.update_data = AsyncMock()
+ _wire_batcher_for_test(prisma_client)
proxy_logging_obj = MagicMock()
proxy_logging_obj.service_logging_obj = MagicMock()
diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py
index 8a47c78db05..0b683745369 100644
--- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py
+++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py
@@ -92,6 +92,37 @@ class MockLiteLLMEndUserTable:
return self._find_many_results
+class MockBatcher:
+ """Captures per-row update calls and exposes them after commit().
+
+ Mirrors prisma's `db.batch_()` ergonomics enough that the reset job's
+ narrow-write helpers (`_write_key_reset_updates` et al) can run against
+ the mock and the test can assert on what would have been written.
+ """
+
+ def __init__(self):
+ self.calls: List[Dict[str, Any]] = []
+ self.committed: bool = False
+
+ class _Table:
+ def __init__(_self, table_name: str, outer: "MockBatcher"):
+ _self._table_name = table_name
+ _self._outer = outer
+
+ def update(_self, where, data):
+ _self._outer.calls.append(
+ {"table": _self._table_name, "where": where, "data": data}
+ )
+
+ self.litellm_verificationtoken = _Table("key", self)
+ self.litellm_usertable = _Table("user", self)
+ self.litellm_teamtable = _Table("team", self)
+
+ async def commit(self):
+ self.committed = True
+ return self.calls
+
+
class MockDB:
def __init__(self):
self.litellm_teammembership = MockLiteLLMTeamMembership()
@@ -99,6 +130,19 @@ class MockDB:
self.litellm_endusertable = MockLiteLLMEndUserTable()
self.litellm_organizationtable = MockLiteLLMOrganizationTable()
self.litellm_tagtable = MockLiteLLMTagTable()
+ self.batch_calls: List[Dict[str, Any]] = []
+
+ def batch_(self):
+ batcher = MockBatcher()
+ # Aggregate calls across all batches so tests can assert on cumulative writes.
+ original_commit = batcher.commit
+
+ async def _record_and_commit():
+ self.batch_calls.extend(batcher.calls)
+ return await original_commit()
+
+ batcher.commit = _record_and_commit # type: ignore[assignment]
+ return batcher
class MockPrismaClient:
@@ -205,6 +249,7 @@ def test_reset_budget_for_key(reset_budget_job, mock_prisma_client):
"budget_duration": "30d",
"budget_reset_at": now,
"id": "test-key-1",
+ "token": "tok-key-1",
},
)
@@ -213,11 +258,16 @@ def test_reset_budget_for_key(reset_budget_job, mock_prisma_client):
# Run the test
asyncio.run(reset_budget_job.reset_budget_for_litellm_keys())
- # Verify results
- assert len(mock_prisma_client.updated_data["key"]) == 1
- updated_key = mock_prisma_client.updated_data["key"][0]
- assert updated_key.spend == 0.0
- assert updated_key.budget_reset_at > now
+ # The reset writes only {spend, budget_reset_at} per row via batch_().
+ # Full-row writes would re-detonate the Prisma DataError on rows carrying
+ # object_permission_id / budget_limits (see #27730).
+ key_writes = [c for c in mock_prisma_client.db.batch_calls if c["table"] == "key"]
+ assert len(key_writes) == 1
+ write = key_writes[0]
+ assert write["where"] == {"token": "tok-key-1"}
+ assert write["data"]["spend"] == 0
+ assert write["data"]["budget_reset_at"] > now
+ assert set(write["data"].keys()) == {"spend", "budget_reset_at"}
def test_reset_budget_for_user(reset_budget_job, mock_prisma_client):
@@ -231,6 +281,7 @@ def test_reset_budget_for_user(reset_budget_job, mock_prisma_client):
"budget_duration": "7d",
"budget_reset_at": now,
"id": "test-user-1",
+ "user_id": "uid-1",
},
)
@@ -239,11 +290,13 @@ def test_reset_budget_for_user(reset_budget_job, mock_prisma_client):
# Run the test
asyncio.run(reset_budget_job.reset_budget_for_litellm_users())
- # Verify results
- assert len(mock_prisma_client.updated_data["user"]) == 1
- updated_user = mock_prisma_client.updated_data["user"][0]
- assert updated_user.spend == 0.0
- assert updated_user.budget_reset_at > now
+ user_writes = [c for c in mock_prisma_client.db.batch_calls if c["table"] == "user"]
+ assert len(user_writes) == 1
+ write = user_writes[0]
+ assert write["where"] == {"user_id": "uid-1"}
+ assert write["data"]["spend"] == 0
+ assert write["data"]["budget_reset_at"] > now
+ assert set(write["data"].keys()) == {"spend", "budget_reset_at"}
def test_reset_budget_for_team(reset_budget_job, mock_prisma_client):
@@ -257,6 +310,7 @@ def test_reset_budget_for_team(reset_budget_job, mock_prisma_client):
"budget_duration": "1mo",
"budget_reset_at": now,
"id": "test-team-1",
+ "team_id": "tid-1",
},
)
@@ -265,11 +319,13 @@ def test_reset_budget_for_team(reset_budget_job, mock_prisma_client):
# Run the test
asyncio.run(reset_budget_job.reset_budget_for_litellm_teams())
- # Verify results
- assert len(mock_prisma_client.updated_data["team"]) == 1
- updated_team = mock_prisma_client.updated_data["team"][0]
- assert updated_team.spend == 0.0
- assert updated_team.budget_reset_at > now
+ team_writes = [c for c in mock_prisma_client.db.batch_calls if c["table"] == "team"]
+ assert len(team_writes) == 1
+ write = team_writes[0]
+ assert write["where"] == {"team_id": "tid-1"}
+ assert write["data"]["spend"] == 0
+ assert write["data"]["budget_reset_at"] > now
+ assert set(write["data"].keys()) == {"spend", "budget_reset_at"}
def test_reset_budget_for_enduser(reset_budget_job, mock_prisma_client):
@@ -324,6 +380,7 @@ def test_reset_budget_all(reset_budget_job, mock_prisma_client):
"budget_duration": "30d",
"budget_reset_at": now,
"id": "test-key-1",
+ "token": "tok-all-1",
},
)
@@ -335,6 +392,7 @@ def test_reset_budget_all(reset_budget_job, mock_prisma_client):
"budget_duration": "7d",
"budget_reset_at": now,
"id": "test-user-1",
+ "user_id": "uid-all-1",
},
)
@@ -346,6 +404,7 @@ def test_reset_budget_all(reset_budget_job, mock_prisma_client):
"budget_duration": "1mo",
"budget_reset_at": now,
"id": "test-team-1",
+ "team_id": "tid-all-1",
},
)
@@ -379,17 +438,22 @@ def test_reset_budget_all(reset_budget_job, mock_prisma_client):
# Run the test
asyncio.run(reset_budget_job.reset_budget())
- # Verify results
- assert len(mock_prisma_client.updated_data["key"]) == 1
- assert len(mock_prisma_client.updated_data["user"]) == 1
- assert len(mock_prisma_client.updated_data["team"]) == 1
+ # key/user/team rows are written via batch_()..update — verify each
+ # one fired exactly once with the narrow {spend, budget_reset_at} payload.
+ for table_name, where in [
+ ("key", {"token": "tok-all-1"}),
+ ("user", {"user_id": "uid-all-1"}),
+ ("team", {"team_id": "tid-all-1"}),
+ ]:
+ writes = [c for c in mock_prisma_client.db.batch_calls if c["table"] == table_name]
+ assert len(writes) == 1, f"expected 1 {table_name} write, got {len(writes)}"
+ assert writes[0]["where"] == where
+ assert writes[0]["data"]["spend"] == 0
+ assert set(writes[0]["data"].keys()) == {"spend", "budget_reset_at"}
+
+ # Enduser + budget rows still go through update_data (not narrowed; different path).
assert len(mock_prisma_client.updated_data["enduser"]) == 1
assert len(mock_prisma_client.updated_data["budget"]) == 1
-
- # Check that all spends were reset to 0
- assert mock_prisma_client.updated_data["key"][0].spend == 0.0
- assert mock_prisma_client.updated_data["user"][0].spend == 0.0
- assert mock_prisma_client.updated_data["team"][0].spend == 0.0
assert mock_prisma_client.updated_data["enduser"][0].spend == 0.0
@@ -1399,6 +1463,105 @@ def test_reset_budget_for_teams_invalidates_redis_counter(
)
+def test_reset_does_not_zero_counter_when_db_write_fails(monkeypatch):
+ """
+ Regression for #27730 (the bypass-half).
+
+ If the DB write inside the reset job raises (e.g. Prisma DataError on a
+ row carrying object_permission_id or budget_limits), the Redis spend
+ counter MUST NOT be zeroed — that would let get_current_spend admit
+ requests past the cap while the DB row still holds the over-budget
+ spend.
+
+ Pre-fix: _reset_budget_common pre-zeroed the counter before the DB
+ write attempt, opening the bypass window.
+ Post-fix: counter invalidation lives in the caller, AFTER the DB write
+ commits. If the write raises, the post-write invalidation never runs.
+ """
+ counter_cache = _make_counter_invalidation_job(monkeypatch)
+
+ now = datetime.now(timezone.utc)
+ prisma_client = MagicMock()
+
+ matching_key = type(
+ "Key",
+ (),
+ {
+ "spend": 100.0,
+ "budget_duration": "30d",
+ "budget_reset_at": now - timedelta(seconds=1),
+ "token": "sk-failing",
+ },
+ )
+
+ # get_data returns one key needing reset; the batched DB write then explodes.
+ async def fake_get_data(table_name, query_type, **kwargs):
+ if table_name == "key":
+ return [matching_key]
+ return []
+
+ prisma_client.get_data = fake_get_data
+
+ batcher = MagicMock()
+ batcher.litellm_verificationtoken.update = MagicMock()
+
+ async def failing_commit():
+ raise RuntimeError("simulated Prisma DataError on update")
+
+ batcher.commit = failing_commit
+ prisma_client.db.batch_ = MagicMock(return_value=batcher)
+
+ job = ResetBudgetJob(
+ proxy_logging_obj=MockProxyLogging(), prisma_client=prisma_client
+ )
+
+ asyncio.run(job.reset_budget_for_litellm_keys())
+
+ # CRITICAL: counter invalidation must NOT have been called at all —
+ # the DB write raised before the post-write invalidation loop. Using
+ # assert_not_called() instead of iterating call_args_list, because the
+ # latter is vacuously true when the list is empty (would pass even if
+ # the bypass were re-introduced via a different code path).
+ counter_cache.in_memory_cache.set_cache.assert_not_called()
+
+
+def test_reset_budget_for_keys_writes_only_spend_and_reset_at(reset_budget_job, mock_prisma_client):
+ """
+ Regression for #27730 (the trigger-half).
+
+ The reset job must write only {spend, budget_reset_at} per row — never
+ the full key object. Sending the full object via the old update_data
+ batcher path made Prisma reject any row carrying object_permission_id
+ or budget_limits (both became non-NULL on UI-created keys after v1.84.0).
+ """
+ now = datetime.now(timezone.utc)
+ key_with_problematic_fields = type(
+ "LiteLLM_VerificationToken",
+ (),
+ {
+ "spend": 50.0,
+ "budget_duration": "30d",
+ "budget_reset_at": now,
+ "token": "sk-problematic",
+ "object_permission_id": "perm-abc", # would be rejected on update
+ "budget_limits": [{"max_budget": 5}], # would be rejected on update
+ "metadata": {"some": "thing"},
+ },
+ )
+ mock_prisma_client.data["key"] = [key_with_problematic_fields]
+
+ asyncio.run(reset_budget_job.reset_budget_for_litellm_keys())
+
+ key_writes = [c for c in mock_prisma_client.db.batch_calls if c["table"] == "key"]
+ assert len(key_writes) == 1
+ payload_keys = set(key_writes[0]["data"].keys())
+ assert payload_keys == {"spend", "budget_reset_at"}, (
+ f"reset payload must not include any field besides spend / budget_reset_at, "
+ f"got: {payload_keys}. Any extra field (object_permission_id, budget_limits, etc.) "
+ f"trips Prisma DataError and detonates the whole batch."
+ )
+
+
def test_reset_budget_for_keys_linked_to_budgets_invalidates_redis_counter(monkeypatch):
"""Resetting keys via budget tier must clear each linked key's counter."""
counter_cache = _make_counter_invalidation_job(monkeypatch)
From a9cc6ed68cd77b133bd252fd9d2cea5453d5dea7 Mon Sep 17 00:00:00 2001
From: ryan-crabbe-berri
Test email body
", + ) + + mock_async_client.post.assert_called_once() + request_body = mock_async_client.post.call_args[1]["json"] + assert request_body["from"] == "alerts@my-verified-domain.com" + + +@pytest.mark.asyncio +async def test_send_email_falls_back_to_argument_when_override_unset(mock_env_vars): + """When RESEND_FROM_EMAIL is unset, the caller-supplied from_email is used.""" + logger = ResendEmailLogger() + + mock_response = mock.Mock(spec=Response) + mock_response.status_code = 200 + mock_response.json.return_value = {"id": "test_email_id"} + mock_response.raise_for_status.return_value = None + + mock_async_client = mock.AsyncMock() + mock_async_client.post.return_value = mock_response + logger.async_httpx_client = mock_async_client + + await logger.send_email( + from_email="notifications@alerts.litellm.ai", + to_email=["recipient@example.com"], + subject="Test Subject", + html_body="Test email body
", + ) + + mock_async_client.post.assert_called_once() + request_body = mock_async_client.post.call_args[1]["json"] + assert request_body["from"] == "notifications@alerts.litellm.ai" From f11c12d1574fb51acd958bc7173a74e994aacf19 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 30 May 2026 11:26:24 -0700 Subject: [PATCH 26/44] Revert "chore(tests): migrate Bedrock CI to AWS account 941277531214 (#28728)" (#29326) This reverts the Bedrock CI account migration (#28728). The original account (888602223428) was put under an AWS security restriction after a leaked key and has since been reactivated, while the replacement account (941277531214) lacks access to several models the suites exercise (legacy Bedrock Claude 3 models, Cohere, Nova Canvas image gen, Bedrock batch inference, and flagship Opus). Pointing CI back at the reactivated account restores that coverage. This is the exact inverse of #28728: all hardcoded 941277531214 references go back to 888602223428 (provisioned/imported-model ARNs, AgentCore runtime ARNs and their suffixes, batch execution role ARN, and the example proxy config), the S3 buckets revert to litellm-proxy and load-testing-oct, the guardrail IDs revert to wf0hkdb5x07f and ff6ujrregl1q, the SageMaker endpoint and Knowledge Base revert to their original ids, and the live-call tests go back to the legacy model strings. The grid_spec fail_reason workaround for the unentitled Opus cells is dropped while keeping the unrelated bedrock_effort_ceiling field added after the migration. The CircleCI AWS_ACCESS_KEY_ID / AWS_SECRET_ACCESS_KEY env vars still point at 941277531214 and must be set to the reactivated account's fresh credentials separately via the CircleCI API; AWS_REGION_NAME stays us-west-2. --- .../bedrock/chat/agentcore/transformation.py | 6 +- .../example_config_yaml/oai_misc_config.yaml | 4 +- .../example_config_yaml/otel_test_config.yaml | 2 +- .../test_a2a_completion_bridge.py | 2 +- .../test_bedrock_files_and_batches.py | 8 +-- .../test_bedrock_guardrails.py | 16 ++--- .../test_bedrock_image_gen_unit_tests.py | 35 +++-------- .../image_gen_tests/test_image_generation.py | 58 +------------------ .../test_litellm_overhead.py | 4 +- .../reasoning_effort_grid/grid_spec.py | 8 +-- .../test_reasoning_effort_grid.py | 4 +- .../llm_translation/test_bedrock_agentcore.py | 24 ++++---- .../test_bedrock_completion.py | 20 +++---- tests/local_testing/test_completion.py | 21 ++++--- .../test_function_call_parsing.py | 3 +- tests/local_testing/test_function_calling.py | 9 +-- tests/local_testing/test_sagemaker.py | 28 ++++----- tests/local_testing/test_streaming.py | 8 +-- .../test_amazing_s3_logs.py | 8 +-- .../test_bedrock_knowledgebase_hook.py | 27 ++++----- .../test_agentcore_transformation.py | 10 ++-- tests/test_openai_endpoints.py | 2 +- .../test_bedrock_vector_store.py | 14 ++--- 23 files changed, 118 insertions(+), 203 deletions(-) diff --git a/litellm/llms/bedrock/chat/agentcore/transformation.py b/litellm/llms/bedrock/chat/agentcore/transformation.py index 9b9b96aae04..44ba1ce3c86 100644 --- a/litellm/llms/bedrock/chat/agentcore/transformation.py +++ b/litellm/llms/bedrock/chat/agentcore/transformation.py @@ -157,8 +157,8 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): def _get_agent_runtime_arn(self, model: str) -> str: """ Extract ARN from model string - model = "agentcore/arn:aws:bedrock-agentcore:us-west-2:941277531214:runtime/hosted_agent_r9jvp-Rq79QFC2fp" - returns: "arn:aws:bedrock-agentcore:us-west-2:941277531214:runtime/hosted_agent_r9jvp-Rq79QFC2fp" + model = "agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC" + returns: "arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC" """ parts = model.split("/", 1) if len(parts) != 2 or parts[0] != "agentcore": @@ -170,7 +170,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): def _extract_region_from_arn(self, arn: str) -> str: """ Extract region from ARN - arn:aws:bedrock-agentcore:us-west-2:941277531214:runtime/hosted_agent_r9jvp-Rq79QFC2fp + arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC returns: us-west-2 """ parts = arn.split(":") diff --git a/litellm/proxy/example_config_yaml/oai_misc_config.yaml b/litellm/proxy/example_config_yaml/oai_misc_config.yaml index 0b647de8a08..16cc69c19a5 100644 --- a/litellm/proxy/example_config_yaml/oai_misc_config.yaml +++ b/litellm/proxy/example_config_yaml/oai_misc_config.yaml @@ -23,11 +23,11 @@ model_list: model: bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0 ######################################################### ########## batch specific params ######################## - s3_bucket_name: litellm-proxy-941277531214 + s3_bucket_name: litellm-proxy s3_region_name: us-west-2 s3_access_key_id: os.environ/AWS_ACCESS_KEY_ID s3_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY - aws_batch_role_arn: arn:aws:iam::941277531214:role/service-role/AmazonBedrockExecutionRoleForAgents_BB9HNW6V4CV + aws_batch_role_arn: arn:aws:iam::888602223428:role/service-role/AmazonBedrockExecutionRoleForAgents_BB9HNW6V4CV model_info: mode: batch diff --git a/litellm/proxy/example_config_yaml/otel_test_config.yaml b/litellm/proxy/example_config_yaml/otel_test_config.yaml index 9c7937efba9..c05e2b1b5df 100644 --- a/litellm/proxy/example_config_yaml/otel_test_config.yaml +++ b/litellm/proxy/example_config_yaml/otel_test_config.yaml @@ -55,7 +55,7 @@ guardrails: litellm_params: guardrail: bedrock # supported values: "bedrock", "lakera" mode: "during_call" - guardrailIdentifier: 4w3d1di3snt5 + guardrailIdentifier: ff6ujrregl1q guardrailVersion: "DRAFT" - guardrail_name: "custom-pre-guard" litellm_params: diff --git a/tests/agent_tests/local_only_agent_tests/test_a2a_completion_bridge.py b/tests/agent_tests/local_only_agent_tests/test_a2a_completion_bridge.py index d3d582f29e4..4369bb800af 100644 --- a/tests/agent_tests/local_only_agent_tests/test_a2a_completion_bridge.py +++ b/tests/agent_tests/local_only_agent_tests/test_a2a_completion_bridge.py @@ -168,7 +168,7 @@ async def test_a2a_completion_bridge_bedrock_agentcore(): litellm._turn_on_debug() # Bedrock AgentCore ARN (streaming-capable runtime) - agentcore_arn = "arn:aws:bedrock-agentcore:us-west-2:941277531214:runtime/hosted_agent_r9jvp-Rq79QFC2fp" + agentcore_arn = "arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC" send_message_payload = { "message": { diff --git a/tests/batches_tests/test_bedrock_files_and_batches.py b/tests/batches_tests/test_bedrock_files_and_batches.py index 97c0802ec99..5148ea4db91 100644 --- a/tests/batches_tests/test_bedrock_files_and_batches.py +++ b/tests/batches_tests/test_bedrock_files_and_batches.py @@ -38,7 +38,7 @@ async def test_async_create_file(): file=open(file_path, "rb"), purpose="batch", custom_llm_provider="bedrock", - s3_bucket_name="litellm-proxy-941277531214", + s3_bucket_name="litellm-proxy", ) @@ -55,7 +55,7 @@ async def test_async_file_and_batch(): file=open(file_path, "rb"), purpose="batch", custom_llm_provider="bedrock", - s3_bucket_name="litellm-proxy-941277531214", + s3_bucket_name="litellm-proxy", ) print("CREATED FILE RESPONSE=", file_obj) @@ -70,7 +70,7 @@ async def test_async_file_and_batch(): # bedrock specific params ######################################################### model="us.anthropic.claude-haiku-4-5-20251001-v1:0", - aws_batch_role_arn="arn:aws:iam::941277531214:role/service-role/AmazonBedrockExecutionRoleForAgents_BB9HNW6V4CV", + aws_batch_role_arn="arn:aws:iam::888602223428:role/service-role/AmazonBedrockExecutionRoleForAgents_BB9HNW6V4CV", ) print("CREATED BATCH RESPONSE=", create_batch_response) @@ -129,7 +129,7 @@ async def test_mock_bedrock_file_url_mapping(): ), purpose="batch", custom_llm_provider="bedrock", - s3_bucket_name="litellm-proxy-941277531214", + s3_bucket_name="litellm-proxy", ) print(f"PUT URL: {captured_put_url}") diff --git a/tests/guardrails_tests/test_bedrock_guardrails.py b/tests/guardrails_tests/test_bedrock_guardrails.py index ea50fe08ae0..6e78a8c4284 100644 --- a/tests/guardrails_tests/test_bedrock_guardrails.py +++ b/tests/guardrails_tests/test_bedrock_guardrails.py @@ -20,7 +20,7 @@ async def test_bedrock_guardrails_pii_masking(): mock_user_api_key_dict = UserAPIKeyAuth() guardrail = BedrockGuardrail( - guardrailIdentifier="zgkmukebruil", + guardrailIdentifier="wf0hkdb5x07f", guardrailVersion="DRAFT", ) @@ -60,7 +60,7 @@ async def test_bedrock_guardrails_pii_masking_content_list(): mock_user_api_key_dict = UserAPIKeyAuth() guardrail = BedrockGuardrail( - guardrailIdentifier="zgkmukebruil", + guardrailIdentifier="wf0hkdb5x07f", guardrailVersion="DRAFT", ) @@ -115,7 +115,7 @@ async def test_bedrock_guardrails_block_messages_api(): mock_user_api_key_dict = UserAPIKeyAuth() guardrail = BedrockGuardrail( - guardrailIdentifier="4w3d1di3snt5", + guardrailIdentifier="ff6ujrregl1q", guardrailVersion="DRAFT", ) @@ -166,7 +166,7 @@ async def test_bedrock_guardrails_block_responses_api(): mock_user_api_key_dict = UserAPIKeyAuth() guardrail = BedrockGuardrail( - guardrailIdentifier="4w3d1di3snt5", + guardrailIdentifier="ff6ujrregl1q", guardrailVersion="DRAFT", ) @@ -211,7 +211,7 @@ async def test_bedrock_guardrails_with_streaming(): ) guardrail = BedrockGuardrail( - guardrailIdentifier="4w3d1di3snt5", + guardrailIdentifier="ff6ujrregl1q", guardrailVersion="DRAFT", supported_event_hooks=[GuardrailEventHooks.post_call], guardrail_name="bedrock-post-guard", @@ -255,7 +255,7 @@ async def test_bedrock_guardrails_with_streaming_no_violation(): ) guardrail = BedrockGuardrail( - guardrailIdentifier="4w3d1di3snt5", + guardrailIdentifier="ff6ujrregl1q", guardrailVersion="DRAFT", supported_event_hooks=[GuardrailEventHooks.post_call], guardrail_name="bedrock-post-guard", @@ -299,7 +299,7 @@ async def test_bedrock_guardrails_streaming_request_body_mock(): # Create the guardrail guardrail = BedrockGuardrail( - guardrailIdentifier="zgkmukebruil", + guardrailIdentifier="wf0hkdb5x07f", guardrailVersion="DRAFT", supported_event_hooks=[GuardrailEventHooks.post_call], guardrail_name="bedrock-post-guard", @@ -382,7 +382,7 @@ async def test_bedrock_guardrail_aws_param_persistence(): from litellm.types.guardrails import GuardrailEventHooks guardrail = BedrockGuardrail( - guardrailIdentifier="zgkmukebruil", + guardrailIdentifier="wf0hkdb5x07f", guardrailVersion="DRAFT", aws_access_key_id="test-access-key", aws_secret_access_key="test-secret-key", diff --git a/tests/image_gen_tests/test_bedrock_image_gen_unit_tests.py b/tests/image_gen_tests/test_bedrock_image_gen_unit_tests.py index 181691b730d..36ae9e1df67 100644 --- a/tests/image_gen_tests/test_bedrock_image_gen_unit_tests.py +++ b/tests/image_gen_tests/test_bedrock_image_gen_unit_tests.py @@ -1,4 +1,3 @@ -import json import logging import os import sys @@ -45,9 +44,6 @@ from litellm.llms.bedrock.image_generation.image_handler import ( ) from litellm.llms.bedrock.common_utils import BedrockError -# Base64 placeholder used for mocked Bedrock image responses (a 1x1 PNG). -_MOCK_BEDROCK_IMAGE_B64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" - @pytest.mark.parametrize( "model,expected", @@ -532,34 +528,17 @@ def test_backward_compatibility_regular_nova_model(): def test_amazon_titan_image_gen(): - """Test Amazon Titan image generation with cost tracking. - - The Bedrock CI account is not entitled to amazon.titan-image-generator, so - the network call is mocked and only the transform + cost-tracking path is - exercised. - """ - from litellm.llms.custom_httpx.http_handler import HTTPHandler + """Test Amazon Titan image generation with cost tracking.""" + from litellm import image_generation # Use v2 as v1 has reached end of life model_id = "bedrock/amazon.titan-image-generator-v2:0" - mock_payload = {"images": [_MOCK_BEDROCK_IMAGE_B64]} - mock_response = MagicMock() - mock_response.status_code = 200 - mock_response.json.return_value = mock_payload - mock_response.text = json.dumps(mock_payload) - mock_response.headers = {} - - client = HTTPHandler() - with patch.object(client, "post", return_value=mock_response): - response = litellm.image_generation( - model=model_id, - prompt="A serene mountain landscape at sunset with a lake reflection", - aws_region_name="us-east-1", - aws_access_key_id="fake-access-key-id", - aws_secret_access_key="fake-secret-access-key", - client=client, - ) + response = litellm.image_generation( + model=model_id, + prompt="A serene mountain landscape at sunset with a lake reflection", + aws_region_name="us-east-1", + ) print(f"response cost: {response._hidden_params['response_cost']}") diff --git a/tests/image_gen_tests/test_image_generation.py b/tests/image_gen_tests/test_image_generation.py index 23a94ef389a..873777189c9 100644 --- a/tests/image_gen_tests/test_image_generation.py +++ b/tests/image_gen_tests/test_image_generation.py @@ -7,6 +7,7 @@ import sys import traceback from unittest.mock import AsyncMock, MagicMock, patch + sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path @@ -135,51 +136,6 @@ class TestVertexAIGeminiImageGeneration(BaseImageGenTest): } -# Base64 placeholder used for mocked Bedrock image responses (a 1x1 PNG). -_MOCK_BEDROCK_IMAGE_B64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" - - -async def _assert_mocked_bedrock_image_generation(call_args: dict) -> None: - """Run ``aimage_generation`` with the Bedrock HTTP call mocked. - - The CI account is not entitled to Nova Canvas, so the network call is - replaced with a canned Bedrock response. This keeps the request transform, - response transform, and cost-tracking path under test without live access. - """ - mock_payload = {"images": [_MOCK_BEDROCK_IMAGE_B64]} - mock_response = MagicMock() - mock_response.status_code = 200 - mock_response.json.return_value = mock_payload - mock_response.text = json.dumps(mock_payload) - mock_response.headers = {} - - custom_logger = TestCustomLogger() - litellm.logging_callback_manager._reset_all_callbacks() - litellm.callbacks = [custom_logger] - - with patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", - new_callable=AsyncMock, - return_value=mock_response, - ): - response = await litellm.aimage_generation( - **call_args, - prompt="A image of a otter", - aws_access_key_id="fake-access-key-id", - aws_secret_access_key="fake-secret-access-key", - ) - - await asyncio.sleep(1) - - assert custom_logger.standard_logging_payload is not None - assert custom_logger.standard_logging_payload["response_cost"] is not None - assert custom_logger.standard_logging_payload["response_cost"] > 0 - assert response.data is not None - for d in response.data: - assert isinstance(d, Image) - assert d.b64_json is not None or d.url is not None - - class TestBedrockNovaCanvasTextToImage(BaseImageGenTest): def get_base_image_generation_call_args(self) -> dict: litellm.in_memory_llm_clients_cache = InMemoryCache() @@ -192,12 +148,6 @@ class TestBedrockNovaCanvasTextToImage(BaseImageGenTest): "aws_region_name": "us-east-1", } - @pytest.mark.asyncio(scope="module") - async def test_basic_image_generation(self): - await _assert_mocked_bedrock_image_generation( - self.get_base_image_generation_call_args() - ) - class TestBedrockNovaCanvasColorGuidedGeneration(BaseImageGenTest): def get_base_image_generation_call_args(self) -> dict: @@ -212,12 +162,6 @@ class TestBedrockNovaCanvasColorGuidedGeneration(BaseImageGenTest): "aws_region_name": "us-east-1", } - @pytest.mark.asyncio(scope="module") - async def test_basic_image_generation(self): - await _assert_mocked_bedrock_image_generation( - self.get_base_image_generation_call_args() - ) - class TestOpenAIGPTImage1(BaseImageGenTest): def get_base_image_generation_call_args(self) -> dict: diff --git a/tests/litellm_utils_tests/test_litellm_overhead.py b/tests/litellm_utils_tests/test_litellm_overhead.py index 60ee849f8eb..3a428e9d588 100644 --- a/tests/litellm_utils_tests/test_litellm_overhead.py +++ b/tests/litellm_utils_tests/test_litellm_overhead.py @@ -82,7 +82,7 @@ async def _vertex_ai_mocks(): "bedrock/mistral.mistral-7b-instruct-v0:2", "openai/gpt-4o", "openai/self_hosted", - "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + "bedrock/anthropic.claude-3-5-haiku-20241022-v1:0", "vertex_ai/gemini-1.5-flash", ], ) @@ -147,7 +147,7 @@ async def test_litellm_overhead_non_streaming(model): [ "bedrock/mistral.mistral-7b-instruct-v0:2", "openai/gpt-4o", - "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + "bedrock/anthropic.claude-3-5-haiku-20241022-v1:0", "openai/self_hosted", ], ) diff --git a/tests/llm_translation/reasoning_effort_grid/grid_spec.py b/tests/llm_translation/reasoning_effort_grid/grid_spec.py index 2f9735274c6..0d709584bde 100644 --- a/tests/llm_translation/reasoning_effort_grid/grid_spec.py +++ b/tests/llm_translation/reasoning_effort_grid/grid_spec.py @@ -1,6 +1,7 @@ from dataclasses import dataclass, field from typing import Dict, FrozenSet, List, Optional, Tuple + OMIT = object() @@ -21,7 +22,6 @@ class ModelEntry: extra_params: Tuple[Tuple[str, str], ...] = field(default_factory=tuple) required_env: FrozenSet[str] = field(default_factory=frozenset) caps: FrozenSet[str] = field(default_factory=frozenset) - fail_reason: Optional[str] = None bedrock_effort_ceiling: Optional[str] = None def params(self) -> Dict[str, str]: @@ -234,12 +234,6 @@ BEDROCK_CONVERSE_MODELS: Tuple[ModelEntry, ...] = ( extra_params=(("aws_region_name", "us-east-1"),), required_env=_BEDROCK_REQ, caps=_CAPS_OPUS_4_7, - fail_reason=( - "claude-opus-4-7 is not entitled on the Bedrock CI account " - "941277531214 (model access requires an AWS Sales request, not " - "self-serve); this cell fails on purpose so it stays loud in CI — " - "remove this fail_reason once access is granted" - ), ), ModelEntry( alias="bedrock-claude-opus-4-6", diff --git a/tests/llm_translation/reasoning_effort_grid/test_reasoning_effort_grid.py b/tests/llm_translation/reasoning_effort_grid/test_reasoning_effort_grid.py index e0b6290ad77..28e2e402d67 100644 --- a/tests/llm_translation/reasoning_effort_grid/test_reasoning_effort_grid.py +++ b/tests/llm_translation/reasoning_effort_grid/test_reasoning_effort_grid.py @@ -15,6 +15,7 @@ from .grid_spec import ( all_cells, ) + _PROMPT_MESSAGES: List[Dict[str, str]] = [ {"role": "user", "content": "Step by step, calculate 47 * 53. Show your work."} ] @@ -167,9 +168,6 @@ async def test_reasoning_effort_grid( if skip_reason: pytest.skip(skip_reason) - if model.fail_reason: - pytest.xfail(model.fail_reason) - if route_name == "bedrock_invoke_messages": status, exc = await _call_messages(model, effort) else: diff --git a/tests/llm_translation/test_bedrock_agentcore.py b/tests/llm_translation/test_bedrock_agentcore.py index 95a814e97e4..40774cf3d60 100644 --- a/tests/llm_translation/test_bedrock_agentcore.py +++ b/tests/llm_translation/test_bedrock_agentcore.py @@ -19,8 +19,8 @@ import httpx @pytest.mark.parametrize( "model", [ - "bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:941277531214:runtime/hosted_agent_13sf6-4046UzHSwy", # non-streaming invocation - "bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:941277531214:runtime/hosted_agent_r9jvp-Rq79QFC2fp", # streaming invocation + "bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_13sf6-cALnp38iZD", # non-streaming invocation + "bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC", # streaming invocation ], ) def test_bedrock_agentcore_basic(model): @@ -44,7 +44,7 @@ def test_bedrock_agentcore_basic(model): @pytest.mark.parametrize( "model", [ - "bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:941277531214:runtime/hosted_agent_13sf6-4046UzHSwy", # streaming invocation + "bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_13sf6-cALnp38iZD", # streaming invocation ], ) async def test_bedrock_agentcore_with_streaming(model): @@ -54,7 +54,7 @@ async def test_bedrock_agentcore_with_streaming(model): print("running streming test for model=", model) # litellm._turn_on_debug() response = await litellm.acompletion( - model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:941277531214:runtime/hosted_agent_r9jvp-Rq79QFC2fp", + model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC", messages=[ { "role": "user", @@ -82,7 +82,7 @@ def test_bedrock_agentcore_with_custom_params(): with patch.object(client, "post", return_value=MagicMock()) as mock_post: try: response = litellm.completion( - model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:941277531214:runtime/hosted_agent_r9jvp-Rq79QFC2fp", + model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC", messages=[ { "role": "user", @@ -105,7 +105,7 @@ def test_bedrock_agentcore_with_custom_params(): url = call_kwargs["url"] print(f"URL: {url}") assert ( - "/runtimes/arn%3Aaws%3Abedrock-agentcore%3Aus-west-2%3A941277531214%3Aruntime%2Fhosted_agent_r9jvp-Rq79QFC2fp/invocations" + "/runtimes/arn%3Aaws%3Abedrock-agentcore%3Aus-west-2%3A888602223428%3Aruntime%2Fhosted_agent_r9jvp-3ySZuRHjLC/invocations" in url ) assert "qualifier=DEFAULT" in url @@ -150,7 +150,7 @@ def test_bedrock_agentcore_with_runtime_user_id(): with patch.object(client, "post", return_value=MagicMock()) as mock_post: try: response = litellm.completion( - model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:941277531214:runtime/hosted_agent_r9jvp-Rq79QFC2fp", + model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC", messages=[ { "role": "user", @@ -189,7 +189,7 @@ def test_bedrock_agentcore_with_session_and_user(): with patch.object(client, "post", return_value=MagicMock()) as mock_post: try: response = litellm.completion( - model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:941277531214:runtime/hosted_agent_r9jvp-Rq79QFC2fp", + model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC", messages=[ { "role": "user", @@ -234,7 +234,7 @@ def test_bedrock_agentcore_with_api_key_bearer_token(): with patch.object(client, "post", return_value=MagicMock()) as mock_post: try: response = litellm.completion( - model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:941277531214:runtime/hosted_agent_r9jvp-Rq79QFC2fp", + model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC", messages=[ { "role": "user", @@ -282,7 +282,7 @@ def test_bedrock_agentcore_with_all_parameters(): with patch.object(client, "post", return_value=MagicMock()) as mock_post: try: response = litellm.completion( - model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:941277531214:runtime/hosted_agent_r9jvp-Rq79QFC2fp", + model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC", messages=[ { "role": "user", @@ -350,7 +350,7 @@ def test_bedrock_agentcore_without_api_key_uses_sigv4(): with patch.object(client, "post", return_value=MagicMock()) as mock_post: try: response = litellm.completion( - model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:941277531214:runtime/hosted_agent_r9jvp-Rq79QFC2fp", + model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC", messages=[ { "role": "user", @@ -625,7 +625,7 @@ def test_agentcore_synchronous_non_streaming_response(): with patch.object(client, "post", return_value=mock_response) as mock_post: # Make a synchronous (non-streaming) completion call response = litellm.completion( - model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:941277531214:runtime/hosted_agent_r9jvp-Rq79QFC2fp", + model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC", messages=[ { "role": "user", diff --git a/tests/llm_translation/test_bedrock_completion.py b/tests/llm_translation/test_bedrock_completion.py index a3b4a010f60..aecd7bc699a 100644 --- a/tests/llm_translation/test_bedrock_completion.py +++ b/tests/llm_translation/test_bedrock_completion.py @@ -115,7 +115,7 @@ def test_completion_bedrock_guardrails(streaming): ], max_tokens=10, guardrailConfig={ - "guardrailIdentifier": "4w3d1di3snt5", + "guardrailIdentifier": "ff6ujrregl1q", "guardrailVersion": "DRAFT", "trace": "enabled", }, @@ -144,7 +144,7 @@ def test_completion_bedrock_guardrails(streaming): stream=True, max_tokens=10, guardrailConfig={ - "guardrailIdentifier": "4w3d1di3snt5", + "guardrailIdentifier": "ff6ujrregl1q", "guardrailVersion": "DRAFT", "trace": "enabled", }, @@ -475,7 +475,7 @@ def test_bedrock_claude_3(image_url): ], } response: ModelResponse = completion( - model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", + model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0", num_retries=3, **data, ) # type: ignore @@ -498,7 +498,7 @@ def test_bedrock_claude_3(image_url): @pytest.mark.parametrize( "model", [ - "us.anthropic.claude-sonnet-4-5-20250929-v1:0", + "anthropic.claude-3-sonnet-20240229-v1:0", # "meta.llama3-70b-instruct-v1:0", # "anthropic.claude-v2", # "mistral.mixtral-8x7b-instruct-v0:1", @@ -537,7 +537,7 @@ def test_bedrock_stop_value(stop, model): @pytest.mark.parametrize( "model", [ - "us.anthropic.claude-sonnet-4-5-20250929-v1:0", + "anthropic.claude-3-sonnet-20240229-v1:0", "mistral.mixtral-8x7b-instruct-v0:1", ], ) @@ -602,7 +602,7 @@ def test_bedrock_claude_3_tool_calling(): } ] response: ModelResponse = completion( - model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", + model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0", messages=messages, tools=tools, tool_choice="auto", @@ -630,7 +630,7 @@ def test_bedrock_claude_3_tool_calling(): ) # In the second response, Claude should deduce answer from tool results second_response = completion( - model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", + model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0", messages=messages, tools=tools, tool_choice="auto", @@ -737,7 +737,7 @@ def test_bedrock_ptu(): from openai.types.chat import ChatCompletion model_id = ( - "arn:aws:bedrock:us-west-2:941277531214:provisioned-model/8fxff74qyhs3" + "arn:aws:bedrock:us-west-2:888602223428:provisioned-model/8fxff74qyhs3" ) try: response = litellm.completion( @@ -752,7 +752,7 @@ def test_bedrock_ptu(): assert "url" in mock_client_post.call_args.kwargs assert ( mock_client_post.call_args.kwargs["url"] - == "https://bedrock-runtime.us-west-2.amazonaws.com/model/arn%3Aaws%3Abedrock%3Aus-west-2%3A941277531214%3Aprovisioned-model%2F8fxff74qyhs3/converse" + == "https://bedrock-runtime.us-west-2.amazonaws.com/model/arn%3Aaws%3Abedrock%3Aus-west-2%3A888602223428%3Aprovisioned-model%2F8fxff74qyhs3/converse" ) mock_client_post.assert_called_once() @@ -2327,7 +2327,7 @@ def test_bedrock_cross_region_inference(monkeypatch): def test_bedrock_empty_content_real_call(): completion( - model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", + model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0", messages=[ { "role": "user", diff --git a/tests/local_testing/test_completion.py b/tests/local_testing/test_completion.py index c7abdb5f493..cce6d33e799 100644 --- a/tests/local_testing/test_completion.py +++ b/tests/local_testing/test_completion.py @@ -299,10 +299,7 @@ def test_completion_claude_3(): @pytest.mark.parametrize( "model", - [ - "anthropic/claude-sonnet-4-5-20250929", - "us.anthropic.claude-sonnet-4-5-20250929-v1:0", - ], + ["anthropic/claude-sonnet-4-5-20250929", "anthropic.claude-3-sonnet-20240229-v1:0"], ) def test_completion_claude_3_function_call(model): litellm.set_verbose = True @@ -388,7 +385,7 @@ def test_completion_claude_3_function_call(model): [ ("gpt-3.5-turbo", None, None), ("claude-sonnet-4-5-20250929", None, None), - ("us.anthropic.claude-sonnet-4-5-20250929-v1:0", None, None), + ("anthropic.claude-3-sonnet-20240229-v1:0", None, None), # ( # "azure_ai/command-r-plus", # os.getenv("AZURE_COHERE_API_KEY"), @@ -1581,7 +1578,7 @@ def test_completion_openai(): [ # ("gpt-4o-2024-08-06", None), # ("azure/gpt-4.1-mini", None), - ("bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", None), + ("bedrock/anthropic.claude-3-sonnet-20240229-v1:0", None), # ("azure/gpt-4o-new-test", "2024-08-01-preview"), ], ) @@ -1669,13 +1666,15 @@ def custom_callback( ################################################# - print(f""" + print( + f""" Model: {model}, Messages: {messages}, User: {user}, Seed: {kwargs["seed"]}, temperature: {kwargs["temperature"]}, - """) + """ + ) assert kwargs["user"] == "ishaans app" assert kwargs["model"] == "gpt-3.5-turbo-1106" @@ -2700,7 +2699,7 @@ def test_bedrock_deepseek_custom_prompt_dict(): def test_bedrock_deepseek_known_tokenizer_config(monkeypatch): model = ( - "deepseek_r1/arn:aws:bedrock:us-west-2:941277531214:imported-model/bnnr6463ejgf" + "deepseek_r1/arn:aws:bedrock:us-west-2:888602223428:imported-model/bnnr6463ejgf" ) from litellm.llms.custom_httpx.http_handler import HTTPHandler from unittest.mock import Mock @@ -2915,8 +2914,8 @@ def response_format_tests(response: litellm.ModelResponse): "model", [ "bedrock/mistral.mistral-large-2407-v1:0", - "us.anthropic.claude-haiku-4-5-20251001-v1:0", - "us.anthropic.claude-sonnet-4-5-20250929-v1:0", + "bedrock/cohere.command-r-plus-v1:0", + "anthropic.claude-3-sonnet-20240229-v1:0", "mistral.mistral-7b-instruct-v0:2", "meta.llama3-8b-instruct-v1:0", ], diff --git a/tests/local_testing/test_function_call_parsing.py b/tests/local_testing/test_function_call_parsing.py index 2453571f1c4..f9582fcc574 100644 --- a/tests/local_testing/test_function_call_parsing.py +++ b/tests/local_testing/test_function_call_parsing.py @@ -142,8 +142,7 @@ def trade(model_name: str) -> List[Trade]: # type: ignore @pytest.mark.parametrize( - "model", - ["claude-haiku-4-5-20251001", "us.anthropic.claude-haiku-4-5-20251001-v1:0"], + "model", ["claude-haiku-4-5-20251001", "anthropic.claude-3-haiku-20240307-v1:0"] ) @pytest.mark.flaky(retries=6, delay=10) def test_function_call_parsing(model): diff --git a/tests/local_testing/test_function_calling.py b/tests/local_testing/test_function_calling.py index 1cad7d1421e..3c7e004b62e 100644 --- a/tests/local_testing/test_function_calling.py +++ b/tests/local_testing/test_function_calling.py @@ -49,7 +49,7 @@ def get_current_weather(location, unit="fahrenheit"): "mistral/mistral-large-latest", "claude-haiku-4-5-20251001", "gemini/gemini-2.5-flash-lite", - "us.anthropic.claude-sonnet-4-5-20250929-v1:0", + "anthropic.claude-3-sonnet-20240229-v1:0", ], ) @pytest.mark.flaky(retries=3, delay=1) @@ -267,6 +267,7 @@ def test_aaparallel_function_call_with_anthropic_thinking(model): from litellm.types.utils import ChatCompletionMessageToolCall, Function, Message + _PARALLEL_TOOL_HISTORY_MESSAGES = [ { "role": "user", @@ -302,7 +303,7 @@ _PARALLEL_TOOL_HISTORY_MESSAGES = [ [ # Bedrock Converse still requires modify_params to inject the dummy tool. ( - "us.anthropic.claude-sonnet-4-5-20250929-v1:0", + "anthropic.claude-3-sonnet-20240229-v1:0", _PARALLEL_TOOL_HISTORY_MESSAGES, True, ), @@ -313,7 +314,7 @@ _PARALLEL_TOOL_HISTORY_MESSAGES = [ False, ), ( - "us.anthropic.claude-sonnet-4-5-20250929-v1:0", + "anthropic.claude-3-sonnet-20240229-v1:0", [ { "role": "user", @@ -578,7 +579,7 @@ def test_groq_parallel_function_call(): @pytest.mark.parametrize( "model", [ - "bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", + "bedrock/anthropic.claude-3-sonnet-20240229-v1:0", ], ) def test_passing_tool_result_as_list(model): diff --git a/tests/local_testing/test_sagemaker.py b/tests/local_testing/test_sagemaker.py index fdc8347c36a..d4c5a5a857f 100644 --- a/tests/local_testing/test_sagemaker.py +++ b/tests/local_testing/test_sagemaker.py @@ -57,7 +57,7 @@ async def test_completion_sagemaker(sync_mode): print("testing sagemaker") if sync_mode is True: response = litellm.completion( - model="sagemaker/litellm-ci-textgen", + model="sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614", messages=[ {"role": "user", "content": "hi"}, ], @@ -67,7 +67,7 @@ async def test_completion_sagemaker(sync_mode): ) else: response = await litellm.acompletion( - model="sagemaker/litellm-ci-textgen", + model="sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614", messages=[ {"role": "user", "content": "hi"}, ], @@ -158,7 +158,7 @@ async def test_completion_sagemaker_messages_api(sync_mode): "model", [ # "sagemaker_chat/huggingface-pytorch-tgi-inference-2024-08-23-15-48-59-245", - "sagemaker/litellm-ci-textgen", + "sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614", ], ) # @pytest.mark.flaky(retries=3, delay=1) @@ -218,7 +218,7 @@ async def test_completion_sagemaker_stream(sync_mode, model): "model", [ # "sagemaker_chat/huggingface-pytorch-tgi-inference-2024-08-23-15-48-59-245", - "sagemaker/litellm-ci-textgen", + "sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614", ], ) async def test_completion_sagemaker_streaming_bad_request(sync_mode, model): @@ -256,7 +256,7 @@ async def test_acompletion_sagemaker_non_stream(): "id": "cmpl-mockid", "object": "text_completion", "created": 1629800000, - "model": "sagemaker/litellm-ci-textgen", + "model": "sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614", "choices": [ { "text": "This is a mock response from SageMaker.", @@ -282,7 +282,7 @@ async def test_acompletion_sagemaker_non_stream(): ) as mock_post: # Act: Call the litellm.acompletion function response = await litellm.acompletion( - model="sagemaker/litellm-ci-textgen", + model="sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614", messages=[ {"role": "user", "content": "hi"}, ], @@ -302,7 +302,7 @@ async def test_acompletion_sagemaker_non_stream(): assert args_to_sagemaker == expected_payload assert ( kwargs["url"] - == "https://runtime.sagemaker.us-west-2.amazonaws.com/endpoints/litellm-ci-textgen/invocations" + == "https://runtime.sagemaker.us-west-2.amazonaws.com/endpoints/jumpstart-dft-hf-textgeneration1-mp-20240815-185614/invocations" ) @@ -316,7 +316,7 @@ async def test_completion_sagemaker_non_stream(): "id": "cmpl-mockid", "object": "text_completion", "created": 1629800000, - "model": "sagemaker/litellm-ci-textgen", + "model": "sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614", "choices": [ { "text": "This is a mock response from SageMaker.", @@ -342,7 +342,7 @@ async def test_completion_sagemaker_non_stream(): ) as mock_post: # Act: Call the litellm.acompletion function response = litellm.completion( - model="sagemaker/litellm-ci-textgen", + model="sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614", messages=[ {"role": "user", "content": "hi"}, ], @@ -362,7 +362,7 @@ async def test_completion_sagemaker_non_stream(): assert args_to_sagemaker == expected_payload assert ( kwargs["url"] - == "https://runtime.sagemaker.us-west-2.amazonaws.com/endpoints/litellm-ci-textgen/invocations" + == "https://runtime.sagemaker.us-west-2.amazonaws.com/endpoints/jumpstart-dft-hf-textgeneration1-mp-20240815-185614/invocations" ) @@ -377,7 +377,7 @@ async def test_completion_sagemaker_prompt_template_non_stream(): "id": "cmpl-mockid", "object": "text_completion", "created": 1629800000, - "model": "sagemaker/litellm-ci-textgen", + "model": "sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614", "choices": [ { "text": "This is a mock response from SageMaker.", @@ -433,7 +433,7 @@ async def test_completion_sagemaker_non_stream_with_aws_params(): "id": "cmpl-mockid", "object": "text_completion", "created": 1629800000, - "model": "sagemaker/litellm-ci-textgen", + "model": "sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614", "choices": [ { "text": "This is a mock response from SageMaker.", @@ -459,7 +459,7 @@ async def test_completion_sagemaker_non_stream_with_aws_params(): ) as mock_post: # Act: Call the litellm.acompletion function response = litellm.completion( - model="sagemaker/litellm-ci-textgen", + model="sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614", messages=[ {"role": "user", "content": "hi"}, ], @@ -482,5 +482,5 @@ async def test_completion_sagemaker_non_stream_with_aws_params(): assert args_to_sagemaker == expected_payload assert ( kwargs["url"] - == "https://runtime.sagemaker.us-west-5.amazonaws.com/endpoints/litellm-ci-textgen/invocations" + == "https://runtime.sagemaker.us-west-5.amazonaws.com/endpoints/jumpstart-dft-hf-textgeneration1-mp-20240815-185614/invocations" ) diff --git a/tests/local_testing/test_streaming.py b/tests/local_testing/test_streaming.py index eb153404a44..10f351714e1 100644 --- a/tests/local_testing/test_streaming.py +++ b/tests/local_testing/test_streaming.py @@ -1174,7 +1174,7 @@ async def test_completion_replicate_llama3_streaming(sync_mode): [ # ["bedrock/ai21.jamba-instruct-v1:0", "us-east-1"], # ["bedrock/cohere.command-r-plus-v1:0", None], - ["us.anthropic.claude-sonnet-4-5-20250929-v1:0", None], + ["anthropic.claude-3-sonnet-20240229-v1:0", None], # ["mistral.mistral-7b-instruct-v0:2", None], # ["meta.llama3-8b-instruct-v1:0", None], ], @@ -1246,7 +1246,7 @@ def test_bedrock_claude_3_streaming(): try: litellm.set_verbose = True response: ModelResponse = completion( # type: ignore - model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", + model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0", messages=messages, max_tokens=10, # type: ignore stream=True, @@ -1276,7 +1276,7 @@ def test_bedrock_claude_3_streaming(): "model", [ "claude-haiku-4-5-20251001", - "us.anthropic.claude-haiku-4-5-20251001-v1:0", # bedrock + "cohere.command-r-plus-v1:0", # bedrock "gpt-3.5-turbo", ], ) @@ -3500,7 +3500,7 @@ def test_unit_test_perplexity_citations_chunk(): [ "gpt-3.5-turbo", "claude-sonnet-4-5-20250929", - "us.anthropic.claude-sonnet-4-5-20250929-v1:0", + "anthropic.claude-3-sonnet-20240229-v1:0", # "vertex_ai/claude-3-5-sonnet@20240620", ], ) diff --git a/tests/logging_callback_tests/test_amazing_s3_logs.py b/tests/logging_callback_tests/test_amazing_s3_logs.py index e6291a94049..dab2a0cc0b9 100644 --- a/tests/logging_callback_tests/test_amazing_s3_logs.py +++ b/tests/logging_callback_tests/test_amazing_s3_logs.py @@ -27,7 +27,7 @@ async def test_basic_s3_logging(sync_mode, streaming): verbose_logger.setLevel(level=logging.DEBUG) litellm.success_callback = ["s3"] litellm.s3_callback_params = { - "s3_bucket_name": "load-testing-oct-941277531214", + "s3_bucket_name": "load-testing-oct", "s3_aws_secret_access_key": "os.environ/AWS_SECRET_ACCESS_KEY", "s3_aws_access_key_id": "os.environ/AWS_ACCESS_KEY_ID", "s3_region_name": "us-west-2", @@ -64,14 +64,14 @@ async def test_basic_s3_logging(sync_mode, streaming): await asyncio.sleep(2) print(f"response: {response}") - total_objects, all_s3_keys = list_all_s3_objects("load-testing-oct-941277531214") + total_objects, all_s3_keys = list_all_s3_objects("load-testing-oct") # assert that atlest one key has response.id in it assert any(response_id in key for key in all_s3_keys) s3 = boto3.client("s3") # delete all objects for key in all_s3_keys: - s3.delete_object(Bucket="load-testing-oct-941277531214", Key=key) + s3.delete_object(Bucket="load-testing-oct", Key=key) @pytest.mark.asyncio @@ -82,7 +82,7 @@ async def test_basic_s3_v2_logging(streaming): from litellm.integrations.s3_v2 import S3Logger litellm.s3_callback_params = { - "s3_bucket_name": "load-testing-oct-941277531214", + "s3_bucket_name": "load-testing-oct", "s3_aws_secret_access_key": "test-secret", "s3_aws_access_key_id": "test-key", "s3_region_name": "us-west-2", diff --git a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py index 0d4405094b5..d6d0652ed77 100644 --- a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py +++ b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py @@ -2,6 +2,7 @@ import io import os import sys + sys.path.insert(0, os.path.abspath("../..")) import asyncio @@ -66,7 +67,7 @@ def setup_vector_store_registry(): litellm.vector_store_registry = VectorStoreRegistry( vector_stores=[ LiteLLM_ManagedVectorStore( - vector_store_id="LCYXFBR2TU", custom_llm_provider="bedrock" + vector_store_id="T37J8R4WTM", custom_llm_provider="bedrock" ) ] ) @@ -110,7 +111,7 @@ async def test_e2e_bedrock_knowledgebase_retrieval_with_completion( response = await litellm.acompletion( model="anthropic/claude-3.5-sonnet", messages=[{"role": "user", "content": "what is litellm?"}], - vector_store_ids=["LCYXFBR2TU"], + vector_store_ids=["T37J8R4WTM"], client=client, ) except Exception as e: @@ -151,7 +152,7 @@ async def test_e2e_bedrock_knowledgebase_retrieval_with_llm_api_call( response = await litellm.acompletion( model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", messages=[{"role": "user", "content": "what is litellm?"}], - vector_store_ids=["LCYXFBR2TU"], + vector_store_ids=["T37J8R4WTM"], client=async_client, ) print("OPENAI RESPONSE:", json.dumps(dict(response), indent=4, default=str)) @@ -195,7 +196,7 @@ async def test_e2e_bedrock_knowledgebase_retrieval_with_llm_api_call_streaming( response = await litellm.acompletion( model=f"anthropic/{os.environ.get('CI_CD_DEFAULT_ANTHROPIC_MODEL', 'claude-haiku-4-5-20251001')}", messages=[{"role": "user", "content": "what is litellm?"}], - vector_store_ids=["LCYXFBR2TU"], + vector_store_ids=["T37J8R4WTM"], stream=True, client=async_client, ) @@ -254,7 +255,7 @@ async def test_e2e_bedrock_knowledgebase_retrieval_with_llm_api_call_with_tools( model=f"anthropic/{os.environ.get('CI_CD_DEFAULT_ANTHROPIC_MODEL', 'claude-haiku-4-5-20251001')}", messages=[{"role": "user", "content": "what is litellm?"}], max_tokens=10, - tools=[{"type": "file_search", "vector_store_ids": ["LCYXFBR2TU"]}], + tools=[{"type": "file_search", "vector_store_ids": ["T37J8R4WTM"]}], ) assert response is not None @@ -278,7 +279,7 @@ async def test_e2e_bedrock_knowledgebase_retrieval_with_llm_api_call_with_tools_ tools=[ { "type": "file_search", - "vector_store_ids": ["LCYXFBR2TU"], + "vector_store_ids": ["T37J8R4WTM"], "filters": { "key": "user_id", "value": "fake-user-id", @@ -386,7 +387,7 @@ async def test_bedrock_kb_request_body_has_transformed_filters( tools=[ { "type": "file_search", - "vector_store_ids": ["LCYXFBR2TU"], + "vector_store_ids": ["T37J8R4WTM"], "filters": { "key": "user_id", "value": "fake-user-id", @@ -460,7 +461,7 @@ async def test_openai_with_knowledge_base_mock_openai(setup_vector_store_registr await litellm.acompletion( model="gpt-5.5", messages=[{"role": "user", "content": "what is litellm?"}], - vector_store_ids=["LCYXFBR2TU"], + vector_store_ids=["T37J8R4WTM"], client=client, ) except Exception as e: @@ -536,7 +537,7 @@ async def test_openai_with_vector_store_ids_in_tool_call_mock_openai( await litellm.acompletion( model="gpt-5.5", messages=[{"role": "user", "content": "what is litellm?"}], - tools=[{"type": "file_search", "vector_store_ids": ["LCYXFBR2TU"]}], + tools=[{"type": "file_search", "vector_store_ids": ["T37J8R4WTM"]}], client=client, ) except Exception as e: @@ -610,7 +611,7 @@ async def test_openai_with_mixed_tool_call_mock_openai(setup_vector_store_regist model="gpt-5.5", messages=[{"role": "user", "content": "what is litellm?"}], tools=[ - {"type": "file_search", "vector_store_ids": ["LCYXFBR2TU"]}, + {"type": "file_search", "vector_store_ids": ["T37J8R4WTM"]}, {"type": "file_search", "vector_store_ids": ["unknownVS"]}, ], client=client, @@ -644,7 +645,7 @@ async def test_openai_with_mixed_tool_call_mock_openai(setup_vector_store_regist # model="gpt-5.5", # messages=[{"role": "user", "content": "what is litellm?"}], # vector_store_ids = [ -# "LCYXFBR2TU" +# "T37J8R4WTM" # ], # ) @@ -666,7 +667,7 @@ async def test_openai_with_mixed_tool_call_mock_openai(setup_vector_store_regist # # expect the vector store request metadata object to have the correct values # vector_store_request_metadata = standard_logging_vector_store_request_metadata[0] -# assert vector_store_request_metadata.get("vector_store_id") == "LCYXFBR2TU" +# assert vector_store_request_metadata.get("vector_store_id") == "T37J8R4WTM" # assert vector_store_request_metadata.get("query") == "what is litellm?" # assert vector_store_request_metadata.get("custom_llm_provider") == "bedrock" @@ -722,7 +723,7 @@ async def test_e2e_bedrock_knowledgebase_retrieval_without_vector_store_registry response = await litellm.acompletion( model="anthropic/claude-3.5-sonnet", messages=[{"role": "user", "content": "what is litellm?"}], - vector_store_ids=["LCYXFBR2TU"], + vector_store_ids=["T37J8R4WTM"], client=client, ) except Exception as e: diff --git a/tests/test_litellm/llms/bedrock/chat/agentcore/test_agentcore_transformation.py b/tests/test_litellm/llms/bedrock/chat/agentcore/test_agentcore_transformation.py index 3287061d37e..64b43b15dcd 100644 --- a/tests/test_litellm/llms/bedrock/chat/agentcore/test_agentcore_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/agentcore/test_agentcore_transformation.py @@ -76,7 +76,7 @@ class TestAgentCoreAcceptHeader: with patch.object(client, "post", return_value=MagicMock()) as mock_post: try: litellm.completion( - model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:941277531214:runtime/test_runtime", + model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/test_runtime", messages=[{"role": "user", "content": "test"}], api_key="test-jwt-token", client=client, @@ -281,7 +281,7 @@ class TestAgentCoreStreamingJsonFallback: with patch.object(client, "post", return_value=mock_response): response = litellm.completion( - model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:941277531214:runtime/test_agent", + model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/test_agent", messages=[{"role": "user", "content": "test"}], stream=True, client=client, @@ -318,7 +318,7 @@ class TestAgentCoreStreamingJsonFallback: client, "post", new_callable=AsyncMock, return_value=mock_response ): response = await litellm.acompletion( - model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:941277531214:runtime/test_agent", + model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/test_agent", messages=[{"role": "user", "content": "test"}], stream=True, client=client, @@ -353,7 +353,7 @@ class TestAgentCoreStreamingJsonFallback: Exception, match="Failed to read/parse JSON response body" ): litellm.completion( - model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:941277531214:runtime/test_agent", + model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/test_agent", messages=[{"role": "user", "content": "test"}], stream=True, client=client, @@ -383,7 +383,7 @@ class TestAgentCoreStreamingJsonFallback: Exception, match="Failed to read/parse JSON response body" ): await litellm.acompletion( - model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:941277531214:runtime/test_agent", + model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/test_agent", messages=[{"role": "user", "content": "test"}], stream=True, client=client, diff --git a/tests/test_openai_endpoints.py b/tests/test_openai_endpoints.py index 29875a04413..e898b88a556 100644 --- a/tests/test_openai_endpoints.py +++ b/tests/test_openai_endpoints.py @@ -446,7 +446,7 @@ async def test_chat_completion_anthropic_structured_output(): client = AsyncOpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000") res = await client.beta.chat.completions.parse( - model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", + model="bedrock/us.anthropic.claude-3-sonnet-20240229-v1:0", messages=messages, response_format=EventsList, timeout=60, diff --git a/tests/vector_store_tests/test_bedrock_vector_store.py b/tests/vector_store_tests/test_bedrock_vector_store.py index 47e73e61c59..d8af1c7188b 100644 --- a/tests/vector_store_tests/test_bedrock_vector_store.py +++ b/tests/vector_store_tests/test_bedrock_vector_store.py @@ -22,7 +22,7 @@ class TestBedrockVectorStore(BaseVectorStoreTest): def get_base_request_args(self): return { - "vector_store_id": "LCYXFBR2TU", + "vector_store_id": "T37J8R4WTM", "custom_llm_provider": "bedrock", "query": "what happens after we add a model", } @@ -106,7 +106,7 @@ async def test_bedrock_search_with_router(): _router = Router(model_list=[]) search_response = await _router.avector_store_search( query="what happens after we add a model", - vector_store_id="LCYXFBR2TU", + vector_store_id="T37J8R4WTM", custom_llm_provider="bedrock", ) print(search_response) @@ -150,7 +150,7 @@ async def test_bedrock_search_with_credentials_managed_registry(): # Create vector store with credential reference vector_store = LiteLLM_ManagedVectorStore( - vector_store_id="LCYXFBR2TU", + vector_store_id="T37J8R4WTM", custom_llm_provider="bedrock", created_at=datetime.now(timezone.utc), updated_at=datetime.now(timezone.utc), @@ -162,7 +162,7 @@ async def test_bedrock_search_with_credentials_managed_registry(): litellm.vector_store_registry = registry # Verify credentials can be retrieved from registry - retrieved_credentials = registry.get_credentials_for_vector_store("LCYXFBR2TU") + retrieved_credentials = registry.get_credentials_for_vector_store("T37J8R4WTM") assert retrieved_credentials, "Should retrieve credentials from registry" assert retrieved_credentials.get("aws_access_key_id") == "test_access_key" assert retrieved_credentials.get("aws_secret_access_key") == "test_secret_key" @@ -194,7 +194,7 @@ async def test_bedrock_search_with_credentials_managed_registry(): search_response = await _router.avector_store_search( query="what happens after we add a model", - vector_store_id="LCYXFBR2TU", + vector_store_id="T37J8R4WTM", custom_llm_provider="bedrock", ) @@ -203,7 +203,7 @@ async def test_bedrock_search_with_credentials_managed_registry(): call_kwargs = mock_handler.call_args[1] # Verify that the credential accessor was called with the correct vector store ID - mock_get_creds.assert_called_with("LCYXFBR2TU") + mock_get_creds.assert_called_with("T37J8R4WTM") # Verify the credentials were injected into the search call litellm_params = call_kwargs.get("litellm_params", {}) @@ -224,7 +224,7 @@ async def test_bedrock_search_with_credentials_managed_registry(): assert search_response["data"][0]["id"] == "test_result" print( - f"✅ Test passed: Credential accessor was called with vector store ID: LCYXFBR2TU" + f"✅ Test passed: Credential accessor was called with vector store ID: T37J8R4WTM" ) print(f"✅ Retrieved credentials: {retrieved_credentials}") print(f"✅ Credentials were injected into search call") From ace3c65ab355108817d1d826fe1f23fd26a03434 Mon Sep 17 00:00:00 2001 From: Shivam Rawat