Merge remote-tracking branch 'origin/main' into litellm_openai_ws_frame_model_gate

# Conflicts:
#	litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
#	litellm/types/passthrough_endpoints/pass_through_endpoints.py

.agents/skills/database-access/SKILL.md

@@ -1,54 +1,73 @@
 ---
 name: database-access
-description: Add database queries, extract existing Prisma calls into repositories, or explicitly port queries to Rust and SQLx without expanding migration scope
+description: Organize database repositories and proxy dependency wiring, manage connection ownership, or port Prisma operations to Rust and SQLx without expanding migration scope
 metadata:
   internal: true
 ---
 
 # Database access
 
+## Situations
+
+- Adding a query or an operation to a repository, read [adding a query](references/python-repository.md#adding-a-query)
+- Moving existing Prisma calls behind a repository, read [moving existing access](references/python-repository.md#moving-existing-access-into-a-repository)
+- Porting a repository from Prisma to Rust and SQLx, read [porting to SQLx](references/python-repository.md#porting-a-repository-to-sqlx)
+- Injecting a repository into a route, read [Python and FastAPI](references/python-fastapi.md) or [Rust and Axum](references/rust-axum.md)
+
 ## Adding new database access
 
 - Put Python database access behind an operation in `litellm/repositories`
-- Give callers typed inputs and outputs, not Prisma clients, models, or query dictionaries
-- Keep query execution and backend selection inside the repository
+- Give each repository its own FastAPI provider, so a route builds only the repositories it uses
+- A provider borrows existing handles and performs no database I/O during construction
+- Define the operations a route calls as a Protocol in `litellm/models/<feature>.py`, next to its domain types
+  - The Protocol takes and returns domain types, never Prisma models, raw rows or serialized payloads
+  - Parsing and serialization of stored values live in the repository
+  - Add a Protocol when a repository gets a provider, not up front for every repository
+- Keep query execution inside the repository, with one implementation per repository and no backend switch or fallback
+- A repository is either all Prisma or all native, never a mix of both
 - Prisma is fine for the implementation, adding a query does not require a Rust port
 - Native Rust callers use `litellm-rust/crates/db` directly
 
-## Asked to move existing access behind a repository
+## Framework wiring
 
-- Identify the calls being moved and their transaction boundary
-- Move the existing Prisma implementation into the relevant repository
-- Replace those direct database calls with repository operations
-- Preserve return values, errors, and transaction behavior, with focused regression tests
-- Keep this an extraction, not a backend migration or unrelated cleanup
+- Python routes receive a per-repository dependency such as `UserBannerRepositoryDep` and call its operations, such as `banners.get_banner()`
+- `DatabaseDep` lives in `litellm/proxy/deps/database.py`, repository providers in `litellm/proxy/deps/repositories.py`
+  - A provider returns the feature's Protocol, so routes never see which backend implements it
+  - Dependencies used by one feature only, like the banner audit binding, stay next to that feature's routes
+- Repositories never import FastAPI, so jobs and helpers can use them without a request
+- Use named FastAPI providers with `Annotated[..., Depends(...)]`, ordinary helpers and jobs receive explicit arguments
+- Preserve each route's no-database behavior, authorization precedence and background-task scheduling
+- Read [Python and FastAPI](references/python-fastapi.md) for provider examples, lifespan state and route verification
+- Read [Rust and Axum](references/rust-axum.md) when adding database access to Rust handlers
 
-## An equivalent Rust or SQLx query already exists
+## Connection ownership
 
-- Leave it and the current backend selection unchanged unless integration is explicitly requested
-- Its existence alone does not justify switching a Prisma caller to Rust
-- Repository extraction can keep using Prisma even when the Rust query is available
+- Keep connection setup and shutdown separate from repositories, which own no pools
+- Reuse existing handles, never connect per request or disconnect from a repository operation
+- Give each pool one owner for setup and shutdown, borrowed handles never close it
+- Finish requests, jobs and queued database work before closing connections
+- Clean up resources opened before a startup failure, without a second disconnect fallback
+- Add a backend to the existing lifecycle owner and pass its handle only to repositories that need it
+- Preserve cancellation, timeout, error and transaction semantics across the native bridge
 
-## Asked to port a Prisma query to Rust or SQLx
+## Changing ownership or reloading resources
 
-- Identify the existing inputs, outputs, errors, and transaction boundary
-- Check `litellm-rust/crates/db` for an equivalent query and reuse it if its behavior matches
-- Otherwise follow `litellm-rust/crates/db/AGENTS.md`
-  - Put static SQL in `queries/` and compile-check it with SQLx
-  - Keep typed inputs and reusable row projections in `src/`
-  - Refresh SQLx query metadata after SQL changes
-- Keep Python conversion in the Python bridge, not in `litellm-rust/crates/db`
-- When caller integration is requested, expose the Rust operation through the bridge and call it from the repository
-  - A query-only port can land without switching callers
-- Verify behavior against real PostgreSQL before switching callers
-  - Check relevant null handling, conflicts, ordering, errors, and transaction semantics
-- Remove the old implementation only when no callers need it
+- Inventory initialization, reload, tests, background jobs and shutdown together
+- Move every writer when changing authoritative storage, copying a replaceable global into state is insufficient
+- A frozen repository does not make its clients immutable
+- Decide whether each request captures one resource generation or reads it again during processing
+- Build and validate replacements before publishing them, preserving the old generation on failure
+- Serialize reloads and coordinate publication with shutdown
+- Drain streams and background users before retiring old resources, retaining every generation grows resource use
+- Define multi-worker coordination separately from an in-process swap
+- Choose mechanisms such as `ArcSwap` only after defining publication, capture and retirement
 
 ## Migration context and boundaries
 
-- Long-term direction is to centralize Python database access in repositories, then replace their Prisma implementations incrementally with Rust APIs
-- Mixed Prisma and Rust implementations are expected during the transition, unrelated existing calls can stay as they are
-- Both backends may have separate pools through PgBouncer, size them together and verify pooling-mode compatibility
+- Long-term direction is to centralize Python database access in repositories, then replace their Prisma implementations with Rust APIs one repository at a time
+- Prisma and native repositories coexist during the transition, unrelated existing calls can stay as they are
+- Prisma and SQLx keep separate client pools to PgBouncer, which pools connections to PostgreSQL
+  - Size both layers together and preserve reader/writer routing and pooling-mode compatibility
 - PgBouncer does not share transactions between Prisma and SQLx, keep each atomic unit of work on one backend and transaction
 - Never retry a failed Rust write through Prisma when it may already have committed
 - Prisma owns schema migrations today, eventual SQLx migration ownership and Prisma removal are separate planned steps, not part of a query port

.agents/skills/database-access/references/python-fastapi.md

@@ -0,0 +1,101 @@
+# Python and FastAPI
+
+## Existing pattern
+
+- `UserBannerStore` in `litellm/models/user_banner.py` is the Protocol the banner routes call
+- `get_database` in `litellm/proxy/deps/database.py` reads the lifespan-owned handle from request state
+- `get_user_banner_repository` in `litellm/proxy/deps/repositories.py` builds the banner repository from that handle
+- Choose each repository's replica routing once in its provider, never at call sites
+- A handler receives `UserBannerStore | None`, never a Prisma client or a concrete repository
+- `UserBannerRepository` implements `UserBannerStore` structurally, keeping the existing Prisma table, serialization and configuration-sync behavior
+- The banner audit provider binds the legacy audit callable at request time, without performing audit I/O
+- A native repository can later implement the same Protocol without changing handler signatures, see [porting a repository to SQLx](python-repository.md#porting-a-repository-to-sqlx)
+
+## Dependency injection
+
+- Add one provider and one `Annotated` alias per repository in `litellm/proxy/deps/repositories.py`
+- Annotate the provider and alias with the feature's Protocol, not the concrete repository
+- A route depends only on the repositories it calls, so a request never builds repositories it does not use
+- Do not build repositories once at lifespan, a reconnect can replace the Prisma client they wrap
+
+```python
+class UserBannerStore(Protocol):
+    async def get_banner(self) -> UserBanner: ...
+
+    async def save_banner(self, banner: UserBanner) -> None: ...
+```
+
+```python
+async def get_user_banner_repository(database: DatabaseDep) -> UserBannerStore | None:
+    return None if database is None else UserBannerRepository(database.prisma)
+
+
+UserBannerRepositoryDep: TypeAlias = Annotated[UserBannerStore | None, Depends(get_user_banner_repository)]
+```
+
+- FastAPI caches each provider within a request, so routes and sub-dependencies share one instance
+- Ordinary helpers and jobs receive explicit arguments, never unresolved `Depends` defaults
+- A route still reading the `prisma_client` global or building a repository inline switches to the repository's `...RepositoryDep`, keeping its no-database behavior and error precedence
+
+```python
+async def get_user_banner(banners: UserBannerRepositoryDep) -> UserBanner:
+    if banners is None:
+        return UserBanner()
+    return await banners.get_banner()
+```
+
+## Lifespan state
+
+- Prisma startup owns connection setup, authentication refresh, health checks and migrations
+- `Database` holds the initialized Prisma client and the `NativeDatabase` handle, without holding fields per feature
+  - The native handle opens no pool yet, its `connect` and `close` are noops
+- The proxy lifespan builds the handle with `Database.open`, registers `Database.aclose` on its exit stack and publishes it in request state
+  - Callers only open and close `Database`, which closes Prisma and the native handle together, and `open` closes Prisma itself when the native connect fails
+- Starlette shallow-copies lifespan state into request state, which is distinct from `app.state`
+- The database module never imports a repository, legacy globals borrow its handle without owning shutdown
+- Drain requests, jobs, logging, spend and audit work before database teardown
+- The exit stack closes the handle on exit, including startup failures after the client is created
+- When adding native resources, initialize them in the serving process, respecting runtime fork restrictions
+- Await native shutdown instead of depending on Python garbage collection
+- Apply the [connection ownership rules](../SKILL.md#connection-ownership) when changing the owner
+
+## Timing and errors
+
+- FastAPI resolves dependencies before entering the handler
+- Handler exception handling cannot catch failures raised during dependency resolution
+- A universal required-database dependency can override authorization and validation behavior
+- Banner GET without a database returns the default disabled banner
+- Banner PATCH checks the admin role before returning the existing missing-database error
+- Invalid stored banner data falls back to the default, but storage failures still propagate
+- Banner writes finish before the success audit is queued as a FastAPI background task
+- Queue post-response work with `BackgroundTasks`, not bare `asyncio.create_task`, so the in-flight drain waits for it before the database closes
+- Do not move background work into an awaited provider merely to remove an import
+
+## Verification
+
+- Exercise real FastAPI handlers using `dependency_overrides` for external boundaries
+- Keep production parsing and persistence code real
+- Restore overrides after each test
+- Verify complete response bodies, status codes and OpenAPI against the base
+- Include combined failure conditions to pin error precedence
+- Check saved values, fresh revisions, audit contents and failed-write behavior
+- Check lifespan state with a running lifespan, not only a TestClient that skips startup
+- Verify providers do not connect, query, disconnect or audit eagerly
+- Use real PostgreSQL and live proxy requests when changing persistence behavior
+
+## Scope
+
+- Preserve direct Python call sites or migrate them explicitly
+- Inventory imports, state writers, side effects and error ordering before migrating a feature
+- Preserve configuration sync and serialization, and state which legacy reads remain
+- Do not migrate every route in a large module just to establish the pattern
+- Do not invent providers for resources without consumers
+
+## Framework references
+
+- [Lifespan events](https://fastapi.tiangolo.com/advanced/events/) and [Starlette lifespan state](https://starlette.dev/lifespan/)
+- [`SessionDep`-style `Annotated` aliases](https://fastapi.tiangolo.com/tutorial/sql-databases/#create-a-session-dependency), which `DatabaseDep` and `UserBannerRepositoryDep` follow
+- [Dependency composition and caching](https://fastapi.tiangolo.com/tutorial/dependencies/sub-dependencies/)
+- [Background tasks and dependencies with `yield`](https://fastapi.tiangolo.com/advanced/advanced-dependencies/#background-tasks-and-dependencies-with-yield-technical-details)
+- [Dependency overrides](https://fastapi.tiangolo.com/advanced/testing-dependencies/)
+- [Testing lifespan](https://fastapi.tiangolo.com/advanced/testing-events/)

.agents/skills/database-access/references/python-repository.md

@@ -0,0 +1,119 @@
+# Python repositories
+
+## Adding a query
+
+Use this when a feature needs a read or write that no repository offers yet, including one more operation on an existing repository
+
+### Before starting
+
+- Search `litellm/repositories` for a repository that already owns the table, and extend it instead of adding a second one
+- Check `litellm/models/<feature>.py` for an existing Protocol and domain types
+- Decide the route's behavior when no database is configured, and where that check sits relative to authorization and validation
+
+### Steps
+
+1. Domain types and Protocol in `litellm/models/<feature>.py`
+   - The Protocol takes and returns domain types, never Prisma models, raw rows or serialized payloads
+   - Add the Protocol together with the provider, not for repositories nothing injects
+2. Repository in `litellm/repositories/<feature>_repository.py`
+   - Implement the operation with Prisma, a Rust port is not required
+   - Compose the Prisma client and type the tables it touches with `TableActions` from `prisma_protocols.py`, never inherit `PrismaTableRepository`
+   - Parse and serialize stored values here, so callers only see domain types
+   - Never import FastAPI, connect, disconnect or choose replicas per call
+3. Provider and `Annotated` alias in `litellm/proxy/deps/repositories.py`, see [dependency injection](python-fastapi.md#dependency-injection)
+4. The route takes `<Feature>RepositoryDep`, handles `None` with the behavior decided above, and calls the operation
+5. Jobs and helpers that need the operation receive the repository as an explicit argument, built by their caller from the handle it already holds
+
+### Done when
+
+- Handler tests drive the real route with `dependency_overrides` and the real repository code, see [verification](python-fastapi.md#verification)
+- Tests cover the no-database path, error precedence, and what was actually stored
+- Persistence changes are checked against real PostgreSQL through a live proxy
+
+### Out of scope
+
+- Porting the query to Rust, see [porting a repository to SQLx](#porting-a-repository-to-sqlx)
+- Moving unrelated direct Prisma calls in the same module
+
+## Moving existing access into a repository
+
+Use this when code calls Prisma directly, through `prisma_client.db.<table>` or a helper, and the calls should live behind a repository. The backend stays Prisma
+
+### Before starting
+
+- List every call being moved, its callers, and the transaction or batch it runs in
+- List the module's imports, global state writers, side effects such as audit and config sync, and error ordering
+- Record current return values, raised errors and response bodies, these are what the change must preserve
+
+### Steps
+
+1. Move the existing Prisma code into the relevant repository in `litellm/repositories` without rewriting the queries
+2. Add domain types and a Protocol in `litellm/models/<feature>.py` when a route will inject the repository
+3. Replace each direct call with the repository operation, through a provider for routes and explicit arguments for jobs and helpers, see [Python and FastAPI](python-fastapi.md)
+4. Calls sharing a transaction move into the same repository or unit of work in one change
+5. State which legacy reads remain and why, when the change does not move all of them
+
+### Done when
+
+- Focused regression tests pin the preserved return values, errors and transaction behavior
+- Response bodies, status codes and OpenAPI match the base
+- Configuration sync and serialization behave the same, including invalid stored data
+
+### Out of scope
+
+- Switching the backend, even when an equivalent Rust query exists
+- Converting other routes in the module just to establish the pattern
+- Cleanups unrelated to the moved calls
+
+## Porting a repository to SQLx
+
+Status: TBD. No repository has been ported yet. `database.native` is a `NativeDatabase` from `crates/python-bridge/src/db.rs` whose `connect` and `close` are noops that open no pool, and the PyO3 repository surface does not exist. A real pool takes typed writer/reader settings and a token provider, never the Prisma `DATABASE_URL`, which carries Prisma-only params and rotating IAM tokens. The constraints below are settled, the steps are a sketch to revise after the first port lands
+
+Use this only when a port is explicitly requested. Repository extraction does not require one
+
+### When not to port
+
+- An equivalent Rust or SQLx query existing is not a reason to switch a Prisma caller
+- Leave existing backend selection unchanged unless integration is explicitly requested
+- Never port part of a repository, it is either all Prisma or all native
+
+### Before starting
+
+- Identify the repository's inputs, outputs, errors and transaction boundary
+- List repositories that share a transaction with it, they move together
+- Check `litellm-rust/crates/db` for an equivalent query and reuse it when its behavior matches
+
+### Steps
+
+1. Rust queries, following `litellm-rust/crates/db/AGENTS.md`
+   - Static SQL in `queries/`, compile-checked with SQLx, typed inputs and row projections in `src/`
+   - Dynamic filters and sort keys become closed Rust enums that select static SQL, every value is bound
+   - Refresh SQLx query metadata after SQL changes
+2. PyO3 surface exposing typed operations, never SQL text, table names or column names from Python. TBD: where it lives and how the handle reaches it
+3. Native repository implementing the same Protocol in `litellm/repositories`
+   - Converts typed PyO3 results to domain types, preserving defaults, invalid-data handling and storage errors
+   - Python conversion stays in the bridge, not in `crates/db`
+4. Switch the provider in `litellm/proxy/deps/repositories.py`, the Protocol and routes stay unchanged
+
+```diff
+ async def get_user_banner_repository(database: DatabaseDep) -> UserBannerStore | None:
+-    return None if database is None else UserBannerRepository(database.prisma)
++    return None if database is None else NativeUserBannerRepository(database.native)
+```
+
+5. Delete the Prisma repository in the same change and switch every caller, including jobs and helpers
+
+### Done when
+
+- Behavior is verified against real PostgreSQL, covering null handling, conflicts, ordering, errors and transaction semantics
+- `DATABASE_URL=... cargo test -p litellm-db --locked -- --include-ignored` passes against a Prisma-migrated database
+- Live proxy requests show unchanged responses. TBD: a Python-side real-PostgreSQL test harness
+
+### Rollback
+
+- Roll back the release, never retry an uncertain native write through Prisma
+
+### Out of scope
+
+- SQLx owning schema migrations, and removing Prisma
+- Converting unrelated routes to dependency injection

.agents/skills/database-access/references/rust-axum.md

@@ -0,0 +1,4 @@
+# Rust and Axum
+
+- This should use same `crates/db` what is exposed to Python world.
+- Details are TBD.

.github/workflows/test-unit.yml

@@ -577,6 +577,7 @@ jobs:
               tests/unit/proxy/rerank_endpoints
               tests/unit/proxy/realtime_endpoints
               tests/unit/proxy/ui_crud_endpoints
+              tests/unit/proxy/deps
               tests/unit/proxy/config_resolvers
               tests/unit/proxy/utils
             python-version: "3.12"

AGENTS.md

@@ -90,6 +90,8 @@ Do not put names of customers or customer company names in code, PR descriptions
 
 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
 
+For database operations, connection ownership and proxy dependency wiring, follow [database-access](.agents/skills/database-access/SKILL.md)
+
 Prisma migrations apply synchronously at proxy boot, before it serves traffic, so a migration must only change schema, never rewrite rows. No `UPDATE`, `DELETE` or `MERGE`, and no `INSERT ... SELECT`: on a spend-log-sized table any of those is minutes of downtime plus a doubled heap that plain autovacuum won't give back. `tests/code_coverage_tests/check_migrations_no_data_rewrites.py` enforces this. When a rewrite is genuinely bounded and has to ship inside the migration, mark the statement `-- data-migration-ok: <what bounds it>`
 
 Follow these coding conventions for new/updated code (a three-line fix in a legacy file shouldn't trigger huge drive-by refactors):

CONTRIBUTING.md

@@ -276,6 +276,8 @@ When the provider reports its own usage or spend numbers, compare the proxy's nu
 
 Provider work (a new provider, a new parameter, a new usage or cost field) is the most common kind of contribution, and it is where PRs get reworked most often. Follow these rules
 
+**Preserve zero prices.** In [`model_prices_and_context_window.json`](model_prices_and_context_window.json), an explicit `0.0` means that the priced operation is free. A missing field or `null` means no rate is specified, so the relevant calculator may use a fallback. For example, when no batch input rate is configured, batch cache writes use `cache_creation_input_token_cost` if set, otherwise the regular input rate, then apply the batch discount. A zero ordinary input rate does not make separately priced cache reads or writes free. Test `rate is None` before substituting a fallback, since `row.get("cache_creation_input_token_cost") or fallback` would discard a valid zero. Speech keeps character pricing ahead of token pricing even when the character rate is zero. If neither rate is configured, speech cost calculation raises an error. Its model lookup preserves unset rates instead of treating compatibility defaults as explicit zero prices. Perplexity keeps a configured legacy search rate ahead of the current search rate, including zero. When registering custom prices, `None` leaves an existing rate unchanged and `0.0` overrides it
+
 **Find the existing machinery before you add plumbing.** LiteLLM already has generic paths that move data between the endpoints, the provider transformations, and cost tracking. A second, hand-rolled path for the same data is a common reason a provider PR gets rewritten before merge. [`ARCHITECTURE.md`](ARCHITECTURE.md) shows where the layers live, and the decision records in [`docs/adrs/`](docs/adrs/AGENTS.md) explain why they are shaped that way. For example, [the ADR on provider usage fields](docs/adrs/adr_001_provider_usage_fields_ride_as_usage_extras.md) covers how provider-specific usage fields reach cost tracking. If you cannot find a mechanism for what you need, say so in the PR or ask in [#pr-review on Slack](https://join.slack.com/t/litellmossslack/shared_invite/zt-3o7nkuyfr-p_kbNJj8taRfXGgQI1~YyA) before you build one
 
 **Bill a new provider usage field from the usage object.** Leave the field on the usage object as an extra and read it in that provider's `cost_calculator.py`. A provider without one gets a new `cost_calculator.py` and an arm for it in `cost_per_token` (`litellm/cost_calculator.py`), which picks the calculator by provider. If the provider's transform builds `Usage` from named fields, as the Anthropic, Gemini, and Bedrock Converse ones do, pass the new field there too. Streamed requests rebuild usage in `ChunkProcessor.calculate_usage` (`litellm/litellm_core_utils/streaming_chunk_builder_utils.py`), which does not copy extras yet the way the Responses-to-chat bridge does (minus a short exclusion list), so until it does, bill the field on streams by naming it there, as Anthropic's `inference_geo` is. A provider that reports its own charge for the request, as OpenRouter does in `usage.cost`, is billed at that charge from `_hidden_params` before any calculator runs, unless the deployment sets its own prices. When the field counts a built-in tool such as web search, add the provider to `get_cost_for_web_search_request` in `litellm/llms/__init__.py`, and add a gate to `litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py` only when no existing gate fires. Then show the field's cost in your proof on every endpoint it reaches, streaming included

enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py

@@ -10,6 +10,7 @@ from dataclasses import replace as dataclasses_replace
 from datetime import datetime, timedelta
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Final, List, Literal, Optional, Protocol, Tuple, cast
+from pydantic import TypeAdapter, ValidationError
 
 from litellm_enterprise.proxy.common_utils.managed_object_attribution import (
     ManagedObjectCreatorAttribution,
@@ -50,6 +51,7 @@ def _error_names_resource(error: Exception, resource_id: str) -> bool:
 CHECK_BATCH_COST_USER_AGENT = "LiteLLM Proxy/CheckBatchCost"
 RECENT_BATCH_POLL_WINDOW: Final = timedelta(hours=48)
 GIVE_UP_POLLING_AFTER: Final = timedelta(days=2 * MANAGED_OBJECT_STALENESS_CUTOFF_DAYS)
+_FILE_OBJECT_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(dict[str, object])
 
 
 PROVIDER_TERMINAL_BATCH_STATUSES: Final[Tuple[str, ...]] = (
@@ -550,6 +552,10 @@ class CheckBatchCost:
         llm_provider: str,
         bare_model_name: str,
     ) -> Optional[Tuple[str, str]]:
+        stored_model_id: Final = self._get_stored_model_id(job)
+        if stored_model_id is not None and self._batch_deployment_exists(stored_model_id):
+            return stored_model_id, job.unified_object_id
+
         deployment_id = self._get_deployment_id_for_bare_model(bare_model_name, llm_provider)
         if deployment_id is None:
             verbose_proxy_logger.info(
@@ -662,23 +668,36 @@ class CheckBatchCost:
         )
 
     @staticmethod
-    def _get_input_file_id(job: "ManagedObjectRow") -> Optional[str]:
-        from litellm.types.utils import LiteLLMBatch
-
+    def _parse_file_object(job: "ManagedObjectRow") -> Optional[dict[str, object]]:
         file_object: Final = job.file_object
         try:
-            parsed_file_object: Final[object] = (
+            return _FILE_OBJECT_ADAPTER.validate_python(
                 json.loads(file_object) if isinstance(file_object, str) else file_object
             )
-        except (json.JSONDecodeError, ValueError):
+        except (json.JSONDecodeError, ValueError, ValidationError):
             return None
-        if not isinstance(parsed_file_object, dict):
+
+    @classmethod
+    def _get_input_file_id(cls, job: "ManagedObjectRow") -> Optional[str]:
+        from litellm.types.utils import LiteLLMBatch
+
+        parsed_file_object: Final = cls._parse_file_object(job)
+        if parsed_file_object is None:
             return None
         try:
             return LiteLLMBatch.model_validate(parsed_file_object).input_file_id
         except Exception:
             return None
 
+    @classmethod
+    def _get_stored_model_id(cls, job: "ManagedObjectRow") -> Optional[str]:
+        """The deployment that created the batch, written into file_object by the managed
+        files hook at create time. Older rows have no such key and fall back to matching
+        the bare model name against the router."""
+        parsed_file_object: Final = cls._parse_file_object(job)
+        model_id: Final = parsed_file_object.get("model_id") if parsed_file_object else None
+        return model_id if isinstance(model_id, str) and model_id else None
+
     async def _track_completed_batch_cost(
         self,
         job: "ManagedObjectRow",

enterprise/litellm_enterprise/proxy/hooks/managed_files.py

@@ -21,6 +21,7 @@ from typing import (
     TypedDict,
     Union,
     cast,
+    runtime_checkable,
 )
 from uuid import NAMESPACE_URL, uuid5
 
@@ -374,6 +375,12 @@ class _ManagedResourcePrismaClient(Protocol):
     def db(self) -> _ManagedResourceDatabase: ...
 
 
+@runtime_checkable
+class _WriterRoutedDatabase(Protocol):
+    @property
+    def writer(self) -> _ManagedResourceDatabase: ...
+
+
 class _SchedulerWithJobLookup(Protocol):
     def get_job(self, job_id: str) -> object: ...
 
@@ -479,6 +486,11 @@ def _managed_object_table(prisma_client: _ManagedResourcePrismaClient) -> _Manag
     return prisma_client.db.litellm_managedobjecttable
 
 
+def _writer_database(db: _ManagedResourceDatabase) -> _ManagedResourceDatabase:
+    """Reads that feed a write must see the create that just happened; a lagging read replica may not"""
+    return db.writer if isinstance(db, _WriterRoutedDatabase) else db
+
+
 _JSON_VALUE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
 _POSTGRES_UNSTORABLE_CHARACTER: Final = "\x00"
 
@@ -512,6 +524,25 @@ def _without_model_output(
     )
 
 
+def _with_model_id(
+    file_object: LiteLLMBatch | LiteLLMFineTuningJob | ResponsesAPIResponse,
+    model_id: str | None,
+) -> LiteLLMBatch | LiteLLMFineTuningJob | ResponsesAPIResponse:
+    return file_object if model_id is None else file_object.model_copy(update={"model_id": model_id})
+
+
+async def _stored_model_id(managed_object_table: _ManagedObjectTableActions, unified_object_id: str) -> str | None:
+    row: Final = await managed_object_table.find_first(where={"unified_object_id": unified_object_id})
+    if row is None:
+        return None
+    try:
+        parsed: Final = _JSON_VALUE_ADAPTER.validate_python(_decode_json_blob(row.file_object))
+    except (json.JSONDecodeError, ValidationError):
+        return None
+    model_id: Final = parsed.get("model_id") if isinstance(parsed, dict) else None
+    return model_id if isinstance(model_id, str) and model_id else None
+
+
 def _storage_metadata_of(file_object: OpenAIFileObject | None) -> Mapping[str, str]:
     hidden_params: Final = cast(  # cast-ok: _hidden_params is an untyped attribute the upload path sets
         "Mapping[str, object]", getattr(file_object, "_hidden_params", None) or {}
@@ -697,9 +728,17 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         persist_attribution: bool = False,
         create_if_missing: bool = True,
         batch_processed: bool = False,
+        creating_model_id: str | None = None,
     ) -> None:
         """Persist a managed object row, caching it and upserting it in the DB.
 
+        creating_model_id is the router deployment that created the object. It is kept
+        inside the stored file_object JSON as `model_id`, because the raw provider id an
+        unmanaged batch stores as unified_object_id says nothing about which deployment
+        made it, and the cost poller needs that deployment to read the result back. A
+        later observe of the row (create_if_missing=False) keeps the stored value rather
+        than overwriting it with whichever deployment served the retrieve.
+
         persist_attribution is set only by the batch and background response creates,
         the callers that can speak for the creator; it gates the api_key, request_tags, and,
         for a response, end_user_id columns that the cost poll jobs bill against, so a later poll or
@@ -720,7 +759,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         It is written only in the upsert create branch.
         """
         verbose_logger.info(f"Storing LiteLLM Managed {file_purpose} object with id={unified_object_id} in cache")
-        stored_object: Final = _without_model_output(file_object)
+        managed_database: Final = _managed_database(self.prisma_client)
+        managed_object_table: Final = managed_database.litellm_managedobjecttable
+        preserved_model_id: Final = (
+            creating_model_id
+            if creating_model_id is not None or create_if_missing
+            else await _stored_model_id(_writer_database(managed_database).litellm_managedobjecttable, unified_object_id)
+        )
+        stored_object: Final = _with_model_id(_without_model_output(file_object), preserved_model_id)
         litellm_managed_object = LiteLLM_ManagedObjectTable(
             unified_object_id=unified_object_id,
             model_object_id=model_object_id,
@@ -752,7 +798,6 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
             "status": file_object.status,
             "updated_by": user_api_key_dict.user_id,
         }
-        managed_object_table: Final = _managed_object_table(self.prisma_client)
         if not create_if_missing:
             await managed_object_table.update_many(
                 where={"unified_object_id": unified_object_id},
@@ -2170,6 +2215,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
                 request_tags=request_tags_from_metadata(request_metadata if isinstance(request_metadata, dict) else {}),
                 persist_attribution=is_batch_create,
                 create_if_missing=is_batch_create,
+                creating_model_id=model_id if is_batch_create else None,
             )
             if not is_batch_create:
                 await self.store_batch_output_file_ownership(

litellm-rust/crates/python-bridge/src/db.rs

@@ -0,0 +1,19 @@
+use litellm_host_python::ready_future;
+use pyo3::{prelude::*, types::PyNone};
+
+#[pyclass(frozen)]
+pub(crate) struct NativeDatabase {}
+
+// TODO: take typed writer/reader settings and a token provider, then open a sqlx PgPool, once the first
+// repository moves to native. The Prisma DATABASE_URL carries Prisma-only params and rotating IAM tokens
+#[pymethods]
+impl NativeDatabase {
+    #[staticmethod]
+    fn connect(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
+        ready_future(py, Bound::new(py, NativeDatabase {})?.as_any())
+    }
+
+    fn close<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+        ready_future(py, PyNone::get(py).as_any())
+    }
+}

litellm-rust/crates/python-bridge/src/lib.rs

@@ -2,6 +2,7 @@ mod cache;
 mod callable;
 mod coercion;
 mod credentials;
+mod db;
 mod diagnostics;
 mod errors;
 mod execution;
@@ -21,6 +22,8 @@ mod _native {
     use litellm_host_python::{ForkedAfterNativeRuntimeStarted, ProcessReservedForForking};
     use pyo3::{prelude::*, types::PyModule};
 
+    #[pymodule_export]
+    use crate::db::NativeDatabase;
     #[cfg(feature = "panic-test")]
     #[pymodule_export]
     use crate::diagnostics::_panic_for_test;
@@ -112,6 +115,7 @@ mod tests {
                 "NativeDiagnosticProcessor",
                 "NativeClickHouseSpendConfig",
                 "NativeClickHouseSpendStorage",
+                "NativeDatabase",
                 "trace_encode_error",
                 "bundled_model_cost_map",
                 "TokenCounter",

litellm/_lazy_imports_registry.py

@@ -1469,11 +1469,11 @@ _UTILS_MODULE_IMPORT_MAP: Final = {
     "AsyncHTTPHandler": ("litellm.llms.custom_httpx.http_handler", "AsyncHTTPHandler"),
     "HTTPHandler": ("litellm.llms.custom_httpx.http_handler", "HTTPHandler"),
     "get_num_retries_from_retry_policy": (
-        "litellm.router_utils.get_retry_from_policy",
+        "litellm.litellm_core_utils.get_retry_from_policy",
         "get_num_retries_from_retry_policy",
     ),
     "reset_retry_policy": (
-        "litellm.router_utils.get_retry_from_policy",
+        "litellm.litellm_core_utils.get_retry_from_policy",
         "reset_retry_policy",
     ),
     "get_coroutine_checker": (

litellm/a2a_protocol/litellm_completion_bridge/handler.py

@@ -22,7 +22,7 @@ from litellm.a2a_protocol.litellm_completion_bridge.transformation import (
     A2AStreamingContext,
 )
 from litellm.a2a_protocol.providers.config_manager import A2AProviderConfigManager
-from litellm.constants import A2A_CHAT_COMPLETION_BRIDGE_PROVIDERS
+from litellm.constants import A2A_CHAT_COMPLETION_BRIDGE_PROVIDERS, A2A_COMPLETION_BRIDGE_WORKING_MESSAGE
 from litellm.interactions.agents.utils import merge_agent_headers
 from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
 from litellm.types.utils import ModelResponse
@@ -268,7 +268,7 @@ class A2ACompletionBridgeHandler:
             ctx=ctx,
             state="working",
             final=False,
-            message_text="Processing request...",
+            message_text=A2A_COMPLETION_BRIDGE_WORKING_MESSAGE,
         )
         yield working_event
 

litellm/a2a_protocol/main.py

@@ -191,6 +191,23 @@ def _set_litellm_params_on_logging_obj(
     }
 
 
+async def _record_send_message_billing(
+    kwargs: Mapping[str, object],
+    request: "SendMessageRequest",
+    response_dict: Mapping[str, object],
+    litellm_params: Mapping[str, object],
+    agent_id: str | None,
+) -> None:
+    prompt_tokens, completion_tokens, _ = await asyncify(A2ARequestUtils.calculate_usage_from_request_response)(
+        request=request,
+        response_dict=response_dict,
+    )
+    _set_usage_on_logging_obj(kwargs=kwargs, prompt_tokens=prompt_tokens, completion_tokens=completion_tokens)
+    if "error" not in response_dict:
+        _set_litellm_params_on_logging_obj(kwargs=kwargs, litellm_params=litellm_params)
+    _set_agent_id_on_logging_obj(kwargs=kwargs, agent_id=agent_id)
+
+
 def _get_a2a_model_info(a2a_client: "A2AClientType", kwargs: Mapping[str, object]) -> str:
     """
     Extract agent info and set model/custom_llm_provider for cost tracking.
@@ -483,13 +500,21 @@ async def asend_message(
     if custom_llm_provider:
         if request is None:
             raise ValueError("request is required for completion bridge")
-        return await _send_message_via_completion_bridge(
+        bridge_response: Final = await _send_message_via_completion_bridge(
             request=request,
             custom_llm_provider=custom_llm_provider,
             api_base=api_base,
             litellm_params=litellm_params,
             agent_extra_headers=agent_extra_headers,
         )
+        await _record_send_message_billing(
+            kwargs=kwargs,
+            request=request,
+            response_dict=bridge_response.model_dump(mode="json", exclude_none=True),
+            litellm_params=litellm_params,
+            agent_id=agent_id,
+        )
+        return bridge_response
 
     # Standard A2A client flow
     if request is None:
@@ -538,30 +563,14 @@ async def asend_message(
     # Wrap in LiteLLM response type for _hidden_params support
     response: Final = LiteLLMSendMessageResponse.from_a2a_response(a2a_response, request_id=str(request.id))
 
-    # Calculate token usage from request and response
-    response_dict: Final[dict[str, object]] = a2a_response.root.model_dump(mode="json", exclude_none=True)
-    (
-        prompt_tokens,
-        completion_tokens,
-        _,
-    ) = await asyncify(A2ARequestUtils.calculate_usage_from_request_response)(
-        request=request,
-        response_dict=response_dict,
-    )
-
-    # Set usage on logging obj for standard logging payload
-    _set_usage_on_logging_obj(
+    await _record_send_message_billing(
         kwargs=kwargs,
-        prompt_tokens=prompt_tokens,
-        completion_tokens=completion_tokens,
+        request=request,
+        response_dict=a2a_response.root.model_dump(mode="json", exclude_none=True),
+        litellm_params=litellm_params,
+        agent_id=agent_id,
     )
 
-    # Merge agent pricing params into the logging obj so cost is calculated
-    _set_litellm_params_on_logging_obj(kwargs=kwargs, litellm_params=litellm_params)
-
-    # Set agent_id on logging obj for SpendLogs tracking
-    _set_agent_id_on_logging_obj(kwargs=kwargs, agent_id=agent_id)
-
     return response
 
 
@@ -646,6 +655,7 @@ async def asend_message_streaming(
     metadata: dict[str, object] | None = None,
     proxy_server_request: dict[str, object] | None = None,
     agent_extra_headers: dict[str, str] | None = None,
+    agent_name: str | None = None,
     **kwargs: object,
 ) -> AsyncIterator["SendStreamingMessageResponse | Mapping[str, object]"]:
     """
@@ -661,6 +671,7 @@ async def asend_message_streaming(
         agent_id: Optional agent ID for tracking in SpendLogs
         metadata: Optional metadata dict (contains user_api_key, user_id, team_id, etc.)
         proxy_server_request: Optional proxy server request data
+        agent_name: Agent name logged as a2a_agent/<agent_name> for completion bridge agents
 
     Yields:
         SendStreamingMessageResponse chunks from the agent
@@ -685,6 +696,7 @@ async def asend_message_streaming(
             print(chunk)
         ```
     """
+    called_at: Final = datetime.datetime.now()
     litellm_params = litellm_params or {}
     custom_llm_provider: Final = litellm_params.get("custom_llm_provider")
 
@@ -705,20 +717,47 @@ async def asend_message_streaming(
             request.params.model_dump(mode="json") if hasattr(request.params, "model_dump") else dict(request.params)
         )
 
-        async for chunk in A2ACompletionBridgeHandler.handle_streaming(
+        bridge_stream: Final = A2ACompletionBridgeHandler.handle_streaming(
             request_id=str(request.id),
             params=params,
             litellm_params=litellm_params,
             api_base=api_base,
             agent_extra_headers=agent_extra_headers,
-        ):
-            yield chunk
+        )
+        caller_logging_obj: Final = kwargs.get("litellm_logging_obj")
+        bridge_logging_obj: Final = (
+            caller_logging_obj
+            if isinstance(caller_logging_obj, Logging)
+            else _build_streaming_logging_obj(
+                request=request,
+                agent_name=agent_name or "unknown",
+                agent_id=agent_id,
+                litellm_params=litellm_params,
+                metadata=metadata,
+                proxy_server_request=proxy_server_request,
+                start_time=called_at,
+            )
+        )
+        _set_litellm_params_on_logging_obj(kwargs=kwargs, litellm_params=litellm_params)
+        _set_agent_id_on_logging_obj(kwargs=kwargs, agent_id=agent_id)
+
+        bridge_iterator: Final = A2AStreamingIterator(
+            stream=bridge_stream,
+            request=request,
+            logging_obj=bridge_logging_obj,
+            agent_name=agent_name or "unknown",
+        )
+        try:
+            async for chunk in bridge_iterator:
+                yield chunk
+        except (GeneratorExit, asyncio.CancelledError):
+            await bridge_iterator.abill_on_disconnect()
+            raise
         return
 
     if request is None:
         raise ValueError("request is required")
 
-    called_at: Final = datetime.datetime.now()
     _raw_logging_obj: Final = kwargs.get("litellm_logging_obj")
     logging_obj: Logging | None = _raw_logging_obj if isinstance(_raw_logging_obj, Logging) else None
 
@@ -742,12 +781,12 @@ async def asend_message_streaming(
 
     assert a2a_client is not None
 
-    agent_name: Final = _get_a2a_model_info(a2a_client, kwargs)
+    card_agent_name: Final = _get_a2a_model_info(a2a_client, kwargs)
 
     if logging_obj is None:
         logging_obj = _build_streaming_logging_obj(
             request=request,
-            agent_name=agent_name,
+            agent_name=card_agent_name,
             agent_id=agent_id,
             litellm_params=litellm_params,
             metadata=metadata,
@@ -755,7 +794,7 @@ async def asend_message_streaming(
             start_time=called_at,
         )
 
-    verbose_logger.info("A2A send_message_streaming request_id=%s, agent=%s", request.id, agent_name)
+    verbose_logger.info("A2A send_message_streaming request_id=%s, agent=%s", request.id, card_agent_name)
 
     agent_card: Final = _get_a2a_client_agent_card(a2a_client)
     card_url: Final = get_agent_card_url(agent_card) if agent_card else None
@@ -766,18 +805,23 @@ async def asend_message_streaming(
         agent_card=agent_card,
         card_url=card_url,
         api_base=api_base,
-        agent_name=agent_name,
+        agent_name=card_agent_name,
     )
 
     _set_agent_id_on_logging_obj(kwargs=kwargs, agent_id=agent_id)
 
-    async for chunk in A2AStreamingIterator(
+    url_agent_iterator: Final = A2AStreamingIterator(
         stream=stream,
         request=request,
         logging_obj=logging_obj,
-        agent_name=agent_name,
-    ):
-        yield chunk
+        agent_name=card_agent_name,
+    )
+    try:
+        async for chunk in url_agent_iterator:
+            yield chunk
+    except (GeneratorExit, asyncio.CancelledError):
+        await url_agent_iterator.abill_on_disconnect()
+        raise
 
 
 async def create_a2a_client(

litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py

@@ -515,6 +515,8 @@ class PydanticAITransformation:
                             "contextId": context_id,
                             "kind": "artifact-update",
                             "taskId": task_id,
+                            "append": i > 0,
+                            "lastChunk": is_last_chunk,
                             "artifact": {
                                 "artifactId": artifact_id,
                                 "parts": [

litellm/a2a_protocol/providers/watsonx_orchestrate/transformation.py

@@ -206,6 +206,8 @@ class WatsonxOrchestrateTransformation:
                     "contextId": context_id,
                     "kind": "artifact-update",
                     "taskId": task_id,
+                    "append": i > 0,
+                    "lastChunk": is_last,
                     "artifact": {
                         "artifactId": artifact_id,
                         "parts": [{"kind": "text", "text": chunk_text}],

litellm/a2a_protocol/streaming_iterator.py

@@ -3,16 +3,18 @@ A2A Streaming Iterator with token tracking and logging support.
 """
 
 import asyncio
-from collections.abc import AsyncIterator
+from collections.abc import AsyncIterator, Mapping
 from datetime import datetime
 from typing import TYPE_CHECKING, Final
 
+import anyio
 from pydantic import TypeAdapter
 
 import litellm
 from litellm._logging import verbose_logger
 from litellm.a2a_protocol.cost_calculator import A2ACostCalculator
 from litellm.a2a_protocol.utils import A2ARequestUtils
+from litellm.constants import A2A_COMPLETION_BRIDGE_WORKING_MESSAGE
 from litellm.litellm_core_utils.asyncify import asyncify
 from litellm.litellm_core_utils.core_helpers import as_str_mapping
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
@@ -23,6 +25,23 @@ if TYPE_CHECKING:
 _OPAQUE_PAYLOAD: Final = TypeAdapter(object)
 
 
+def _chunk_payload(chunk: "SendStreamingMessageResponse | Mapping[str, object]") -> Mapping[str, object]:
+    if isinstance(chunk, Mapping):
+        return chunk
+    return chunk.model_dump(mode="json", exclude_none=True) if hasattr(chunk, "model_dump") else {}
+
+
+def _artifact_update(payload: Mapping[str, object]) -> tuple[str, bool] | None:
+    result: Final = as_str_mapping(payload.get("result"))
+    if result is None or result.get("kind") != "artifact-update":
+        return None
+    artifact: Final = as_str_mapping(result.get("artifact"))
+    artifact_id: Final = artifact.get("artifactId") if artifact is not None else None
+    if not isinstance(artifact_id, str):
+        return None
+    return artifact_id, result.get("append") is True
+
+
 class A2AStreamingIterator:
     """
     Async iterator for A2A streaming responses with token tracking.
@@ -32,7 +51,7 @@ class A2AStreamingIterator:
 
     def __init__(
         self,
-        stream: AsyncIterator["SendStreamingMessageResponse"],
+        stream: AsyncIterator["SendStreamingMessageResponse | Mapping[str, object]"],
         request: "SendStreamingMessageRequest",
         logging_obj: LiteLLMLoggingObj,
         agent_name: str = "unknown",
@@ -44,14 +63,16 @@ class A2AStreamingIterator:
         self.start_time = logging_obj.start_time
 
         # Collect chunks for token counting
-        self.chunks: list[SendStreamingMessageResponse] = []
-        self.collected_text_parts: list[str] = []
-        self.final_chunk: SendStreamingMessageResponse | None = None
+        self.chunks: list[SendStreamingMessageResponse | Mapping[str, object]] = []
+        self.latest_message_text = ""
+        self.artifact_texts: Mapping[str, str] = {}
+        self.saw_error = False
+        self.final_chunk: SendStreamingMessageResponse | Mapping[str, object] | None = None
 
     def __aiter__(self):
         return self
 
-    async def __anext__(self) -> "SendStreamingMessageResponse":
+    async def __anext__(self) -> "SendStreamingMessageResponse | Mapping[str, object]":
         try:
             chunk: Final = await self.stream.__anext__()
             if self.logging_obj.completion_start_time is None:
@@ -59,38 +80,54 @@ class A2AStreamingIterator:
 
             # Store chunk
             self.chunks.append(chunk)
+            payload: Final = _chunk_payload(chunk)
+            self.saw_error = self.saw_error or payload.get("error") is not None
 
             # Extract text from chunk for token counting
-            self._collect_text_from_chunk(chunk)
+            self._collect_text_from_chunk(payload)
 
             # Check if this is the final chunk (completed status)
-            if self._is_completed_chunk(chunk):
+            if self._is_completed_chunk(payload):
                 self.final_chunk = chunk
 
             return chunk
 
         except StopAsyncIteration:
-            # Stream ended - handle logging
-            if self.final_chunk is None and self.chunks:
-                self.final_chunk = self.chunks[-1]
             await self._handle_stream_complete()
             raise
 
-    def _collect_text_from_chunk(self, chunk: "SendStreamingMessageResponse") -> None:
-        """Extract text from a streaming chunk and add to collected parts."""
+    async def abill_on_disconnect(self) -> None:
+        if not self.chunks:
+            return
+        with anyio.CancelScope(shield=True):
+            await self._handle_stream_complete()
+
+    def _collect_text_from_chunk(self, payload: Mapping[str, object]) -> None:
+        """Record the reply text a streaming chunk carries."""
         try:
-            chunk_dict: Final = chunk.model_dump(mode="json", exclude_none=True) if hasattr(chunk, "model_dump") else {}
-            text: Final = A2ARequestUtils.extract_text_from_response(chunk_dict)
-            if text:
-                self.collected_text_parts.append(text)
+            text: Final = A2ARequestUtils.extract_text_from_stream_event(payload)
+            if not text:
+                return
+            artifact_update: Final = _artifact_update(payload)
+            if artifact_update is None and text == A2A_COMPLETION_BRIDGE_WORKING_MESSAGE:
+                return
+            if artifact_update is None:
+                self.latest_message_text = text
+                return
+            artifact_id, append = artifact_update
+            artifact_text: Final = self.artifact_texts.get(artifact_id, "") + text if append else text
+            self.artifact_texts = {**self.artifact_texts, artifact_id: artifact_text}
         except Exception:
             verbose_logger.debug("Failed to extract text from A2A streaming chunk")
 
-    def _is_completed_chunk(self, chunk: "SendStreamingMessageResponse") -> bool:
+    def _reply_text(self) -> str:
+        artifact_text: Final = " ".join(text for text in self.artifact_texts.values() if text)
+        return artifact_text or self.latest_message_text
+
+    def _is_completed_chunk(self, payload: Mapping[str, object]) -> bool:
         """Check if chunk indicates stream completion."""
         try:
-            chunk_dict: Final = chunk.model_dump(mode="json", exclude_none=True) if hasattr(chunk, "model_dump") else {}
-            result: Final = chunk_dict.get("result", {})
+            result: Final = payload.get("result", {})
             if isinstance(result, dict):
                 status: Final = result.get("status", {})
                 if isinstance(status, dict):
@@ -101,6 +138,8 @@ class A2AStreamingIterator:
 
     async def _handle_stream_complete(self) -> None:
         """Handle logging and token counting when stream completes."""
+        if self.final_chunk is None and self.chunks:
+            self.final_chunk = self.chunks[-1]
         try:
             end_time: Final = datetime.now()
 
@@ -109,9 +148,7 @@ class A2AStreamingIterator:
             input_text: Final = A2ARequestUtils.extract_text_from_message(input_message)
             prompt_tokens: Final = await asyncify(A2ARequestUtils.count_tokens)(input_text)
 
-            # Use the last (most complete) text from chunks
-            output_text: Final = self.collected_text_parts[-1] if self.collected_text_parts else ""
-            completion_tokens: Final = await asyncify(A2ARequestUtils.count_tokens)(output_text)
+            completion_tokens: Final = await asyncify(A2ARequestUtils.count_tokens)(self._reply_text())
 
             total_tokens: Final = prompt_tokens + completion_tokens
 
@@ -128,7 +165,7 @@ class A2AStreamingIterator:
             self.logging_obj.model_call_details["stream"] = False
 
             # Calculate cost using A2ACostCalculator
-            response_cost: Final = A2ACostCalculator.calculate_a2a_cost(self.logging_obj)
+            response_cost: Final = 0.0 if self.saw_error else A2ACostCalculator.calculate_a2a_cost(self.logging_obj)
             self.logging_obj.model_call_details["response_cost"] = response_cost
 
             # Build result for logging
@@ -167,9 +204,7 @@ class A2AStreamingIterator:
         # Add final chunk result if available
         if self.final_chunk:
             try:
-                chunk_dict: Final = as_str_mapping(
-                    _OPAQUE_PAYLOAD.validate_python(self.final_chunk.model_dump(mode="json", exclude_none=True))
-                )
+                chunk_dict: Final = as_str_mapping(_OPAQUE_PAYLOAD.validate_python(_chunk_payload(self.final_chunk)))
                 if chunk_dict is not None:
                     result["result"] = chunk_dict.get("result", {})
             except Exception:

litellm/a2a_protocol/utils.py

@@ -6,13 +6,17 @@ import hashlib
 from collections.abc import Mapping
 from typing import TYPE_CHECKING, Final, Protocol
 
+from pydantic import TypeAdapter
+
 import litellm
 from litellm._logging import verbose_logger
-from litellm.litellm_core_utils.core_helpers import has_attributes
+from litellm.litellm_core_utils.core_helpers import as_str_mapping, has_attributes
 
 if TYPE_CHECKING:
     from a2a.types import SendMessageRequest, SendStreamingMessageRequest
 
+_OBJECT_ITEMS: Final = TypeAdapter(tuple[object, ...])
+
 
 class _SupportsModelDump(Protocol):
     def model_dump(self, *, mode: str) -> object: ...
@@ -34,6 +38,9 @@ class A2ARequestUtils:
         """
         if message is None:
             return ""
+        # a2a-sdk Part RootModels keep kind/text on part.root, so read pydantic messages as dicts
+        if has_attributes(message, _SupportsModelDump, "model_dump"):
+            return A2ARequestUtils.extract_text_from_message(message.model_dump(mode="json"))
 
         # Handle both dict and object access
         if isinstance(message, dict):
@@ -63,8 +70,8 @@ class A2ARequestUtils:
         Returns:
             Text from response message parts
         """
-        result: Final = response_dict.get("result", {})
-        if not isinstance(result, dict):
+        result: Final = as_str_mapping(response_dict.get("result"))
+        if result is None:
             return ""
 
         # Direct message format (A2A spec): detect by explicit kind tag only.
@@ -72,10 +79,36 @@ class A2ARequestUtils:
         # type that happens to include a "parts" field.
         if result.get("kind") == "message":
             return A2ARequestUtils.extract_text_from_message(result)
+        if result.get("kind") == "task":
+            return A2ARequestUtils.extract_text_from_task(result)
 
         message: Final = result.get("message", {})
         return A2ARequestUtils.extract_text_from_message(message)
 
+    @staticmethod
+    def extract_text_from_task(task: Mapping[str, object]) -> str:
+        artifacts: Final = task.get("artifacts")
+        artifact_items: Final = _OBJECT_ITEMS.validate_python(artifacts) if isinstance(artifacts, list) else ()
+        artifact_texts: Final = tuple(
+            A2ARequestUtils.extract_text_from_message(artifact) for artifact in artifact_items
+        )
+        if any(artifact_texts):
+            return " ".join(text for text in artifact_texts if text)
+        status: Final = as_str_mapping(task.get("status"))
+        return A2ARequestUtils.extract_text_from_message(status.get("message")) if status is not None else ""
+
+    @staticmethod
+    def extract_text_from_stream_event(event: Mapping[str, object]) -> str:
+        result: Final = as_str_mapping(event.get("result"))
+        if result is None:
+            return ""
+        if result.get("kind") == "artifact-update":
+            return A2ARequestUtils.extract_text_from_message(result.get("artifact"))
+        if result.get("kind") == "status-update":
+            status: Final = as_str_mapping(result.get("status"))
+            return A2ARequestUtils.extract_text_from_message(status.get("message")) if status is not None else ""
+        return A2ARequestUtils.extract_text_from_response(event)
+
     @staticmethod
     def get_input_message_from_request(
         request: "SendMessageRequest | SendStreamingMessageRequest",
@@ -128,17 +161,8 @@ class A2ARequestUtils:
         Returns:
             Tuple of (prompt_tokens, completion_tokens, total_tokens)
         """
-        # Count input tokens. Dump the message to a dict first so extraction hits
-        # the dict branch — request-side parts are a2a-sdk Part RootModels whose
-        # kind/text live on part.root, which the object branch cannot read. This
-        # mirrors how the response side already works (it operates on model_dump).
         input_message: Final = A2ARequestUtils.get_input_message_from_request(request)
-        input_text: Final = A2ARequestUtils.extract_text_from_message(
-            input_message.model_dump(mode="json")
-            if has_attributes(input_message, _SupportsModelDump, "model_dump")
-            else input_message
-        )
-        prompt_tokens: Final = A2ARequestUtils.count_tokens(input_text)
+        prompt_tokens: Final = A2ARequestUtils.count_tokens(A2ARequestUtils.extract_text_from_message(input_message))
 
         # Count output tokens
         output_text: Final = A2ARequestUtils.extract_text_from_response(response_dict)

litellm/constants.py

@@ -902,6 +902,7 @@ PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO: Final = frozenset(
 
 SUBSCRIPTION_BACKED_PROVIDERS: Final = frozenset({"chatgpt"})
 A2A_CHAT_COMPLETION_BRIDGE_PROVIDERS: Final = frozenset({"databricks_agent"})
+A2A_COMPLETION_BRIDGE_WORKING_MESSAGE: Final = "Processing request..."
 
 LITELLM_EMBEDDING_PROVIDERS_SUPPORTING_INPUT_ARRAY_OF_TOKENS: Final = [
     "openai",
@@ -2046,6 +2047,7 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES: Final = [
 SPECIAL_LITELLM_AUTH_TOKEN: Final = ["ui-token"]
 DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL = get_env_int("DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL", 60)
 AUTH_OBJECTS_TARGET: Final = "auth_objects"
+HEALTH_CHECKS_TARGET: Final = "health_checks"
 DEFAULT_ACCESS_GROUP_CACHE_TTL: Final = get_env_int("DEFAULT_ACCESS_GROUP_CACHE_TTL", 600)
 SPEND_LOG_KEY_METADATA_CACHE_TTL: Final = 600
 SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL: Final = 30

litellm/cost_calculator.py

@@ -623,7 +623,11 @@ def cost_per_token(
         )
         if lyria_generation_cost is not None:
             return 0.0, lyria_generation_cost
-        speech_model_info = litellm.get_model_info(model=model_without_prefix, custom_llm_provider=custom_llm_provider)
+        speech_model_info: Final = cached_get_model_info_helper(
+            model=model_without_prefix,
+            custom_llm_provider=custom_llm_provider,
+            preserve_unset_costs=True,
+        )
         cost_metric: Final = select_cost_metric_for_model(speech_model_info)
         prompt_cost: float = 0.0
         completion_cost: float = 0.0
@@ -3084,7 +3088,7 @@ def batch_cost_calculator(
             batch_rates.cache_creation,
             batch_rates.cache_creation_rates_by_ttl,
         )
-    elif input_cost_per_token:
+    elif input_cost_per_token is not None:
         details: Final = parse_prompt_tokens_details(usage)
         cache_read_tokens: Final = details["cache_hit_tokens"]
         cache_creation_tokens: Final = details["cache_creation_tokens"]
@@ -3100,7 +3104,10 @@ def batch_cost_calculator(
         cache_read_cost_key: Final = get_service_tier_cost_key("cache_read_input_token_cost", None)
         total_prompt_cost += calculate_cost_component(model_info, cache_read_cost_key, cache_read_tokens) / 2
 
-        cache_creation_cost: Final = model_info.get("cache_creation_input_token_cost") or input_cost_per_token
+        configured_cache_creation_cost: Final = model_info.get("cache_creation_input_token_cost")
+        cache_creation_cost: Final = (
+            input_cost_per_token if configured_cache_creation_cost is None else configured_cache_creation_cost
+        )
         total_prompt_cost += cache_creation_tokens * cache_creation_cost / 2
     text_rate: Final = _batch_or_half(batch_rates.output, output_cost_per_token)
     image_rate: Final = _batch_or_half(

litellm/decisions/call.py

@@ -218,7 +218,7 @@ def supported_decisions_providers() -> tuple[str, ...]:
     )
 
 
-def _provider_config(model: str, custom_llm_provider: str) -> BaseDecisionsConfig:
+def get_decisions_provider_config(model: str, custom_llm_provider: str) -> BaseDecisionsConfig:
     provider: Final = next((member for member in LlmProviders if member.value == custom_llm_provider), None)
     provider_config: Final = (
         None if provider is None else ProviderConfigManager.get_provider_decisions_config(model, provider)
@@ -408,7 +408,7 @@ def _provider_ir_request(
     return supported
 
 
-def _outbound_headers(
+def decisions_outbound_headers(
     extra_headers: Mapping[str, str] | None, custom_llm_provider: str, kwargs: Mapping[str, object]
 ) -> Mapping[str, str]:
     forwarded: Final = ProviderSpecificHeaderUtils.get_provider_specific_headers(
@@ -466,7 +466,7 @@ def prepare_call(
         api_base=call_api_base,
         api_key=api_key,
     )
-    provider_config: Final = _provider_config(upstream_model, provider)
+    provider_config: Final = get_decisions_provider_config(upstream_model, provider)
     if not upstream_model:
         raise litellm.BadRequestError(
             message="A model name is required for the Decisions API",
@@ -478,7 +478,7 @@ def prepare_call(
     except ValueError as error:
         raise litellm.BadRequestError(message=str(error), model=model, llm_provider=provider) from error
     try:
-        headers: Final = _outbound_headers(extra_headers, provider, kwargs)
+        headers: Final = decisions_outbound_headers(extra_headers, provider, kwargs)
         request_safety_identifier: Final = _request_safety_identifier(safety_identifier, provider_config, kwargs)
         request: Final = _fields_without_unsupported_media(
             fields, provider_config, canonical_model, provider, kwargs

litellm/integrations/custom_guardrail.py

@@ -71,6 +71,8 @@ GUARDRAIL_SESSIONS_TARGET: Final = "guardrail_sessions"
 # field to suppress a guardrail on the direct-SDK path that never reaches the
 # proxy's metadata sanitizer.
 _PRE_CALL_EXECUTED_TOKEN: Final = secrets.token_hex(16)
+_PRE_CALL_RAISED_TOKEN: Final = secrets.token_hex(16)
+_DEPLOYMENT_HOOK_RAN_TOKEN: Final = secrets.token_hex(16)
 
 _GUARDRAIL_BLOCK_STATUS_CODES: Final = frozenset({400, 403, 422})
 
@@ -260,6 +262,25 @@ def _request_is_streaming(data: object, event_type: GuardrailEventHooks | None =
     )
 
 
+def _str_collection(value: object) -> frozenset[str]:
+    if not isinstance(value, (list, tuple, set, frozenset)):
+        return frozenset()
+    items: Final = cast("Iterable[object]", value)  # cast-ok: a metadata collection of guardrail names
+    return frozenset(item for item in items if isinstance(item, str))
+
+
+def _request_carries_marker(data: Mapping[str, object], marker: str | None) -> bool:
+    if marker is None:
+        return False
+    for meta_key in ("metadata", "litellm_metadata"):
+        meta = data.get(meta_key)
+        if isinstance(meta, dict):
+            executed = meta.get(PRE_CALL_EXECUTED_GUARDRAILS_KEY)
+            if isinstance(executed, list) and marker in executed:
+                return True
+    return False
+
+
 _REQUEST_FIELDS: Final = TypeAdapter(Mapping[str, object])
 
 
@@ -869,13 +890,16 @@ class CustomGuardrail(CustomLogger):
 
         return False
 
-    def _pre_call_marker(self) -> str | None:
+    def _marker(self, token: str) -> str | None:
         name: Final = self.guardrail_name
         if not name:
             return None
-        return f"{_PRE_CALL_EXECUTED_TOKEN}:{name}"
+        return f"{token}:{name}"
 
-    def mark_pre_call_hook_ran(self, data: dict[str, object]) -> None:
+    def _pre_call_marker(self, *, raised: bool = False) -> str | None:
+        return self._marker(_PRE_CALL_RAISED_TOKEN if raised else _PRE_CALL_EXECUTED_TOKEN)
+
+    def mark_pre_call_hook_ran(self, data: dict[str, object], *, raised: bool = False) -> None:
         """
         Record that this guardrail's ``async_pre_call_hook`` already ran for this
         request, so the deployment-level hook does not run it a second time.
@@ -884,8 +908,11 @@ class CustomGuardrail(CustomLogger):
         router later spreads a deployment's model-level ``guardrails`` into the
         top-level request kwargs, which would otherwise re-trigger the same hook
         from ``async_pre_call_deployment_hook``.
+
+        With ``raised``, records instead that the hook raised as a policy pipeline
+        step the pipeline then moved past. The deployment-level hook still runs it.
         """
-        marker: Final = self._pre_call_marker()
+        marker: Final = self._pre_call_marker(raised=raised)
         if marker is None:
             return
         for meta_key in ("metadata", "litellm_metadata"):
@@ -900,17 +927,57 @@ class CustomGuardrail(CustomLogger):
                 return
         data["metadata"] = {PRE_CALL_EXECUTED_GUARDRAILS_KEY: [marker]}
 
-    def _pre_call_hook_already_ran(self, data: dict[str, object]) -> bool:
-        marker: Final = self._pre_call_marker()
-        if marker is None:
-            return False
-        for meta_key in ("metadata", "litellm_metadata"):
-            meta = data.get(meta_key)
-            if isinstance(meta, dict):
-                executed = meta.get(PRE_CALL_EXECUTED_GUARDRAILS_KEY)
-                if isinstance(executed, list) and marker in executed:
-                    return True
-        return False
+    def pre_call_hook_ran_or_raised(self, data: Mapping[str, object]) -> bool:
+        """Whether this guardrail's ``async_pre_call_hook`` ran for this request, whatever it decided."""
+        return self._pre_call_hook_already_ran(data) or self._pre_call_hook_already_ran(data, raised=True)
+
+    def _pre_call_hook_already_ran(self, data: Mapping[str, object], *, raised: bool = False) -> bool:
+        return _request_carries_marker(data, self._pre_call_marker(raised=raised))
+
+    def deployment_hook_ran(self, data: Mapping[str, object]) -> bool:
+        """Whether a deployment-level hook ran this guardrail as a policy pipeline step of this request."""
+        return _request_carries_marker(data, self._marker(_DEPLOYMENT_HOOK_RAN_TOKEN))
+
+    def _record_deployment_hook_run(
+        self,
+        data: dict[str, object],  # mutable-ok: the request payload the run is recorded on
+    ) -> None:
+        """
+        Record that a deployment-level hook ran this guardrail when it is a policy pipeline step of
+        this request, and name it in the applied guardrails when the request's metadata lists it.
+        """
+        name: Final = self.guardrail_name
+        marker: Final = self._marker(_DEPLOYMENT_HOOK_RAN_TOKEN)
+        buckets: Final = tuple(
+            cast("dict[str, object]", bucket)  # cast-ok: metadata buckets are str-keyed
+            for bucket in (data.get("metadata"), data.get("litellm_metadata"))
+            if isinstance(bucket, dict)
+        )
+        policy_state: Final = next(
+            (bucket for bucket in buckets if name in _str_collection(bucket.get("_pipeline_managed_guardrails"))),
+            None,
+        )
+        if marker is None or policy_state is None:
+            return
+        recorded: Final = policy_state.get(PRE_CALL_EXECUTED_GUARDRAILS_KEY)
+        if not isinstance(recorded, list):
+            policy_state[PRE_CALL_EXECUTED_GUARDRAILS_KEY] = [marker]
+        elif marker not in recorded:
+            cast("list[object]", recorded).append(marker)  # cast-ok: markers are only compared as strings
+        if not any(name in _str_collection(bucket.get("guardrails")) for bucket in buckets):
+            return
+        named_lists: Final = tuple(
+            cast("list[object]", bucket["applied_guardrails"])  # cast-ok: entries are only compared with names
+            for bucket in buckets
+            if isinstance(bucket.get("applied_guardrails"), list)
+        )
+        for applied in named_lists:
+            if name not in applied:
+                applied.append(name)
+        if named_lists:
+            return
+        metadata_bucket: Final = cast("dict[str, object]", get_or_create_metadata_bucket(data)[1])  # cast-ok: str keys
+        metadata_bucket["applied_guardrails"] = [name]
 
     @classmethod
     def uses_apply_guardrail_interface(cls) -> bool:
@@ -959,12 +1026,15 @@ class CustomGuardrail(CustomLogger):
             target: Final = self._deployment_hook_target()
             if target is not self:
                 kwargs["guardrail_to_apply"] = self
-            result: Final = await target.async_pre_call_hook(
-                user_api_key_dict=_user_api_key_auth_from_request(kwargs),
-                cache=dc,
-                data=kwargs,
-                call_type="completion" if call_type == CallTypes.completion else "acompletion",
-            )
+            try:
+                result: Final = await target.async_pre_call_hook(
+                    user_api_key_dict=_user_api_key_auth_from_request(kwargs),
+                    cache=dc,
+                    data=kwargs,
+                    call_type="completion" if call_type == CallTypes.completion else "acompletion",
+                )
+            finally:
+                self._record_deployment_hook_run(kwargs)
 
             if result is not None and isinstance(result, dict):
                 result_messages: Final = result.get("messages")
@@ -1048,6 +1118,9 @@ class CustomGuardrail(CustomLogger):
         finally:
             if target is not self:
                 request_data.pop("guardrail_to_apply", None)
+            self._record_deployment_hook_run(
+                cast("dict[str, object]", request_data)  # cast-ok: proxy request payloads are str-keyed
+            )
 
         if not self._is_valid_response_type(result):
             return None

litellm/integrations/otel/routing.py

@@ -21,6 +21,11 @@ _ROUTING_SCALAR_FIELDS: Final = frozenset(
         "classifier_cost",
         "classifier_failure_reason",
         "classifier_error_type",
+        "classifier_input_budget_tokens",
+        "classifier_input_tokens",
+        "classifier_input_count_method",
+        "classifier_input_truncated",
+        "classifier_input_budget_status",
         "classifier_confidence",
         "classifier_primary_rule",
         "classifier_capability_boundary",

litellm/integrations/prometheus.py

@@ -66,7 +66,7 @@ from litellm.repositories.organization_repository import OrganizationRepository
 from litellm.repositories.table_repositories import EndUserRepository
 from litellm.repositories.team_repository import TeamRepository
 from litellm.repositories.user_repository import UserRepository
-from litellm.router_utils.router_callbacks.track_deployment_metrics import find_deployment_metadata
+from litellm.router_api import deployment_metadata
 from litellm.types.guardrails import GuardrailEventHooks
 from litellm.types.integrations.prometheus import *
 from litellm.types.integrations.prometheus import (
@@ -3022,8 +3022,8 @@ class PrometheusLogger(CustomLogger):
         ).get("custom_llm_provider")
         if isinstance(provider, str) and provider:
             return provider
-        deployment_metadata: Final = find_deployment_metadata(litellm_params)
-        return _deployment_model_provider(deployment_metadata.get("deployment") if deployment_metadata else None)
+        picked_deployment: Final = deployment_metadata(litellm_params)
+        return _deployment_model_provider(picked_deployment.get("deployment") if picked_deployment else None)
 
     @staticmethod
     def _extract_api_provider_from_exception(exception: Exception) -> str | None:
@@ -5209,7 +5209,7 @@ def _tag_matches_wildcard_configured_pattern(tags: Sequence[str], configured_tag
     """
     import re
 
-    from litellm.router_utils.pattern_match_deployments import PatternMatchRouter
+    from litellm.litellm_core_utils.pattern_match_deployments import PatternMatchRouter
 
     pattern_router: Final = PatternMatchRouter()
     regex_pattern: Final = pattern_router.pattern_to_regex(configured_tag)

litellm/litellm_core_utils/fallback_utils.py

@@ -3,13 +3,13 @@ from typing import Final
 import litellm
 from litellm._logging import verbose_logger
 from litellm._uuid import uuid
+from litellm.litellm_core_utils.add_retry_fallback_headers import (
+    add_fallback_headers_to_response,
+)
 from litellm.litellm_core_utils.core_helpers import (
     filter_internal_params,
     safe_deep_copy,
 )
-from litellm.router_utils.add_retry_fallback_headers import (
-    add_fallback_headers_to_response,
-)
 
 from .asyncify import run_async_function
 

litellm/litellm_core_utils/health_check_helpers.py

@@ -119,7 +119,7 @@ def _is_route_segment(custom_llm_provider: str, segment: str, model_after_segmen
 
 def wildcard_health_check_models(wildcard: str, custom_llm_provider: str) -> tuple[str, ...]:
     from litellm.litellm_core_utils.llm_request_utils import rank_chat_models_from_llm_provider
-    from litellm.router_utils.pattern_match_deployments import PatternMatchRouter
+    from litellm.litellm_core_utils.pattern_match_deployments import PatternMatchRouter
 
     provider_prefix: Final = f"{custom_llm_provider}/"
     segments: Final = wildcard.removeprefix(provider_prefix).split("/")

litellm/litellm_core_utils/llm_cost_calc/utils.py

@@ -28,6 +28,7 @@ from litellm.types.utils import (
     DataResidency,
     ImageResponse,
     ModelInfo,
+    ModelInfoBase,
     PassthroughCallTypes,
     PromptTokensDetailsWrapper,
     ServiceTier,
@@ -197,15 +198,15 @@ def get_billable_input_tokens(usage: Usage) -> int:
 
 
 def select_cost_metric_for_model(
-    model_info: ModelInfo,
+    model_info: ModelInfoBase,
 ) -> Literal["cost_per_character", "cost_per_token"]:
     """
     Select 'cost_per_character' if model_info has 'input_cost_per_character'
     Select 'cost_per_token' if model_info has 'input_cost_per_token'
     """
-    if model_info.get("input_cost_per_character"):
+    if model_info.get("input_cost_per_character") is not None:
         return "cost_per_character"
-    elif model_info.get("input_cost_per_token"):
+    elif model_info.get("input_cost_per_token") is not None:
         return "cost_per_token"
     else:
         raise ValueError(

litellm/litellm_core_utils/ptu_shares.py

@@ -12,13 +12,13 @@ from datetime import date
 from typing import Final, Generic, TypeAlias, TypeVar
 
 from litellm.litellm_core_utils.ptu_pricing import parsed_ptu_shares, ptu_terms
+from litellm.litellm_core_utils.routing_common_utils import get_request_team_id, team_may_use_deployment
 from litellm.llms.azure.ptu_capacity import (
     PTUCapacity,
     deployment_ptu_capacity,
     is_azure_deployment,
     is_str_keyed_mapping,
 )
-from litellm.router_utils.common_utils import get_request_team_id, team_may_use_deployment
 
 _DeploymentT: Final = TypeVar("_DeploymentT", bound=Mapping[str, object])
 WildcardRoute: TypeAlias = Callable[[str], Sequence[Mapping[str, object]]]

litellm/litellm_core_utils/token_counter.py

@@ -56,9 +56,9 @@ from litellm.litellm_core_utils.tokenizer import HuggingFaceTokenizer, OpenAIEnc
 
 if TYPE_CHECKING:
     from litellm.litellm_core_utils.tokenizer import Encoding, HuggingFace
+from litellm.litellm_core_utils.deployment_pricing import deployment_pricing_models
 from litellm.litellm_core_utils.url_utils import safe_get
 from litellm.llms.custom_httpx.http_handler import get_httpx_client, header_value
-from litellm.router_utils.deployment_pricing import deployment_pricing_models
 from litellm.rust_bridge.tokenizer import get_encoding
 from litellm.types.llms.anthropic import (
     AnthropicContentParamSource,

litellm/llms/anthropic/messages/native/handler.py

@@ -19,6 +19,13 @@ from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheCont
 from litellm.litellm_core_utils.core_helpers import as_str_mapping
 from litellm.litellm_core_utils.exception_mapping_utils import exception_type
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+from litellm.litellm_core_utils.reasoning_effort_capability import (
+    RequestedReasoningEffort,
+    deployment_reasoning_effort_levels,
+    keeps_effort,
+    model_group_from_metadata,
+    resolve_deployment_reasoning_effort,
+)
 from litellm.llms.anthropic.common_utils import (
     AnthropicModelInfo,
     strip_provider_specific_fields_from_anthropic_messages,
@@ -31,13 +38,6 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException
 from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
 from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
 from litellm.responses.mcp.tool_precheck import may_name_litellm_mcp_gateway
-from litellm.router_utils.reasoning_effort_capability import (
-    RequestedReasoningEffort,
-    deployment_reasoning_effort_levels,
-    keeps_effort,
-    model_group_from_metadata,
-    resolve_deployment_reasoning_effort,
-)
 from litellm.types.llms.anthropic import AllAnthropicToolsValues
 from litellm.types.llms.anthropic_messages.anthropic_request import AnthropicMetadata
 from litellm.types.llms.anthropic_messages.anthropic_response import (

litellm/llms/anthropic/messages/native/interceptors/advisor.py

@@ -22,7 +22,7 @@ import litellm
 import litellm.constants as _c
 from litellm.litellm_core_utils.url_utils import validate_url
 from litellm.llms.anthropic.common_utils import strip_advisor_blocks_from_messages
-from litellm.router_utils.cooldown_handlers import mark_advisor_orchestration_failure
+from litellm.router_api import mark_advisor_orchestration_failure
 from litellm.types.llms.anthropic import ANTHROPIC_ADVISOR_TOOL_TYPE
 from litellm.types.llms.anthropic_messages.anthropic_response import (
     AnthropicMessagesResponse,

litellm/llms/anthropic/messages/utils.py

@@ -214,7 +214,7 @@ def normalize_reasoning_effort_value(
     if chain is None:
         return effort
 
-    from litellm.router_utils.reasoning_effort_capability import (
+    from litellm.litellm_core_utils.reasoning_effort_capability import (
         deployment_rounds_reasoning_effort,
         resolve_supported_reasoning_efforts,
     )

litellm/llms/azure_ai/decisions/transformation.py

@@ -12,7 +12,7 @@ from litellm.llms.azure_ai.common_utils import (
     AZURE_ENTRA_LITELLM_PARAM_KEYS,
     get_azure_ai_auth_headers,
 )
-from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig
+from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig, DecisionsInputLimits
 
 _FOUNDRY_ROUTE_SUFFIX: Final = re.compile(r"(/api/projects/[^/]+)?(/openai/v1|/openai|/models|/v1)?$")
 _AZURE_AI_CREDENTIAL_PARAM_KEYS: Final = AZURE_ENTRA_LITELLM_PARAM_KEYS | frozenset({"azure_ad_token_provider"})
@@ -26,6 +26,16 @@ class AzureAIDecisionsConfig(BaseDecisionsConfig):
     api_base_env = ("AZURE_AI_API_BASE",)
     api_key_required = False
 
+    def get_input_limits(
+        self, model: str, custom_llm_provider: str, headers: Mapping[str, str]
+    ) -> DecisionsInputLimits:
+        limits: Final = super().get_input_limits(model, custom_llm_provider, headers)
+        return (
+            limits
+            if limits.max_request_tokens is not None
+            else super().get_input_limits("Microsoft-Decision-1", custom_llm_provider, headers)
+        )
+
     def missing_api_base_message(self, custom_llm_provider: str) -> str:
         return "Missing AZURE_AI_API_BASE - set AZURE_AI_API_BASE or pass api_base, e.g. https://<resource>.services.ai.azure.com"
 

litellm/llms/base_llm/decisions/transformation.py

@@ -1,7 +1,7 @@
 import json
 from abc import ABC
 from collections.abc import Callable, Mapping, Sequence
-from dataclasses import replace
+from dataclasses import dataclass, replace
 from itertools import chain
 from types import MappingProxyType
 from typing import Final
@@ -511,6 +511,12 @@ def _is_decisions_model(model: str, custom_llm_provider: str) -> bool:
     return is_decisions_model_mode(model_mode)
 
 
+@dataclass(frozen=True, slots=True)
+class DecisionsInputLimits:
+    max_request_tokens: int | None = None
+    max_question_tokens: int | None = None
+
+
 class BaseDecisionsConfig(ABC):
     path: str = "/v1/systemone"
     api_key_env: tuple[str, ...] = ()
@@ -533,6 +539,19 @@ class BaseDecisionsConfig(ABC):
     def provider_reported_cost(self, response: DecisionsIRResponse) -> float | None:
         return None
 
+    def get_input_limits(
+        self, model: str, custom_llm_provider: str, headers: Mapping[str, str]
+    ) -> DecisionsInputLimits:
+        try:
+            limit: Final = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider).get(
+                "max_input_tokens"
+            )
+        except Exception:  # noqa: BLE001  # get_model_info raises bare Exception for an unmapped model
+            return DecisionsInputLimits()
+        return DecisionsInputLimits(
+            max_request_tokens=limit if isinstance(limit, int) and not isinstance(limit, bool) and limit > 0 else None
+        )
+
     def missing_api_base_message(self, custom_llm_provider: str) -> str:
         return f"api_base is required for Decisions provider '{custom_llm_provider}'"
 

litellm/llms/clf_ai_gateway/chat/transformation.py

@@ -4,8 +4,8 @@ import litellm
 from litellm._logging import verbose_logger
 from litellm.exceptions import UnsupportedParamsError
 from litellm.litellm_core_utils.core_helpers import as_str_mapping
+from litellm.litellm_core_utils.reasoning_effort_capability import declared_reasoning_efforts_for_model
 from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig
-from litellm.router_utils.reasoning_effort_capability import declared_reasoning_efforts_for_model
 from litellm.secret_managers.main import get_secret_str
 
 SUPPORTED_OPENAI_PARAMS: Final = (

litellm/llms/compaction.py

@@ -9,7 +9,7 @@ from pydantic import TypeAdapter
 from litellm.types.llms.openai import AllMessageValues
 
 if TYPE_CHECKING:
-    from litellm.router import Router
+    from litellm.router_api import RoutingContext
 
 CompactionProtocol: TypeAlias = Literal["chat", "messages"]
 _MAPPING: Final = TypeAdapter(Mapping[str, object])
@@ -32,7 +32,9 @@ def get_native_compaction_provider(params: Mapping[str, object]) -> NativeCompac
     return compaction if compaction.supports_native_compaction(params) else None
 
 
-async def dispatch(router: Router, protocol: CompactionProtocol, payload: Mapping[str, object]) -> Mapping[str, object]:
+async def dispatch(
+    router: RoutingContext, protocol: CompactionProtocol, payload: Mapping[str, object]
+) -> Mapping[str, object]:
     if protocol == "messages":
         return _MAPPING.validate_python(
             await router.aanthropic_messages(custom_llm_provider=None, client=None, **payload)
@@ -45,4 +47,4 @@ async def dispatch(router: Router, protocol: CompactionProtocol, payload: Mappin
             {key: value for key, value in payload.items() if key not in ("model", "messages", "stream")}
         ),
     )
-    return _MAPPING.validate_python(response.model_dump())
+    return _MAPPING.validate_python(response.model_dump())  # pyright: ignore[reportAttributeAccessIssue]  # stream=False: acompletion returns a ModelResponse

litellm/llms/custom_httpx/aiohttp_transport.py

@@ -110,6 +110,15 @@ def map_aiohttp_exceptions() -> Generator[None, None, None]:
         raise mapped_exc(message) from exc
 
 
+def _unread_payload(content: object) -> bytes:
+    """Bytes the upstream sent before the connection dropped that aiohttp buffered but never handed out, since its
+    reads raise the stored payload error ahead of the buffer."""
+    if not isinstance(content, aiohttp.StreamReader):
+        return b""
+    buffered: Final = b"".join(content._buffer)  # pyright: ignore[reportPrivateUsage]  # no public read drains it
+    return buffered[content._buffer_offset :]  # pyright: ignore[reportPrivateUsage]  # see above
+
+
 def _release_transport_anchor(_transport: object) -> None:
     """Release a transport anchored to a response stream. See ``_anchor_transport_to``."""
 
@@ -144,6 +153,12 @@ class AiohttpResponseStream(httpx.AsyncByteStream):
             raise httpx.ReadError(str(e)) from e
         except aiohttp.http_exceptions.TransferEncodingError as e:
             raise httpx.ReadError(str(e)) from e
+        except aiohttp.ClientPayloadError:
+            unread: Final = _unread_payload(self._aiohttp_response.content)
+            if unread:
+                yield unread
+            with map_aiohttp_exceptions():
+                raise
         except Exception:
             # For other exceptions, use the normal mapping
             with map_aiohttp_exceptions():

litellm/llms/mistral/chat/transformation.py

@@ -25,16 +25,16 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
     strip_litellm_internal_message_fields,
     strip_none_values_from_message,
 )
+from litellm.litellm_core_utils.reasoning_effort_capability import (
+    declared_reasoning_efforts_for_model,
+    nearest_declared_reasoning_effort,
+)
 from litellm.litellm_core_utils.reasoning_effort_utils import OPAQUE_VALUE, warn_dropped_reasoning
 from litellm.litellm_core_utils.url_utils import ensure_api_route
 from litellm.llms.openai.chat.gpt_transformation import (
     OpenAIChatCompletionStreamingHandler,
     OpenAIGPTConfig,
 )
-from litellm.router_utils.reasoning_effort_capability import (
-    declared_reasoning_efforts_for_model,
-    nearest_declared_reasoning_effort,
-)
 from litellm.secret_managers.main import get_secret_str
 from litellm.types.llms.mistral import (
     MistralCacheKeyParams,

litellm/llms/oci/chat/transformation.py

@@ -21,6 +21,7 @@ from litellm._logging import verbose_logger
 from litellm.constants import DEFAULT_OCI_CHAT_MAX_TOKENS
 from litellm.litellm_core_utils.logging_utils import track_llm_api_timing
 from litellm.litellm_core_utils.prompt_templates.common_utils import messages_with_file_parts_as_image_urls
+from litellm.litellm_core_utils.reasoning_effort_capability import declared_reasoning_efforts_for_model
 from litellm.litellm_core_utils.reasoning_effort_utils import warn_dropped_reasoning
 from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
 from litellm.litellm_core_utils.url_utils import join_api_base
@@ -54,7 +55,6 @@ from litellm.llms.oci.common_utils import (
     sign_oci_request,
     validate_oci_environment,
 )
-from litellm.router_utils.reasoning_effort_capability import declared_reasoning_efforts_for_model
 from litellm.types.llms.oci import (
     CohereChatRequest,
     OCIChatRequestPayload,

litellm/llms/perplexity/cost_calculator.py

@@ -100,8 +100,9 @@ def cost_per_token(model: str, usage: Usage, current_time: datetime | None = Non
         num_search_queries = getattr(usage.prompt_tokens_details, "web_search_requests", 0) or 0
 
     # Check both possible keys for search cost (legacy and current)
-    search_cost_value: Final = model_info.get("search_queries_cost_per_query") or model_info.get(
-        "search_context_cost_per_query"
+    legacy_search_cost: Final = model_info.get("search_queries_cost_per_query")
+    search_cost_value: Final = (
+        model_info.get("search_context_cost_per_query") if legacy_search_cost is None else legacy_search_cost
     )
     if num_search_queries > 0 and search_cost_value is not None:
         # Handle both dict and float formats

litellm/llms/perplexity/search/transformation.py

@@ -84,9 +84,9 @@ class PerplexitySearchConfig(BaseSearchConfig):
     def transform_search_request(
         self,
         query: str | list[str],
-        optional_params: dict,
-        **kwargs,
-    ) -> dict:
+        optional_params: dict[str, object],
+        **kwargs: object,
+    ) -> dict[str, object]:
         """
         Transform Search request to Perplexity API format.
 
@@ -98,11 +98,7 @@ class PerplexitySearchConfig(BaseSearchConfig):
 
         Args:
             query: Search query (string or list of strings)
-            optional_params: Optional parameters for the request
-                - max_results: Maximum number of search results (1-20)
-                - search_domain_filter: List of domains to filter (max 20)
-                - max_tokens_per_page: Max tokens per page (default 1024)
-                - country: Country code filter (e.g., 'US', 'GB', 'DE')
+            optional_params: Provider parameters, forwarded unchanged when not None.
 
         Returns:
             Dict with typed request data following PerplexitySearchRequest spec
@@ -111,24 +107,10 @@ class PerplexitySearchConfig(BaseSearchConfig):
             "query": query,
         }
 
-        # Add optional parameters following Perplexity API spec (only if not None)
-        max_results: Final = optional_params.get("max_results")
-        if max_results is not None:
-            request_data["max_results"] = max_results
-
-        search_domain_filter: Final = optional_params.get("search_domain_filter")
-        if search_domain_filter is not None:
-            request_data["search_domain_filter"] = search_domain_filter
-
-        max_tokens_per_page: Final = optional_params.get("max_tokens_per_page")
-        if max_tokens_per_page is not None:
-            request_data["max_tokens_per_page"] = max_tokens_per_page
-
-        country: Final = optional_params.get("country")
-        if country is not None:
-            request_data["country"] = country
-
-        return dict(request_data)
+        return {
+            **{key: value for key, value in optional_params.items() if value is not None},
+            **request_data,
+        }
 
     def transform_search_response(
         self,

litellm/llms/together_ai/chat/transformation.py

@@ -18,7 +18,7 @@ import litellm
 from litellm._logging import verbose_logger
 from litellm.exceptions import UnsupportedParamsError
 from litellm.litellm_core_utils.prompt_templates.common_utils import strip_litellm_response_only_fields
-from litellm.router_utils.reasoning_effort_capability import declared_reasoning_efforts_for_model
+from litellm.litellm_core_utils.reasoning_effort_capability import declared_reasoning_efforts_for_model
 from litellm.types.llms.openai import AllMessageValues
 from litellm.utils import supports_function_calling, supports_reasoning, supports_response_schema
 

litellm/llms/typesafe/decisions/transformation.py

@@ -6,7 +6,7 @@ from types import MappingProxyType
 from typing import Final, Literal
 
 from litellm.constants import TYPESAFE_DANGER_ZONE_HEADER_PREFIX
-from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig, image_media_type
+from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig, DecisionsInputLimits, image_media_type
 from litellm.types.decisions import (
     DecisionsAsset,
     DecisionsIRMessages,
@@ -19,6 +19,7 @@ from litellm.types.decisions import (
 )
 
 _IMAGES_HEADER: Final = f"{TYPESAFE_DANGER_ZONE_HEADER_PREFIX}images"
+_LONG_CONTEXT_HEADER: Final = f"{TYPESAFE_DANGER_ZONE_HEADER_PREFIX}long-context"
 _ASSET_FORMATS: Final[Mapping[str, Literal["png", "jpg"]]] = MappingProxyType(
     {"image/png": "png", "image/jpeg": "jpg", "image/jpg": "jpg"}
 )
@@ -84,6 +85,16 @@ class TypeSafeDecisionsConfig(BaseDecisionsConfig):
     def get_default_api_base(self) -> str | None:
         return "https://api.typesafe.ai"
 
+    def get_input_limits(
+        self, model: str, custom_llm_provider: str, headers: Mapping[str, str]
+    ) -> DecisionsInputLimits:
+        long_context_values: Final = tuple(
+            value.strip().lower() for name, value in headers.items() if name.lower() == _LONG_CONTEXT_HEADER
+        )
+        if long_context_values and all(value == "true" for value in long_context_values):
+            return DecisionsInputLimits(max_request_tokens=262144, max_question_tokens=262144)
+        return DecisionsInputLimits(max_request_tokens=65536, max_question_tokens=32768)
+
     def supports_images_in(
         self, decision_input: DecisionsIRState | DecisionsIRMessages, headers: Mapping[str, str]
     ) -> bool:

litellm/main.py

@@ -113,6 +113,11 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
     messages_without_pdf_file_parts,
 )
 from litellm.litellm_core_utils.provider_affinity import add_provider_affinity_header
+from litellm.litellm_core_utils.reasoning_effort_capability import (
+    model_group_from_metadata,
+    reasoning_effort_needs_allowing,
+    resolve_chat_reasoning,
+)
 from litellm.litellm_core_utils.request_timeout_resolver import (
     get_configured_request_timeout,
 )
@@ -144,11 +149,6 @@ from litellm.llms.vertex_ai.common_utils import (
 )
 from litellm.realtime_api.main import realtime_health_check
 from litellm.responses.mcp.tool_precheck import may_name_litellm_mcp_gateway
-from litellm.router_utils.reasoning_effort_capability import (
-    model_group_from_metadata,
-    reasoning_effort_needs_allowing,
-    resolve_chat_reasoning,
-)
 from litellm.secret_managers.main import get_secret_bool, get_secret_str
 from litellm.text_completions.streaming_iterator import TextCompletionStreamWrapper
 from litellm.types.completion import (

litellm/model_prices_and_context_window_backup.json

@@ -23078,6 +23078,24 @@
         "output_cost_per_token_batches": 1e-06,
         "source": "https://developers.openai.com/api/docs/pricing"
     },
+    "deepgram/aura": {
+        "input_cost_per_character": 1.5e-05,
+        "litellm_provider": "deepgram",
+        "metadata": {
+            "calculation": "$0.0150/1000 characters = $0.000015 per character",
+            "note": "Aura-1 text-to-speech voices (aura-<voice>-<lang>), pay as you go"
+        },
+        "source": "https://deepgram.com/pricing"
+    },
+    "deepgram/aura-2": {
+        "input_cost_per_character": 3e-05,
+        "litellm_provider": "deepgram",
+        "metadata": {
+            "calculation": "$0.030/1000 characters = $0.00003 per character",
+            "note": "Aura-2 text-to-speech voices (aura-2-<voice>-<lang>), pay as you go"
+        },
+        "source": "https://deepgram.com/pricing"
+    },
     "deepgram/base": {
         "input_cost_per_second": 0.00020833,
         "litellm_provider": "deepgram",
@@ -25284,6 +25302,38 @@
         "cache_creation_input_token_cost": 1e-06,
         "prompt_cache_min_tokens": 2048
     },
+    "eu.anthropic.claude-haiku-4-5-20251001-v1:0": {
+        "cache_creation_input_token_cost": 1.375e-06,
+        "cache_creation_input_token_cost_above_1hr": 2.2e-06,
+        "cache_read_input_token_cost": 1.1e-07,
+        "input_cost_per_token": 1.1e-06,
+        "deprecation_date": "2026-10-15",
+        "litellm_provider": "bedrock_converse",
+        "supports_tool_search": true,
+        "max_input_tokens": 200000,
+        "max_output_tokens": 64000,
+        "mode": "chat",
+        "output_cost_per_token": 5.5e-06,
+        "source": "https://aws.amazon.com/bedrock/pricing/",
+        "supports_assistant_prefill": true,
+        "supports_computer_use": true,
+        "supports_function_calling": true,
+        "supports_pdf_input": true,
+        "supports_prompt_caching": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true,
+        "supports_vision": true,
+        "supports_native_structured_output": true,
+        "supports_parallel_tool_use_config": true,
+        "prompt_cache_min_tokens": 4096,
+        "input_cost_per_token_batches": 5.5e-07,
+        "output_cost_per_token_batches": 2.75e-06,
+        "token_count_model": {
+            "model": "anthropic.claude-haiku-4-5-20251001-v1:0",
+            "litellm_provider": "bedrock"
+        }
+    },
     "eu.anthropic.claude-3-5-sonnet-20240620-v1:0": {
         "input_cost_per_token": 3e-06,
         "litellm_provider": "bedrock",
@@ -31224,6 +31274,9 @@
         "output_cost_per_image_token": 3e-05,
         "output_cost_per_reasoning_token": 2.5e-06,
         "output_cost_per_token": 2.5e-06,
+        "output_cost_per_token_batches": 1.25e-06,
+        "output_cost_per_token_flex": 1.25e-06,
+        "output_cost_per_token_priority": 4.5e-06,
         "rpm": 100000,
         "source": "https://ai.google.dev/gemini-api/docs/pricing",
         "supported_endpoints": [
@@ -31404,6 +31457,7 @@
         "mode": "image_generation",
         "output_cost_per_image": 0.0336,
         "output_cost_per_image_token": 3e-05,
+        "output_cost_per_image_token_batches": 1.5e-05,
         "output_cost_per_token": 1.5e-06,
         "output_cost_per_token_batches": 7.5e-07,
         "rpm": 1000,
@@ -59744,12 +59798,13 @@
         "cache_read_input_token_cost": 1.375e-06,
         "output_cost_per_token": 8.25e-05,
         "litellm_provider": "bedrock_mantle",
-        "max_input_tokens": 1050000,
+        "max_input_tokens": 922000,
         "max_output_tokens": 128000,
         "mode": "responses",
         "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-cyber.html",
         "use_openai_responses_path": true,
         "supported_endpoints": [
+            "/v1/chat/completions",
             "/v1/responses"
         ],
         "supported_modalities": [
@@ -75428,6 +75483,7 @@
         "mode": "image_generation",
         "output_cost_per_image": 0.0336,
         "output_cost_per_image_token": 3e-05,
+        "output_cost_per_image_token_batches": 1.5e-05,
         "output_cost_per_token": 7.5e-06,
         "output_cost_per_token_batches": 3.75e-06,
         "rpm": 1000.0,

litellm/models/deployment_warning.py

@@ -0,0 +1,24 @@
+from collections.abc import Mapping
+from typing import Literal, Protocol
+
+from pydantic import AwareDatetime, ConfigDict, Field
+
+from litellm.types.llms.base import LiteLLMBaseModel
+
+
+class AdminIgnore(LiteLLMBaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    source: Literal["admin"] = "admin"
+    ignored_by: str
+    ignored_at: AwareDatetime
+    ignored_until: AwareDatetime | None = Field(
+        default=None,
+        description="When the snooze ends. Null means the warning is ignored until an admin restores it.",
+    )
+
+
+class DeploymentWarningIgnoreStore(Protocol):
+    async def get_ignores(self) -> Mapping[str, AdminIgnore]: ...
+
+    async def save_ignores(self, ignores: Mapping[str, AdminIgnore]) -> None: ...

litellm/models/gateway_activity.py

@@ -0,0 +1,31 @@
+from typing import Literal, Protocol, TypeAlias
+
+from pydantic import ConfigDict
+
+from litellm.types.llms.base import LiteLLMBaseModel
+
+
+class GatewayRouteActivity(LiteLLMBaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    kind: Literal["route"]
+    date: str
+    category: str
+    route: str
+    successful_requests: int
+    failed_requests: int
+
+
+class GatewayStatusCodeActivity(LiteLLMBaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    kind: Literal["status"]
+    status_code: int
+    failed_requests: int
+
+
+GatewayActivity: TypeAlias = GatewayRouteActivity | GatewayStatusCodeActivity
+
+
+class GatewayActivityStore(Protocol):
+    async def get_activity(self, start_date: str, end_date: str) -> tuple[GatewayActivity, ...]: ...

litellm/models/user_banner.py

@@ -0,0 +1,48 @@
+from typing import Final, Literal, Protocol
+
+from pydantic import ConfigDict, Field, model_validator
+
+from litellm.types.llms.base import LiteLLMBaseModel
+
+USER_BANNER_MAX_MESSAGE_LENGTH: Final = 4000
+UserBannerSeverity = Literal["info", "warning", "error"]
+
+
+class UserBannerUpdate(LiteLLMBaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    enabled: bool = Field(
+        default=False,
+        description="If true, the banner is shown to all authenticated dashboard users.",
+    )
+    message: str = Field(
+        default="",
+        max_length=USER_BANNER_MAX_MESSAGE_LENGTH,
+        description="Banner text shown to dashboard users. Markdown is supported.",
+    )
+    severity: UserBannerSeverity = Field(
+        default="info",
+        description="Visual style of the banner.",
+    )
+
+    @model_validator(mode="after")
+    def _require_message_when_enabled(self) -> "UserBannerUpdate":
+        if self.enabled and not self.message.strip():
+            raise ValueError("message must be non-empty when the banner is enabled")
+        return self
+
+
+class UserBanner(UserBannerUpdate):
+    revision: str = Field(
+        default="",
+        description=(
+            "Server-stamped opaque publish identity; a fresh value is generated on every "
+            "update so clients re-surface dismissed banners on republish."
+        ),
+    )
+
+
+class UserBannerStore(Protocol):
+    async def get_banner(self) -> UserBanner: ...
+
+    async def save_banner(self, banner: UserBanner) -> None: ...

litellm/proxy/_lazy_openapi_snapshot.json

@@ -58015,7 +58015,7 @@
     "paths": {
       "/toolset/{toolset_name}/mcp": {
         "delete": {
-          "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset/<name>/mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.",
+          "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset/<name>/mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: API keys other than an admin's key with no\nteam must hold the toolset through one of their grant sources, or the\nrequest will be rejected with a 403.",
           "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_delete",
           "parameters": [
             {
@@ -58054,7 +58054,7 @@
           ]
         },
         "get": {
-          "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset/<name>/mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.",
+          "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset/<name>/mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: API keys other than an admin's key with no\nteam must hold the toolset through one of their grant sources, or the\nrequest will be rejected with a 403.",
           "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_get",
           "parameters": [
             {
@@ -58093,7 +58093,7 @@
           ]
         },
         "head": {
-          "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset/<name>/mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.",
+          "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset/<name>/mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: API keys other than an admin's key with no\nteam must hold the toolset through one of their grant sources, or the\nrequest will be rejected with a 403.",
           "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_head",
           "parameters": [
             {
@@ -58132,7 +58132,7 @@
           ]
         },
         "options": {
-          "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset/<name>/mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.",
+          "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset/<name>/mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: API keys other than an admin's key with no\nteam must hold the toolset through one of their grant sources, or the\nrequest will be rejected with a 403.",
           "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_options",
           "parameters": [
             {
@@ -58171,7 +58171,7 @@
           ]
         },
         "patch": {
-          "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset/<name>/mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.",
+          "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset/<name>/mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: API keys other than an admin's key with no\nteam must hold the toolset through one of their grant sources, or the\nrequest will be rejected with a 403.",
           "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_patch",
           "parameters": [
             {
@@ -58210,7 +58210,7 @@
           ]
         },
         "post": {
-          "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset/<name>/mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.",
+          "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset/<name>/mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: API keys other than an admin's key with no\nteam must hold the toolset through one of their grant sources, or the\nrequest will be rejected with a 403.",
           "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_post",
           "parameters": [
             {
@@ -58249,7 +58249,7 @@
           ]
         },
         "put": {
-          "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset/<name>/mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.",
+          "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset/<name>/mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: API keys other than an admin's key with no\nteam must hold the toolset through one of their grant sources, or the\nrequest will be rejected with a 403.",
           "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_put",
           "parameters": [
             {

litellm/proxy/_types.py

@@ -124,6 +124,7 @@ from .types_utils.utils import get_instance_fn, validate_custom_validate_return_
 if TYPE_CHECKING:
     from opentelemetry.trace import Span as _Span
 
+    from litellm.repositories.database import Database
     from litellm.tracing import TraceReceiver
 
     Span = _Span | Any
@@ -133,6 +134,7 @@ else:
 
 class ProxyLifespanState(TypedDict):
     tracing_receiver: ReadOnly["TraceReceiver | None"]
+    database: ReadOnly["Database | None"]
 
 
 class ReconcileOutcome(NamedTuple):

litellm/proxy/agent_endpoints/a2a_endpoints.py

@@ -431,6 +431,7 @@ async def _handle_stream_message(
     proxy_server_request: dict[str, object] | None = None,
     *,
     agent_extra_headers: dict[str, str] | None = None,
+    agent_name: str | None = None,
     user_api_key_dict: UserAPIKeyAuth | None = None,
     request_data: dict[str, object] | None = None,
     proxy_logging_obj: ProxyLogging | None = None,
@@ -504,6 +505,7 @@ async def _handle_stream_message(
                 metadata=metadata,
                 proxy_server_request=proxy_server_request,
                 agent_extra_headers=agent_extra_headers,
+                agent_name=agent_name,
             )
 
             if (
@@ -932,6 +934,7 @@ async def invoke_agent_a2a(
                 metadata=data.get("metadata", {}),
                 proxy_server_request=data.get("proxy_server_request"),
                 agent_extra_headers=agent_extra_headers,
+                agent_name=str(agent_name),
                 user_api_key_dict=user_api_key_auth_for_call,
                 request_data=data,
                 proxy_logging_obj=proxy_logging_obj,

litellm/proxy/auth/auth_checks.py

@@ -167,6 +167,7 @@ from litellm.repositories.table_repositories import (
 from litellm.repositories.team_repository import TeamRepository
 from litellm.repositories.user_repository import UserRepository
 from litellm.router import Router
+from litellm.router_api import RegistryReader
 from litellm.types.proxy.auth.auth_checks import UserNotFoundError
 from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget
 from litellm.types.router import DeploymentTypedDict
@@ -525,7 +526,7 @@ def _log_budget_lookup_failure(entity: str, error: Exception) -> None:
     )
 
 
-def _get_router_zero_cost_cache(llm_router: Router) -> dict[str, bool] | None:
+def _get_router_zero_cost_cache(llm_router: RegistryReader) -> dict[str, bool] | None:
     """
     Return the router's per-instance zero-cost cache, or ``None`` for objects
     that don't expose one (e.g. ``MagicMock`` stand-ins in unit tests).
@@ -541,7 +542,7 @@ def _get_router_zero_cost_cache(llm_router: Router) -> dict[str, bool] | None:
     return cache if isinstance(cache, dict) else None
 
 
-def is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None) -> bool:
+def is_model_cost_zero(model: str | list[str] | None, llm_router: RegistryReader | None) -> bool:
     """
     Check if a model has zero cost (no configured pricing).
 
@@ -667,7 +668,7 @@ _NO_MODEL_INFO: Final[Mapping[str, object]] = MappingProxyType({})
 _TEAM_GRANT_RELATIONS: Final[Mapping[str, object]] = MappingProxyType({"litellm_model_table": True})
 
 
-def _deployments_served_for(model: str, llm_router: "Router") -> Sequence[DeploymentTypedDict]:
+def _deployments_served_for(model: str, llm_router: RegistryReader) -> Sequence[DeploymentTypedDict]:
     """
     The deployments a request to ``model`` routes to, for the gates that waive budget checks.
 
@@ -4842,7 +4843,7 @@ async def get_skill_names_from_access_groups(
 
 def _resolve_all_team_model_sentinel_for_auth_check(
     models: list[str],
-    llm_router: Router | None,
+    llm_router: RegistryReader | None,
     team_id: str | None,
 ) -> list[str]:
     if SpecialModelNames.all_team_models.value not in models or team_id is None or llm_router is None:
@@ -4856,7 +4857,7 @@ def _resolve_all_team_model_sentinel_for_auth_check(
 
 def check_model_access_helper(
     model: str,
-    llm_router: Router | None,
+    llm_router: RegistryReader | None,
     models: list[str],
     team_model_aliases: dict[str, str] | None = None,
     team_id: str | None = None,
@@ -4907,7 +4908,7 @@ _check_model_access_helper: Final = check_model_access_helper
 
 def _owning_team_names_for_team_deployment(
     model: str,
-    llm_router: Router | None,
+    llm_router: RegistryReader | None,
     team_id: str | None,
     team_model_aliases: Mapping[str, str] | None,
 ) -> tuple[str, ...]:
@@ -4925,7 +4926,7 @@ def _owning_team_names_for_team_deployment(
 
 def _can_object_call_model(
     model: str | list[str],
-    llm_router: Router | None,
+    llm_router: RegistryReader | None,
     models: list[str],
     team_model_aliases: dict[str, str] | None = None,
     team_id: str | None = None,
@@ -5052,7 +5053,7 @@ def _resolve_team_alias(
 
 
 def _live_team_alias_target(
-    model: str, team_model_aliases: Mapping[str, str], team_id: str | None, llm_router: Router | None
+    model: str, team_model_aliases: Mapping[str, str], team_id: str | None, llm_router: RegistryReader | None
 ) -> str:
     target: Final = team_model_aliases.get(model)
     if target is None:
@@ -5466,7 +5467,7 @@ async def can_key_call_model(
     model: str | list[str],
     llm_model_list: Sequence[object] | None,
     valid_token: UserAPIKeyAuth,
-    llm_router: litellm.Router | None,
+    llm_router: RegistryReader | None,
     prisma_client: DatabaseClient | None = None,
 ) -> Literal[True]:
     """
@@ -5545,7 +5546,7 @@ async def can_key_call_resolved_model(
     model: str,
     llm_model_list: Sequence[object] | None,
     valid_token: UserAPIKeyAuth,
-    llm_router: litellm.Router | None,
+    llm_router: RegistryReader | None,
 ) -> None:
     from litellm.proxy.proxy_server import (
         prisma_client,
@@ -5693,7 +5694,7 @@ def can_org_access_model(
 async def can_team_access_model(
     model: str | list[str],
     team_object: LiteLLM_TeamTable | None,
-    llm_router: Router | None,
+    llm_router: RegistryReader | None,
     team_model_aliases: dict[str, str] | None = None,
     key_model_aliases: Mapping[str, str] | None = None,
     prisma_client: DatabaseClient | None = None,
@@ -5796,7 +5797,7 @@ async def key_access_group_grants_model(
     model: str | list[str],
     valid_token: UserAPIKeyAuth | None,
     team_object: LiteLLM_TeamTable | None,
-    llm_router: Router | None,
+    llm_router: RegistryReader | None,
     key_model_aliases: Mapping[str, str] | None = None,
 ) -> bool:
     """
@@ -5832,7 +5833,7 @@ _key_access_group_grants_model: Final = key_access_group_grants_model
 def can_project_access_model(
     model: str | list[str],
     project_object: LiteLLM_ProjectTable,
-    llm_router: Router | None,
+    llm_router: RegistryReader | None,
     key_model_aliases: Mapping[str, str] | None = None,
 ) -> Literal[True]:
     """
@@ -5853,7 +5854,7 @@ def can_project_access_model(
 def can_customer_access_model(
     model: str | list[str],
     end_user_object: LiteLLM_EndUserTable,
-    llm_router: Router | None,
+    llm_router: RegistryReader | None,
     valid_token: UserAPIKeyAuth | None,
 ) -> Literal[True]:
     team_model_aliases: Final = team_model_aliases_for_auth_check(valid_token) if valid_token is not None else None
@@ -6016,7 +6017,7 @@ async def _team_grants_personal_key_model(
     model: str,
     team_object: LiteLLM_TeamTableCachedObj,
     valid_token: UserAPIKeyAuth,
-    llm_router: Router | None,
+    llm_router: RegistryReader | None,
     prisma_client: PrismaClient | None,
     user_api_key_cache: UserApiKeyCache,
     proxy_logging_obj: ProxyLogging,
@@ -6054,7 +6055,7 @@ async def _loaded_teams_allow_model(
     model: str,
     teams: tuple[LiteLLM_TeamTableCachedObj | None, ...],
     valid_token: UserAPIKeyAuth,
-    llm_router: Router | None,
+    llm_router: RegistryReader | None,
     prisma_client: PrismaClient | None,
     user_api_key_cache: UserApiKeyCache,
     proxy_logging_obj: ProxyLogging,
@@ -6320,7 +6321,7 @@ async def can_personal_key_call_model_via_teams(
     model: str | Sequence[str],
     user_object: LiteLLM_UserTable | None,
     valid_token: UserAPIKeyAuth,
-    llm_router: Router | None,
+    llm_router: RegistryReader | None,
     prisma_client: PrismaClient | None,
     user_api_key_cache: UserApiKeyCache,
     proxy_logging_obj: ProxyLogging,
@@ -7112,7 +7113,7 @@ async def check_team_member_model_access(
     model: str | list[str],
     team_object: LiteLLM_TeamTable,
     valid_token: UserAPIKeyAuth,
-    llm_router: Router | None,
+    llm_router: RegistryReader | None,
     prisma_client: Optional["PrismaClient"],
     user_api_key_cache: UserApiKeyCache,
     proxy_logging_obj: ProxyLogging,

litellm/proxy/auth/fallback_budget.py

@@ -44,7 +44,7 @@ from litellm.proxy.auth.auth_checks import (  # noqa: F401  # legacy module expo
     _is_model_cost_zero,  # pyright: ignore[reportPrivateUsage,reportUnusedImport]  # backwards-compatible package export
     is_model_cost_zero,  # pyright: ignore[reportPrivateUsage]  # the zero-cost predicate the auth-time budget checks use; no public equivalent
 )
-from litellm.router import Router
+from litellm.router_api import RegistryReader, RoutingContext
 from litellm.types.llms.base import LiteLLMBaseModel
 
 
@@ -118,7 +118,7 @@ async def is_token_within_budget_for_model(
     *,
     model: str,
     valid_token: UserAPIKeyAuth,
-    llm_router: Router,
+    llm_router: RegistryReader,
     team_member_model_budget_limiter: _TeamMemberModelBudgetLimiter | None = None,
 ) -> bool:
     """
@@ -187,19 +187,19 @@ class RouterFallbackBudgetCheck:
     is_enforced: Callable[[], bool]
     team_member_model_budget_limiter: _TeamMemberModelBudgetLimiter | None = None
 
-    async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], llm_router: Router) -> bool:
+    async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], ctx: RoutingContext) -> bool:
         if not self.is_enforced():
             return True
         valid_token: Final = _user_api_key_auth_from_request(request_kwargs)
         if valid_token is None:
             return True
-        budget_models: Final = tuple(dict.fromkeys((model, llm_router.registry().resolve(model))))
+        budget_models: Final = tuple(dict.fromkeys((model, ctx.registry().resolve(model))))
         for budget_model in budget_models:
             try:
                 if not await is_token_within_budget_for_model(
                     model=budget_model,
                     valid_token=valid_token,
-                    llm_router=llm_router,
+                    llm_router=ctx,
                     team_member_model_budget_limiter=self.team_member_model_budget_limiter,
                 ):
                     return False

litellm/proxy/auth/fallback_model_access.py

@@ -18,7 +18,7 @@ from pydantic import ValidationError
 from litellm._logging import verbose_proxy_logger
 from litellm.proxy._types import ProxyException, UserAPIKeyAuth
 from litellm.proxy.auth.auth_checks import can_key_call_resolved_model, can_token_call_search_tool
-from litellm.router import Router
+from litellm.router_api import RegistryReader, RoutingContext
 from litellm.search import asearch
 from litellm.types.llms.base import LiteLLMBaseModel
 
@@ -31,7 +31,7 @@ class _FallbackAccessSettings(LiteLLMBaseModel):
     enforce_fallback_model_access: bool = True
 
 
-async def is_model_authorized_for_token(*, model: str, valid_token: UserAPIKeyAuth, llm_router: Router) -> bool:
+async def is_model_authorized_for_token(*, model: str, valid_token: UserAPIKeyAuth, llm_router: RegistryReader) -> bool:
     try:
         await can_key_call_resolved_model(
             model=model,
@@ -98,7 +98,7 @@ class RouterFallbackAccessCheck:
 
     is_enforced: Callable[[], bool]
 
-    async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], llm_router: Router) -> bool:
+    async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], ctx: RoutingContext) -> bool:
         valid_token: Final = _user_api_key_auth_from_request(request_kwargs)
         if valid_token is None:
             return True
@@ -106,7 +106,7 @@ class RouterFallbackAccessCheck:
             return await is_search_tool_authorized_for_token(search_tool_name=model, valid_token=valid_token)
         if not self.is_enforced():
             return True
-        return await is_model_authorized_for_token(model=model, valid_token=valid_token, llm_router=llm_router)
+        return await is_model_authorized_for_token(model=model, valid_token=valid_token, llm_router=ctx)
 
 
 router_fallback_access_check: Final = RouterFallbackAccessCheck(is_enforced=_enforced_by_general_settings)

litellm/proxy/auth/model_checks.py

@@ -14,13 +14,13 @@ from typing import Final, TypeVar
 import litellm
 from litellm._logging import verbose_proxy_logger
 from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
+from litellm.litellm_core_utils.pattern_match_deployments import PatternMatchRouter
 from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping
 from litellm.proxy._types import SpecialModelNames, UserAPIKeyAuth
 from litellm.repositories.object_permission_repository import ObjectPermissionRepository
 from litellm.responses.main import strip_responses_routing_prefix
 from litellm.router import Router
-from litellm.router_utils.fallback_event_handlers import get_fallback_model_group
-from litellm.router_utils.pattern_match_deployments import PatternMatchRouter
+from litellm.router_api import fallback_chain
 from litellm.types.router import CredentialLiteLLMParams, LiteLLM_Params
 from litellm.types.utils import LlmProviders
 from litellm.utils import ProviderConfigManager, get_valid_models, potential_model_names_in_lookup_order
@@ -708,7 +708,7 @@ def get_all_fallbacks(
 
     try:
         # Use existing function to get fallback model group
-        fallback_model_group, _ = get_fallback_model_group(fallbacks=fallbacks_config, model_group=model)
+        fallback_model_group, _ = fallback_chain(fallbacks=fallbacks_config, model_group=model)
 
         if fallback_model_group is None:
             return []

litellm/proxy/auth/user_api_key_auth.py

@@ -40,6 +40,7 @@ from litellm.integrations.otel.model.config import is_otel_v2_enabled
 from litellm.integrations.otel.runtime import phase_event, phase_span, seed_request_identity
 from litellm.litellm_core_utils.dd_tracing import tracer
 from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value
+from litellm.litellm_core_utils.routing_common_utils import resolve_model_group_alias, resolve_served_model
 from litellm.models.end_user import LiteLLM_EndUserTable
 from litellm.proxy._types import *
 from litellm.proxy.agent_endpoints.auth.agent_caller import agent_caller_from_headers
@@ -182,13 +183,7 @@ from litellm.proxy.utils import (
 from litellm.repositories.table_repositories import JWTKeyMappingRepository, TeamMembershipRepository
 from litellm.repositories.user_repository import UserRepository
 from litellm.repositories.verification_token_repository import VerificationTokenRepository
-from litellm.router_utils.common_utils import resolve_model_group_alias, resolve_served_model
-from litellm.router_utils.fallback_event_handlers import (
-    PROVIDER_SCOPED_RESOURCE_KEYS,
-    fallback_lookup_groups,
-    fallbacks_disabled_for_request,
-    get_fallback_model_group_for_lookup_groups,
-)
+from litellm.router_api import fallbacks_disabled, references_provider_scoped_resource, request_fallback_chain
 from litellm.secret_managers.main import get_secret_bool
 from litellm.types.services import ServiceTypes
 
@@ -811,11 +806,7 @@ async def _check_team_member_model_budget_for_model(
             model=model_name,
         )
     except litellm.BudgetExceededError:
-        if (
-            request_data.get("model") != model_name
-            or llm_router is None
-            or fallbacks_disabled_for_request(kwargs=request_data)
-        ):
+        if request_data.get("model") != model_name or llm_router is None or fallbacks_disabled(request_data):
             raise
         fallback_succeeded: Final = await _try_team_member_model_budget_fallback(
             user_id=user_id,
@@ -857,9 +848,9 @@ async def _try_team_member_model_budget_fallback(
     user_api_key_cache: UserApiKeyCache | None = None,
     proxy_logging_obj: ProxyLogging | None = None,
 ) -> bool:
-    if request.headers.get("content-type", "").lower().startswith("multipart/form-data") or any(
-        request_data.get(key) for key in PROVIDER_SCOPED_RESOURCE_KEYS
-    ):
+    if request.headers.get("content-type", "").lower().startswith(
+        "multipart/form-data"
+    ) or references_provider_scoped_resource(request_data):
         return False
 
     from litellm.proxy import proxy_server
@@ -896,10 +887,7 @@ async def _try_team_member_model_budget_fallback(
     if router_fallbacks is None:
         return False
 
-    fallback_models, _ = get_fallback_model_group_for_lookup_groups(
-        fallbacks=list(router_fallbacks),
-        lookup_groups=fallback_lookup_groups(kwargs=request_data, model_group=model_name),
-    )
+    fallback_models, _ = request_fallback_chain(list(router_fallbacks), request_data, model_name)
     fallback_targets: Final[Sequence[object]] = fallback_models or ()
     fallback_candidates: Final[tuple[str, ...]] = tuple(
         model for target in fallback_targets if (model := _budget_fallback_target_model(target)) is not None

litellm/proxy/bug_report_config.py

@@ -19,8 +19,8 @@ from litellm.litellm_core_utils.bug_report import (
     build_bug_report,
     build_environment_report,
 )
+from litellm.litellm_core_utils.routing_groups import VALID_ROUTING_STRATEGIES
 from litellm.proxy._types import ConfigGeneralSettings
-from litellm.router_utils.routing_groups import VALID_ROUTING_STRATEGIES
 from litellm.types.caching import LiteLLMCacheType
 from litellm.types.guardrails import GuardrailEventHooks, LitellmParams, SupportedGuardrailIntegrations
 from litellm.types.secret_managers.main import KeyManagementSystem

litellm/proxy/common_request_processing.py

@@ -53,6 +53,7 @@ from litellm.constants import (
 )
 from litellm.integrations.custom_guardrail import CustomGuardrail
 from litellm.integrations.otel.runtime import phase_event
+from litellm.litellm_core_utils.add_retry_fallback_headers import get_hidden_params_dict
 from litellm.litellm_core_utils.bug_report import (
     allowlisted,
     bug_report_notice,
@@ -81,6 +82,7 @@ from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import (
 from litellm.litellm_core_utils.llm_response_utils.get_headers import (
     get_response_headers,
 )
+from litellm.litellm_core_utils.routing_common_utils import resolve_model_group_alias, resolve_served_model
 from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 from litellm.litellm_core_utils.served_output_texts import (
     record_served_output_texts,
@@ -146,8 +148,6 @@ from litellm.responses.streaming_iterator import partial_responses_incomplete_ev
 from litellm.router import Router
 from litellm.router_api import RoutingContext
 from litellm.router_strategy.auto_router.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX
-from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict
-from litellm.router_utils.common_utils import resolve_model_group_alias, resolve_served_model
 from litellm.types.extraction import ExtractionResponse
 from litellm.types.guardrails import GuardrailEventHooks
 from litellm.types.router import RouterRateLimitError
@@ -2565,7 +2565,7 @@ class ProxyBaseLLMRequestProcessing:
         aliased_model: Final = self.data.get("model")
         if llm_router is not None and route_type in ROUTER_SCHEDULED_ROUTE_TYPES and isinstance(aliased_model, str):
             self.data.update(
-                llm_router.priority_header_update(
+                llm_router.health().priority_header_fields(
                     model=aliased_model, kwargs=self.data, header=request.headers.get(SCHEDULER_PRIORITY_HEADER)
                 )
             )
@@ -2831,12 +2831,9 @@ class ProxyBaseLLMRequestProcessing:
 
     @staticmethod
     def _resolve_fallback_models(model: str, fallbacks: list) -> list[str] | None:
-        from litellm.router_utils.fallback_event_handlers import get_fallback_model_group
+        from litellm.router_api import fallback_chain
 
-        fallback_model_group, generic_fallback_idx = get_fallback_model_group(
-            fallbacks=fallbacks,
-            model_group=model,
-        )
+        fallback_model_group, generic_fallback_idx = fallback_chain(fallbacks, model)
         if fallback_model_group is not None:
             return fallback_model_group
         return fallbacks[generic_fallback_idx]["*"] if generic_fallback_idx is not None else None
@@ -4190,8 +4187,8 @@ class ProxyBaseLLMRequestProcessing:
         HiddenParamsAsyncIteratorWrapper, so class sniffing runs on the
         unwrapped inner iterator.
         """
+        from litellm.litellm_core_utils.add_retry_fallback_headers import HiddenParamsAsyncIteratorWrapper
         from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
-        from litellm.router_utils.add_retry_fallback_headers import HiddenParamsAsyncIteratorWrapper
 
         unwrapped: Final = response.inner if isinstance(response, HiddenParamsAsyncIteratorWrapper) else response
 

litellm/proxy/common_utils/cache_aware_routing.py

@@ -16,6 +16,7 @@ from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_f
 from litellm.llms.anthropic.cache_aware_routing import AnthropicCacheRouting, TokenCounter
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.auth.auth_checks import can_key_call_resolved_model
+from litellm.router_deprecated import renamed_router_param
 from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig
 from litellm.types.llms.base import LiteLLMBaseModel
 from litellm.types.router import Deployment, PreRoutingHookResponse
@@ -23,6 +24,7 @@ from litellm.types.utils import StandardLoggingRoutingDecision
 
 if TYPE_CHECKING:
     from litellm.router import Router
+    from litellm.router_api import RoutingContext
 
 _MESSAGES: Final = TypeAdapter(list[Mapping[str, object]])
 _MAPPING: Final = TypeAdapter(Mapping[str, object])
@@ -231,7 +233,7 @@ async def select_cached_model(
 
 async def _choose_cached_model(
     *,
-    router: Router,
+    ctx: RoutingContext,
     config: ComplexityRouterConfig,
     params_for_model: Callable[[str, str], Mapping[str, object]],
     response: PreRoutingHookResponse | None,
@@ -249,8 +251,10 @@ async def _choose_cached_model(
     )
     from litellm.router_strategy.complexity_router.context_compaction import compaction_pending
 
+    router: Final = proxy_server.llm_router
     if (
-        proxy_server.llm_router is not router
+        router is None
+        or not ctx.belongs_to(router)
         or router.registry().plugins("routing")
         or has_request_transforms()
         or compaction_pending(request_kwargs)
@@ -318,9 +322,10 @@ async def _choose_cached_model(
     )
 
 
+@renamed_router_param("router")
 async def choose_cached_model(
     *,
-    router: Router,
+    ctx: RoutingContext,
     config: ComplexityRouterConfig,
     params_for_model: Callable[[str, str], Mapping[str, object]],
     response: PreRoutingHookResponse | None,
@@ -332,7 +337,7 @@ async def choose_cached_model(
     try:
         return await asyncio.wait_for(
             _choose_cached_model(
-                router=router,
+                ctx=ctx,
                 config=config,
                 params_for_model=params_for_model,
                 response=response,

litellm/proxy/common_utils/credential_hydration.py

@@ -12,10 +12,10 @@ from types import MappingProxyType
 from typing import Final, cast  # noqa: TID251  # narrows untyped credential values at the decrypt boundary
 
 import litellm
+from litellm.litellm_core_utils.clientside_credential_handler import clientside_credential_keys
 from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper
 from litellm.proxy.utils import PrismaClient
 from litellm.repositories.credentials_repository import CredentialsRepository
-from litellm.router_utils.clientside_credential_handler import clientside_credential_keys
 from litellm.types.router import (
     GenericLiteLLMParams,
     server_owned_wif_fields_present,

litellm/proxy/common_utils/model_listing_utils.py

@@ -21,7 +21,7 @@ from pydantic import TypeAdapter, ValidationError
 
 import litellm
 from litellm.constants import CLAUDE_CODE_CLIENT, CLAUDE_CODE_PICKER_PATTERN, GATEWAY_CLIENT_HEADER
-from litellm.router_utils.common_utils import resolve_model_group_alias
+from litellm.litellm_core_utils.routing_common_utils import resolve_model_group_alias
 
 if TYPE_CHECKING:
     from litellm.router import Router

litellm/proxy/db/accounting_ledger.py

@@ -31,6 +31,13 @@ class AccountingConflictError(ValueError):
     pass
 
 
+@dataclass(frozen=True, slots=True)
+class UsageWrite:
+    event: InboundUsageEvent | OutboundUsageEvent
+    envelope: str
+    targets: tuple[DeliveryTarget, ...] = ()
+
+
 _Transaction: TypeAlias = AbstractAsyncContextManager["Prisma"]
 
 
@@ -107,6 +114,35 @@ SELECT event_id FROM "LiteLLM_AccountingEvent" WHERE kind = 'cost' AND source_ev
 _LOCK_EVENT: Final = """
 SELECT payload_hash FROM "LiteLLM_AccountingEvent" WHERE event_id = $1 FOR UPDATE
 """
+_INSERT_USAGE_BATCH: Final = """
+INSERT INTO "LiteLLM_AccountingEvent"
+    (event_id, kind, schema_version, payload_hash, occurred_at, payload)
+SELECT event_id, kind::"AccountingEventKind", 1, payload_hash, occurred_at::timestamp, payload
+FROM jsonb_to_recordset($1::jsonb) AS records(
+    event_id text, kind text, payload_hash text, occurred_at text, payload jsonb)
+ORDER BY event_id
+ON CONFLICT (event_id) DO NOTHING RETURNING event_id
+"""
+_LOCK_USAGE_BATCH: Final = """
+SELECT event_id, payload_hash FROM "LiteLLM_AccountingEvent"
+WHERE event_id IN (SELECT value FROM jsonb_array_elements_text($1::jsonb))
+ORDER BY event_id FOR UPDATE
+"""
+_LOAD_USAGE_BATCH: Final = """
+SELECT event.payload, ARRAY(
+    SELECT target::text FROM "LiteLLM_AccountingDelivery" delivery
+    WHERE delivery.event_id = event.event_id ORDER BY target::text
+) AS targets
+FROM "LiteLLM_AccountingEvent" event
+WHERE event.event_id IN (SELECT value FROM jsonb_array_elements_text($1::jsonb))
+    AND event.kind = 'outbound_usage' AND event.schema_version = 1
+"""
+_INSERT_USAGE_DELIVERIES: Final = """
+INSERT INTO "LiteLLM_AccountingDelivery" (event_id, target, version)
+SELECT records.event_id, destination.value::"AccountingDeliveryTarget", 1
+FROM jsonb_to_recordset($1::jsonb) AS records(event_id text, targets jsonb)
+CROSS JOIN LATERAL jsonb_array_elements_text(records.targets) AS destination
+"""
 _COMPLETED_FAMILY: Final = """
 root.source_event_id IS NULL AND root.kind IN ('inbound_usage', 'outbound_usage')
 AND root.recorded_at < $1::timestamp
@@ -199,6 +235,53 @@ class AccountingLedger:
 
         return cls(transaction)
 
+    async def insert_usage_batch(self, writes: tuple[UsageWrite, ...]) -> None:
+        records: Final = tuple(_prepare(write.event, write.targets, write.envelope, 1, 1) for write in writes)
+        if not records:
+            return
+        unique: Final = {record.event_id: record for record in records}
+        if any(record.payload_hash != unique[record.event_id].payload_hash for record in records):
+            raise AccountingConflictError("Batch contains conflicting copies of a usage identity")
+        batch: Final = _json(
+            tuple(
+                {
+                    "event_id": record.event_id,
+                    "kind": record.kind.value,
+                    "payload_hash": record.payload_hash,
+                    "occurred_at": record.occurred_at.astimezone(timezone.utc).isoformat(),
+                    "payload": _JSON.validate_json(record.payload),
+                }
+                for record in unique.values()
+            )
+        )
+        async with self._transaction() as db:
+            inserted: Final = _STRINGS.validate_python(await db.query_raw(_INSERT_USAGE_BATCH, batch))
+            accepted: Final = _STRINGS.validate_python(await db.query_raw(_LOCK_USAGE_BATCH, _json(tuple(unique))))
+            if len(accepted) != len(unique) or any(
+                row["payload_hash"] != unique[row["event_id"]].payload_hash for row in accepted
+            ):
+                raise AccountingConflictError("Usage identity conflicts with accepted content or delivery set")
+            deliveries: Final = _json(
+                tuple({"event_id": row["event_id"], "targets": unique[row["event_id"]].targets} for row in inserted)
+            )
+            await db.execute_raw(_INSERT_USAGE_DELIVERIES, deliveries)
+
+    async def load_usage_batch(self, event_ids: tuple[str, ...]) -> tuple[UsageWrite, ...]:
+        if not event_ids:
+            return ()
+        async with self._transaction() as db:
+            rows: Final = TypeAdapter(tuple[dict[str, JsonValue], ...]).validate_python(
+                await db.query_raw(_LOAD_USAGE_BATCH, _json(event_ids))
+            )
+        return tuple(
+            UsageWrite(
+                OutboundUsageEvent.model_validate(_JSON.validate_python(row["payload"])["event"]),
+                _json(row["payload"]),
+                TypeAdapter(tuple[DeliveryTarget, ...]).validate_python(row["targets"]),
+            )
+            for row in rows
+        )
+
     async def insert_event(
         self,
         event: AccountingEvent,

litellm/proxy/db/exception_handler.py

@@ -72,6 +72,22 @@ def _exception_types(*candidates: object) -> tuple[type[BaseException], ...]:
     return tuple(c for c in candidates if isinstance(c, type) and issubclass(c, BaseException))
 
 
+def _prisma_error_code(e: Exception) -> str | None:
+    import prisma.engine.errors
+
+    if isinstance(e, _exception_types(prisma.errors.DataError)):
+        code: Final[object] = getattr(e, "code", None)
+        return code if isinstance(code, str) else None
+    if not isinstance(e, _exception_types(prisma.engine.errors.EngineRequestError)):
+        return None
+    try:
+        body: Final = _DATABASE_ERROR_META.validate_json(str(e).partition(": ")[2])
+    except ValidationError:
+        return None
+    engine_code: Final = body.get("error_code")
+    return engine_code if isinstance(engine_code, str) else None
+
+
 class PrismaDBExceptionHandler:
     """
     Class to handle DB Exceptions or Connection Errors
@@ -206,6 +222,8 @@ class PrismaDBExceptionHandler:
             return True
         if PrismaDBExceptionHandler.is_database_capacity_error(e):
             return False
+        if _prisma_error_code(e) == "P1017":
+            return True
         if isinstance(e, _exception_types(prisma.errors.PrismaError)):
             error_message: Final = str(e).lower()
             connection_keywords: Final = (

litellm/proxy/db/routing_prisma_wrapper.py

@@ -4,6 +4,7 @@ to a writer client. Used when DATABASE_URL_READ_REPLICA is configured;
 otherwise PrismaClient uses the writer-only PrismaWrapper directly.
 """
 
+import asyncio
 import os
 from collections.abc import Callable
 from datetime import timedelta
@@ -233,8 +234,13 @@ class RoutingPrismaWrapper:
         await self._reader.start_token_refresh_task()
 
     async def stop_token_refresh_task(self) -> None:
-        await self._writer.stop_token_refresh_task()
-        await self._reader.stop_token_refresh_task()
+        results: Final = await asyncio.gather(
+            self._writer.stop_token_refresh_task(),
+            self._reader.stop_token_refresh_task(),
+            return_exceptions=True,
+        )
+        if (error := next((result for result in results if isinstance(result, BaseException)), None)) is not None:
+            raise error
 
     async def recreate_prisma_client(
         self,

litellm/proxy/deps/database.py

@@ -0,0 +1,22 @@
+from typing import Annotated, Final, TypeAlias
+
+from fastapi import Depends, Request
+
+from litellm.repositories.database import Database
+
+_MISSING_DATABASE: Final = object()
+
+
+class DatabaseStateMissingError(RuntimeError):
+    def __init__(self) -> None:
+        super().__init__("The proxy lifespan did not publish a Database in request state")
+
+
+async def get_database(request: Request) -> Database | None:
+    database: Final[object] = getattr(request.state, "database", _MISSING_DATABASE)
+    if database is None or isinstance(database, Database):
+        return database
+    raise DatabaseStateMissingError()
+
+
+DatabaseDep: TypeAlias = Annotated[Database | None, Depends(get_database)]

litellm/proxy/deps/repositories.py

@@ -0,0 +1,36 @@
+from typing import Annotated, TypeAlias
+
+from fastapi import Depends
+
+from litellm.models.deployment_warning import DeploymentWarningIgnoreStore
+from litellm.models.gateway_activity import GatewayActivityStore
+from litellm.models.user_banner import UserBannerStore
+from litellm.proxy.deps.database import DatabaseDep
+from litellm.repositories.deployment_warning_ignore_repository import DeploymentWarningIgnoreRepository
+from litellm.repositories.gateway_activity_repository import GatewayActivityRepository
+from litellm.repositories.user_banner_repository import UserBannerRepository
+
+
+async def get_user_banner_repository(database: DatabaseDep) -> UserBannerStore | None:
+    return None if database is None else UserBannerRepository(database.prisma)
+
+
+UserBannerRepositoryDep: TypeAlias = Annotated[UserBannerStore | None, Depends(get_user_banner_repository)]
+
+
+async def get_deployment_warning_ignore_repository(database: DatabaseDep) -> DeploymentWarningIgnoreStore | None:
+    return None if database is None else DeploymentWarningIgnoreRepository(database.prisma, use_writer=True)
+
+
+DeploymentWarningIgnoreRepositoryDep: TypeAlias = Annotated[
+    DeploymentWarningIgnoreStore | None, Depends(get_deployment_warning_ignore_repository)
+]
+
+
+async def get_gateway_activity_repository(database: DatabaseDep) -> GatewayActivityStore | None:
+    return None if database is None else GatewayActivityRepository(database.prisma)
+
+
+GatewayActivityRepositoryDep: TypeAlias = Annotated[
+    GatewayActivityStore | None, Depends(get_gateway_activity_repository)
+]

litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py

@@ -29,6 +29,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
     has_tool_with_name,
 )
 from litellm.litellm_core_utils.provider_affinity import get_caller_session_id
+from litellm.litellm_core_utils.routing_common_utils import request_user_api_key_auth
 from litellm.llms.custom_httpx.http_handler import (
     get_async_httpx_client,  # pyright: ignore[reportUnknownVariableType]
     httpxSpecialProvider,
@@ -40,7 +41,6 @@ from litellm.proxy.guardrails.guardrail_hooks.content_text import (
     merge_rewritten_text_parts,
 )
 from litellm.proxy.spend_tracking.compression_savings import HEADROOM_GUARDRAIL_PROVIDER
-from litellm.router_utils.common_utils import request_user_api_key_auth
 from litellm.secret_managers.main import get_secret_str
 from litellm.types.guardrails import GuardrailEventHooks, Mode
 from litellm.types.integrations.custom_logger import (

litellm/proxy/guardrails/guardrail_hooks/semantic_guard/route_loader.py

@@ -14,6 +14,7 @@ from pydantic import TypeAdapter
 
 from litellm._logging import verbose_logger
 from litellm.constants import DEFAULT_SEMANTIC_GUARD_SIMILARITY_THRESHOLD
+from litellm.router_api import RoutingContext
 
 if TYPE_CHECKING:
     from semantic_router.routers import SemanticRouter
@@ -134,7 +135,7 @@ class SemanticGuardRouteLoader:
         )
 
         encoder: Final = LiteLLMRouterEncoder(
-            litellm_router_instance=litellm_router,
+            ctx=RoutingContext(litellm_router),
             model_name=embedding_model,
             score_threshold=global_threshold,
         )

litellm/proxy/health_check.py

@@ -27,12 +27,12 @@ from litellm.constants import (
     HEALTH_CHECK_TIMEOUT_SECONDS,
 )
 from litellm.litellm_core_utils.health_check_helpers import native_health_check_mode
+from litellm.litellm_core_utils.routing_common_utils import deployment_model_id_pin, is_model_capability_field
 from litellm.router_strategy.auto_router.auto_router_model_naming import (
     StrategyRouterDependency,
     classify_strategy_router_model,
     strategy_router_dependencies,
 )
-from litellm.router_utils.common_utils import deployment_model_id_pin, is_model_capability_field
 from litellm.types.utils import secret_bearing_wif_litellm_params, server_owned_wif_litellm_params
 
 # Provider routing and server-owned federation or OAuth token-exchange fields. Allowed for proxy

litellm/proxy/health_check_utils/shared_health_check_manager.py

@@ -13,10 +13,10 @@ from litellm.caching.redis_cache import RedisCache
 from litellm.constants import (
     DEFAULT_SHARED_HEALTH_CHECK_LOCK_TTL,
     DEFAULT_SHARED_HEALTH_CHECK_TTL,
+    HEALTH_CHECKS_TARGET,
 )
 from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 from litellm.proxy.health_check import perform_health_check
-from litellm.router_utils.health_state_cache import HEALTH_CHECKS_TARGET
 
 if TYPE_CHECKING:
     from litellm.router import Router

litellm/proxy/health_endpoints/_health_endpoints.py

@@ -26,7 +26,12 @@ from litellm.integrations.SlackAlerting.ms_teams import (
     build_ms_teams_payload,
     get_ms_teams_webhook_url,
 )
+from litellm.litellm_core_utils.clientside_credential_handler import (
+    _ADMIN_CONFIG_FIELDS_TO_CLEAR_ON_BASE_OVERRIDE,  # pyright: ignore[reportPrivateUsage]  # one canonical list, shared with the router path
+    clientside_credential_keys,
+)
 from litellm.litellm_core_utils.custom_logger_registry import CustomLoggerRegistry
+from litellm.litellm_core_utils.routing_common_utils import resolve_model_group_alias
 from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
 from litellm.proxy._types import (
     AlertType,
@@ -94,11 +99,6 @@ from litellm.proxy.middleware.in_flight_requests_middleware import (
 from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager
 from litellm.router import Router
 from litellm.router_api import deployment_usable_by_team
-from litellm.router_utils.clientside_credential_handler import (
-    _ADMIN_CONFIG_FIELDS_TO_CLEAR_ON_BASE_OVERRIDE,  # pyright: ignore[reportPrivateUsage]  # one canonical list, shared with the router path
-    clientside_credential_keys,
-)
-from litellm.router_utils.common_utils import resolve_model_group_alias
 from litellm.secret_managers.main import get_secret_bool
 from litellm.types.proxy.litellm_pre_call_utils import RedactedDict, SecretFields
 from litellm.types.router import reject_server_owned_wif_fields

litellm/proxy/health_endpoints/deployment_warning_endpoints.py

@@ -1,4 +1,3 @@
-import json
 import logging
 from collections.abc import Mapping
 from dataclasses import dataclass
@@ -7,12 +6,14 @@ from types import MappingProxyType
 from typing import Annotated, Final, Literal, Protocol, TypeAlias
 
 from fastapi import APIRouter, Depends, HTTPException, status
-from pydantic import AwareDatetime, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError, model_validator
+from pydantic import ConfigDict, Field, JsonValue, TypeAdapter, model_validator
 
-from litellm._logging import verbose_logger, verbose_proxy_logger
+from litellm._logging import verbose_logger
 from litellm.constants import LITELLM_PROXY_ADMIN_NAME
+from litellm.models.deployment_warning import AdminIgnore, DeploymentWarningIgnoreStore
 from litellm.proxy._types import AUDIT_ACTIONS, LitellmTableNames, LitellmUserRoles, UserAPIKeyAuth
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+from litellm.proxy.deps.repositories import DeploymentWarningIgnoreRepositoryDep
 from litellm.proxy.health_endpoints._health_endpoints import (
     DISABLE_NO_REDIS_WARNING_ENV_VAR,
     env_credential_login_enabled,
@@ -22,7 +23,6 @@ from litellm.proxy.health_endpoints._health_endpoints import (
 )
 from litellm.repositories.deployment_warning_ignore_repository import (
     DEPLOYMENT_WARNING_IGNORES_ROW_ID,
-    DeploymentWarningIgnoreRepository,
 )
 from litellm.types.llms.base import LiteLLMBaseModel
 
@@ -50,18 +50,6 @@ class WarningCondition:
     env_var_ignore: str | None = None
 
 
-class AdminIgnore(LiteLLMBaseModel):
-    model_config = ConfigDict(frozen=True)
-
-    source: Literal["admin"] = "admin"
-    ignored_by: str
-    ignored_at: AwareDatetime
-    ignored_until: AwareDatetime | None = Field(
-        default=None,
-        description="When the snooze ends. Null means the warning is ignored until an admin restores it.",
-    )
-
-
 class EnvVarIgnore(LiteLLMBaseModel):
     model_config = ConfigDict(frozen=True)
 
@@ -109,12 +97,6 @@ class DeploymentWarningIgnoreResponse(LiteLLMBaseModel):
     ignore: AdminIgnore | None
 
 
-class DeploymentWarningIgnoreStore(Protocol):
-    async def get_raw_ignores(self) -> object: ...
-
-    async def upsert_ignores(self, payload: str) -> None: ...
-
-
 class ConfigAuditLogger(Protocol):
     async def __call__(
         self,
@@ -127,7 +109,6 @@ class ConfigAuditLogger(Protocol):
     ) -> None: ...
 
 
-_IGNORES_ADAPTER: Final[TypeAdapter[Mapping[str, AdminIgnore]]] = TypeAdapter(Mapping[str, AdminIgnore])
 _JSON_VALUE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
 _AUDIT_ENTRIES_ADAPTER: Final[TypeAdapter[tuple[DeploymentWarningIgnoreResponse, ...]]] = TypeAdapter(
     tuple[DeploymentWarningIgnoreResponse, ...]
@@ -185,18 +166,8 @@ def is_known_warning_id(warning_id: str) -> bool:
         return False
 
 
-def _parse_stored_ignores(raw: object) -> Mapping[str, AdminIgnore]:
-    if raw is None:
-        return MappingProxyType({})
-    try:
-        return MappingProxyType(_IGNORES_ADAPTER.validate_python(json.loads(raw) if isinstance(raw, str) else raw))
-    except (json.JSONDecodeError, ValidationError):
-        verbose_proxy_logger.warning("Ignoring unreadable %s row", DEPLOYMENT_WARNING_IGNORES_ROW_ID)
-        return MappingProxyType({})
-
-
 async def _read_unexpired_ignores(store: DeploymentWarningIgnoreStore, now: datetime) -> Mapping[str, AdminIgnore]:
-    stored: Final = _parse_stored_ignores(await store.get_raw_ignores())
+    stored: Final = await store.get_ignores()
     return MappingProxyType(
         {
             warning_id: ignore
@@ -225,14 +196,6 @@ def utc_now() -> datetime:
     return datetime.now(timezone.utc)
 
 
-def deployment_warning_ignore_store() -> DeploymentWarningIgnoreStore | None:
-    from litellm.proxy.proxy_server import prisma_client
-
-    if prisma_client is None:
-        return None
-    return DeploymentWarningIgnoreRepository(prisma_client, use_writer=True)
-
-
 def config_audit_logger() -> ConfigAuditLogger:
     from litellm.proxy.proxy_server import create_config_audit_log
 
@@ -281,7 +244,7 @@ async def current_warning_conditions(now: Annotated[datetime, Depends(utc_now)])
 async def get_deployment_warnings(
     _: Annotated[UserAPIKeyAuth, Depends(_deployment_warning_reader)],
     conditions: Annotated[tuple[WarningCondition, ...], Depends(current_warning_conditions)],
-    store: Annotated[DeploymentWarningIgnoreStore | None, Depends(deployment_warning_ignore_store)],
+    store: DeploymentWarningIgnoreRepositoryDep,
     now: Annotated[datetime, Depends(utc_now)],
 ) -> DeploymentWarningsResponse:
     """
@@ -309,7 +272,7 @@ async def update_deployment_warning_ignore(
     warning_id: str,
     update: DeploymentWarningIgnoreUpdate,
     user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(_deployment_warning_admin)],
-    store: Annotated[DeploymentWarningIgnoreStore | None, Depends(deployment_warning_ignore_store)],
+    store: DeploymentWarningIgnoreRepositoryDep,
     audit_logger: Annotated[ConfigAuditLogger, Depends(config_audit_logger)],
     now: Annotated[datetime, Depends(utc_now)],
 ) -> DeploymentWarningIgnoreResponse:
@@ -338,7 +301,7 @@ async def update_deployment_warning_ignore(
     others: Final = {other_id: other for other_id, other in before.items() if other_id != warning_id}
     after: Final = MappingProxyType(others if ignore is None else {**others, warning_id: ignore})
 
-    await store.upsert_ignores(_IGNORES_ADAPTER.dump_json(dict(after)).decode())
+    await store.save_ignores(after)
     await audit_logger(
         param_name=DEPLOYMENT_WARNING_IGNORES_ROW_ID,
         action="updated",

litellm/proxy/hooks/dynamic_rate_limiter_v3.py

@@ -15,6 +15,10 @@ from litellm._internal_context import with_service_target
 from litellm._logging import verbose_proxy_logger
 from litellm.caching.caching import DualCache
 from litellm.integrations.custom_logger import CustomLogger
+from litellm.litellm_core_utils.add_retry_fallback_headers import (
+    ensure_response_additional_headers,
+    response_has_hidden_params,
+)
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.common_utils.proxy_rate_limit_error import (
     ProxyRateLimitError,
@@ -34,10 +38,6 @@ from litellm.proxy.hooks.rate_limiter_utils import (
     resolve_llm_provider_for_rate_limit,
 )
 from litellm.proxy.utils import InternalUsageCache
-from litellm.router_utils.add_retry_fallback_headers import (
-    ensure_response_additional_headers,
-    response_has_hidden_params,
-)
 from litellm.types.router import ModelGroupInfo
 from litellm.types.utils import CallTypesLiteral, LLMResponseTypes
 

litellm/proxy/hooks/model_max_budget_limiter.py

@@ -204,7 +204,7 @@ def _team_public_names_by_group(router: "Router", team_id: str | None) -> Mappin
     groups_by_name: Final = MappingProxyType(
         {
             name: model_access_groups_of(router=router, model=name, team_id=team_id)
-            for name in router.team_public_model_names_of(team_id)
+            for name in router.registry().names(kind="team_public", team_id=team_id)
         }
     )
     return MappingProxyType(

litellm/proxy/hooks/parallel_request_limiter_v3.py

@@ -49,6 +49,10 @@ from litellm.constants import (
 )
 from litellm.cost_calculator import get_usage_object
 from litellm.integrations.custom_logger import CustomLogger
+from litellm.litellm_core_utils.add_retry_fallback_headers import (
+    ensure_response_additional_headers,
+    response_has_hidden_params,
+)
 from litellm.litellm_core_utils.audio_utils.utils import estimate_audio_tokens
 from litellm.litellm_core_utils.core_helpers import as_str_mapping
 from litellm.litellm_core_utils.prompt_templates.common_utils import (
@@ -56,6 +60,8 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
     split_media_for_tpm,
 )
 from litellm.litellm_core_utils.ptu_pricing import is_ptu_cost_attribution_enabled
+from litellm.litellm_core_utils.ptu_shares import PTUTeamCeiling
+from litellm.litellm_core_utils.routing_common_utils import resolve_model_group_alias
 from litellm.litellm_core_utils.token_counter import estimate_image_tokens, offload_token_count
 from litellm.llms.azure.ptu_capacity import normalized_tokens
 from litellm.proxy._types import UserAPIKeyAuth
@@ -86,12 +92,6 @@ from litellm.proxy.spend_tracking.input_tokens import (
     openai_chat_pdf_file_tokens,
 )
 from litellm.responses.litellm_completion_transformation.compaction_items import billable_input_since_newest_compaction
-from litellm.router_utils.add_retry_fallback_headers import (
-    ensure_response_additional_headers,
-    response_has_hidden_params,
-)
-from litellm.router_utils.common_utils import resolve_model_group_alias
-from litellm.router_utils.ptu_shares import PTUTeamCeiling
 from litellm.types.caching import RedisPipelineIncrementOperation
 from litellm.types.extraction import ExtractionResponse
 from litellm.types.llms.openai import BaseLiteLLMOpenAIResponseObject, ResponseAPIUsage
@@ -3561,9 +3561,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
         Returns True if any deployment has failures in the current minute.
         """
         from litellm.proxy.proxy_server import llm_router
-        from litellm.router_utils.router_callbacks.track_deployment_metrics import (
-            get_deployment_failures_for_current_minute,
-        )
 
         if llm_router is None:
             return False
@@ -3581,10 +3578,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
                     continue
 
                 # Use router's existing failure tracking
-                failure_count = get_deployment_failures_for_current_minute(
-                    litellm_router_instance=llm_router,
-                    deployment_id=deployment_id,
-                )
+                failure_count = llm_router.health().failures_this_minute(deployment_id)
 
                 if failure_count > DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE:
                     verbose_proxy_logger.debug(

litellm/proxy/litellm_pre_call_utils.py

@@ -112,6 +112,7 @@ _EXPLICIT_SESSION_HEADERS: Final = frozenset({"x-litellm-trace-id", "x-litellm-s
 # ``session-id``/``thread-id``; builds before the codex-api split sent
 # ``session_id``/``conversation_id``. Ordered session before thread.
 _CODEX_SESSION_ID_HEADERS: Final = ("session-id", "session_id", "thread-id", "conversation_id")
+_BEDROCK_SESSION_ID_METADATA_KEYS: Final = ("sessionId", "threadId", "conversationId")
 # Session-id values must be non-empty strings of alphanumerics, hyphens, or underscores
 # (covers UUIDs and most common session-id formats).
 _SESSION_ID_VALUE_RE: Final = re.compile(r"^[a-zA-Z0-9_\-]{8,}$")
@@ -1055,6 +1056,22 @@ def _get_anthropic_session_id_from_metadata(metadata: object) -> str | None:
     return session_id
 
 
+def _get_bedrock_session_id_from_request_metadata(
+    normalized_headers: Mapping[str, str], request_body: Mapping[str, object]
+) -> str | None:
+    """Read the conversation id a Bedrock SDK client tags its calls with: Converse's ``requestMetadata``
+    body field or InvokeModel's ``X-Amzn-Bedrock-Request-Metadata`` JSON header. Only the Bedrock
+    pass-through route nests the client's body under ``data``."""
+    if request_body.get("custom_llm_provider") != "bedrock" or not isinstance(request_body.get("data"), Mapping):
+        return None
+    header: Final = normalized_headers.get("x-amzn-bedrock-request-metadata")
+    body_tags: Final = _request_fields(_request_fields(request_body.get("data")).get("requestMetadata"))
+    header_tags: Final = _request_fields(safe_json_loads(header) if header else None)
+    tag_sources: Final = itertools.product((body_tags, header_tags), _BEDROCK_SESSION_ID_METADATA_KEYS)
+    values: Final = (tags.get(key) for tags, key in tag_sources)
+    return next((value for value in values if isinstance(value, str) and value), None)
+
+
 def _is_llm_inference_route(request: Request) -> bool:
     route: Final = get_request_route(request)
     return RouteChecks.is_llm_api_route(route=route) and not RouteChecks.check_route_access(
@@ -1967,6 +1984,14 @@ class LiteLLMProxyRequestSetup:
                     body_metadata["user_id"] = session_id
                 verbose_proxy_logger.debug("Extracted session_id from Anthropic metadata.user_id")
 
+        if "litellm_session_id" not in data:
+            bedrock_session_id: Final = _get_bedrock_session_id_from_request_metadata(normalized_headers, request_body)
+            metadata_session_id: Final = _caller_trace_field(request_body, _metadata_variable_name, "session_id")
+            if bedrock_session_id and metadata_session_id in (None, bedrock_session_id):
+                metadata_from_headers["session_id"] = bedrock_session_id
+                data["litellm_session_id"] = bedrock_session_id  # rebind-ok: data is an out-param
+                verbose_proxy_logger.debug("Extracted session_id from Bedrock request metadata")
+
         # Last-resort fallback: the W3C standards for trace/session propagation
         # (https://www.w3.org/TR/trace-context/, https://www.w3.org/TR/baggage/).
         # Lower priority than everything above - only fires when neither the
@@ -3266,7 +3291,7 @@ def _fallback_edges(
     and its provider-stripped suffix, plus the router matcher run over the fuzzy
     entries that have not fired yet (tracked in ``fired_fuzzy`` by
     ``(list_index, entry_index)``), so each entry's targets are expanded once."""
-    from litellm.router_utils.fallback_event_handlers import get_fallback_model_group
+    from litellm.router_api import fallback_chain
 
     stripped: Final = model_group.split("/", 1)[1] if "/" in model_group else None
     targets: Final[set[str]] = set()  # mutable-ok: accumulates fallback targets
@@ -3277,7 +3302,7 @@ def _fallback_edges(
         pending = tuple(index for index in range(len(fuzzy)) if (list_index, index) not in fired_fuzzy)
         if not pending:
             continue
-        resolved = get_fallback_model_group(fallbacks=[fuzzy[index] for index in pending], model_group=model_group)[0]
+        resolved = fallback_chain([fuzzy[index] for index in pending], model_group)[0]
         values = resolved if isinstance(resolved, list) else ([resolved] if resolved is not None else [])
         targets.update(name for value in values if (name := _fallback_target_name(value)) is not None)
         if resolved is not None:
@@ -3464,7 +3489,7 @@ async def _resolve_user_provider_credentials_for_request(
             status_code=400,
             detail=(f"fallback lists in the request body cannot exceed {_FALLBACK_DISCOVERY_LIMIT} entries"),
         )
-    from litellm.router_utils.common_utils import resolve_model_group_alias
+    from litellm.litellm_core_utils.routing_common_utils import resolve_model_group_alias
 
     # Discovery must see the group the router actually routes to, so include the
     # post-alias name from every alias map that may still rewrite data["model"]

litellm/proxy/management_endpoints/auto_router_endpoints.py

@@ -60,6 +60,7 @@ from litellm.repositories.table_repositories import PrismaTableRepository
 from litellm.repositories.team_repository import TeamRepository
 from litellm.repositories.user_repository import UserRepository
 from litellm.repositories.verification_token_repository import VerificationTokenRepository
+from litellm.router_api import RoutingContext
 from litellm.router_strategy.auto_router.auto_router_model_naming import (
     StrategyRouterDependencyRole,
     classify_strategy_router_model,
@@ -614,7 +615,7 @@ async def preview_auto_router_routing(
     try:
         complexity_strategy: Final = ComplexityRouter(
             model_name=resolved.router_name,
-            litellm_router_instance=llm_router,
+            ctx=RoutingContext(llm_router),
             complexity_router_config=resolved.complexity_router_config.model_dump(exclude_none=True),
             default_model=resolved.complexity_router_config.resolve_default_model(resolved.default_model),
             derive_savings_baseline=False,

litellm/proxy/management_endpoints/gateway_request_endpoints.py

@@ -13,15 +13,15 @@ and the endpoint is restricted to proxy admin roles.
 
 from collections.abc import Sequence
 from datetime import datetime, timedelta, timezone
-from typing import Annotated, Final, Literal
+from typing import Annotated, Final
 
 from fastapi import APIRouter, Depends, HTTPException, Query
-from pydantic import ConfigDict, Field, TypeAdapter
 
 from litellm._logging import verbose_proxy_logger
+from litellm.models.gateway_activity import GatewayRouteActivity, GatewayStatusCodeActivity
 from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
-from litellm.types.llms.base import LiteLLMBaseModel
+from litellm.proxy.deps.repositories import GatewayActivityRepositoryDep
 from litellm.types.proxy.gateway_requests import (
     GatewayRequestActivityResponse,
     GatewayRequestBreakdownEntry,
@@ -33,50 +33,6 @@ router: Final = APIRouter()
 
 _DEFAULT_LOOKBACK_DAYS: Final = 30
 
-_ACTIVITY_SQL: Final = """
-    SELECT
-        'route' AS kind, date, category, route,
-        SUM(successful_requests)::bigint AS successful_requests,
-        SUM(failed_requests)::bigint AS failed_requests,
-        NULL::integer AS status_code
-    FROM "LiteLLM_DailyGatewayRequests"
-    WHERE date >= $1 AND date <= $2
-    GROUP BY date, category, route
-    UNION ALL
-    SELECT
-        'status' AS kind, NULL, NULL, NULL,
-        0::bigint,
-        SUM(failed_requests)::bigint,
-        status_code
-    FROM "LiteLLM_DailyGatewayFailedRequests"
-    WHERE date >= $1 AND date <= $2
-    GROUP BY status_code
-"""
-
-
-class _AggregateRow(LiteLLMBaseModel):
-    """Validates one query_raw row so the handler works with typed values, not Any."""
-
-    kind: Literal["route"]
-    date: str
-    category: str
-    route: str
-    successful_requests: int
-    failed_requests: int
-
-
-class _StatusCodeAggregateRow(LiteLLMBaseModel):
-    model_config = ConfigDict(frozen=True)
-
-    kind: Literal["status"]
-    status_code: int
-    failed_requests: int
-
-
-_ACTIVITY_ROWS_ADAPTER: Final = TypeAdapter(
-    tuple[Annotated[_AggregateRow | _StatusCodeAggregateRow, Field(discriminator="kind")], ...]
-)
-
 
 def _default_range() -> tuple[str, str]:
     end: Final = datetime.now(timezone.utc)
@@ -84,7 +40,7 @@ def _default_range() -> tuple[str, str]:
     return start.strftime("%Y-%m-%d"), end.strftime("%Y-%m-%d")
 
 
-def _fold_by_date(rows: Sequence[_AggregateRow]) -> tuple[GatewayRequestDailyEntry, ...]:
+def _fold_by_date(rows: Sequence[GatewayRouteActivity]) -> tuple[GatewayRequestDailyEntry, ...]:
     dates: Final = sorted(frozenset(row.date for row in rows))
     return tuple(
         GatewayRequestDailyEntry(
@@ -96,7 +52,7 @@ def _fold_by_date(rows: Sequence[_AggregateRow]) -> tuple[GatewayRequestDailyEnt
     )
 
 
-def _fold_by_route(rows: Sequence[_AggregateRow]) -> tuple[GatewayRequestBreakdownEntry, ...]:
+def _fold_by_route(rows: Sequence[GatewayRouteActivity]) -> tuple[GatewayRequestBreakdownEntry, ...]:
     pairs: Final = sorted(frozenset((row.category, row.route) for row in rows))
     entries: Final = tuple(
         GatewayRequestBreakdownEntry(
@@ -113,7 +69,7 @@ def _fold_by_route(rows: Sequence[_AggregateRow]) -> tuple[GatewayRequestBreakdo
 
 
 def _fold_by_status_code(
-    rows: Sequence[_StatusCodeAggregateRow],
+    rows: Sequence[GatewayStatusCodeActivity],
 ) -> tuple[GatewayRequestStatusCodeEntry, ...]:
     return tuple(
         GatewayRequestStatusCodeEntry(status_code=row.status_code, failed_requests=row.failed_requests)
@@ -128,6 +84,7 @@ def _fold_by_status_code(
 )
 async def get_gateway_daily_activity(
     user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
+    activity: GatewayActivityRepositoryDep,
     start_date: str | None = Query(default=None, description="Start date in YYYY-MM-DD format"),
     end_date: str | None = Query(default=None, description="End date in YYYY-MM-DD format"),
 ) -> GatewayRequestActivityResponse:
@@ -137,8 +94,6 @@ async def get_gateway_daily_activity(
     Deployment-wide: the underlying table has no per-key or per-user dimension,
     so this is admin-only.
     """
-    from litellm.proxy.proxy_server import prisma_client
-
     if user_api_key_dict.user_role not in (
         LitellmUserRoles.PROXY_ADMIN,
         LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
@@ -148,22 +103,15 @@ async def get_gateway_daily_activity(
             detail="Only proxy admin roles can view gateway request counts across the deployment",
         )
 
-    if prisma_client is None:
+    if activity is None:
         raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
 
     default_start, default_end = _default_range()
     selected_start: Final = start_date or default_start
     selected_end: Final = end_date or default_end
-    raw_rows: Final = await prisma_client.db.query_raw(  # pyright: ignore[reportAny]  # untyped prisma client
-        _ACTIVITY_SQL,
-        selected_start,
-        selected_end,
-    )
-    activity_rows: Final = _ACTIVITY_ROWS_ADAPTER.validate_python(raw_rows or ())
-    rows: Final[tuple[_AggregateRow, ...]] = tuple(row for row in activity_rows if isinstance(row, _AggregateRow))
-    status_code_rows: Final[tuple[_StatusCodeAggregateRow, ...]] = tuple(
-        row for row in activity_rows if isinstance(row, _StatusCodeAggregateRow)
-    )
+    activity_rows: Final = await activity.get_activity(selected_start, selected_end)
+    rows: Final = tuple(row for row in activity_rows if isinstance(row, GatewayRouteActivity))
+    status_code_rows: Final = tuple(row for row in activity_rows if isinstance(row, GatewayStatusCodeActivity))
     verbose_proxy_logger.debug("/gateway/daily/activity - aggregated %d rows", len(rows))
 
     return GatewayRequestActivityResponse(

litellm/proxy/management_endpoints/model_management_endpoints.py

@@ -1515,14 +1515,10 @@ def update_db_model(
 
     # convert to prisma compatible format
 
-    for key, value in merged_model_info.items():
-        if isinstance(value, datetime.datetime):
-            merged_model_info[key] = value.isoformat()
-
     prisma_compatible_model_dict: Final = PrismaCompatibleUpdateDBModel(
         model_name=merged_model_name,
         litellm_params=json.dumps(merged_litellm_params),
-        model_info=json.dumps(merged_model_info),
+        model_info=TypeAdapter(dict[str, object]).dump_json(merged_model_info).decode(),
     )
 
     if updated_patch.blocked is not None:

litellm/proxy/management_endpoints/ptu_consumption.py

@@ -12,14 +12,14 @@ from types import MappingProxyType
 from typing import Final
 
 from litellm.litellm_core_utils.ptu_pricing import is_ptu_cost_attribution_enabled
-from litellm.llms.azure.ptu_capacity import PTUCapacity, normalized_tokens, ptu_hours
-from litellm.router import Router
-from litellm.router_utils.common_utils import resolve_model_group_alias
-from litellm.router_utils.ptu_shares import (
+from litellm.litellm_core_utils.ptu_shares import (
     model_group_ptu_capacity,
     routed_deployments,
     team_servable_deployments,
 )
+from litellm.litellm_core_utils.routing_common_utils import resolve_model_group_alias
+from litellm.llms.azure.ptu_capacity import PTUCapacity, normalized_tokens, ptu_hours
+from litellm.router import Router
 from litellm.types.proxy.management_endpoints.common_daily_activity import (
     DailySpendData,
     MetricWithMetadata,

litellm/proxy/management_endpoints/types.py

@@ -93,5 +93,6 @@ def get_litellm_user_role(role_str: object) -> LitellmUserRoles | None:
 
 class CustomOpenID(OpenID):
     team_ids: list[str]
+    team_id_to_upsert: str | None = None
     user_role: LitellmUserRoles | None = None
     extra_fields: dict[str, object] | None = None

litellm/proxy/management_endpoints/ui_sso.py

@@ -87,7 +87,7 @@ from litellm.proxy._types import (
     TeamMemberAddRequest,
     UserAPIKeyAuth,
 )
-from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken, get_user_object
+from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken, get_team_object, get_user_object
 from litellm.proxy.auth.auth_utils import (  # noqa: F401  # legacy module exports
     _get_request_ip_address,  # pyright: ignore[reportPrivateUsage,reportUnusedImport]  # backwards-compatible package export
     get_request_ip_address,
@@ -1266,6 +1266,7 @@ def generic_response_convertor(
     sso_jwt_handler: JWTHandler | None = None,
     role_mappings: Optional["RoleMappings"] = None,
     team_mappings: Optional["TeamMappings"] = None,
+    jwt_auth_enabled: bool = False,
 ) -> CustomOpenID:
     generic_user_id_attribute_name: Final = os.getenv("GENERIC_USER_ID_ATTRIBUTE", "preferred_username")
     generic_user_display_name_attribute_name: Final = os.getenv("GENERIC_USER_DISPLAY_NAME_ATTRIBUTE", "sub")
@@ -1286,14 +1287,15 @@ def generic_response_convertor(
         generic_user_email_attribute_name,
     )
 
-    all_teams: Final = []
+    token: Final = cast(dict[str, object], response)  # cast-ok: the IdP userinfo JSON object fastapi-sso hands back
+    all_teams: Final = list[str]()
     if sso_jwt_handler is not None:
-        team_ids = sso_jwt_handler.get_all_jwt_team_ids(cast(dict, response))
+        team_ids = sso_jwt_handler.get_all_jwt_team_ids(token)
         all_teams.extend(team_ids)
 
     if team_mappings is not None and team_mappings.team_ids_jwt_field is not None:
         team_ids_from_db_mapping: Final[list[str] | None] = get_nested_value(
-            data=cast(dict, response),
+            data=token,
             key_path=team_mappings.team_ids_jwt_field,
             default=[],
         )
@@ -1305,8 +1307,17 @@ def generic_response_convertor(
                 team_ids_from_db_mapping,
             )
     else:
-        team_ids = jwt_handler.get_all_jwt_team_ids(cast(dict, response))
+        team_ids = jwt_handler.get_all_jwt_team_ids(token)
         all_teams.extend(team_ids)
+    team_id_jwt_field: Final = jwt_handler.litellm_jwtauth.team_id_jwt_field
+    upsert_candidate: Final[object] = (
+        get_nested_value(data=token, key_path=team_id_jwt_field)
+        if jwt_auth_enabled and jwt_handler.litellm_jwtauth.team_id_upsert and team_id_jwt_field is not None
+        else None
+    )
+    team_id_to_upsert: Final = (
+        upsert_candidate if isinstance(upsert_candidate, str) and upsert_candidate.strip() else None
+    )
 
     # Determine user role based on role_mappings if available
     # Only apply role_mappings for GENERIC SSO provider
@@ -1375,6 +1386,7 @@ def generic_response_convertor(
         last_name=get_nested_value(response, generic_user_last_name_attribute_name),
         provider=get_nested_value(response, generic_provider_attribute_name),
         team_ids=all_teams,
+        team_id_to_upsert=team_id_to_upsert,
         user_role=user_role,
         extra_fields=extra_fields,
     )
@@ -1614,6 +1626,8 @@ async def get_generic_sso_response(
     # make generic sso provider
     from fastapi_sso.sso.base import DiscoveryDocument
 
+    from litellm.proxy import proxy_server
+
     received_response: dict | None = None
     sso_assertion: SSOIdentityAssertion | None = None
 
@@ -1665,6 +1679,7 @@ async def get_generic_sso_response(
             sso_jwt_handler=sso_jwt_handler,
             role_mappings=role_mappings,
             team_mappings=team_mappings,
+            jwt_auth_enabled=proxy_server.general_settings.get("enable_jwt_auth", False) is True,
         )
 
     SSOProvider: Final = create_generic_sso_provider(
@@ -1884,6 +1899,29 @@ async def find_existing_team_ids(team_ids: Sequence[str]) -> frozenset[str] | No
     return frozenset(row.team_id for row in rows)
 
 
+async def create_sso_team(team_id: str) -> None:
+    """Create the ``team_id_jwt_field`` team on UI SSO login the way the API JWT path does under ``team_id_upsert``."""
+    from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
+
+    if prisma_client is None:
+        return
+    await get_team_object(
+        team_id=team_id,
+        prisma_client=prisma_client,
+        user_api_key_cache=user_api_key_cache,
+        check_db_only=True,
+        team_id_upsert=True,
+    )
+
+
+async def _create_sso_team_or_warn(team_id: str, create_team: Callable[[str], Awaitable[object]]) -> None:
+    """A concurrent login may create the team first, so a failed create is not fatal: the join still runs."""
+    try:
+        await create_team(team_id)
+    except Exception as e:  # noqa: BLE001  # fail-safe: a failed team create must not fail the login or skip the join
+        verbose_proxy_logger.warning("SSO login could not create team %s, trying to join it anyway: %s", team_id, e)
+
+
 async def add_missing_team_member(
     user_info: NewUserResponse | LiteLLM_UserTable,
     sso_teams: Sequence[str],
@@ -1892,21 +1930,30 @@ async def add_missing_team_member(
         [str, NewUserResponse | LiteLLM_UserTable], Awaitable[object]
     ] = create_team_member_add_task,
     existing_team_ids: Callable[[Sequence[str]], Awaitable[frozenset[str] | None]] = find_existing_team_ids,
+    team_id_to_upsert: str | None = None,
+    create_team: Callable[[str], Awaitable[object]] = create_sso_team,
 ) -> None:
     """
     - Get missing teams (diff b/w user_info.team_ids and sso_teams)
     - Drop the IdP groups that are not teams with one lookup, since group claims often carry many of them
+    - Join the ``team_id_to_upsert`` team too, creating it when it is missing (``team_id_upsert``), the way the API JWT path does;
+      it stays out of ``result.team_ids`` so restricted_sso_group and the access-token group fallback see only the claimed groups
     - Add missing user to missing teams, one at a time in sso_teams order, so user.teams keeps the claim's order
     """
     held_teams: Final = frozenset(user_info.teams or ())
-    missing_teams: Final = tuple(team_id for team_id in dict.fromkeys(sso_teams) if team_id not in held_teams)
+    claimed_teams: Final = sso_teams if team_id_to_upsert is None else (*sso_teams, team_id_to_upsert)
+    missing_teams: Final = tuple(team_id for team_id in dict.fromkeys(claimed_teams) if team_id not in held_teams)
     if not missing_teams:
         return
     existing: Final = await existing_team_ids(missing_teams)
     teams_to_add: Final = (
-        missing_teams if existing is None else tuple(team_id for team_id in missing_teams if team_id in existing)
+        missing_teams
+        if existing is None
+        else tuple(team_id for team_id in missing_teams if team_id in existing or team_id == team_id_to_upsert)
     )
     for team_id in teams_to_add:
+        if team_id == team_id_to_upsert and (existing is None or team_id not in existing):
+            await _create_sso_team_or_warn(team_id, create_team)
         await add_team_member(team_id, user_info)
 
 
@@ -3663,7 +3710,11 @@ class SSOAuthenticationHandler:
             verbose_proxy_logger.debug("User not found in LiteLLM DB, skipping team member addition")
             return
         sso_teams: Final = result.team_ids if isinstance(result, CustomOpenID) else []
-        await add_missing_team_member(user_info=user_info, sso_teams=sso_teams)
+        await add_missing_team_member(
+            user_info=user_info,
+            sso_teams=sso_teams,
+            team_id_to_upsert=result.team_id_to_upsert if isinstance(result, CustomOpenID) else None,
+        )
 
     @staticmethod
     def verify_user_in_restricted_sso_group(

litellm/proxy/management_helpers/model_quota_policy.py

@@ -9,10 +9,10 @@ from pydantic import BaseModel, ConfigDict, Field, field_validator
 import litellm
 from litellm.exceptions import BadRequestError
 from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider, is_openai_compatible_endpoint
+from litellm.litellm_core_utils.pattern_match_deployments import PatternMatchRouter, PatternUtils
+from litellm.litellm_core_utils.routing_common_utils import provider_for_generic_call
 from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping
 from litellm.llms.openai_like.json_loader import JSONProviderRegistry
-from litellm.router_utils.common_utils import provider_for_generic_call
-from litellm.router_utils.pattern_match_deployments import PatternMatchRouter, PatternUtils
 from litellm.secret_managers.main import get_secret_str
 
 if TYPE_CHECKING:

litellm/proxy/management_helpers/ptu_share_guard.py

@@ -5,8 +5,8 @@ from typing import Final
 
 from fastapi import HTTPException
 
+from litellm.litellm_core_utils.ptu_shares import deployments_shared_with
 from litellm.router import Router
-from litellm.router_utils.ptu_shares import deployments_shared_with
 
 
 def raise_if_teams_hold_ptu_shares(team_ids: Sequence[str], llm_router: Router | None) -> None:

litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py

@@ -438,6 +438,14 @@ def is_mcp_admitted_user_subject(user_api_key_auth: UserAPIKeyAuth | None) -> bo
 _is_mcp_admitted_user_subject: Final = is_mcp_admitted_user_subject
 
 
+def has_team_unbound_mcp_admin_view(user_api_key_auth: UserAPIKeyAuth | None) -> bool:
+    return (
+        user_api_key_auth is not None
+        and user_api_key_has_admin_view(user_api_key_auth)
+        and (not user_api_key_auth.team_id or user_api_key_auth.team_id == UI_TEAM_ID)
+    )
+
+
 def _gateway_dcr_challenge_target(
     route: str,
     mcp_servers: list[str] | None,
@@ -4290,19 +4298,19 @@ class MCPRequestHandler:
         """Whether this principal's admin-view role grants the unscoped MCP resolution, whatever
         credential carries it (admin key, dashboard session, or OAuth-admitted session subject).
 
-        Two bounds disqualify, one per ownership of the row. A CREDENTIAL's explicit
-        ``object_permission.mcp_servers`` scope wins even for admins, including the empty list. An
-        admitted subject's object_permission is the user's own row, whose ``mcp_servers`` column is
-        [] by DB default, so for that shape the row binds through the entitlement ceiling instead
-        (any non-empty entitlement, or an unresolved one, disqualifies), exactly as
-        ``operator_open_server_ids`` reads the same row. The one owner of this predicate: the
-        server-axis registry resolution in ``get_allowed_mcp_servers`` and the tools-axis open
-        channel in ``resolve_admitted_subject_tools`` both consult it, so the two axes cannot
-        disagree."""
+        A credential on a real team never qualifies, a virtual key or a ``lite login`` session alike: its team
+        bounds its MCP access whoever owns it, as it bounds the credential's model access. Two more bounds disqualify, one per ownership of the row. A CREDENTIAL's
+        explicit ``object_permission.mcp_servers`` scope wins even for admins, including the empty list. An
+        admitted subject's object_permission is the user's own row, whose ``mcp_servers`` column is [] by DB
+        default, so for that shape the row binds through the entitlement ceiling instead (any non-empty
+        entitlement, or an unresolved one, disqualifies), exactly as ``operator_open_server_ids`` reads the same
+        row. The one owner of this predicate: the server-axis registry resolution in ``get_allowed_mcp_servers``
+        and the tools-axis open channel in ``resolve_admitted_subject_tools`` both consult it, so the two axes
+        cannot disagree."""
         if (
             user_api_key_auth is None
             or user_api_key_auth.mcp_explicit_grants_only
-            or not user_api_key_has_admin_view(user_api_key_auth)
+            or not has_team_unbound_mcp_admin_view(user_api_key_auth)
         ):
             return False
         object_permission: Final = user_api_key_auth.object_permission

litellm/proxy/mcp_server/mcp_server_manager.py

@@ -3705,7 +3705,7 @@ class MCPServerManager:
 
         Priority:
         1. If object_permission.mcp_servers is explicitly set, use it (even for admins)
-        2. If admin and no object_permission, return all servers
+        2. If admin, not bound to a real team, and no object_permission, return all servers
         3. Otherwise, use standard permission checks
         """
         if managed_agent_policy(user_api_key_auth) is not None:

litellm/proxy/mcp_server/operations.py

@@ -17,6 +17,7 @@ import anyio
 from fastapi import HTTPException
 from mcp import MCPError, ReadResourceResult, Resource
 from mcp.types import (
+    METHOD_NOT_FOUND,
     CallToolRequest,
     CallToolRequestParams,
     CallToolResult,
@@ -222,6 +223,7 @@ __all__ = (
     "mcp_get_prompt",
     "mcp_read_resource",
     "raise_denied_scoped_mcp_access",
+    "raise_generic_mcp_access_denial",
 )
 
 
@@ -762,8 +764,14 @@ async def raise_denied_scoped_mcp_access(
                 )
             }
             raise HTTPException(status_code=403, detail=group_denial)
+    raise_generic_mcp_access_denial(requested_names)
+
+
+def raise_generic_mcp_access_denial(requested_names: Sequence[str]) -> NoReturn:
     generic_denial: Final[_McpDeniedDetail] = {
-        "error": f"The key is not allowed to access the requested MCP servers: {', '.join(requested_names)}"
+        "error": (
+            f"The key is not allowed to access the requested MCP servers: {', '.join(dict.fromkeys(requested_names))}"
+        )
     }
     raise HTTPException(status_code=403, detail=generic_denial)
 
@@ -2806,6 +2814,58 @@ async def _fire_mcp_tool_call_logging(
     return result
 
 
+async def log_mcp_sdk_rejection(
+    *,
+    method: str,
+    request_id: str | int | None,
+    error_code: int,
+    start_time: datetime,
+    user_api_key_auth: UserAPIKeyAuth,
+) -> None:
+    from litellm.proxy.proxy_server import proxy_logging_obj
+    from litellm.utils import Rules, function_setup
+
+    message: Final = "Method not found" if error_code == METHOD_NOT_FOUND else "Invalid request parameters"
+    rejection: Final = {
+        "method": "".join(char for char in method[:256] if char.isprintable()),
+        "request_id": (
+            "".join(char for char in str(request_id)[:128] if char.isprintable()) if request_id is not None else None
+        ),
+        "error_code": error_code,
+        "error_message": message,
+        "stage": "sdk_dispatch",
+        "route": user_api_key_auth.request_route or "/mcp",
+    }
+    metadata: Final = {
+        **LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_auth),
+        "spend_logs_metadata": {"mcp_sdk_rejection": rejection},
+    }
+    logging_obj, request_data = function_setup(
+        original_function=CallTypes.call_mcp_tool.value,
+        rules_obj=Rules(),
+        start_time=start_time,
+        model="MCP: request_rejected",
+        litellm_call_id=str(uuid.uuid4()),
+        metadata=metadata,
+    )
+    logging_obj.update_environment_variables(  # pyright: ignore[reportUnknownMemberType]  # untyped legacy signature
+        litellm_params={"metadata": metadata},
+        optional_params={},
+        user=user_api_key_auth.user_id,
+    )
+    safe_error: Final = HTTPException(status_code=400, detail=message)
+    end_time: Final = datetime.now()  # noqa: DTZ005  # matches the failure logging start time
+    await asyncio.to_thread(logging_obj.failure_handler, safe_error, "", start_time, end_time)
+    await logging_obj.async_failure_handler(safe_error, "", start_time, end_time)
+    await proxy_logging_obj.post_call_failure_hook(  # pyright: ignore[reportUnknownMemberType]  # legacy signature
+        request_data={**request_data, "litellm_logging_obj": logging_obj},
+        original_exception=safe_error,
+        user_api_key_dict=user_api_key_auth,
+        route="/mcp",
+        traceback_str="",
+    )
+
+
 async def fire_mcp_tool_call_failure_logging(
     logging_obj: LiteLLMLoggingObj | None,
     exception: Exception,

litellm/proxy/mcp_server/semantic_tool_filter.py

@@ -14,6 +14,7 @@ from litellm.litellm_core_utils.exception_mapping_utils import ExceptionCheckers
 from litellm.proxy.mcp_server.catalog import catalog_operation, global_manager
 from litellm.proxy.mcp_server.faults import iter_exception_tree
 from litellm.proxy.mcp_server.utils import MCP_TOOL_PREFIX_SEPARATOR
+from litellm.router_api import RoutingContext
 
 if TYPE_CHECKING:
     from semantic_router.routers import SemanticRouter
@@ -172,7 +173,7 @@ class SemanticMCPToolFilter:
             self.tool_router = SemanticRouter(
                 routes=routes,
                 encoder=LiteLLMRouterEncoder(
-                    litellm_router_instance=self.router_instance,
+                    ctx=RoutingContext(self.router_instance),
                     model_name=self.embedding_model,
                     score_threshold=self.similarity_threshold,
                 ),
@@ -247,7 +248,7 @@ class SemanticMCPToolFilter:
                 router: Final = SemanticRouter(
                     routes=[],
                     encoder=LiteLLMRouterEncoder(
-                        litellm_router_instance=self.router_instance,
+                        ctx=RoutingContext(self.router_instance),
                         model_name=self.embedding_model,
                         score_threshold=self.similarity_threshold,
                     ),

litellm/proxy/mcp_server/server.py

@@ -14,6 +14,7 @@ import time
 import types
 from collections import Counter
 from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterable, Mapping, Sequence, Set
+from datetime import datetime
 from typing import TYPE_CHECKING, Final, Literal, NoReturn, Protocol
 
 import httpx
@@ -519,6 +520,7 @@ if MCP_AVAILABLE:
         "mcp_get_prompt",
         "mcp_read_resource",
         "raise_denied_scoped_mcp_access",
+        "raise_generic_mcp_access_denial",
         "redact_mcp_resource_url",
     )
     from mcp.server import Server
@@ -530,13 +532,15 @@ if MCP_AVAILABLE:
     )
     from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser
     from mcp.server.auth.provider import AccessToken
-    from mcp.server.context import ServerRequestContext
+    from mcp.server.context import CallNext, HandlerResult, ServerRequestContext
     from mcp.server.lowlevel.server import NotificationOptions
     from mcp.server.lowlevel.server import Server as MCPProtocolServer
     from mcp.server.models import InitializationOptions
     from mcp.shared.exceptions import MCPError
     from mcp.shared.inbound import InboundLadderRejection, classify_inbound_request, find_duplicated_routing_header
     from mcp.types import (
+        INVALID_PARAMS,
+        METHOD_NOT_FOUND,
         CallToolRequest,
         GetPromptRequest,
         JSONRPCRequest,
@@ -546,6 +550,7 @@ if MCP_AVAILABLE:
         ListToolsRequest,
         ReadResourceRequest,
     )
+    from mcp_types.methods import CLIENT_REQUESTS, SPEC_CLIENT_METHODS
 
     from litellm.proxy.mcp_server import operations
     from litellm.proxy.mcp_server.contracts import OperationContext
@@ -643,7 +648,56 @@ if MCP_AVAILABLE:
         name=LITELLM_MCP_SERVER_NAME,
         version=LITELLM_MCP_SERVER_VERSION,
     )
-    server.middleware.extend((finish_scope_response, GatewayVersionPolicy()))
+
+    def _sdk_rejection_code(ctx: ServerRequestContext, error: MCPError | ValidationError) -> int | None:
+        if ctx.request_id is None:
+            return None
+        entry: Final = server.get_request_handler(ctx.method)
+        if entry is None:
+            if ctx.method == "initialize":
+                return None
+            if isinstance(error, ValidationError):
+                return INVALID_PARAMS
+            if error.code == METHOD_NOT_FOUND:
+                return METHOD_NOT_FOUND
+            return None
+        if not isinstance(error, ValidationError):
+            unavailable_for_version: Final = (
+                ctx.method in SPEC_CLIENT_METHODS and (ctx.method, ctx.protocol_version) not in CLIENT_REQUESTS
+            )
+            return METHOD_NOT_FOUND if error.code == METHOD_NOT_FOUND and unavailable_for_version else None
+        try:
+            entry.params_type.model_validate({} if ctx.params is None else ctx.params, by_name=False)
+        except ValidationError:
+            return INVALID_PARAMS
+        return None
+
+    async def _log_sdk_rejection(ctx: ServerRequestContext, call_next: CallNext) -> HandlerResult:
+        start_time: Final = datetime.now()  # noqa: DTZ005  # failure logging uses naive timestamps
+        try:
+            return await call_next(ctx)
+        except (MCPError, ValidationError) as error:
+            try:
+                code: Final = _sdk_rejection_code(ctx, error)
+                if code is not None:
+                    token: Final = active_mcp_request_ctx_var.set(ctx)
+                    try:
+                        auth, *_ = await get_or_extract_auth_context()
+                    finally:
+                        active_mcp_request_ctx_var.reset(token)
+                    if auth is not None:
+                        await operations.log_mcp_sdk_rejection(
+                            method=ctx.method,
+                            request_id=ctx.request_id,
+                            error_code=code,
+                            start_time=start_time,
+                            user_api_key_auth=auth,
+                        )
+            except Exception:  # noqa: BLE001  # logging must not replace the SDK response or expose rejected input
+                verbose_logger.warning("Failed to log MCP SDK rejection")
+            raise
+
+    server.middleware.extend((finish_scope_response, GatewayVersionPolicy(), _log_sdk_rejection))
     server.create_initialization_options = types.MethodType(_gateway_create_initialization_options, server)
     sse: Final[SseServerTransport] = SseServerTransport("/sse/messages")
 
@@ -1144,6 +1198,7 @@ if MCP_AVAILABLE:
         apply_display_name_overrides,
         filter_tools_by_allowed_tools,
         raise_denied_scoped_mcp_access,
+        raise_generic_mcp_access_denial,
     )
 
     @contextlib.asynccontextmanager
@@ -1771,10 +1826,10 @@ if MCP_AVAILABLE:
         row capping it.
 
         Raises HTTPException(403) unless the principal holds toolset_id through one
-        of its grant sources. Admins always pass.
+        of its grant sources. Admins pass unless the key is bound to a real team.
         """
         from litellm.proxy._types import LiteLLM_ObjectPermissionTable
-        from litellm.proxy.management_endpoints.common_utils import user_api_key_has_admin_view
+        from litellm.proxy.mcp_server.auth.user_api_key_auth_mcp import has_team_unbound_mcp_admin_view
 
         # A key scoped to no MCP servers opts out of every MCP path. Enforce it
         # here too, since toolset scoping replaces mcp_servers and would otherwise
@@ -1788,7 +1843,7 @@ if MCP_AVAILABLE:
             )
 
         acting: Final = await acting_user(user_api_key_auth)
-        is_admin: Final = user_api_key_has_admin_view(acting)
+        is_admin: Final = has_team_unbound_mcp_admin_view(acting)
         if not is_admin and toolset_id not in await granted(acting):
             raise HTTPException(
                 status_code=403,

litellm/proxy/openai_files_endpoints/files_endpoints.py

@@ -1786,7 +1786,7 @@ def _routed_model_name_on_provider(
 def _stored_model_name_routes_to_provider(
     llm_router: Router, custom_llm_provider: str, provider_model_names: frozenset[str], model_name: str
 ) -> bool:
-    routing_group: Final = llm_router.get_routing_group(model_name)
+    routing_group: Final = llm_router.registry().routing_group(model_name)
     routed_model_names: Final = (
         tuple(routing_group.models) if routing_group is not None else (llm_router.registry().resolve(model_name),)
     )