Merge upstream litellm_internal_staging into litellm_redis_atomic_increment_ttl

Test-file only: upstream added pool-timeout breaker tests next to the spy tests
of this branch. Both sets are kept, production code is unchanged

Co-authored-by: songkuan-zheng <252822057+songkuan-zheng@users.noreply.github.com>
Co-authored-by: songkuan-zheng <songkuan-zheng@users.noreply.github.com>

.gitignore

@@ -147,3 +147,6 @@ crash.*.log
 
 ui/litellm-dashboard/out/
 litellm.log
+
+.coverage-rust
+coverage-rust.xml

enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py

@@ -780,7 +780,10 @@ async def update_project(
 
         # Handle budget updates
         budget_fields = LiteLLM_BudgetTable.model_fields.keys()
-        budget_updates = {k: v for k, v in update_data.items() if k in budget_fields}
+        budget_updates = {
+            **{k: v for k, v in update_data.items() if k in budget_fields},
+            **({"max_budget": None} if "max_budget" in data.model_fields_set and data.max_budget is None else {}),
+        }
 
         if budget_updates and existing_project.budget_id:
             # Update existing budget

enterprise/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "litellm-enterprise"
-version = "0.1.66"
+version = "0.1.67"
 description = "Package for LiteLLM Enterprise features"
 readme = "README.md"
 requires-python = ">=3.9"
@@ -26,7 +26,7 @@ required-version = ">=0.10.9"
 module-root = ""
 
 [tool.commitizen]
-version = "0.1.66"
+version = "0.1.67"
 version_files = [
     "pyproject.toml:^version",
     "../pyproject.toml:litellm-enterprise==",

litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py

@@ -37,11 +37,13 @@ raised it above the deploy default keeps that larger budget for deploy unless
 the deploy override says otherwise.
 """
 
+import importlib.util
 import math
 import os
 import shutil
 import signal
 import subprocess
+import sys
 from collections.abc import Mapping, Sequence
 from dataclasses import dataclass
 from pathlib import Path
@@ -64,6 +66,7 @@ DEFAULT_PRISMA_BOOTSTRAP_TIMEOUT = 600.0
 DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT = 600.0
 
 BOOTSTRAP_ARG = "--version"
+PRISMA_CONSOLE_SCRIPT = "prisma"
 
 
 @dataclass(frozen=True)
@@ -184,6 +187,28 @@ def _kill_process_group(process: "subprocess.Popen[str]") -> None:
         return
 
 
+def prisma_cli_available() -> bool:
+    """Whether some way of running the Prisma CLI exists: the console script on PATH or the importable package."""
+    if shutil.which(PRISMA_CONSOLE_SCRIPT) is not None:
+        return True
+    return importlib.util.find_spec(PRISMA_CONSOLE_SCRIPT) is not None
+
+
+def resolve_prisma_argv(argv: Sequence[str]) -> tuple[str, ...]:
+    """Route a bare ``prisma`` command through ``python -m prisma`` when the console script is not on PATH.
+
+    The console script and ``python -m prisma`` are the same entry point, but
+    only the module form survives an interpreter whose ``bin`` directory is
+    missing from PATH, which is how the proxy gets started under launchers and
+    init systems. Any other executable name is left untouched.
+    """
+    if not argv or argv[0] != PRISMA_CONSOLE_SCRIPT:
+        return tuple(argv)
+    if shutil.which(PRISMA_CONSOLE_SCRIPT) is not None:
+        return tuple(argv)
+    return (sys.executable, "-m", PRISMA_CONSOLE_SCRIPT, *argv[1:])
+
+
 def run_prisma(
     argv: Sequence[str],
     *,
@@ -200,7 +225,7 @@ def run_prisma(
     text unless ``stdout``/``stderr`` say otherwise.
     """
     with subprocess.Popen(
-        argv,
+        resolve_prisma_argv(argv),
         env=env,
         stdout=stdout,
         stderr=stderr,

litellm-proxy-extras/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "litellm-proxy-extras"
-version = "0.4.96"
+version = "0.4.97"
 description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
 readme = "README.md"
 requires-python = ">=3.9"
@@ -26,7 +26,7 @@ required-version = ">=0.10.9"
 module-root = ""
 
 [tool.commitizen]
-version = "0.4.96"
+version = "0.4.97"
 version_files = [
     "pyproject.toml:^version",
     "../pyproject.toml:litellm-proxy-extras==",

litellm-rust/ADDING_A_PROVIDER.md

@@ -1,29 +0,0 @@
-# Adding a provider / route to litellm-rust
-
-Everything for a route lives in `crates/core/src/<route>/`; `crates/core/src/messages` is the reference. A host (the axum gateway, the Python bridge) only calls the route's entrypoint.
-
-1. **Entrypoint** — `mod.rs`: `pub async fn <route>(request) -> CoreResult<Response>`, the Rust equivalent of `litellm.<route>()`, plus a `<route>_stream` variant when the route streams. It is the only thing a host touches.
-2. **Transform contract** — `transformation.rs`: a `…ProviderConfig` trait (URL build + request/response transforms) with types in `types.rs`.
-3. **Provider config** — `crates/core/src/providers/<provider>/<route>/transformation.rs`: implement that trait as a `const <PROVIDER>_<ROUTE>_CONFIG`, mirroring the Python provider tree. Add parity unit tests.
-4. **Prepare + handler** — `prepare.rs` resolves provider/model, credentials, auth headers, and URL, then transforms the request; `handler.rs` performs the provider call through the shared client in `client.rs` and transforms the response.
-
-## Coding standards
-
-Before writing new logic, look for an existing base to extend. When a change is
-“the same behavior for one more provider/endpoint/integration”, the codebase
-almost always already has a shared abstraction for it (for example, provider
-`BaseConfig` transformation classes in `litellm/llms/base_llm/`, shared
-helpers in `litellm_core_utils/`, typed request/response models, or factory
-functions). Find it first with a search, then add the new variant by inheriting
-from or composing that base, overriding only what genuinely differs (model
-name, parameter mapping, or auth).
-
-Never copy an existing implementation and edit it in place, and never hand-roll
-a parallel version of logic a base already provides. If you catch yourself
-writing a second copy of a pattern that exists twice already, stop and extract a
-base instead: put the shared shape in one place and make both call sites thin
-variants of it. The test for a good abstraction is that adding the next provider
-is a few declarative lines, not a new file of duplicated flow. Only diverge from
-the base when behavior is genuinely different, and say so explicitly in the PR.
-
-**Calling:** hosts invoke the core entrypoint — the Python bridge and the `ai-gateway` route service both call `litellm_core::messages::messages`. Never add a provider handler to `ai-gateway`. Register new modules in `lib.rs` / `mod.rs`, then run the commands under "Checks" in [CLAUDE.md](CLAUDE.md).

litellm-rust/AGENTS.md

@@ -1,45 +0,0 @@
-# AGENTS.md
-
-litellm-rust has six crates. A crate is a layer or shared foundation, not a route. Routes (ocr, realtime, chat) and providers (mistral, openai) are modules inside the layers.
-
-## Crates
-
-| Crate | Role |
-|-------|------|
-| litellm-core | The LiteLLM SDK in Rust. One public entrypoint per top-level call (`messages::messages()`), owning types, transforms, provider resolution, auth, and the provider HTTP call. Call it, get a typed response. |
-| litellm-token-counter | Standalone input token counting shared by host integrations without pulling in the full SDK. |
-| litellm-config | Config-loading boundary. Returns resolved core deployment data and optionally delegates loading to Python. |
-| litellm-ai-gateway | The axum server (behind the `server` feature) plus the WebSocket hosts. Translates HTTP/WS to core entrypoints; owns no provider logic and no handlers. |
-| litellm-python-interop | Domain-neutral PyO3 foundation for GIL handling and typed Python/Serde conversion. |
-| litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. Owns API registration, domain wiring, and Python exception mapping. |
-
-Dependency direction is acyclic: `litellm-config` depends on `litellm-core`, the gateway depends on both, and `litellm-python-bridge` depends on the domain layers, `litellm-token-counter`, and `litellm-python-interop`. The token counter and interop foundations depend on no LiteLLM domain crate.
-
-## Where a route lives
-
-A top-level LiteLLM call is a module under `crates/core/src/<route>/`, shaped like `messages`:
-
-```
-core/src/messages/
-  mod.rs             # pub async fn messages(..) -> CoreResult<..>  (+ messages_stream for SSE)
-  types.rs           # request/response types, MessagesRequest
-  transformation.rs  # the provider template trait
-  prepare.rs         # provider resolution, auth headers, URL
-  handler.rs         # the provider call
-  client.rs          # the shared reqwest client
-```
-
-Handlers never live in `ai-gateway`. `ocr`, `audio_transcription`, and `realtime` are still hosted there from before this rule; they move to `core` as they are touched.
-
-Adding a crate: default to a module. A new crate requires a real trigger: separate artifact (binary/cdylib), proc-macro, shared foundation, or publishable standalone. A new provider or route is none of these.
-
-Adding a crate fails crates/core/tests/workspace_crate_allowlist.rs until you update its allowlist and this file — intentional.
-
-## Style
-
-All Rust in `litellm-rust/` follows the official Rust Style Guide:
-https://doc.rust-lang.org/style-guide/
-
-`rustfmt` implements its formatting by default, so run `cargo fmt` before committing; CI gates every PR on `cargo fmt --check`. Do not hand-format against rustfmt or add a `rustfmt.toml` that diverges from the default style.
-
-Beyond formatting, follow the guide's naming and idiom conventions rustfmt cannot auto-apply: `snake_case` items/functions/modules, `UpperCamelCase` types/traits/variants, `SCREAMING_SNAKE_CASE` constants/statics (acronyms as one word, e.g. `HttpClient`), and the import grouping and item ordering it prescribes. See CLAUDE.md for the detailed version.

litellm-rust/CLAUDE.md

@@ -1,189 +0,0 @@
-# CLAUDE.md
-
-This file defines the rules for Rust work in LiteLLM.
-
-## Provider Coding Standards
-
-Before writing new logic, look for an existing base to extend. When a change is
-“the same behavior for one more provider/endpoint/integration”, the codebase
-almost always already has a shared abstraction for it (for example, provider
-`BaseConfig` transformation classes in `litellm/llms/base_llm/`, shared
-helpers in `litellm_core_utils/`, typed request/response models, or factory
-functions). Find it first with a search, then add the new variant by inheriting
-from or composing that base, overriding only what genuinely differs (model
-name, parameter mapping, or auth).
-
-Never copy an existing implementation and edit it in place, and never hand-roll
-a parallel version of logic a base already provides. If you catch yourself
-writing a second copy of a pattern that exists twice already, stop and extract a
-base instead: put the shared shape in one place and make both call sites thin
-variants of it. The test for a good abstraction is that adding the next provider
-is a few declarative lines, not a new file of duplicated flow. Only diverge from
-the base when behavior is genuinely different, and say so explicitly in the PR.
-
-## Crates (see AGENTS.md)
-
-`litellm-core` **is** the LiteLLM SDK in Rust: it makes the LLM call.
-`litellm-config` is the config-loading boundary and returns resolved core types.
-`litellm-ai-gateway` is an HTTP/WebSocket server in front of it, and
-`litellm-python-bridge` exposes it to the Python SDK. `litellm-python-interop`
-holds domain-neutral PyO3 primitives shared by Python-facing Rust code. A crate
-is a layer or shared foundation, not a route; add modules, not crates.
-
-## Core Boundary
-
-`litellm-core` owns the whole call. The Rust equivalent of `litellm.messages()`
-is `litellm_core::messages::messages(request).await`: you call it, it does the
-provider call, and you get a typed non-streaming response back.
-
-Route-level Rust structure mirrors LiteLLM's Python responsibilities:
-- `core/src/<route>/` owns the route end to end: the public entrypoint fn named
-  after the route in `mod.rs`, the request/response types (`types.rs`), the
-  provider template trait (`transformation.rs`), the provider/auth/URL
-  resolution (`prepare.rs`), the HTTP client (`client.rs`), and the handler that
-  performs the call (`handler.rs`). `core/src/messages` is the reference.
-- `core/src/providers/<provider>/<route>/transformation.rs` owns the
-  provider-specific transform. For Anthropic Messages, this means
-  `core/src/providers/anthropic/messages/transformation.rs`.
-- Handlers live in `core`, never in a host. `ai-gateway` must not contain a
-  route handler that talks to a provider; its axum route reads the HTTP request,
-  picks a deployment, and calls the `core` entrypoint. `python-bridge` marshals
-  Python objects and calls the same entrypoint.
-
-Streaming keeps the same shape: the route entrypoint has a `<route>_stream`
-variant in `core` that returns the upstream response so a host can splice it to
-its own caller; the host still owns no provider logic.
-
-Call-hook and lifecycle instrumentation, including phase timing, usage
-accumulation, and callback payload construction, always lives in `core`.
-Hosts feed observed events into core and dispatch the completed payloads through
-their I/O logger; hosts must not own callback orchestration.
-
-Allowed in `core`:
-- The public entrypoint for a top-level LiteLLM call
-- Request/response transforms and stream chunk normalization
-- Provider resolution, auth header construction, and URL building
-- The provider HTTP call itself, through a shared reused client with connect and
-  request timeouts
-- Shared data types and validation errors
-- Deterministic token/cost helper logic
-
-Not allowed in `core`:
-- Serving HTTP: axum routes, extractors, and transport concerns stay in the host
-- Filesystem access
-- Database access
-- Config file reading and rollout state
-- Logging callbacks, spend writes, or custom callbacks
-- Global mutable runtime state
-
-Env reads in `core` are limited to credential fallback inside a route's
-`prepare.rs` (the `env_lookup` closure), mirroring what the Python SDK does when
-no key is passed. Everything else config-shaped is resolved by the host and
-passed in.
-
-Routes still hosted in `ai-gateway` (`ocr`, `audio_transcription`, `realtime`)
-predate this rule and are being moved into `core` route modules; do not add new
-ones there, and prefer moving one when you touch it.
-
-Python owns rollout state and fallback while Rust is being introduced. Rust
-paths must be off by default until parity tests prove equivalence with Python.
-A new provider/route may instead be implemented rust-only with no Python
-reference; then the Python interface is a thin dispatch that calls Rust with no
-fallback, and you state the rust-only choice explicitly in the PR. Either way
-the Python side stays minimal (it only marshals inputs and calls the Rust
-interface), never add a per-route feature flag, and never push provider
-dispatch into `litellm/main.py`; put it in a thin dispatch class under
-`litellm/llms/<provider>/<route>/`.
-
-## Production Bar
-
-Rust code in this workspace is held to a strict parity and robustness bar from
-the first PR:
-
-- Correctness parity is proven with tests. Do not rely on README claims or
-  manual inspection for a port that mirrors Python behavior.
-- Every provider transform must have unit tests for supported-parameter
-  filtering, request body shape, response normalization, missing/null fields,
-  and bad-input errors.
-- When Rust is exposed through Python, add Python tests that prove disabled,
-  enabled, and unavailable-bridge fallback behavior.
-- Avoid panics on user/provider input. Return typed errors and let the host map
-  them to Python exceptions or HTTP responses.
-- OCR handles documents that often contain personal data. Do not log document
-  contents, base64 payloads, provider response bodies, or secrets.
-- Error messages must be useful but data-minimized. Truncate or sanitize any
-  upstream body before it crosses a host boundary.
-- Treat empty or whitespace-only credentials, URLs, and config values as absent
-  at the host/config resolution layer.
-- Preserve Python output shape intentionally. If a field is always serialized as
-  `null` for Python parity, leave a short comment explaining that parity choice.
-
-## Network I/O Rules
-
-These rules apply to every module that executes network I/O, whether it is a
-`core` route handler or a host such as `ai-gateway`:
-
-- Set connect and full-request timeouts. No unbounded waits.
-- Reuse HTTP clients; do not construct clients per request.
-- Prefer rustls TLS for portable Python wheels and Linux images unless there is
-  a documented reason not to.
-- Add request IDs and structured tracing at the host layer, without logging OCR
-  document contents or secrets.
-- Do not echo raw upstream response bodies to callers. Sanitize and bound them.
-- Avoid `expect`/`unwrap` in server startup and request paths unless the panic is
-  impossible by construction and documented.
-
-## Rust Style Guide
-
-All Rust in `litellm-rust/` follows the official Rust Style Guide:
-https://doc.rust-lang.org/style-guide/
-
-`rustfmt` implements the guide's formatting rules by default, so the mechanical
-side is enforced for you: run `cargo fmt` before committing and CI gates every
-PR on `cargo fmt --check` (see Checks). Do not hand-format against rustfmt or add
-a `rustfmt.toml` that diverges from the default style; the default style *is* the
-guide.
-
-The guide also covers conventions rustfmt cannot auto-apply; follow these too:
-- Naming: `snake_case` for items, functions, and modules; `UpperCamelCase` for
-  types, traits, and enum variants; `SCREAMING_SNAKE_CASE` for constants and
-  statics; acronyms count as one word (`HttpClient`, not `HTTPClient`).
-- Ordering and grouping the guide prescribes: imports grouped std / external /
-  crate-local, derives before other attributes, and consistent item order.
-- Idioms the guide recommends over the formatter fighting you (e.g. prefer
-  restructuring an over-long expression rather than forcing an awkward wrap).
-
-## Constants
-
-Magic numbers and fixed strings go in a crate-level `constants.rs`, never
-hardcoded inline — the Rust mirror of Python's `litellm/constants.py`.
-
-- Each crate that needs them has `src/constants.rs` (declared `mod constants;`);
-  import from it (`use crate::constants::...`). Don't scatter `const` values at
-  the top of feature modules.
-- An env-overridable tunable still lives in `constants.rs` as its `DEFAULT_*`
-  value; the env read (with fallback to that default) happens at the host/config
-  resolution layer, not in `core`/`providers`.
-- Exception: a value that is purely local to one function and has no meaning
-  elsewhere may stay inline, but prefer `constants.rs` when in doubt.
-
-## Checks
-
-Run these before pushing Rust changes. The same checks run in GitHub Actions
-for changes under `litellm-rust/`.
-
-```bash
-cd litellm-rust
-cargo fmt --check
-cargo clippy --workspace --all-targets -- -D warnings
-cargo clippy -p litellm-core --all-targets --features bedrock-auth -- -D warnings
-# the ai-gateway binary + server code is behind the `server` feature
-cargo clippy -p litellm-ai-gateway --all-targets --all-features -- -D warnings
-cargo test --workspace
-cargo test -p litellm-core --features bedrock-auth
-# the `auth`, `routes`, `state` and `realtime` tests only exist under `server`
-cargo test -p litellm-ai-gateway --features server
-```
-
-When a Rust path is exposed through Python, add Python parity tests that compare
-the existing Python output with the Rust-backed output.

litellm-rust/Cargo.lock

@@ -1948,12 +1948,17 @@ dependencies = [
  "azure_core",
  "azure_identity",
  "base64 0.22.1",
+ "bytes",
  "data-url",
+ "futures-util",
  "gcp_auth",
+ "mime_guess",
  "moka",
  "rand 0.8.7",
  "reqwest 0.12.28",
  "rstest",
+ "rustls 0.23.42",
+ "rustls-native-certs",
  "serde",
  "serde_json",
  "serde_path_to_error",
@@ -1962,6 +1967,7 @@ dependencies = [
  "subtle",
  "thiserror 2.0.19",
  "tokio",
+ "tokio-tungstenite",
  "tracing",
  "tracing-subscriber",
  "url",
@@ -1974,12 +1980,12 @@ version = "0.1.0"
 dependencies = [
  "criterion",
  "futures-util",
- "litellm-ai-gateway",
  "litellm-core",
  "litellm-python-interop",
  "litellm-token-counter",
  "pyo3",
  "pyo3-async-runtimes",
+ "rstest",
  "serde",
  "serde_json",
  "tokio",

litellm-rust/Cargo.toml

@@ -16,6 +16,7 @@ license = "MIT"
 repository = "https://github.com/BerriAI/litellm"
 
 [workspace.dependencies]
+bytes = "1"
 tracing = "0.1"
 tracing-subscriber = { version = "0.3", default-features = false, features = ["registry", "std"] }
 litellm-core = { path = "crates/core" }

litellm-rust/README.md

@@ -1,56 +0,0 @@
-# LiteLLM Rust
-
-This workspace contains the staged Rust implementation for LiteLLM.
-
-`litellm-core` is the LiteLLM SDK in Rust: one entrypoint per top-level call
-that makes the LLM call and hands back a typed response, the same shape as
-`litellm.messages()` in Python.
-
-```rust
-let response = litellm_core::messages::messages(MessagesRequest {
-    model: "claude-sonnet-4-5",
-    body,
-    api_key: Some(key),
-    ..
-})
-.await?;
-```
-
-Python continues to own configuration, retries, routing policy, logging,
-callbacks, spend tracking, and customer plugins until each Rust path has parity
-coverage and production evidence.
-
-## Crates
-
-| Crate | Role |
-|-------|------|
-| litellm-core | The SDK. Per-route entrypoints (`messages::messages()`), types, provider transforms (modules under `providers/`), provider resolution, auth, the provider HTTP call, and the router. |
-| litellm-config | Config-loading boundary. Returns resolved deployments and optionally delegates loading to Python. |
-| litellm-ai-gateway | The axum server (behind the `server` feature) and WebSocket hosts. Translates HTTP/WS to core entrypoints; no provider handlers. |
-| litellm-python-interop | Domain-neutral PyO3 foundation for GIL handling and typed Python/Serde conversion. |
-| litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. Owns API registration, domain wiring, and Python exception mapping. |
-
-Dependency direction is acyclic: config depends on core, the gateway depends on config and core, and the Python bridge depends on the domain layers and Python interop.
-
-## Layout
-
-```text
-crates/
-  core/           The SDK: route modules + provider transforms.
-    src/messages/   mod.rs (entrypoint), types, transformation, prepare, handler, client
-    src/providers/anthropic/messages/transformation.rs
-  config/         Config loading and resolved deployments.
-  ai-gateway/     Axum server + WebSocket hosts; calls core entrypoints.
-  python-interop/ Domain-neutral PyO3 conversion and GIL primitives.
-  python-bridge/  PyO3 API adapter for Python LiteLLM.
-```
-
-The folder shape follows the Python provider tree:
-`core/src/providers/<provider>/<route>/transformation.rs`. The bridge exposes one
-function per top-level route, mirroring the core entrypoints.
-
-## Checks
-
-Run the commands under "Checks" in [CLAUDE.md](CLAUDE.md) before pushing Rust
-changes. That list is the single source of truth and matches what GitHub Actions
-runs for changes under `litellm-rust/`.

litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md

@@ -1,53 +0,0 @@
-# Provider coding standards (litellm-rust)
-
-Rules for adding or changing an LLM provider/route in `litellm-rust`. `messages` (`core/src/messages`, `ANTHROPIC_MESSAGES_CONFIG`) is the reference: a route is a `core` module with a public entrypoint that makes the call and returns a typed response.
-
-## Provider resolution
-
-1. Always resolve the provider/model first with `get_custom_llm_provider` (`core/src/routing_utils/provider.rs`). Nothing downstream may branch on a raw model string.
-2. Model/provider is resolved once, in `prepare.rs`, and passed down as typed fields. Don't re-resolve or re-parse it in transforms or handlers.
-
-## Transforms and the base config
-
-3. Every route defines a base config trait with `transform_request` + `transform_response` (+ `complete_url`, `supported_params`), living in `core/src/<route>/transformation.rs` (e.g. `AnthropicMessagesProviderConfig`, mirroring `OcrProviderConfig`).
-4. Each provider implements that trait as a `const <PROVIDER>_<ROUTE>_CONFIG` in `core/src/providers/<provider>/<route>/transformation.rs`, mirroring the Python provider tree.
-5. Individual configs implement only the request/response transforms. Shared behavior (param filtering, defaults) stays as trait default methods so future providers inherit existing logic instead of reimplementing it.
-6. Prefer composition: a provider that extends another reuses the base trait's defaults or wraps another config; don't copy transform bodies between providers.
-
-## Boundaries
-
-7. Layers never cross: `core` = the call itself (entrypoint, types, transforms, provider resolution, auth headers, provider HTTP, lifecycle hooks); `ai-gateway` = serving HTTP/WS (routing, extractors, auth of *our* callers, streaming to the client); `python-bridge` = thin PyO3 adapter. Hosts call the core entrypoint; they never build a provider request.
-8. Generic/route files contain zero provider-specific branches. A provider is one module under `core/src/providers/<provider>/<route>/`; a route is a module, never a new crate.
-9. Route entry point stays thin: `core::<route>::<route>()` -> `prepare_*` -> handler (or `CallLifecycle::run_request`, which owns the pre_call -> during_call -> provider call -> success/failure order and phase timing). Axum handlers validate and delegate to a service that calls the entrypoint; no business logic in them.
-10. Constants (URLs, env-var names, API versions, error messages) live in a crate `constants.rs`, never inline. Config-shaped env reads happen at the host/config layer with the `DEFAULT_*` fallback defined in `constants.rs`; the only env read in `core` is the credential fallback in a route's `prepare.rs`.
-
-## Types and errors
-
-11. Typed contracts only: no bare `serde_json::Value` / `String` / `Vec<String>` as a transform input or output. Parse wire bytes into typed structs/enums at the host edge; a `type` discriminator is a typed field, not a raw string.
-12. Model failures as values: return typed `CoreError`, don't panic. No `unwrap`/`expect`/`panic!` on user or provider input.
-13. No mutation: build values in one shot (comprehensions/iterators, `collect`), prefer immutable bindings and owned typed structs over seeding-and-mutating.
-14. Early returns over deep nesting; small focused files over god modules.
-15. Preserve Python output shape intentionally. If a field is always serialized as `null` for parity, keep it and pin it with a test.
-
-## Safety and data minimization
-
-16. Never log request/response bodies, base64 payloads, document contents, or secrets. Truncate and bound any upstream body before it crosses a host boundary.
-17. Treat empty/whitespace credentials, URLs, and config values as absent at the host resolution layer.
-18. Network I/O sets connect + request timeouts (no unbounded waits), reuses a shared HTTP client, and prefers rustls TLS.
-
-## Tests and rollout
-
-19. Every provider transform ships tests for: supported-param filtering, request body shape, response normalization, missing/null fields, bad input, and `*_match_python` fixture parity.
-20. Lifecycle/hook tests cover hook order, success + failure callback payloads, pre-call guardrail blocking before any provider I/O, during-call body mutation, and provider-error mapping.
-21. When a route has a Python reference implementation, the Rust path stays off by default and behind Python parity tests (disabled / enabled-equals-Python / bridge-unavailable fallback) until parity is proven. A new provider/route may instead be implemented rust-only with no Python reference; then the Python interface is a thin dispatch to Rust with no fallback, and tests cover the rust-backed path plus the unavailable-bridge error. State the rust-only choice explicitly in the PR.
-
-## Python bridge (SDK side)
-
-22. A Python -> Rust bridge keeps the Python side minimal: the Python interface only marshals inputs and calls the Rust interface, with no transform, handler, or business logic. Aim for well under 100 lines of interface code per route; if the Python grows past that, the logic belongs in Rust.
-23. Do not bloat `litellm/main.py`. A route's provider dispatch lives in a thin dispatch class under `litellm/llms/<provider>/<route>/` that calls the Rust bridge; `main.py` only instantiates it and calls its sync/async method.
-24. Do not add new feature flags unless explicitly requested. Reuse the existing LiteLLM Rust rollout mechanism (`litellm.rust`); never introduce a per-route env flag such as `LITELLM_USE_RUST_<ROUTE>`.
-
-## Checks before push
-
-25. Run, and keep green, the commands under "Checks" in `litellm-rust/CLAUDE.md`.
-    That list is the single source of truth and matches what GitHub Actions runs.

litellm-rust/crates/ai-gateway/AGENTS.md

@@ -1,54 +0,0 @@
-# ai-gateway — folder architecture
-
-The Axum server that fronts the Rust gateway. It owns transport + config + auth
-only; deployment selection lives in `core::router`, and the LLM call itself
-(transforms, auth headers, provider HTTP) lives behind a `core` route entrypoint
-such as `litellm_core::messages::messages`. No provider handler lives here.
-
-```
-src/
-  main.rs            # entrypoint: build AppState (router + master key), bind, serve
-  state.rs           # AppState — shared Arc<Router> + master_key
-  auth/              # authentication as an axum extractor — added to handler args
-    mod.rs           #   RequireMasterKey: FromRequestParts, single master key (LITELLM_MASTER_KEY)
-  routes/            # one module per route, all matching the same template
-    AGENTS.md        #   ← the route template (read this before adding a route)
-    mod.rs           #   app(): merges every module's router()
-    health.rs        #   simple route (one file): router() + liveness/readiness
-    realtime/        #   route with logic → axum surface + a no-axum service:
-      mod.rs         #     router() + handler + WS<->events adapter (the axum surface)
-      service.rs     #     business logic (select deployment, call provider) — no axum, testable
-```
-
-## Rules
-
-- **Routes follow one template.** Each route module exposes
-  `pub fn router() -> Router<AppState>`; `routes/mod.rs` only merges them. Simple
-  routes are one file; non-trivial routes are a folder (`handler`/`service`/
-  `transport`). See `routes/AGENTS.md`.
-- **Auth is an extractor.** Add `crate::auth::RequireMasterKey` to a handler's
-  args; it runs during extraction. Never re-implement the check per route.
-- **Handlers are thin.** A handler validates and delegates to its `service`. No
-  business logic, no provider calls, no transforms in handlers.
-- **Services call `core`, they don't reimplement it.** A `service` picks the
-  deployment and calls the `core` route entrypoint. Provider resolution, auth
-  headers, URL building, and the HTTP call are `core`'s job; a service that
-  builds a provider request itself is a bug (`routes/messages/service.rs` is
-  the reference).
-- **State is shared and cheap to clone.** Long-lived handles live behind `Arc` in
-  `state.rs`; read env/config only in `main.rs` when building state.
-
-## Auth (interim)
-
-A single **master key** (`LITELLM_MASTER_KEY`), enforced by the
-`auth::RequireMasterKey` extractor: any caller presenting it as
-`Authorization: Bearer <key>` may invoke the gateway. Fails closed (500) when
-unset; constant-time compare. The server binds `127.0.0.1` by default (`HOST` to
-override). Full per-key auth + budgets/rate-limits are delegated to the Python
-proxy in a later phase. Health routes don't add the extractor (unauthenticated).
-
-## Python interop
-
-Python-backed loading lives in `litellm-config` and is **load-time only**. The
-gateway's `python-config` feature forwards to that crate. The realtime data path
-never takes the GIL.

litellm-rust/crates/ai-gateway/ARCHITECTURE.md

@@ -1,14 +0,0 @@
-# ai-gateway architecture
-
-The Rust ai-gateway does LLM inference (realtime WebSocket). Spend tracking is an
-API callback: it POSTs each finished session to the LiteLLM proxy, which records
-spend and runs the usual callbacks.
-
-```mermaid
-flowchart LR
-  C[client] <--> G[Rust ai-gateway<br/>LLM inference]
-  G <--> O[OpenAI realtime]
-  G -. spend tracking callback .-> P[litellm proxy]
-  F[litellm-config<br/>load-time only] --> G
-  F -. Python backend .-> P
-```

litellm-rust/crates/ai-gateway/Cargo.toml

@@ -13,6 +13,11 @@ name = "litellm-ai-gateway"
 path = "src/main.rs"
 required-features = ["server"]
 
+[[bin]]
+name = "trace-parity-gateway"
+path = "src/bin/trace_parity_gateway.rs"
+required-features = ["trace-parity"]
+
 [dependencies]
 tracing.workspace = true
 litellm-core = { workspace = true, features = ["bedrock-auth"] }

litellm-rust/crates/ai-gateway/benchmarks/realtime/README.md

@@ -1,55 +0,0 @@
-# Realtime gateway benchmark — pool on/off
-
-Measures what the gateway adds over talking to OpenAI's realtime WebSocket
-directly, and what the pre-warmed connection pool removes. See
-`../../src/routes/realtime/README.md` for how the pool works.
-
-## Results
-
-5000 calls / 500 concurrency, gateway at 10 instances, pool ON
-(`REALTIME_POOL_SIZE=64`), upstream OpenAI `gpt-realtime`. Each leg run twice.
-Times in **ms**. Phases per connection: **dial** = TCP+TLS+WS upgrade,
-**session** = upgrade → `session.created` (the phase the pool removes),
-**1st-audio** = `response.create` → first audio delta (OpenAI inference),
-**total** = full wall-clock.
-
-| metric             | Direct OpenAI | Gateway (pool ON) | Overhead (ms) | vs OpenAI  |
-| ------------------ | ------------- | ----------------- | ------------- | ---------- |
-| success rate (%)   | 99.8          | 99.8              | —             | —          |
-| dial p50 (ms)      | 276           | 158               | −118          | **faster** |
-| session p50 (ms)   | 7             | 0                 | −7            | **faster** |
-| 1st-audio p50 (ms) | 440           | 664               | +224          | slower¹    |
-| total p50 (ms)     | 816           | 1010              | +194          | slower¹    |
-| total p95 (ms)     | 2152          | 1970              | −182          | **faster** |
-| total p99 (ms)     | 2692          | 2610              | −82           | **faster** |
-
-The gateway is **faster than direct on 4 of 6 metrics**. The warm pool makes the
-**session phase sub-millisecond** at the median — ~76% of connects hit the pool,
-~70% had session < 1 ms. ¹ The two "slower" rows are not gateway overhead:
-`1st-audio` is OpenAI's own inference time (the gateway only relays it), which ran
-slower during the gateway legs and drags `total p50` with it.
-
-**Pool OFF** (control, `REALTIME_POOL_SIZE=0`): session p50 was **367 ms** — the
-fresh-dial overhead the pool removes.
-
-## Reproduce
-
-The load generator lives in a separate repo:
-**https://github.com/ishaan-berri/litellm-realtime-bench**
-
-```bash
-git clone https://github.com/ishaan-berri/litellm-realtime-bench
-cd litellm-realtime-bench && go build -o wsbench .
-
-# Direct to OpenAI (baseline)
-./wsbench -host api.openai.com -key "$OPENAI_API_KEY" -m gpt-realtime -n 5000 -c 500 -t 60
-
-# Through the gateway — run once with pool ON, once with REALTIME_POOL_SIZE=0
-./wsbench -host <gateway-host> -key "$LITELLM_MASTER_KEY" -m gpt-realtime -n 5000 -c 500 -t 60
-```
-
-Run the gateway with the env stand-in (`OPENAI_REALTIME_MODEL=gpt-realtime`,
-`OPENAI_API_KEY`, `LITELLM_MASTER_KEY`, `REALTIME_POOL_SIZE`, `HOST=0.0.0.0`). At
-500 concurrency over N instances, size the pool to `≈ 500 / N` per instance (64 was
-used here for 10 instances). The bench repo's README covers running 500-concurrency
-legs from a hosted multi-vCPU runner. **Never commit keys — pass them via `-key`.**

litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs

@@ -277,7 +277,7 @@ fn core_error_kind(error: &Error) -> &'static str {
         Error::InvalidProvider(_) => "InvalidProvider",
         Error::InvalidRequest(_) => "InvalidRequest",
         Error::InvalidType { .. } => "InvalidType",
-        Error::MissingField(_) => "MissingField",
+        Error::MissingField(_) | Error::MissingDocumentUrl => "MissingField",
         Error::Http { .. } => "HttpError",
         Error::InvalidResponse(_) => "InvalidResponse",
         Error::Network(_) => "NetworkError",

litellm-rust/crates/ai-gateway/src/bin/trace_parity_gateway.rs

@@ -0,0 +1,40 @@
+use std::io::Read;
+
+use serde::Deserialize;
+use serde_json::Value;
+
+#[derive(Deserialize)]
+struct Input {
+    model_alias: String,
+    provider_model: String,
+    api_base: String,
+    body: Value,
+}
+
+#[tokio::main]
+async fn main() {
+    let mut input = String::new();
+    if let Err(error) = std::io::stdin().read_to_string(&mut input) {
+        fail(error);
+    }
+    let input: Input = match serde_json::from_str(&input) {
+        Ok(input) => input,
+        Err(error) => fail(error),
+    };
+    let result = litellm_ai_gateway::trace_parity::traced_messages_request(
+        input.model_alias,
+        input.provider_model,
+        input.api_base,
+        input.body,
+    )
+    .await;
+    match serde_json::to_string(&result) {
+        Ok(result) => println!("{result}"),
+        Err(error) => fail(error),
+    }
+}
+
+fn fail(error: impl std::fmt::Display) -> ! {
+    eprintln!("{error}");
+    std::process::exit(1)
+}

litellm-rust/crates/ai-gateway/src/io/responses_ws.rs

@@ -1,5 +1,3 @@
-use std::collections::HashMap;
-use std::sync::Arc;
 use std::time::Duration;
 
 use futures_util::stream::{SplitSink, SplitStream};
@@ -10,106 +8,21 @@ use litellm_core::auth::error::MissingCredential;
 use litellm_core::providers::openai::responses::transformation::OPENAI_RESPONSES_WS_CONFIG;
 use litellm_core::responses::types::ResponsesWsEvent;
 use litellm_core::responses::websocket::ResponsesWebSocketProviderConfig;
-use tokio::net::TcpStream;
-use tokio::sync::Mutex;
 use tokio_tungstenite::tungstenite::Message;
 use tokio_tungstenite::tungstenite::client::IntoClientRequest;
 use tokio_tungstenite::tungstenite::http::HeaderValue;
-use tokio_tungstenite::tungstenite::http::header::{AUTHORIZATION, HeaderName};
-use tokio_tungstenite::{MaybeTlsStream, WebSocketStream};
+use tokio_tungstenite::tungstenite::http::header::AUTHORIZATION;
 
-use crate::io::tls::connect_upstream;
+use litellm_core::responses::websocket::{ResponsesUpstreamWs, connect_upstream};
 
 use crate::constants::{
     DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS, DEFAULT_RESPONSES_WS_IDLE_TIMEOUT_SECS,
 };
 
 const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY";
-pub type ResponsesUpstreamWs = WebSocketStream<MaybeTlsStream<TcpStream>>;
 type UpstreamTx = SplitSink<ResponsesUpstreamWs, Message>;
 type UpstreamRx = SplitStream<ResponsesUpstreamWs>;
 
-#[derive(Clone)]
-pub struct ResponsesWebSocketConnection {
-    socket: Arc<Mutex<Option<ResponsesUpstreamWs>>>,
-}
-
-impl ResponsesWebSocketConnection {
-    pub async fn connect_url(
-        url: &str,
-        headers: &HashMap<String, String>,
-        timeout: Option<Duration>,
-    ) -> Result<Self, Error> {
-        let mut request = url
-            .into_client_request()
-            .map_err(|error| Error::Network(error.to_string()))?;
-        for (name, value) in headers {
-            let header_name = name
-                .parse::<HeaderName>()
-                .map_err(|error| Error::InvalidRequest(error.to_string()))?;
-            let header_value = HeaderValue::from_str(value)
-                .map_err(|error| Error::InvalidRequest(error.to_string()))?;
-            request.headers_mut().insert(header_name, header_value);
-        }
-        let connect = connect_upstream(request);
-        let result = match timeout {
-            Some(timeout) => tokio::time::timeout(timeout, connect).await.map_err(|_| {
-                Error::Network("Responses WebSocket connection timed out".to_string())
-            })?,
-            None => connect.await,
-        };
-        let (socket, _) = result.map_err(|error| match *error {
-            tokio_tungstenite::tungstenite::Error::Http(response) => Error::Http {
-                status: response.status().as_u16(),
-                body: String::new(),
-            },
-            other => Error::Network(other.to_string()),
-        })?;
-        Ok(Self {
-            socket: Arc::new(Mutex::new(Some(socket))),
-        })
-    }
-
-    pub async fn send_text(&self, text: String) -> Result<(), Error> {
-        let mut socket = self.socket.lock().await;
-        let Some(socket) = socket.as_mut() else {
-            return Err(Error::Network("Responses WebSocket is closed".to_string()));
-        };
-        socket
-            .send(Message::Text(text))
-            .await
-            .map_err(|error| Error::Network(error.to_string()))
-    }
-
-    pub async fn recv_text(&self) -> Result<Option<String>, Error> {
-        let mut socket_guard = self.socket.lock().await;
-        let Some(socket) = socket_guard.as_mut() else {
-            return Ok(None);
-        };
-        match socket.next().await {
-            Some(Ok(Message::Text(text))) => Ok(Some(text)),
-            Some(Ok(Message::Binary(bytes))) => String::from_utf8(bytes.to_vec())
-                .map(Some)
-                .map_err(|error| Error::InvalidResponse(error.to_string())),
-            Some(Ok(Message::Close(_))) | None => Ok(None),
-            Some(Ok(_)) => Ok(None),
-            Some(Err(error)) => Err(Error::Network(error.to_string())),
-        }
-    }
-
-    pub async fn close(&self) -> Result<(), Error> {
-        let mut socket = self.socket.lock().await;
-        if let Some(socket) = socket.as_mut() {
-            socket
-                .close(None)
-                .await
-                .map_err(|error| Error::Network(error.to_string()))?;
-        }
-        *socket = None;
-        Ok(())
-    }
-}
-
 pub(crate) fn resolve_api_key(api_key: Option<&str>) -> Result<String, Error> {
     api_key
         .map(str::trim)

litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs

@@ -118,7 +118,8 @@ impl IntoResponse for MessagesRouteError {
             | Error::Connect(_)
             | Error::InvalidResponse(_)
             | Error::InvalidType { .. }
-            | Error::MissingField(_) => (
+            | Error::MissingField(_)
+            | Error::MissingDocumentUrl => (
                 StatusCode::BAD_GATEWAY,
                 "messages provider request failed".to_string(),
             ),

litellm-rust/crates/ai-gateway/src/trace_parity.rs

@@ -10,6 +10,7 @@ use litellm_core::router::{Deployment, LiteLLMParams, Router as ModelRouter};
 use serde::Serialize;
 use serde_json::Value;
 use tower::ServiceExt;
+use tracing::instrument::WithSubscriber;
 
 use crate::io::realtime_pool::RealtimePool;
 use crate::routes;
@@ -21,6 +22,38 @@ pub struct GatewayResponse {
     pub body: Value,
 }
 
+#[derive(Debug, Serialize)]
+pub struct TracedGatewayResponse {
+    pub response: Option<GatewayResponse>,
+    pub error: Option<String>,
+    pub trace: Vec<litellm_core::observability::FunctionTraceEvent>,
+}
+
+pub async fn traced_messages_request(
+    model_alias: String,
+    provider_model: String,
+    api_base: String,
+    body: Value,
+) -> TracedGatewayResponse {
+    let trace = litellm_core::observability::FunctionTrace::default();
+    let result = messages_request(model_alias, provider_model, api_base, body)
+        .with_subscriber(trace.dispatcher())
+        .await;
+    let events = trace.events();
+    match result {
+        Ok(response) => TracedGatewayResponse {
+            response: Some(response),
+            error: None,
+            trace: events,
+        },
+        Err(error) => TracedGatewayResponse {
+            response: None,
+            error: Some(error.to_string()),
+            trace: events,
+        },
+    }
+}
+
 pub async fn messages_request(
     model_alias: String,
     provider_model: String,

litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs

@@ -2,10 +2,10 @@
 //! API has to resolve its own crypto provider, in a test binary where nothing
 //! has installed a process-wide one, and has to leave it uninstalled.
 
-use std::collections::HashMap;
 use std::time::Duration;
 
-use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection;
+use futures_util::{sink, stream};
+use litellm_ai_gateway::io::responses_ws::async_responses_websocket;
 use tokio::net::TcpListener;
 
 async fn dead_tls_server() -> u16 {
@@ -30,10 +30,15 @@ async fn dead_tls_server() -> u16 {
 async fn dialing_wss_returns_an_error_instead_of_panicking() {
     let port = dead_tls_server().await;
 
-    let result = ResponsesWebSocketConnection::connect_url(
-        &format!("wss://127.0.0.1:{port}/"),
-        &HashMap::new(),
+    let result = async_responses_websocket(
+        "gpt-5",
+        Some("test-key"),
+        Some(&format!("wss://127.0.0.1:{port}/")),
+        None,
         Some(Duration::from_secs(10)),
+        |_| {},
+        stream::empty(),
+        sink::drain(),
     )
     .await;
 

litellm-rust/crates/core/AGENTS.md

@@ -2,6 +2,6 @@ litellm-core is the LiteLLM SDK in Rust — it makes the LLM call. Each top-leve
 
 A route module owns everything the call needs: types, the provider template trait, provider transforms (under `providers/`), provider/auth/URL resolution, and the handler that performs the HTTP call. Handlers belong here, never in a host crate.
 
-Not here: serving HTTP (axum routes, extractors), config file reading, rollout state, databases, or callback dispatch. Env reads are limited to credential fallback in a route's `prepare.rs`.
+Not here: serving HTTP (axum routes, extractors), config file reading, rollout state, databases, or host-specific callback execution. Core owns lifecycle sequencing and callback payload construction; hosts execute the selected integrations. Env reads are limited to credential fallback in a route's `prepare.rs`.
 
 Routes (messages, ocr, realtime) and providers (anthropic, mistral, openai) are modules, not crates.

litellm-rust/crates/core/CLAUDE.md

@@ -1,66 +0,0 @@
-# CLAUDE.md
-
-Rules for `litellm-rust/crates/core`.
-
-## Responsibility
-
-`core` is the LiteLLM SDK in Rust: it makes the LLM call. Every top-level
-LiteLLM call has a public entrypoint here, named after the route
-(`messages::messages()` is the Rust equivalent of `litellm.messages()`), and
-calling it returns a typed non-streaming response.
-
-Allowed:
-- The public entrypoint for a route, plus its `<route>_stream` variant when the
-  route supports streaming.
-- Provider resolution, auth header construction, URL building, and the provider
-  HTTP call (shared reused client, connect + request timeouts).
-- Shared request/response structs.
-- Typed errors with stable, non-sensitive messages.
-- Deterministic validation helpers.
-- Serialization helpers that intentionally mirror Python output shape.
-- Route templates that match Python base config responsibilities, such as
-  `messages::transformation::AnthropicMessagesProviderConfig`.
-
-Not allowed:
-- Serving HTTP: axum routers, extractors, and other transport concerns.
-- Filesystem, database, or cache access.
-- Config file reading or rollout state; the host resolves those and passes them
-  in. Env reads are limited to credential fallback in a route's `prepare.rs`.
-- Logging callbacks, tracing spans, spend writes, or customer callbacks.
-- Provider-specific branching that belongs in `providers`.
-- Panics for user/provider-controlled input.
-
-## Typed Contracts (core rule)
-
-Trait and function boundaries MUST be strongly typed. No stringly-typed JSON
-(`&str` / `String` / `Vec<String>` / bare `serde_json::Value`) as a transform
-input or output. Parse wire bytes into typed structs/enums at the host edge;
-`core` and `providers` operate only on those types (e.g. `RealtimeEvent`,
-`RealtimeTransformResult`, `OcrRequestData`). A `type`-style discriminator is a
-typed field on a struct, not a raw string threaded through the API.
-
-## Structure
-
-Use route names directly under `src/`: `messages`, `ocr`, future
-`chat_completions`, `embeddings`, and similar top-level LiteLLM calls. Do not
-invent broad names like `engine` for route contracts.
-
-`src/messages` is the reference shape for a route module:
-
-```
-mod.rs             pub async fn messages(..) (+ messages_stream)
-types.rs           request/response types
-transformation.rs  the provider template trait
-prepare.rs         provider resolution, auth headers, URL
-handler.rs         the provider call
-client.rs          the shared reqwest client
-```
-
-## Parity Rules
-
-- Every shared type used by a provider transform needs unit tests for
-  serialization shape.
-- If Python parity requires always emitting a `null` field instead of omitting
-  it, document that in code and pin it with a test.
-- Error enums should preserve enough detail for Python/HTTP hosts to map errors
-  consistently without exposing document contents or upstream bodies.

litellm-rust/crates/core/Cargo.toml

@@ -6,25 +6,27 @@ license.workspace = true
 repository.workspace = true
 autotests = false
 
-[[test]]
-name = "workspace_crate_allowlist"
-path = "tests/workspace_crate_allowlist.rs"
-
 [dependencies]
+bytes.workspace = true
+futures-util.workspace = true
 base64.workspace = true
 azure_core.workspace = true
 azure_identity.workspace = true
 data-url = "0.3.2"
 gcp_auth.workspace = true
 moka.workspace = true
+mime_guess = "2.0.5"
 rand.workspace = true
 reqwest.workspace = true
+rustls.workspace = true
+rustls-native-certs.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 serde_path_to_error = "0.1"
 strum.workspace = true
 subtle.workspace = true
-tokio.workspace = true
+tokio = { workspace = true, features = ["sync"] }
+tokio-tungstenite.workspace = true
 thiserror.workspace = true
 tracing.workspace = true
 tracing-subscriber = { workspace = true, optional = true }

litellm-rust/crates/core/src/auth/credential.rs

@@ -9,6 +9,21 @@ use crate::AuthError;
 
 use super::{ResolvedCredential, SecretValue, TokenProviderHandle};
 
+pub fn credential_index(requested: &str, names: &[String]) -> Option<usize> {
+    names.iter().position(|name| name == requested)
+}
+
+pub fn credential_default_fields<'a>(
+    supplied: &[String],
+    credential_fields: &'a [String],
+) -> Vec<&'a str> {
+    credential_fields
+        .iter()
+        .filter(|name| !supplied.contains(name))
+        .map(String::as_str)
+        .collect()
+}
+
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub enum CredentialFileRef {
     Path(PathBuf),

litellm-rust/crates/core/src/auth/mod.rs

@@ -49,6 +49,7 @@ impl<T> Sourced<T> {
 pub use credential::{
     CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialPlan,
     CredentialPlanResolution, CredentialRef, CredentialResolver, CredentialResolverHandle,
+    credential_default_fields, credential_index,
 };
 pub use http::{CredentialPlacement, RequestAuth};
 pub use policy::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy};

litellm-rust/crates/core/src/call_lifecycle/README.md

@@ -1,167 +0,0 @@
-# Call lifecycle
-
-`litellm_core::call_lifecycle` is the shared execution wrapper for LiteLLM call
-types migrated to Rust. It owns lifecycle ordering, phase timing, and trace
-observer calls. It must not know about OCR, chat, messages, responses,
-completions, provider auth, request transforms, or response normalization.
-
-Call-type modules own their domain behavior. For example, OCR owns document
-payloads, OCR provider transforms, safe document fetch, guardrail payload shape,
-callback payload shape, and provider HTTP execution.
-
-## Runtime order
-
-Every wrapped call runs in this order:
-
-1. `async_pre_call_hook`
-2. `async_during_call_hook`
-3. provider call
-4. `async_log_success_event` or `async_log_failure_event`
-
-`async_pre_call_hook` receives the initial LiteLLM request shape. It is where
-pre-call custom guardrails run.
-
-`async_during_call_hook` converts the initial request into the provider-ready
-request. It is where provider config selection, parameter mapping, auth/header
-resolution, request transforms, and during-call guardrails belong.
-
-The provider call receives only the provider-ready request. It should execute
-I/O and call the provider response transform.
-
-Success and failure callbacks receive `CallLifecycleTiming`. Callback failures
-must not replace the original provider or guardrail result.
-
-## Trace contract
-
-The lifecycle runner records:
-
-- full call start and end time
-- `pre_call` phase timing
-- `during_call` phase timing
-- `provider_call` phase timing
-- `success_callback` phase timing
-- `failure_callback` phase timing
-
-`CallLifecycleObserver` receives phase start and end events. The default
-observer is a no-op. Future OTEL support should implement this observer instead
-of editing OCR, chat, messages, responses, completions, or provider modules.
-
-## Required shape
-
-Each migrated call type should use this folder shape:
-
-```text
-litellm-rust/crates/ai-gateway/src/<call_type>/
-  mod.rs          # thin public entrypoint
-  types.rs        # public request, prepared request, provider request, response types
-  prepare.rs      # model/provider/callback/guardrail setup
-  hooks.rs        # CallLifecycleHooks implementation
-  handler.rs      # provider I/O and response normalization
-  tests.rs        # call-type lifecycle and handler tests
-```
-
-Provider transforms can live in `litellm-rust/crates/core/src/providers/...`.
-Shared call-type helpers can live beside the call type, but generic lifecycle
-code stays in this folder.
-
-## Core API
-
-The prepared request implements `CallLifecycleRequest`:
-
-```rust
-impl CallLifecycleRequest for PreparedMessagesRequest {
-    fn lifecycle_context(&self) -> CallLifecycleContext {
-        CallLifecycleContext::new(
-            "messages",
-            self.model.clone(),
-            self.custom_llm_provider.clone(),
-            self.litellm_call_id.clone(),
-        )
-    }
-}
-```
-
-The call-type hooks implement `CallLifecycleHooks`:
-
-```rust
-impl CallLifecycleHooks<
-    PreparedMessagesRequest,
-    ProviderMessagesRequest,
-    MessagesResponse,
-> for MessagesLifecycleHooks {
-    fn async_pre_call_hook(...) {
-        // run pre-call custom guardrails against the LiteLLM request shape
-    }
-
-    fn async_during_call_hook(...) {
-        // map params, validate env, transform request, run during-call guardrails
-    }
-
-    fn async_log_success_event(...) {
-        // call async_log_success_event on configured custom loggers
-    }
-
-    fn async_log_failure_event(...) {
-        // call async_log_failure_event without swallowing the original error
-    }
-}
-```
-
-The public entrypoint stays thin:
-
-```rust
-pub async fn messages(request: MessagesRequest<'_>) -> CoreResult<MessagesResponse> {
-    let PreparedMessagesCall { request, hooks } = prepare_messages_call(request)?;
-
-    CallLifecycle::default()
-        .run_request(request, &hooks, execute_messages_provider_call)
-        .await
-}
-```
-
-Use `run_request` for new call types. Keep `run` available only for specialized
-tests or existing code that already has a `CallLifecycleContext`.
-
-## Adding a new call type
-
-1. Add `<call_type>/types.rs`
-
-Define the public request accepted by the bridge, the prepared request used by
-the lifecycle runner, and the provider request consumed by the handler.
-
-2. Implement `CallLifecycleRequest`
-
-Return `call_type`, `model`, `custom_llm_provider`, and `litellm_call_id`.
-Do not put provider-specific logic here.
-
-3. Add `<call_type>/prepare.rs`
-
-Resolve model/provider once, generate or preserve `litellm_call_id`, construct
-callback and guardrail runners, and return `Prepared<CallType>Call`.
-
-4. Add `<call_type>/hooks.rs`
-
-Implement `CallLifecycleHooks`. Put pre-call guardrail payload construction,
-provider config selection, param mapping, request transform, during-call
-guardrail payload construction, and callback payload construction here.
-
-5. Add `<call_type>/handler.rs`
-
-Execute the provider request and normalize the provider response. Do not repeat
-provider-specific transforms here; call the provider config.
-
-6. Add tests
-
-Cover hook order, success callback payload, failure callback payload, pre-call
-guardrail blocking before provider I/O, during-call body mutation, and provider
-error mapping.
-
-## Review checklist
-
-- Core lifecycle has no call-type or provider-specific branches
-- Public call-type entrypoint only prepares and calls `run_request`
-- Provider behavior lives behind provider config/transformation code
-- Hook method names map to the Python custom logger and guardrail concepts
-- Phase timing is recorded once in lifecycle, not separately per call type
-- Callback failures never hide the original provider or guardrail error
-- Tests prove the provider socket is not touched when pre-call guardrails block

litellm-rust/crates/core/src/call_lifecycle/host.rs

@@ -0,0 +1,121 @@
+use std::future::Future;
+use std::pin::Pin;
+
+pub enum HostCallStep<O, C> {
+    Host(O),
+    Complete(C),
+}
+
+pub type HostCallFuture<'a, O, C> =
+    Pin<Box<dyn Future<Output = Result<HostCallStep<O, C>, crate::Error>> + Send + 'a>>;
+
+pub trait HostCall: Send + Sync {
+    type Operation: Send + 'static;
+    type Result: Send + 'static;
+    type Complete: Send + 'static;
+
+    fn resume(
+        &mut self,
+        result: Option<Self::Result>,
+    ) -> HostCallFuture<'_, Self::Operation, Self::Complete>;
+
+    fn interrupt(
+        &mut self,
+        failure: HostFailure,
+    ) -> HostCallFuture<'_, Self::Operation, Self::Complete>;
+}
+
+pub enum HostStep<V, S> {
+    Ready(V),
+    Suspend(S),
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum HostPhase {
+    Setup,
+    DeploymentPreCall,
+    Prepare,
+    Execute,
+    ConstructResponse,
+    DeploymentPostCall,
+    Finalize,
+    Success,
+    MapFailure,
+    DeploymentFailure,
+    Failure,
+    AsyncFailure,
+    Complete,
+}
+
+#[derive(Clone, Debug)]
+pub enum HostFailure {
+    Error(crate::Error),
+    Cancelled(crate::Error),
+}
+
+pub struct HostLifecycle {
+    phase: HostPhase,
+    asynchronous: bool,
+}
+
+impl HostLifecycle {
+    pub fn new(asynchronous: bool) -> Self {
+        Self {
+            phase: HostPhase::Setup,
+            asynchronous,
+        }
+    }
+
+    pub fn phase(&self) -> HostPhase {
+        self.phase
+    }
+
+    pub fn accept(&mut self, result: Result<(), HostFailure>) -> Option<crate::Error> {
+        if let Err(failure) = result {
+            if self.phase == HostPhase::DeploymentFailure {
+                self.phase = HostPhase::Failure;
+                return None;
+            }
+            let error = match failure {
+                HostFailure::Cancelled(error) => {
+                    self.phase = HostPhase::Complete;
+                    return Some(error);
+                }
+                HostFailure::Error(error) => error,
+            };
+            match self.phase {
+                HostPhase::Failure | HostPhase::AsyncFailure => {
+                    self.advance();
+                    return None;
+                }
+                HostPhase::Success => self.phase = HostPhase::Complete,
+                HostPhase::Execute | HostPhase::ConstructResponse => {
+                    self.phase = HostPhase::MapFailure;
+                }
+                _ => self.phase = HostPhase::Failure,
+            }
+            return Some(error);
+        }
+        self.advance();
+        None
+    }
+
+    fn advance(&mut self) {
+        self.phase = match self.phase {
+            HostPhase::Setup if self.asynchronous => HostPhase::DeploymentPreCall,
+            HostPhase::Setup | HostPhase::DeploymentPreCall => HostPhase::Prepare,
+            HostPhase::Prepare => HostPhase::Execute,
+            HostPhase::Execute => HostPhase::ConstructResponse,
+            HostPhase::ConstructResponse if self.asynchronous => HostPhase::DeploymentPostCall,
+            HostPhase::ConstructResponse | HostPhase::DeploymentPostCall => HostPhase::Finalize,
+            HostPhase::Finalize => HostPhase::Success,
+            HostPhase::MapFailure if self.asynchronous => HostPhase::DeploymentFailure,
+            HostPhase::MapFailure | HostPhase::DeploymentFailure => HostPhase::Failure,
+            HostPhase::Failure if self.asynchronous => HostPhase::AsyncFailure,
+            HostPhase::Failure
+            | HostPhase::AsyncFailure
+            | HostPhase::Success
+            | HostPhase::Complete => HostPhase::Complete,
+        };
+    }
+}

litellm-rust/crates/core/src/call_lifecycle/mod.rs

@@ -3,6 +3,10 @@ use std::time::{Instant, SystemTime, UNIX_EPOCH};
 
 use crate::Error;
 
+pub mod host;
+#[cfg(test)]
+#[path = "../../tests/host_lifecycle.rs"]
+mod host_tests;
 pub mod types;
 
 pub use types::{

litellm-rust/crates/core/src/constants.rs

@@ -46,9 +46,10 @@ pub const FUNCTION_TRACE_TARGET: &str = "litellm::function_trace";
 
 pub(crate) const MEDIA_CONNECT_TIMEOUT_SECS: u64 = 10;
 
+pub(crate) const OCR_RESPONSE_MAX_BYTES: usize = 64 * 1024 * 1024;
 pub(crate) const OCR_HTTP_TIMEOUT_SECS: u64 = 600;
 pub(crate) const OCR_CONNECT_TIMEOUT_SECS: u64 = 10;
-pub(crate) const OCR_INLINE_MAX_BYTES: usize = 50 * 1024 * 1024;
+pub const OCR_INLINE_MAX_BYTES: usize = 50 * 1024 * 1024;
 pub(crate) const OCR_DOWNLOAD_MAX_BYTES: u64 = 50 * 1024 * 1024;
 pub(crate) const OCR_MAX_FETCH_REDIRECTS: usize = 10;
 pub(crate) const OCR_POLL_TIMEOUT_SECS: u64 = 120;
@@ -63,3 +64,6 @@ pub(crate) const REDUCTO_API_KEY_ENV: &str = "REDUCTO_API_KEY";
 pub(crate) const REDUCTO_ID_PREFIX: &str = "reducto://";
 pub(crate) const AZURE_AI_OCR_PATH: &str = "/providers/mistral/azure/ocr";
 pub(crate) const MISTRAL_OCR_API_BASE: &str = "https://api.mistral.ai/v1";
+
+pub(crate) const COHERE_PARSE_API_BASE: &str = "https://api.cohere.com";
+pub(crate) const COHERE_API_KEY_ENV: &str = "COHERE_API_KEY";

litellm-rust/crates/core/src/error.rs

@@ -1,6 +1,6 @@
 use thiserror::Error as ThisError;
 
-#[derive(Debug, ThisError, PartialEq, Eq)]
+#[derive(Clone, Debug, ThisError, PartialEq, Eq)]
 pub enum Error {
     #[error("expected {expected}, got {actual}")]
     InvalidType {
@@ -9,6 +9,8 @@ pub enum Error {
     },
     #[error("missing required field: {0}")]
     MissingField(&'static str),
+    #[error("Document URL is required")]
+    MissingDocumentUrl,
     #[error("invalid response: {0}")]
     InvalidResponse(String),
     #[error("invalid provider: {0}")]
@@ -52,6 +54,17 @@ pub enum Error {
     Unsupported(&'static str),
 }
 
+impl Error {
+    pub const fn http_status_code(&self) -> Option<u16> {
+        match self {
+            Self::InvalidRequest(_) => Some(400),
+            Self::MissingDocumentUrl => Some(500),
+            Self::Http { status, .. } => Some(*status),
+            _ => None,
+        }
+    }
+}
+
 #[derive(Debug, ThisError)]
 pub(crate) enum MediaError {
     #[error("media URL rejected by network policy")]
@@ -106,6 +119,7 @@ impl From<crate::ocr::error::OcrRequestError> for Error {
     fn from(error: crate::ocr::error::OcrRequestError) -> Self {
         match error {
             crate::ocr::error::OcrRequestError::MissingField(field) => Self::MissingField(field),
+            crate::ocr::error::OcrRequestError::MissingDocumentUrl => Self::MissingDocumentUrl,
             error => Self::InvalidRequest(error.to_string()),
         }
     }

litellm-rust/crates/core/src/ocr/adapters/azure/cohere.rs

@@ -0,0 +1,131 @@
+use super::super::OcrAdapter;
+use crate::Error;
+use crate::ocr::OcrClient;
+use crate::ocr::codecs::cohere::{
+    CohereParams, CohereResponse, transform_request, transform_response, validate_document,
+};
+use crate::ocr::document::{inline_remote_document, validate_inline_document};
+use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError};
+use crate::ocr::prepare::{credential_env, transform_request_body};
+use crate::ocr::registry::OcrProvider;
+use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse};
+use crate::providers::azure_ai::auth::AzureAuthInputs;
+use crate::url_utils::ApiUrl;
+
+const AZURE_AI_API_BASE_ENV: &str = "AZURE_AI_API_BASE";
+
+pub(crate) struct AzureCohereAdapter;
+
+impl OcrAdapter for AzureCohereAdapter {
+    type ProviderResponse = CohereResponse;
+    const PROVIDER: OcrProvider = OcrProvider::AzureAi;
+
+    async fn prepare_request(
+        &self,
+        request: &LiteLLMOcrRequest,
+        client: &OcrClient,
+    ) -> Result<reqwest::Request, OcrError> {
+        let params = super::super::super::wire::decode_request_value::<CohereParams>(
+            serde_json::Value::Object(request.optional_params.clone()),
+            "optional_params",
+        )?;
+        let mut config = AzureAuthInputs::from_sourced_optional_params(
+            &request.optional_params,
+            &request.input_sources,
+        )
+        .map_err(Error::from)?;
+        config.azure_ad_token_provider = request.azure_ad_token_provider.clone();
+        let base = request
+            .connection
+            .api_base
+            .clone()
+            .or_else(|| credential_env(AZURE_AI_API_BASE_ENV))
+            .filter(|base| !base.trim().is_empty())
+            .ok_or_else(|| {
+                Error::Auth(
+                    "Missing Azure AI API Base - Set AZURE_AI_API_BASE or pass api_base".into(),
+                )
+            })?;
+        let headers =
+            super::validate_ai_environment(&request.connection, &config, &credential_env).await?;
+        validate_document(&request.document)?;
+        let remote = request.document.source().starts_with("http://")
+            || request.document.source().starts_with("https://");
+        let document = inline_remote_document(
+            client.document_fetcher(),
+            request.document.clone(),
+            &request.connection,
+        )
+        .await?;
+        let body = transform_request(&request.model, document, params)?;
+        transform_request_body(
+            client,
+            request,
+            &complete_url(&base)?,
+            &headers,
+            !remote,
+            body,
+            |body| {
+                validate_document(&body.document)?;
+                validate_inline_document(&body.document)
+            },
+        )
+        .await
+    }
+
+    fn transform_ocr_response(
+        &self,
+        request: &LiteLLMOcrRequest,
+        response: Self::ProviderResponse,
+    ) -> Result<LiteLLMOcrResponse, OcrResponseError> {
+        transform_response(&request.model, response)
+    }
+}
+
+fn complete_url(base: &str) -> Result<String, OcrError> {
+    let mut url = reqwest::Url::parse(base).map_err(|_| invalid_api_base())?;
+    if !matches!(url.scheme(), "http" | "https") {
+        return Err(invalid_api_base().into());
+    }
+    let path = url.path().trim_end_matches('/').to_string();
+    if path.ends_with("/v2/parse") {
+        url.set_path(&path);
+        return Ok(url.into());
+    }
+    url.set_path(path.strip_suffix("/models").unwrap_or(&path));
+    ApiUrl::parse(url.as_str())
+        .and_then(|url| url.complete_path(&["providers", "cohere", "v2", "parse"]))
+        .map(|url| url.into_string())
+        .map_err(|_| invalid_api_base().into())
+}
+
+fn invalid_api_base() -> OcrRequestError {
+    OcrRequestError::RequestField {
+        path: "api_base".into(),
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn completes_foundry_urls_without_duplicate_paths_and_preserves_queries() {
+        for suffix in [
+            "",
+            "/models",
+            "/providers/cohere/v2",
+            "/providers/cohere/v2/parse",
+        ] {
+            assert_eq!(
+                complete_url(&format!("https://example.com{suffix}?tenant=a")).unwrap(),
+                "https://example.com/providers/cohere/v2/parse?tenant=a"
+            );
+        }
+        assert_eq!(
+            complete_url("https://example.com/v2/parse?tenant=a").unwrap(),
+            "https://example.com/v2/parse?tenant=a"
+        );
+        assert!(complete_url("relative/path").is_err());
+    }
+}

litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/mod.rs

@@ -10,7 +10,6 @@ use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError};
 use crate::ocr::prepare::{credential_env, transform_request_body};
 use crate::ocr::registry::OcrProvider;
 use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection, OcrResponseFormat};
-use crate::ocr::wire::DecodedOcrResponse;
 use crate::providers::azure_ai::auth::AzureAuthInputs;
 use crate::url_utils::ApiUrl;
 
@@ -32,18 +31,19 @@ impl OcrAdapter for AzureDocumentIntelligenceAdapter {
         client: &OcrClient,
     ) -> Result<reqwest::Request, OcrError> {
         let params = map_ocr_params(request)?;
-        let config = AzureAuthInputs::from_sourced_optional_params(
+        let mut config = AzureAuthInputs::from_sourced_optional_params(
             &request.optional_params,
             &request.input_sources,
         )
         .map_err(Error::from)?;
+        config.azure_ad_token_provider = request.azure_ad_token_provider.clone();
         let headers = validate_environment(&request.connection, &config, &credential_env).await?;
         let endpoint = nonblank(request.connection.api_base.clone())
             .or_else(|| nonblank(credential_env(AZURE_DI_ENDPOINT_ENV)))
             .ok_or_else(|| Error::Auth("Missing Azure Document Intelligence API Base - Set AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT or pass api_base".into()))?;
         let url = get_complete_url(&endpoint, &request.model, &params)?;
         let body = document_intelligence::transform_ocr_request(request.document.clone())?;
-        transform_request_body(client, request, &url, &headers, body, |_| Ok(())).await
+        transform_request_body(client, request, &url, &headers, false, body, |_| Ok(())).await
     }
 
     fn transform_ocr_response(
@@ -61,7 +61,7 @@ impl OcrAdapter for AzureDocumentIntelligenceAdapter {
         url: &str,
         headers: &[(String, String)],
         request: &LiteLLMOcrRequest,
-    ) -> Result<DecodedOcrResponse<Self::ProviderResponse>, OcrError> {
+    ) -> Result<crate::ocr::wire::DecodedOcrResponse<Self::ProviderResponse>, OcrError> {
         polling::read_operation_response(
             client.polling_http(),
             response,
@@ -69,6 +69,7 @@ impl OcrAdapter for AzureDocumentIntelligenceAdapter {
             headers,
             &request.connection,
             request.response_format()? == OcrResponseFormat::Native,
+            &request.hooks,
         )
         .await
     }

litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/polling.rs

@@ -1,3 +1,4 @@
+use std::sync::Arc;
 use std::time::Duration;
 
 use reqwest::Url;
@@ -9,6 +10,7 @@ use crate::ocr::codecs::document_intelligence::{
     AzureDocumentIntelligenceOperation, OperationStatus,
 };
 use crate::ocr::error::{OcrError, OcrPollingError, OcrResponseError};
+use crate::ocr::hooks::OcrHooks;
 use crate::ocr::types::OcrConnection;
 use crate::ocr::wire::DecodedOcrResponse;
 
@@ -19,24 +21,33 @@ pub(super) async fn read_operation_response(
     headers: &[(String, String)],
     connection: &OcrConnection,
     native: bool,
+    hooks: &Arc<dyn OcrHooks>,
 ) -> Result<DecodedOcrResponse<AzureDocumentIntelligenceOperation>, OcrError> {
     if response.status() != reqwest::StatusCode::ACCEPTED {
-        return read_json_response(response, native).await;
+        let bytes =
+            crate::ocr::client::read_response_bytes(response, connection.max_response_bytes)
+                .await?;
+        crate::ocr::handler::post_call(hooks, &bytes).await?;
+        return Ok(crate::ocr::wire::decode_response(&bytes, native)?);
     }
     let location = response
         .headers()
         .get("operation-location")
         .and_then(|value| value.to_str().ok())
-        .ok_or(OcrPollingError::PollLocation)?;
+        .ok_or(OcrPollingError::PollLocation)?
+        .to_string();
     let original = Url::parse(original_url).map_err(|_| OcrPollingError::PollOrigin)?;
-    let operation = Url::parse(location).map_err(|_| OcrPollingError::PollOrigin)?;
+    let operation = Url::parse(&location).map_err(|_| OcrPollingError::PollOrigin)?;
     if original.origin() != operation.origin()
         || !operation.username().is_empty()
         || operation.password().is_some()
     {
         return Err(OcrPollingError::PollOrigin.into());
     }
-    poll_operation(http_client, operation, headers, connection, native).await
+    let bytes =
+        crate::ocr::client::read_response_bytes(response, connection.max_response_bytes).await?;
+    crate::ocr::handler::post_call(hooks, &bytes).await?;
+    poll_operation(http_client, operation, headers, connection, native, hooks).await
 }
 
 async fn poll_operation(
@@ -45,6 +56,7 @@ async fn poll_operation(
     headers: &[(String, String)],
     connection: &OcrConnection,
     native: bool,
+    hooks: &Arc<dyn OcrHooks>,
 ) -> Result<DecodedOcrResponse<AzureDocumentIntelligenceOperation>, OcrError> {
     let deadline = Instant::now()
         .checked_add(connection.poll_timeout)
@@ -75,12 +87,19 @@ async fn poll_operation(
             .max(1);
         let decoded = tokio::time::timeout_at(
             deadline,
-            read_json_response::<AzureDocumentIntelligenceOperation>(response, native),
+            read_json_response::<AzureDocumentIntelligenceOperation>(
+                response,
+                native,
+                connection.max_response_bytes,
+            ),
         )
         .await
         .map_err(|_| OcrPollingError::PollTimeout)??;
         match &decoded.data.status {
-            Some(OperationStatus::Succeeded) => return Ok(decoded),
+            Some(OperationStatus::Succeeded) => {
+                crate::ocr::handler::post_call(hooks, decoded.text.as_bytes()).await?;
+                return Ok(decoded);
+            }
             Some(OperationStatus::Running | OperationStatus::NotStarted) => {
                 tokio::time::timeout_at(deadline, tokio::time::sleep(Duration::from_secs(retry)))
                     .await

litellm-rust/crates/core/src/ocr/adapters/azure/mistral.rs

@@ -33,13 +33,16 @@ impl OcrAdapter for AzureMistralAdapter {
             known: params,
             extra_params: _extra_params,
         } = _prepare_ocr_request::<MistralOcrParams>(request)?;
-        let config = AzureAuthInputs::from_sourced_optional_params(
+        let mut config = AzureAuthInputs::from_sourced_optional_params(
             &request.optional_params,
             &request.input_sources,
         )
         .map_err(Error::from)?;
-        let headers = validate_environment(&request.connection, &config, &credential_env).await?;
+        config.azure_ad_token_provider = request.azure_ad_token_provider.clone();
         let url = get_complete_url(request.connection.api_base.as_deref(), &credential_env)?;
+        let headers = validate_environment(&request.connection, &config, &credential_env).await?;
+        let retains_document = !request.document.source().starts_with("http://")
+            && !request.document.source().starts_with("https://");
         let document = inline_remote_document(
             client.document_fetcher(),
             request.document.clone(),
@@ -47,9 +50,15 @@ impl OcrAdapter for AzureMistralAdapter {
         )
         .await?;
         let body = mistral::transform_ocr_request(&request.model, document, &params)?;
-        transform_request_body(client, request, &url, &headers, body, |body| {
-            validate_inline_document(&body.document)
-        })
+        transform_request_body(
+            client,
+            request,
+            &url,
+            &headers,
+            retains_document,
+            body,
+            |body| validate_inline_document(&body.document),
+        )
         .await
     }
 
@@ -83,12 +92,15 @@ fn get_complete_url(
         })
 }
 
-async fn validate_environment(
+pub(in crate::ocr::adapters) async fn validate_environment(
     connection: &OcrConnection,
     config: &AzureAuthInputs,
     env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
 ) -> Result<Vec<(String, String)>, OcrError> {
     if crate::http_utils::has_header(&connection.extra_headers, "authorization") {
+        if config.azure_ad_token_provider.is_some() {
+            super::resolve_entra(config, env_lookup).await?;
+        }
         super::validate_destination(connection, connection.extra_headers_source)?;
         return Ok(connection.extra_headers.clone());
     }

litellm-rust/crates/core/src/ocr/adapters/azure/mod.rs

@@ -1,3 +1,4 @@
+mod cohere;
 mod document_intelligence;
 mod mistral;
 
@@ -10,8 +11,10 @@ use crate::ocr::error::OcrError;
 use crate::ocr::types::OcrConnection;
 use crate::providers::azure_ai::auth::{AzureAuthInputs, AzureAuthService};
 
+pub(crate) use cohere::AzureCohereAdapter;
 pub(crate) use document_intelligence::AzureDocumentIntelligenceAdapter;
 pub(crate) use mistral::AzureMistralAdapter;
+pub(super) use mistral::validate_environment as validate_ai_environment;
 
 async fn resolve_entra(
     config: &AzureAuthInputs,
@@ -22,6 +25,10 @@ async fn resolve_entra(
         .get_or_init(AzureAuthService::default)
         .get_azure_ad_token(config, env_lookup)
         .await
+        .or_else(|error| match error {
+            crate::AuthError::EmptyAzureToken => Ok(None),
+            other => Err(other),
+        })
         .map(|credential| {
             credential.map(|credential| {
                 let source = credential.source();

litellm-rust/crates/core/src/ocr/adapters/cohere.rs

@@ -0,0 +1,123 @@
+use super::OcrAdapter;
+use crate::Error;
+use crate::constants::{COHERE_API_KEY_ENV, COHERE_PARSE_API_BASE};
+use crate::ocr::OcrClient;
+use crate::ocr::codecs::cohere::{
+    CohereParams, CohereResponse, transform_request, transform_response, validate_document,
+};
+use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError};
+use crate::ocr::prepare::{credential_env, transform_request_body};
+use crate::ocr::registry::OcrProvider;
+use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection};
+use crate::url_utils::ApiUrl;
+
+pub(crate) struct CohereAdapter;
+
+impl OcrAdapter for CohereAdapter {
+    type ProviderResponse = CohereResponse;
+    const PROVIDER: OcrProvider = OcrProvider::Cohere;
+
+    async fn prepare_request(
+        &self,
+        request: &LiteLLMOcrRequest,
+        client: &OcrClient,
+    ) -> Result<reqwest::Request, OcrError> {
+        let params = super::super::wire::decode_request_value::<CohereParams>(
+            serde_json::Value::Object(request.optional_params.clone()),
+            "optional_params",
+        )?;
+        let headers = validate_environment(&request.connection, &credential_env)?;
+        let url = complete_url(
+            request
+                .connection
+                .api_base
+                .as_deref()
+                .unwrap_or(COHERE_PARSE_API_BASE),
+        )?;
+        let body = transform_request(&request.model, request.document.clone(), params)?;
+        transform_request_body(client, request, &url, &headers, true, body, |body| {
+            validate_document(&body.document)
+        })
+        .await
+    }
+
+    fn transform_ocr_response(
+        &self,
+        request: &LiteLLMOcrRequest,
+        response: Self::ProviderResponse,
+    ) -> Result<LiteLLMOcrResponse, OcrResponseError> {
+        transform_response(&request.model, response)
+    }
+}
+
+fn complete_url(base: &str) -> Result<String, OcrError> {
+    let parsed = reqwest::Url::parse(base).map_err(|_| invalid_api_base())?;
+    if !matches!(parsed.scheme(), "http" | "https") {
+        return Err(invalid_api_base().into());
+    }
+    ApiUrl::parse(base)
+        .and_then(|url| url.complete_path(&["v2", "parse"]))
+        .map(|url| url.into_string())
+        .map_err(|_| invalid_api_base().into())
+}
+
+fn invalid_api_base() -> OcrRequestError {
+    OcrRequestError::RequestField {
+        path: "api_base".into(),
+    }
+}
+
+fn validate_environment(
+    connection: &OcrConnection,
+    env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
+) -> Result<Vec<(String, String)>, OcrError> {
+    if crate::http_utils::has_header(&connection.extra_headers, "authorization") {
+        return Ok(connection.extra_headers.clone());
+    }
+    let key = connection
+        .api_key
+        .as_deref()
+        .map(str::trim)
+        .filter(|key| !key.is_empty())
+        .map(str::to_string)
+        .or_else(|| env_lookup(COHERE_API_KEY_ENV).filter(|key| !key.trim().is_empty()))
+        .ok_or_else(|| {
+            Error::Auth("Missing COHERE_API_KEY - set it in the environment or pass api_key".into())
+        })?;
+    Ok(
+        std::iter::once(("Authorization".into(), format!("Bearer {key}")))
+            .chain(connection.extra_headers.clone())
+            .collect(),
+    )
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn completes_provider_urls_without_duplicate_paths_and_preserves_queries() {
+        for suffix in ["", "/v2", "/v2/parse"] {
+            assert_eq!(
+                complete_url(&format!("https://example.com{suffix}?tenant=a")).unwrap(),
+                "https://example.com/v2/parse?tenant=a"
+            );
+        }
+    }
+
+    #[test]
+    fn rejects_invalid_urls_and_blank_keys() {
+        assert!(complete_url("relative/path").is_err());
+        assert!(complete_url("ftp://example.com").is_err());
+        assert!(matches!(
+            validate_environment(
+                &OcrConnection {
+                    api_key: Some("  ".into()),
+                    ..Default::default()
+                },
+                &|_| None,
+            ),
+            Err(OcrError::Public(Error::Auth(_)))
+        ));
+    }
+}

litellm-rust/crates/core/src/ocr/adapters/mistral.rs

@@ -33,7 +33,7 @@ impl OcrAdapter for MistralAdapter {
         let url = get_complete_url(request.connection.api_base.as_deref())?;
         let body =
             mistral::transform_ocr_request(&request.model, request.document.clone(), &params)?;
-        transform_request_body(client, request, &url, &headers, body, |_| Ok(())).await
+        transform_request_body(client, request, &url, &headers, true, body, |_| Ok(())).await
     }
 
     fn transform_ocr_response(

litellm-rust/crates/core/src/ocr/adapters/mod.rs

@@ -5,15 +5,16 @@ use serde::de::DeserializeOwned;
 use super::OcrClient;
 use super::error::{OcrError, OcrResponseError};
 use super::registry::OcrProvider;
-use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrResponseFormat};
-use super::wire::DecodedOcrResponse;
+use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse};
 
 mod azure;
+mod cohere;
 mod mistral;
 mod reducto;
 mod vertex;
 
-pub(crate) use azure::{AzureDocumentIntelligenceAdapter, AzureMistralAdapter};
+pub(crate) use azure::{AzureCohereAdapter, AzureDocumentIntelligenceAdapter, AzureMistralAdapter};
+pub(crate) use cohere::CohereAdapter;
 pub(crate) use mistral::MistralAdapter;
 pub(crate) use reducto::{ReductoLegacyAdapter, ReductoV3Adapter};
 pub(crate) use vertex::{VertexDeepSeekAdapter, VertexMistralAdapter};
@@ -55,18 +56,27 @@ pub(crate) trait OcrAdapter: Send + Sync + Sized + 'static {
         _url: &str,
         _headers: &[(String, String)],
         request: &LiteLLMOcrRequest,
-    ) -> impl Future<Output = Result<DecodedOcrResponse<Self::ProviderResponse>, OcrError>> + Send
-    {
-        let retain_native = request
-            .response_format()
-            .map(|format| format == OcrResponseFormat::Native);
-        async move { super::client::read_json_response(response, retain_native?).await }
+    ) -> impl Future<
+        Output = Result<super::wire::DecodedOcrResponse<Self::ProviderResponse>, OcrError>,
+    > + Send {
+        async move {
+            let bytes =
+                super::client::read_response_bytes(response, request.connection.max_response_bytes)
+                    .await?;
+            super::handler::post_call(&request.hooks, &bytes).await?;
+            Ok(super::wire::decode_response(
+                &bytes,
+                request.response_format()? == super::types::OcrResponseFormat::Native,
+            )?)
+        }
     }
 }
 
 macro_rules! for_each_ocr_adapter {
     ($callback:ident) => {
         $callback! {
+            Cohere, $crate::ocr::adapters::CohereAdapter, $crate::ocr::adapters::CohereAdapter, Cohere;
+            AzureCohere, $crate::ocr::adapters::AzureCohereAdapter, $crate::ocr::adapters::AzureCohereAdapter, AzureAi;
             Mistral, $crate::ocr::adapters::MistralAdapter, $crate::ocr::adapters::MistralAdapter, Mistral;
             AzureMistral, $crate::ocr::adapters::AzureMistralAdapter, $crate::ocr::adapters::AzureMistralAdapter, AzureAi;
             AzureDocumentIntelligence, $crate::ocr::adapters::AzureDocumentIntelligenceAdapter, $crate::ocr::adapters::AzureDocumentIntelligenceAdapter, AzureAi;

litellm-rust/crates/core/src/ocr/adapters/reducto/legacy.rs

@@ -27,7 +27,7 @@ impl OcrAdapter for ReductoLegacyAdapter {
         } = _prepare_ocr_request::<ReductoLegacyParams>(request)?;
         let headers = super::validate_environment(&request.connection, &credential_env)?;
         let url = super::get_complete_url(request.connection.api_base.as_deref(), "parse")?;
-        let document = guardrail_document(request, &url).await?;
+        let (document, headers) = guardrail_document(request, &url, &headers).await?;
         let document =
             super::prepare_document(client, document, &request.connection, &headers).await?;
         let body = reducto::transform_legacy_ocr_request(&request.model, document, &params)?;

litellm-rust/crates/core/src/ocr/adapters/reducto/mod.rs

@@ -93,7 +93,7 @@ pub(super) async fn prepare_document(
         .map_err(crate::error::TransportError::from)?;
     let uploaded = crate::ocr::client::read_json_response::<
         crate::ocr::codecs::reducto::ReductoUploadResponse,
-    >(response, false)
+    >(response, false, connection.max_response_bytes)
     .await?
     .data;
     let file_id = uploaded

litellm-rust/crates/core/src/ocr/adapters/reducto/v3.rs

@@ -27,7 +27,7 @@ impl OcrAdapter for ReductoV3Adapter {
         } = _prepare_ocr_request::<ReductoV3Params>(request)?;
         let headers = super::validate_environment(&request.connection, &credential_env)?;
         let url = super::get_complete_url(request.connection.api_base.as_deref(), "parse")?;
-        let document = guardrail_document(request, &url).await?;
+        let (document, headers) = guardrail_document(request, &url, &headers).await?;
         let document =
             super::prepare_document(client, document, &request.connection, &headers).await?;
         let body = reducto::transform_v3_ocr_request(&request.model, document, &params)?;

litellm-rust/crates/core/src/ocr/adapters/vertex/deepseek.rs

@@ -57,9 +57,15 @@ impl OcrAdapter for VertexDeepSeekAdapter {
         let document = request.document.clone();
         let body =
             deepseek::transform_ocr_request(&provider_model(&request.model), document, &params)?;
-        transform_request_body(client, request, &url, &authentication.headers, body, |_| {
-            Ok(())
-        })
+        transform_request_body(
+            client,
+            request,
+            &url,
+            &authentication.headers,
+            false,
+            body,
+            |_| Ok(()),
+        )
         .await
     }
 

litellm-rust/crates/core/src/ocr/adapters/vertex/mistral.rs

@@ -54,6 +54,8 @@ impl OcrAdapter for VertexMistralAdapter {
             &location,
             &request.model,
         )?;
+        let retains_document = !request.document.source().starts_with("http://")
+            && !request.document.source().starts_with("https://");
         let document = inline_remote_document(
             client.document_fetcher(),
             request.document.clone(),
@@ -66,6 +68,7 @@ impl OcrAdapter for VertexMistralAdapter {
             request,
             &url,
             &authentication.headers,
+            retains_document,
             body,
             |body| validate_inline_document(&body.document),
         )

litellm-rust/crates/core/src/ocr/client.rs

@@ -1,10 +1,10 @@
 use std::sync::OnceLock;
 use std::time::Duration;
 
+use bytes::{Bytes, BytesMut};
 use serde::de::DeserializeOwned;
 
-use super::error::OcrError;
-use super::handler::perform_ocr_request;
+use super::error::{OcrError, OcrResponseError};
 use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse};
 use super::wire::{DecodedOcrResponse, decode_response};
 use crate::Error;
@@ -32,6 +32,10 @@ impl OcrClient {
         })
     }
 
+    pub fn shared() -> Result<Self, Error> {
+        shared_client()
+    }
+
     #[tracing::instrument(
         name = "ocr",
         target = "litellm::function_trace",
@@ -39,7 +43,34 @@ impl OcrClient {
         skip_all
     )]
     pub async fn perform(&self, request: LiteLLMOcrRequest) -> Result<LiteLLMOcrResponse, Error> {
-        perform_ocr_request(self, request).await
+        use super::{
+            NativeOutcome, OcrAdmission, OcrCall, OcrCallStep, OcrHookHost, OcrHost,
+            OcrHostOperation, OcrHostResult,
+        };
+
+        let host = OcrHookHost::new(request.hooks.clone());
+        let mut request = Some(request);
+        let NativeOutcome::Completed(mut call) = OcrCall::admit(self.clone(), OcrAdmission::all())
+        else {
+            return Err(Error::InvalidRequest(
+                "native OCR host admission declined".into(),
+            ));
+        };
+        let mut result = None;
+        loop {
+            match call.resume(result.take()).await? {
+                OcrCallStep::Host(OcrHostOperation::ProjectRequest) => {
+                    result = Some(OcrHostResult::Request(Ok((
+                        Box::new(request.take().ok_or_else(|| {
+                            Error::InvalidRequest("OCR request was already projected".into())
+                        })?),
+                        false,
+                    ))))
+                }
+                OcrCallStep::Host(operation) => result = Some(host.invoke(operation).await),
+                OcrCallStep::Complete(response) => return Ok(response),
+            }
+        }
     }
 
     pub(crate) fn provider_http(&self) -> &reqwest::Client {
@@ -77,7 +108,7 @@ fn no_redirect_http() -> Result<reqwest::Client, TransportError> {
         .map_err(TransportError::from)
 }
 
-pub async fn ocr(request: LiteLLMOcrRequest) -> Result<LiteLLMOcrResponse, Error> {
+pub(crate) fn shared_client() -> Result<OcrClient, Error> {
     static CLIENT: OnceLock<Result<OcrClient, TransportError>> = OnceLock::new();
     let client = CLIENT
         .get_or_init(|| {
@@ -88,18 +119,50 @@ pub async fn ocr(request: LiteLLMOcrRequest) -> Result<LiteLLMOcrResponse, Error
                 .and_then(OcrClient::new)
         })
         .clone()?;
-    client.perform(request).await
+    Ok(client)
+}
+
+pub async fn ocr(request: LiteLLMOcrRequest) -> Result<LiteLLMOcrResponse, Error> {
+    shared_client()?.perform(request).await
 }
 
 pub async fn read_json_response<T: DeserializeOwned>(
     response: reqwest::Response,
     native: bool,
+    max_response_bytes: usize,
 ) -> Result<DecodedOcrResponse<T>, OcrError> {
+    let bytes = read_response_bytes(response, max_response_bytes).await?;
+    Ok(decode_response(&bytes, native)?)
+}
+
+pub(crate) async fn read_response_bytes(
+    mut response: reqwest::Response,
+    max_response_bytes: usize,
+) -> Result<Bytes, OcrError> {
     let status = response.status();
-    let bytes = response
-        .bytes()
-        .await
-        .map_err(crate::error::TransportError::from)?;
+    let limit = if status.is_success() {
+        max_response_bytes
+    } else {
+        max_response_bytes.min(4 * (crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS + 1))
+    };
+    if status.is_success()
+        && response
+            .content_length()
+            .is_some_and(|length| length > limit as u64)
+    {
+        return Err(OcrResponseError::TooLarge { limit }.into());
+    }
+    let mut bytes = BytesMut::new();
+    while let Some(chunk) = response.chunk().await.map_err(transport_error)? {
+        let remaining = limit.saturating_sub(bytes.len());
+        if status.is_success() && chunk.len() > remaining {
+            return Err(OcrResponseError::TooLarge { limit }.into());
+        }
+        bytes.extend_from_slice(&chunk[..chunk.len().min(remaining)]);
+        if !status.is_success() && bytes.len() == limit {
+            break;
+        }
+    }
     if !status.is_success() {
         return Err(crate::error::TransportError::Http {
             status: status.as_u16(),
@@ -107,5 +170,41 @@ pub async fn read_json_response<T: DeserializeOwned>(
         }
         .into());
     }
-    Ok(decode_response(&bytes, native)?)
+    Ok(bytes.freeze())
+}
+
+pub(crate) fn transport_error(error: reqwest::Error) -> Error {
+    if error.is_timeout() {
+        return Error::Http {
+            status: 408,
+            body: "OCR request timed out".into(),
+        };
+    }
+    crate::error::TransportError::from(error).into()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn request_timeout_has_an_http_408_status() {
+        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let address = listener.local_addr().unwrap();
+        let server = tokio::spawn(async move {
+            let _connection = listener.accept().await.unwrap();
+            tokio::time::sleep(Duration::from_secs(1)).await;
+        });
+        let error = reqwest::Client::new()
+            .get(format!("http://{address}"))
+            .timeout(Duration::from_millis(10))
+            .send()
+            .await
+            .unwrap_err();
+        assert!(matches!(
+            transport_error(error),
+            Error::Http { status: 408, .. }
+        ));
+        server.abort();
+    }
 }

litellm-rust/crates/core/src/ocr/codecs/cohere.rs

@@ -0,0 +1,254 @@
+use serde::{Deserialize, Serialize};
+use serde_json::{Map, Value, json};
+
+use crate::ocr::document::InlineDocument;
+use crate::ocr::error::{OcrRequestError, OcrResponseError};
+use crate::ocr::types::{LiteLLMOcrResponse, OcrDocument};
+
+#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize)]
+#[serde(rename_all = "lowercase")]
+pub(crate) enum OutputFormat {
+    #[default]
+    Markdown,
+    Blocks,
+}
+
+#[derive(Deserialize)]
+pub(crate) struct CohereParams {
+    #[serde(default)]
+    pub output_format: OutputFormat,
+}
+
+#[derive(Deserialize, Serialize)]
+pub(crate) struct CohereRequest {
+    pub model: String,
+    pub document: OcrDocument,
+    pub output_format: OutputFormat,
+}
+
+pub(crate) fn validate_document(document: &OcrDocument) -> Result<(), OcrRequestError> {
+    let OcrDocument::ImageUrl { image_url, .. } = document else {
+        return Err(OcrRequestError::CohereImageOnly);
+    };
+    if image_url.is_empty() {
+        return Err(OcrRequestError::CohereImageOnly);
+    }
+    if let Some(inline) = InlineDocument::parse(image_url)? {
+        if !inline.mime_type().type_.eq_ignore_ascii_case("image") {
+            return Err(OcrRequestError::CohereImageOnly);
+        }
+        inline.decode(crate::constants::OCR_INLINE_MAX_BYTES)?;
+    }
+    Ok(())
+}
+
+#[derive(Deserialize)]
+pub(crate) struct CohereResponse {
+    #[serde(default)]
+    pages: Vec<CoherePage>,
+    meta: Option<CohereMeta>,
+}
+
+#[derive(Deserialize)]
+struct CoherePage {
+    index: Option<i64>,
+    markdown: Option<CohereMarkdown>,
+    blocks: Option<Vec<Map<String, Value>>>,
+}
+
+#[derive(Deserialize)]
+struct CohereMarkdown {
+    #[serde(default)]
+    content: String,
+    images: Option<Vec<Map<String, Value>>>,
+}
+
+#[derive(Deserialize)]
+struct CohereMeta {
+    billed_units: Option<CohereBilledUnits>,
+}
+
+#[derive(Deserialize)]
+struct CohereBilledUnits {
+    pages: Option<i64>,
+}
+
+pub(crate) fn transform_response(
+    model: &str,
+    response: CohereResponse,
+) -> Result<LiteLLMOcrResponse, OcrResponseError> {
+    let pages_processed = response
+        .meta
+        .and_then(|meta| meta.billed_units)
+        .and_then(|units| units.pages)
+        .map(Ok)
+        .unwrap_or_else(|| {
+            i64::try_from(response.pages.len()).map_err(|_| OcrResponseError::NumericRange("pages"))
+        })?;
+    let pages = response
+        .pages
+        .into_iter()
+        .enumerate()
+        .map(|(position, page)| {
+            let index = page.index.map(Ok).unwrap_or_else(|| {
+                i64::try_from(position).map_err(|_| OcrResponseError::NumericRange("page index"))
+            })?;
+            let (content, images) = page
+                .markdown
+                .map(|markdown| {
+                    let images =
+                        markdown
+                            .images
+                            .filter(|images| !images.is_empty())
+                            .map(|images| {
+                                images
+                                    .into_iter()
+                                    .map(|mut image| {
+                                        if let Some(Value::Object(bbox)) =
+                                            image.get("bounding_box").cloned()
+                                        {
+                                            image.insert("bbox".into(), Value::Object(bbox));
+                                        }
+                                        Value::Object(image)
+                                    })
+                                    .collect::<Vec<_>>()
+                            });
+                    (markdown.content, images)
+                })
+                .unwrap_or_default();
+            let mut normalized = json!({"index": index, "markdown": content, "images": images});
+            if let Some(blocks) = page.blocks {
+                normalized["blocks"] = json!(blocks);
+            }
+            Ok(normalized)
+        })
+        .collect::<Result<Vec<_>, OcrResponseError>>()?;
+    Ok(LiteLLMOcrResponse {
+        pages,
+        model: model.into(),
+        document_annotation: None,
+        usage_info: Some(json!({"pages_processed": pages_processed})),
+        object: "ocr".into(),
+        extra_fields: Map::new(),
+        provider_native_response: None,
+    })
+}
+
+pub(crate) fn transform_request(
+    model: &str,
+    document: OcrDocument,
+    params: CohereParams,
+) -> Result<CohereRequest, OcrRequestError> {
+    validate_document(&document)?;
+    Ok(CohereRequest {
+        model: model.into(),
+        document,
+        output_format: params.output_format,
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn response_normalizes_markdown_images_blocks_and_billed_pages() {
+        let response = serde_json::from_value(json!({
+            "pages": [
+                {
+                    "type":"markdown",
+                    "index":4,
+                    "markdown":{
+                        "content":"receipt",
+                        "images":[{
+                            "id":"image",
+                            "bounding_box":{"top_left_x":1,"bottom_right_x":48},
+                            "bounding_box_normalized":{"top_left_x":0.04,"bottom_right_x":0.15},
+                            "description":"scan",
+                            "category":"logo"
+                        }]
+                    }
+                },
+                {"type":"blocks","blocks":[{"type":"text","text":{"content":"total"}}]}
+            ],
+            "meta":{"api_version":{"version":"2"},"billed_units":{"pages":3}}
+        }))
+        .unwrap();
+        let normalized = transform_response("parse-v5.0", response).unwrap();
+        assert_eq!(normalized.pages[0]["index"], 4);
+        assert_eq!(normalized.pages[0]["markdown"], "receipt");
+        assert_eq!(normalized.pages[0]["images"][0]["bbox"]["top_left_x"], 1);
+        assert_eq!(
+            normalized.pages[0]["images"][0]["bounding_box_normalized"]["bottom_right_x"],
+            0.15
+        );
+        assert_eq!(normalized.pages[0]["images"][0]["description"], "scan");
+        assert_eq!(normalized.pages[0]["images"][0]["category"], "logo");
+        assert_eq!(normalized.pages[1]["index"], 1);
+        assert_eq!(normalized.pages[1]["markdown"], "");
+        assert_eq!(normalized.pages[1]["blocks"][0]["text"]["content"], "total");
+        assert_eq!(normalized.usage_info.unwrap()["pages_processed"], 3);
+    }
+
+    #[test]
+    fn response_defaults_and_invalid_fields() {
+        for value in [
+            json!({}),
+            json!({"meta":null}),
+            json!({"pages":[],"meta":{"billed_units":null}}),
+        ] {
+            let normalized =
+                transform_response("parse", serde_json::from_value(value).unwrap()).unwrap();
+            assert!(normalized.pages.is_empty());
+            assert_eq!(normalized.usage_info.unwrap()["pages_processed"], 0);
+        }
+        for value in [
+            json!({"pages":null}),
+            json!({"pages":[{"markdown":"text"}]}),
+            json!({"pages":[{"index":"bad"}]}),
+        ] {
+            assert!(serde_json::from_value::<CohereResponse>(value).is_err());
+        }
+        let normalized = transform_response(
+            "parse",
+            serde_json::from_value(json!({"pages":[{"markdown":null}]})).unwrap(),
+        )
+        .unwrap();
+        assert_eq!(normalized.usage_info.unwrap()["pages_processed"], 1);
+        assert!(normalized.pages[0]["images"].is_null());
+    }
+
+    #[test]
+    fn request_requires_image_and_supported_output_format() {
+        for value in [
+            json!({"type":"document_url","document_url":"https://example.com/a.pdf"}),
+            json!({"type":"image_url","image_url":""}),
+            json!({"type":"image_url","image_url":"data:application/pdf;base64,YQ=="}),
+        ] {
+            assert_eq!(
+                validate_document(&serde_json::from_value(value).unwrap()),
+                Err(OcrRequestError::CohereImageOnly)
+            );
+        }
+        assert!(serde_json::from_value::<CohereParams>(json!({"output_format":"html"})).is_err());
+        for format in ["markdown", "blocks"] {
+            assert!(
+                serde_json::from_value::<CohereParams>(json!({"output_format":format})).is_ok()
+            );
+        }
+        let request = transform_request(
+            "parse-v5.0",
+            serde_json::from_value(json!({
+                "type":"image_url",
+                "image_url":"https://example.com/image.png"
+            }))
+            .unwrap(),
+            serde_json::from_value(json!({})).unwrap(),
+        )
+        .unwrap();
+        assert_eq!(
+            serde_json::to_value(request).unwrap()["output_format"],
+            "markdown"
+        );
+    }
+}

litellm-rust/crates/core/src/ocr/codecs/deepseek/transformation.rs

@@ -12,13 +12,17 @@ pub(crate) fn transform_ocr_request(
     params: &DeepSeekOcrParams,
 ) -> Result<DeepSeekOcrRequest, OcrRequestError> {
     if document.source().is_empty() {
-        return Err(OcrRequestError::MissingField("document URL"));
+        return Err(OcrRequestError::MissingDocumentUrl);
     }
+    let content = OcrDocument::ImageUrl {
+        image_url: document.source().to_string(),
+        extra_fields: serde_json::Map::new(),
+    };
     Ok(DeepSeekOcrRequest {
         model: provider_model.to_string(),
         messages: vec![DeepSeekOcrMessage {
             role: UserRole::User,
-            content: vec![document],
+            content: vec![content],
         }],
         params: params.clone(),
     })

litellm-rust/crates/core/src/ocr/codecs/document_intelligence/params.rs

@@ -163,6 +163,30 @@ mod tests {
         );
     }
 
+    #[rstest]
+    #[case(json!([0, 1, 2]), Some("1,2,3"))]
+    #[case(json!([2, 0, 0, 1]), Some("1,2,3"))]
+    #[case(json!([]), None)]
+    #[case(json!("3-9"), Some("3-9"))]
+    #[case(json!("1-3, 5"), Some("1-3,5"))]
+    #[case(json!(["1", "3-5"]), Some("1,3-5"))]
+    fn page_mapping_matches_python(#[case] input: Value, #[case] expected: Option<&str>) {
+        assert_eq!(
+            map(json!({"pages": input})).unwrap().pages.as_deref(),
+            expected
+        );
+    }
+
+    #[rstest]
+    #[case(json!("a,b"))]
+    #[case(json!([-1]))]
+    #[case(json!([true, false]))]
+    #[case(json!([1, "2"]))]
+    #[case(json!(5))]
+    fn invalid_page_mapping_matches_python(#[case] input: Value) {
+        assert!(map(json!({"pages": input})).is_err());
+    }
+
     #[rstest]
     #[case(json!(["keyValuePairs"]), "keyValuePairs")]
     #[case(json!(["keyValuePairs", "languages"]), "keyValuePairs,languages")]

litellm-rust/crates/core/src/ocr/codecs/document_intelligence/transformation.rs

@@ -13,7 +13,7 @@ pub(crate) fn transform_ocr_request(
 ) -> Result<DocumentIntelligenceRequest, OcrRequestError> {
     let source = document.source();
     if source.is_empty() {
-        return Err(OcrRequestError::MissingField("document URL"));
+        return Err(OcrRequestError::MissingDocumentUrl);
     }
     Ok(if let Some(document) = InlineDocument::parse(source)? {
         DocumentIntelligenceRequest::Base64Source(
@@ -46,10 +46,7 @@ pub(crate) fn transform_ocr_response(
     let mut extra_fields = Map::new();
     extra_fields.insert("content".into(), option_value(result.content));
     extra_fields.insert("tables".into(), option_value(result.tables));
-    extra_fields.insert(
-        "key_value_pairs".into(),
-        option_value(result.key_value_pairs),
-    );
+    extra_fields.insert("keyValuePairs".into(), option_value(result.key_value_pairs));
     Ok(LiteLLMOcrResponse {
         pages,
         model: model.into(),

litellm-rust/crates/core/src/ocr/codecs/mistral/transformation.rs

@@ -114,6 +114,7 @@ mod tests {
     #[rstest]
     #[case("table_format", json!("html"))]
     #[case("confidence_scores_granularity", json!("word"))]
+    #[case("confidence_scores_granularity", json!("block"))]
     #[case("document_annotation_prompt", json!("extract"))]
     #[case("include_blocks", json!(true))]
     #[case("id", json!("req-123"))]
@@ -133,6 +134,7 @@ mod tests {
 
     #[rstest]
     #[case("pages", json!([0, 2]))]
+    #[case("pages", json!("0,2-4"))]
     #[case("include_image_base64", json!(true))]
     #[case("image_limit", json!(2))]
     #[case("image_min_size", json!(100))]
@@ -196,8 +198,16 @@ mod tests {
     #[rstest]
     fn transform_ocr_response_preserves_blocks_and_confidence_scores() {
         let response: MistralOcrResponse = serde_json::from_value(json!({
-            "pages":[{"index":0,"markdown":"hello","blocks":[{"type":"title"}],"confidence_scores":{"mean":0.99}}],
+            "pages":[{
+                "index":0,
+                "markdown":"hello",
+                "images":[{"id":"img-0","image_base64":"data:image/png;base64,AA=="}],
+                "dimensions":{"width":612,"height":792,"dpi":72},
+                "blocks":[{"type":"title","bbox":{"x":1},"confidence_scores":{"mean":0.98}}],
+                "confidence_scores":{"average_page_confidence_score":0.99,"minimum_page_confidence_score":0.97}
+            }],
             "model":"returned-model",
+            "document_annotation":"{\"language\":\"en\"}",
             "usage_info":{"pages_processed":1}
         }))
         .unwrap();
@@ -205,7 +215,20 @@ mod tests {
             .unwrap()
             .into_json();
         assert_eq!(result["pages"][0]["blocks"][0]["type"], "title");
-        assert_eq!(result["pages"][0]["confidence_scores"]["mean"], 0.99);
+        assert_eq!(result["pages"][0]["blocks"][0]["bbox"]["x"], 1);
+        assert_eq!(
+            result["pages"][0]["blocks"][0]["confidence_scores"]["mean"],
+            0.98
+        );
+        assert_eq!(
+            result["pages"][0]["confidence_scores"]["average_page_confidence_score"],
+            0.99
+        );
+        assert_eq!(result["pages"][0]["images"][0]["id"], "img-0");
+        assert_eq!(result["pages"][0]["dimensions"]["dpi"], 72);
+        assert_eq!(result["model"], "returned-model");
+        assert_eq!(result["document_annotation"], "{\"language\":\"en\"}");
+        assert_eq!(result["usage_info"]["pages_processed"], 1);
     }
 
     #[rstest]

litellm-rust/crates/core/src/ocr/codecs/mistral/types.rs

@@ -3,10 +3,17 @@ use serde_json::{Map, Value};
 
 use crate::ocr::types::OcrDocument;
 
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
+#[serde(untagged)]
+pub(crate) enum MistralOcrPages {
+    Range(String),
+    Indices(Vec<i64>),
+}
+
 #[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
 pub(crate) struct MistralOcrParams {
     #[serde(skip_serializing_if = "Option::is_none")]
-    pub pages: Option<Vec<i64>>,
+    pub pages: Option<MistralOcrPages>,
     #[serde(skip_serializing_if = "Option::is_none")]
     pub include_image_base64: Option<bool>,
     #[serde(skip_serializing_if = "Option::is_none")]

litellm-rust/crates/core/src/ocr/codecs/mod.rs

@@ -1,3 +1,4 @@
+pub(crate) mod cohere;
 pub(crate) mod deepseek;
 pub(crate) mod document_intelligence;
 pub(crate) mod mistral;

litellm-rust/crates/core/src/ocr/document.rs

@@ -2,13 +2,90 @@ use base64::{Engine, engine::general_purpose::STANDARD};
 use data_url::mime::Mime;
 use data_url::{DataUrl, DataUrlError, forgiving_base64::DecodeError};
 use reqwest::Url;
+use serde_json::Map;
 
 use super::error::{OcrError, OcrRequestError, OcrResponseError};
 use super::types::{OcrConnection, OcrDocument};
-use crate::constants::OCR_MAX_FETCH_REDIRECTS;
+use crate::constants::{OCR_INLINE_MAX_BYTES, OCR_MAX_FETCH_REDIRECTS};
 use crate::error::{MediaError, TransportError};
 use crate::media::{DownloadPolicy, MediaFetcher};
 
+pub fn encode_file_document(
+    bytes: &[u8],
+    file_name: Option<&str>,
+    mime_type: Option<&str>,
+) -> Result<OcrDocument, OcrRequestError> {
+    if bytes.is_empty() {
+        return Err(OcrRequestError::EmptyFile);
+    }
+    if bytes.len() > OCR_INLINE_MAX_BYTES {
+        return Err(OcrRequestError::InlineDocumentTooLarge);
+    }
+    if let Some(value) = mime_type
+        && !valid_mime_type(value)
+    {
+        return Err(OcrRequestError::InvalidMimeType(value.into()));
+    }
+    let mime_type = mime_type
+        .map(str::to_string)
+        .or_else(|| file_name.map(|name| mime_type_for_name(name).to_string()))
+        .unwrap_or_else(|| "application/octet-stream".into());
+    let source = format!("data:{mime_type};base64,{}", STANDARD.encode(bytes));
+    Ok(if mime_type.starts_with("image/") {
+        OcrDocument::ImageUrl {
+            image_url: source,
+            extra_fields: Map::new(),
+        }
+    } else {
+        OcrDocument::DocumentUrl {
+            document_url: source,
+            extra_fields: Map::new(),
+        }
+    })
+}
+
+fn valid_mime_type(value: &str) -> bool {
+    let Some((kind, subtype)) = value.split_once('/') else {
+        return false;
+    };
+    !kind.is_empty()
+        && !subtype.is_empty()
+        && kind.chars().chain(subtype.chars()).all(|character| {
+            character.is_alphanumeric() || matches!(character, '.' | '+' | '-' | '_')
+        })
+}
+
+pub fn mime_type_for_name(name: &str) -> &'static str {
+    let extension = std::path::Path::new(name)
+        .extension()
+        .and_then(|value| value.to_str())
+        .unwrap_or_default();
+    match extension.to_ascii_lowercase().as_str() {
+        "pdf" => "application/pdf",
+        "png" => "image/png",
+        "jpg" | "jpeg" => "image/jpeg",
+        "gif" => "image/gif",
+        "webp" => "image/webp",
+        "tiff" | "tif" => "image/tiff",
+        "bmp" => "image/bmp",
+        _ => mime_guess::from_path(name)
+            .first_raw()
+            .unwrap_or("application/octet-stream"),
+    }
+}
+
+pub fn upload_mime_type<'a>(file_name: Option<&str>, content_type: Option<&'a str>) -> &'a str {
+    match content_type
+        .and_then(|value| value.split(';').next())
+        .map(str::trim)
+    {
+        Some(value) if !value.is_empty() && value != "application/octet-stream" => value,
+        _ => file_name
+            .map(mime_type_for_name)
+            .unwrap_or("application/octet-stream"),
+    }
+}
+
 pub(crate) struct InlineDocument<'a>(DataUrl<'a>);
 
 impl<'a> InlineDocument<'a> {
@@ -95,9 +172,11 @@ fn map_media_error(error: MediaError) -> OcrError {
             body: "OCR document download failed".into(),
         }
         .into(),
-        MediaError::Timeout => {
-            TransportError::Network("OCR document download timed out".into()).into()
+        MediaError::Timeout => TransportError::Http {
+            status: 408,
+            body: "OCR document download timed out".into(),
         }
+        .into(),
         MediaError::Transport(error) => error.into(),
     }
 }
@@ -114,6 +193,90 @@ mod tests {
         }
     }
 
+    #[test]
+    fn file_bytes_are_encoded_with_core_owned_mime_policy() {
+        assert_eq!(
+            encode_file_document(b"abc", Some("scan.png"), None).unwrap(),
+            OcrDocument::ImageUrl {
+                image_url: "data:image/png;base64,YWJj".into(),
+                extra_fields: Map::new(),
+            }
+        );
+        assert_eq!(
+            encode_file_document(b"abc", None, Some("application/pdf")).unwrap(),
+            document("data:application/pdf;base64,YWJj")
+        );
+    }
+
+    #[test]
+    fn file_name_mime_mapping_matches_python() {
+        for (name, expected) in [
+            ("document.pdf", "application/pdf"),
+            ("image.png", "image/png"),
+            ("photo.jpg", "image/jpeg"),
+            ("photo.jpeg", "image/jpeg"),
+            ("animation.gif", "image/gif"),
+            ("image.webp", "image/webp"),
+            ("scan.tiff", "image/tiff"),
+            ("scan.tif", "image/tiff"),
+            ("bitmap.bmp", "image/bmp"),
+            ("DOCUMENT.PDF", "application/pdf"),
+            ("IMAGE.PNG", "image/png"),
+            ("file.unknown-extension", "application/octet-stream"),
+        ] {
+            assert_eq!(mime_type_for_name(name), expected);
+        }
+    }
+
+    #[test]
+    fn upload_mime_mapping_matches_python() {
+        assert_eq!(
+            upload_mime_type(Some("report.pdf"), Some("application/octet-stream")),
+            "application/pdf"
+        );
+        assert_eq!(upload_mime_type(Some("image.png"), None), "image/png");
+        assert_eq!(upload_mime_type(None, None), "application/octet-stream");
+        assert_eq!(
+            upload_mime_type(Some("doc.pdf"), Some("application/pdf; charset=utf-8")),
+            "application/pdf"
+        );
+        assert_eq!(
+            upload_mime_type(
+                Some("img.png"),
+                Some("image/png; charset=utf-8; boundary=something")
+            ),
+            "image/png"
+        );
+    }
+
+    #[test]
+    fn file_encoding_enforces_decoded_size_limit() {
+        let bytes = vec![b'a'; OCR_INLINE_MAX_BYTES + 1];
+        assert_eq!(
+            encode_file_document(&bytes, None, None),
+            Err(OcrRequestError::InlineDocumentTooLarge)
+        );
+        let document = encode_file_document(&bytes[..OCR_INLINE_MAX_BYTES], None, None).unwrap();
+        let inline = InlineDocument::parse(document.source()).unwrap().unwrap();
+        assert_eq!(
+            inline.decode(OCR_INLINE_MAX_BYTES).unwrap(),
+            bytes[..OCR_INLINE_MAX_BYTES]
+        );
+    }
+
+    #[test]
+    fn file_encoding_rejects_empty_bytes_and_invalid_explicit_mime() {
+        assert!(encode_file_document(b"", None, None).is_err());
+        for mime in [
+            "text/plain;bad",
+            "text/plain/extra",
+            " text/plain",
+            "text/plain\n",
+        ] {
+            assert!(encode_file_document(b"abc", None, Some(mime)).is_err());
+        }
+    }
+
     #[test]
     fn decodes_data_urls_and_limits_decoded_size() {
         for (source, expected) in [

litellm-rust/crates/core/src/ocr/error.rs

@@ -4,15 +4,27 @@ use crate::error::TransportError;
 
 #[derive(Debug, Clone, PartialEq, Eq, Error)]
 pub enum OcrRequestError {
+    #[error("File is empty or could not be read")]
+    EmptyFile,
+    #[error("Invalid MIME type: {0}")]
+    InvalidMimeType(String),
+    #[error(
+        "Cohere Parse only accepts `image_url` documents; document_url and PDF inputs are not supported"
+    )]
+    CohereImageOnly,
     #[error("Invalid `req_format`. Expected 'native' or 'litellm'.")]
     RequestFormat,
     #[error("invalid OCR request field: {path}")]
     RequestField { path: String },
     #[error("missing required field: {0}")]
     MissingField(&'static str),
+    #[error("Document URL is required")]
+    MissingDocumentUrl,
     #[error("invalid OCR document data URI")]
     InvalidDataUri,
-    #[error("Reducto requires a reducto:// id or a data URI")]
+    #[error(
+        "Reducto requires a reducto:// id or a data URI; plain HTTP URLs are not supported, upload the file first"
+    )]
     ReductoSource,
     #[error("inline OCR document exceeds the size limit")]
     InlineDocumentTooLarge,
@@ -34,6 +46,8 @@ pub enum OcrRequestError {
 
 #[derive(Debug, Clone, PartialEq, Eq, Error)]
 pub enum OcrResponseError {
+    #[error("OCR response exceeds the size limit of {limit} bytes")]
+    TooLarge { limit: usize },
     #[error("invalid OCR response field: {path}")]
     ResponseField { path: String },
     #[error("OCR response is missing non-empty content")]

litellm-rust/crates/core/src/ocr/handler.rs

@@ -1,15 +1,17 @@
 use super::OcrClient;
 use super::adapters::OcrAdapter;
-use super::hooks::OcrLifecycleHooks;
+use super::hooks::{OcrHooks, OcrLifecycleHooks, OcrPostCallRequest};
 use super::registry::OcrAdapterKind;
 use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse};
 use crate::Error;
 use crate::call_lifecycle::{CallLifecycle, CallLifecycleContext};
+use std::sync::Arc;
 
 pub(crate) async fn perform_ocr_request(
     client: &OcrClient,
     request: LiteLLMOcrRequest,
 ) -> Result<LiteLLMOcrResponse, Error> {
+    request.response_format()?;
     let context = CallLifecycleContext::new(
         "ocr",
         request.model.clone(),
@@ -23,27 +25,71 @@ pub(crate) async fn perform_ocr_request(
         hooks: request.hooks.clone(),
         provider_name: context.custom_llm_provider.clone(),
     };
-    CallLifecycle::default().run(context, request, &hooks, |request| async move {
-        macro_rules! execute_selected_adapter {
+    CallLifecycle::default()
+        .run(context, request, &hooks, |request| async move {
+            PreparedOcrCall::prepare(client.clone(), request)
+                .await?
+                .execute()
+                .await?
+                .normalize()
+        })
+        .await
+}
+
+pub(crate) struct PreparedOcrCall {
+    client: OcrClient,
+    request: LiteLLMOcrRequest,
+    http: reqwest::Request,
+}
+
+impl PreparedOcrCall {
+    pub(crate) async fn prepare(
+        client: OcrClient,
+        request: LiteLLMOcrRequest,
+    ) -> Result<Self, Error> {
+        macro_rules! prepare_adapter {
             ($( $variant:ident, $adapter:ty, $instance:expr, $provider:ident; )+) => {
                 match request.adapter {
-                    $( OcrAdapterKind::$variant => execute_ocr_provider_call(client, &$instance, request).await, )+
+                    $( OcrAdapterKind::$variant => $instance.prepare_request(&request, &client).await?, )+
                 }
             };
         }
-        super::adapters::for_each_ocr_adapter!(execute_selected_adapter)
-    }).await
+        let http = super::adapters::for_each_ocr_adapter!(prepare_adapter);
+        Ok(Self {
+            client,
+            request,
+            http,
+        })
+    }
+
+    pub(crate) async fn execute(self) -> Result<OcrProviderResponse, Error> {
+        let url = self.http.url().to_string();
+        let headers = request_headers(&self.http)?;
+        let response = crate::http_utils::http_request(reqwest::RequestBuilder::from_parts(
+            self.client.provider_http().clone(),
+            self.http,
+        ))
+        .await
+        .map_err(super::client::transport_error)?;
+        macro_rules! read_adapter {
+            ($( $variant:ident, $adapter:ty, $instance:expr, $provider:ident; )+) => {
+                match self.request.adapter {
+                    $( OcrAdapterKind::$variant => {
+                        let decoded = $instance.read_response(&self.client, response, &url, &headers, &self.request).await?;
+                        Ok(OcrProviderResponse {
+                            request: self.request,
+                            data: OcrProviderData::$variant(decoded),
+                        })
+                    }, )+
+                }
+            };
+        }
+        super::adapters::for_each_ocr_adapter!(read_adapter)
+    }
 }
 
-#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
-async fn execute_ocr_provider_call<A: OcrAdapter>(
-    client: &OcrClient,
-    adapter: &A,
-    request: LiteLLMOcrRequest,
-) -> Result<LiteLLMOcrResponse, Error> {
-    let provider_request = adapter.prepare_request(&request, client).await?;
-    let url = provider_request.url().to_string();
-    let headers = provider_request
+fn request_headers(request: &reqwest::Request) -> Result<Vec<(String, String)>, Error> {
+    request
         .headers()
         .iter()
         .map(|(name, value)| {
@@ -53,20 +99,41 @@ async fn execute_ocr_provider_call<A: OcrAdapter>(
                 .map_err(|_| super::error::OcrRequestError::RequestField {
                     path: "headers".into(),
                 })
+                .map_err(Error::from)
         })
-        .collect::<Result<Vec<_>, _>>()?;
-    let response = crate::http_utils::http_request(reqwest::RequestBuilder::from_parts(
-        client.provider_http().clone(),
-        provider_request,
-    ))
-    .await
-    .map_err(crate::error::TransportError::from)?;
-    let decoded = adapter
-        .read_response(client, response, &url, &headers, &request)
-        .await?;
-    let response = adapter.transform_ocr_response(&request, decoded.data)?;
-    Ok(LiteLLMOcrResponse {
-        provider_native_response: decoded.native,
-        ..response
-    })
+        .collect()
 }
+
+macro_rules! provider_data {
+    ($( $variant:ident, $adapter:ty, $instance:expr, $provider:ident; )+) => {
+        enum OcrProviderData {
+            $( $variant(super::wire::DecodedOcrResponse<<$adapter as OcrAdapter>::ProviderResponse>), )+
+        }
+
+        impl OcrProviderResponse {
+            pub(crate) fn normalize(self) -> Result<LiteLLMOcrResponse, Error> {
+                match self.data {
+                    $( OcrProviderData::$variant(decoded) => {
+                        let response = $instance.transform_ocr_response(&self.request, decoded.data)?;
+                        Ok(LiteLLMOcrResponse { provider_native_response: decoded.native, ..response })
+                    }, )+
+                }
+            }
+        }
+    };
+}
+
+pub(crate) struct OcrProviderResponse {
+    request: LiteLLMOcrRequest,
+    data: OcrProviderData,
+}
+
+pub(crate) async fn post_call(hooks: &Arc<dyn OcrHooks>, bytes: &[u8]) -> Result<(), Error> {
+    let original_response = serde_json::Value::String(String::from_utf8_lossy(bytes).into_owned());
+    hooks
+        .post_call(OcrPostCallRequest { original_response })
+        .await?;
+    Ok(())
+}
+
+super::adapters::for_each_ocr_adapter!(provider_data);

litellm-rust/crates/core/src/ocr/hooks.rs

@@ -24,11 +24,19 @@ pub struct OcrDuringCallRequest {
     pub model: String,
     pub custom_llm_provider: String,
     pub url: String,
+    pub headers: Vec<(String, String)>,
     pub body: Value,
+    #[serde(skip)]
+    pub retained_fields: Vec<String>,
+}
+
+#[derive(Clone, Debug, Serialize)]
+pub struct OcrPostCallRequest {
+    pub original_response: Value,
 }
 
 pub trait OcrHooks: Send + Sync {
-    fn has_guardrails(&self) -> bool {
+    fn intercepts_requests(&self) -> bool {
         false
     }
     fn pre_call(&self, request: OcrPreCallRequest) -> OcrHookFuture<'_, OcrPreCallRequest> {
@@ -40,6 +48,9 @@ pub trait OcrHooks: Send + Sync {
     ) -> OcrHookFuture<'_, OcrDuringCallRequest> {
         Box::pin(async move { Ok(request) })
     }
+    fn post_call(&self, request: OcrPostCallRequest) -> OcrHookFuture<'_, OcrPostCallRequest> {
+        Box::pin(async move { Ok(request) })
+    }
     fn success<'a>(
         &'a self,
         _context: &'a CallLifecycleContext,
@@ -80,7 +91,7 @@ impl CallLifecycleHooks<LiteLLMOcrRequest, LiteLLMOcrRequest, LiteLLMOcrResponse
         request: LiteLLMOcrRequest,
     ) -> Self::PreCallFuture<'a> {
         Box::pin(async move {
-            if !self.hooks.has_guardrails() {
+            if !self.hooks.intercepts_requests() {
                 return Ok(request);
             }
             let changed = self

litellm-rust/crates/core/src/ocr/lifecycle.rs

@@ -0,0 +1,640 @@
+use std::future::Future;
+use std::pin::Pin;
+use std::sync::Arc;
+
+use tokio::sync::{mpsc, oneshot};
+
+use super::handler::perform_ocr_request;
+use super::hooks::{
+    OcrDuringCallRequest, OcrHookFuture, OcrHooks, OcrLogFuture, OcrPostCallRequest,
+    OcrPreCallRequest,
+};
+use super::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrClient};
+use crate::AuthError;
+use crate::Error;
+use crate::auth::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle};
+use crate::call_lifecycle::host::{
+    HostCall, HostCallFuture, HostCallStep, HostFailure, HostLifecycle, HostPhase,
+};
+use crate::call_lifecycle::{CallLifecycleContext, CallLifecycleTiming};
+
+pub type NativeResult<T> = Result<NativeOutcome<T>, Error>;
+
+#[derive(Debug, PartialEq, Eq)]
+pub enum NativeOutcome<T> {
+    Completed(T),
+    Declined(OcrDecline),
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum OcrDecline {
+    ProviderWorkflow,
+    HostOperations,
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub struct OcrAdmission {
+    pub provider_workflow: bool,
+    pub host_operations: bool,
+    pub asynchronous: bool,
+}
+
+impl OcrAdmission {
+    pub const fn all() -> Self {
+        Self {
+            provider_workflow: true,
+            host_operations: true,
+            asynchronous: false,
+        }
+    }
+}
+
+#[derive(Clone, Debug)]
+pub enum OcrHostOperation {
+    ProjectRequest,
+    Lifecycle(HostPhase),
+    ConstructResponse(Arc<LiteLLMOcrResponse>),
+    MapFailure(Error),
+    Success {
+        context: CallLifecycleContext,
+        response: Arc<LiteLLMOcrResponse>,
+        timing: CallLifecycleTiming,
+    },
+    Failure {
+        context: CallLifecycleContext,
+        error: Error,
+        timing: CallLifecycleTiming,
+    },
+    AcquireAzureAdToken,
+    PreCall(OcrPreCallRequest),
+    DuringCall(OcrDuringCallRequest),
+    PostCall(OcrPostCallRequest),
+}
+
+impl OcrHostOperation {
+    pub const fn phase(&self) -> Option<HostPhase> {
+        match self {
+            Self::Lifecycle(phase) => Some(*phase),
+            Self::Success { .. } => Some(HostPhase::Success),
+            Self::Failure { .. } => Some(HostPhase::Failure),
+            _ => None,
+        }
+    }
+}
+
+pub enum OcrHostResult {
+    Request(Result<(Box<LiteLLMOcrRequest>, bool), Error>),
+    Lifecycle(Result<(), HostFailure>),
+    AzureAdToken(Result<ResolvedCredential, AuthError>),
+    PreCall(Result<OcrPreCallRequest, Error>),
+    DuringCall(Result<OcrDuringCallRequest, Error>),
+    PostCall(Result<OcrPostCallRequest, Error>),
+}
+
+pub type OcrCallStep = HostCallStep<OcrHostOperation, LiteLLMOcrResponse>;
+
+pub struct OcrCall {
+    lifecycle: HostLifecycle,
+    execution: OcrExecution,
+    response: Option<Arc<LiteLLMOcrResponse>>,
+    error: Option<Error>,
+    pending: bool,
+    completed: bool,
+    projecting: bool,
+}
+
+impl OcrCall {
+    pub fn admit(client: OcrClient, admission: OcrAdmission) -> NativeOutcome<Self> {
+        if !admission.provider_workflow {
+            return NativeOutcome::Declined(OcrDecline::ProviderWorkflow);
+        }
+        if !admission.host_operations {
+            return NativeOutcome::Declined(OcrDecline::HostOperations);
+        }
+        NativeOutcome::Completed(Self {
+            lifecycle: HostLifecycle::new(admission.asynchronous),
+            execution: OcrExecution::new(client),
+            response: None,
+            error: None,
+            pending: false,
+            completed: false,
+            projecting: false,
+        })
+    }
+
+    pub async fn resume(&mut self, result: Option<OcrHostResult>) -> Result<OcrCallStep, Error> {
+        if self.completed {
+            return Err(Error::InvalidRequest(
+                "OCR call cannot be resumed after completion".into(),
+            ));
+        }
+        if self.pending != result.is_some() {
+            return Err(Error::InvalidRequest(
+                "OCR host operation result does not match pending state".into(),
+            ));
+        }
+        match &result {
+            Some(OcrHostResult::Lifecycle(Ok(())))
+                if self.lifecycle.phase() == HostPhase::Execute =>
+            {
+                return Err(Error::InvalidRequest(
+                    "OCR provider operation requires a typed result".into(),
+                ));
+            }
+            Some(result)
+                if !matches!(result, OcrHostResult::Lifecycle(_))
+                    && self.lifecycle.phase() != HostPhase::Execute =>
+            {
+                return Err(Error::InvalidRequest(
+                    "unexpected OCR provider operation result".into(),
+                ));
+            }
+            _ => {}
+        }
+        self.pending = false;
+        let provider_result = match result {
+            Some(OcrHostResult::Request(result)) if self.projecting => {
+                self.projecting = false;
+                match result {
+                    Ok((request, azure_ad_token_provider)) => {
+                        self.execution.request = Some(*request);
+                        self.execution.azure_ad_token_provider = azure_ad_token_provider;
+                    }
+                    Err(error) => self.accept(Err(HostFailure::Error(error))),
+                }
+                None
+            }
+            Some(OcrHostResult::Request(_)) => {
+                return Err(Error::InvalidRequest(
+                    "unexpected OCR request projection".into(),
+                ));
+            }
+            Some(OcrHostResult::Lifecycle(result)) => {
+                self.accept(result);
+                None
+            }
+            result => result,
+        };
+        if self.lifecycle.phase() == HostPhase::Execute {
+            if self.execution.request.is_none()
+                && self.execution.execution.is_none()
+                && !self.execution.completed
+            {
+                self.projecting = true;
+                return Ok(self.host_step(OcrHostOperation::ProjectRequest));
+            }
+            match self.execution.resume(provider_result).await {
+                Ok(OcrCallStep::Host(operation)) => return Ok(self.host_step(operation)),
+                Ok(OcrCallStep::Complete(response)) => {
+                    self.response = Some(Arc::new(response));
+                    self.accept(Ok(()));
+                }
+                Err(error) => self.accept(Err(HostFailure::Error(error))),
+            }
+        }
+        if self.error.is_some() {
+            self.execution.stop().await;
+        }
+        let operation = match self.lifecycle.phase() {
+            HostPhase::Complete => {
+                self.completed = true;
+                return match self.error.take() {
+                    Some(error) => Err(error),
+                    None => self
+                        .response
+                        .take()
+                        .map(Arc::unwrap_or_clone)
+                        .map(OcrCallStep::Complete)
+                        .ok_or_else(|| {
+                            Error::InvalidRequest("OCR completed without a response".into())
+                        }),
+                };
+            }
+            HostPhase::ConstructResponse => OcrHostOperation::ConstructResponse(
+                self.response
+                    .as_ref()
+                    .ok_or_else(|| Error::InvalidRequest("missing OCR response".into()))?
+                    .clone(),
+            ),
+            HostPhase::MapFailure => OcrHostOperation::MapFailure(
+                self.error
+                    .as_ref()
+                    .ok_or_else(|| Error::InvalidRequest("missing OCR failure".into()))?
+                    .clone(),
+            ),
+            HostPhase::Success | HostPhase::Failure => {
+                let snapshot = self
+                    .execution
+                    .terminal
+                    .lock()
+                    .unwrap_or_else(|error| error.into_inner())
+                    .clone();
+                match (self.lifecycle.phase(), snapshot) {
+                    (HostPhase::Success, Some((context, timing))) => OcrHostOperation::Success {
+                        context,
+                        response: self
+                            .response
+                            .as_ref()
+                            .ok_or_else(|| Error::InvalidRequest("missing OCR response".into()))?
+                            .clone(),
+                        timing,
+                    },
+                    (HostPhase::Failure, Some((context, timing))) => OcrHostOperation::Failure {
+                        context,
+                        error: self
+                            .error
+                            .as_ref()
+                            .ok_or_else(|| Error::InvalidRequest("missing OCR failure".into()))?
+                            .clone(),
+                        timing,
+                    },
+                    (phase, _) => OcrHostOperation::Lifecycle(phase),
+                }
+            }
+            phase => OcrHostOperation::Lifecycle(phase),
+        };
+        Ok(self.host_step(operation))
+    }
+
+    fn accept(&mut self, result: Result<(), HostFailure>) {
+        let cancelled = matches!(&result, Err(HostFailure::Cancelled(_)));
+        if let Some(error) = self.lifecycle.accept(result) {
+            if cancelled {
+                self.error = Some(error);
+            } else {
+                self.error.get_or_insert(error);
+            }
+            self.execution.cancel();
+        }
+    }
+
+    pub async fn interrupt(&mut self, failure: HostFailure) -> Result<OcrCallStep, Error> {
+        if self.completed {
+            return Err(Error::InvalidRequest(
+                "OCR call cannot be interrupted after completion".into(),
+            ));
+        }
+        self.pending = false;
+        self.accept(Err(failure));
+        self.resume(None).await
+    }
+
+    fn host_step(&mut self, operation: OcrHostOperation) -> OcrCallStep {
+        self.pending = true;
+        OcrCallStep::Host(operation)
+    }
+}
+
+impl HostCall for OcrCall {
+    type Operation = OcrHostOperation;
+    type Result = OcrHostResult;
+    type Complete = LiteLLMOcrResponse;
+
+    fn resume(
+        &mut self,
+        result: Option<Self::Result>,
+    ) -> HostCallFuture<'_, Self::Operation, Self::Complete> {
+        Box::pin(OcrCall::resume(self, result))
+    }
+
+    fn interrupt(
+        &mut self,
+        failure: HostFailure,
+    ) -> HostCallFuture<'_, Self::Operation, Self::Complete> {
+        Box::pin(OcrCall::interrupt(self, failure))
+    }
+}
+
+struct PendingOperation {
+    operation: OcrHostOperation,
+    result: oneshot::Sender<OcrHostResult>,
+}
+
+struct OcrExecution {
+    client: Option<OcrClient>,
+    request: Option<LiteLLMOcrRequest>,
+    operations_tx: mpsc::UnboundedSender<PendingOperation>,
+    operations_rx: mpsc::UnboundedReceiver<PendingOperation>,
+    pending_result: Option<oneshot::Sender<OcrHostResult>>,
+    execution: Option<tokio::task::JoinHandle<Result<LiteLLMOcrResponse, Error>>>,
+    completed: bool,
+    azure_ad_token_provider: bool,
+    terminal: Arc<std::sync::Mutex<Option<(CallLifecycleContext, CallLifecycleTiming)>>>,
+}
+
+impl OcrExecution {
+    fn new(client: OcrClient) -> Self {
+        let (operations_tx, operations_rx) = mpsc::unbounded_channel();
+        Self {
+            client: Some(client),
+            request: None,
+            operations_tx,
+            operations_rx,
+            pending_result: None,
+            execution: None,
+            completed: false,
+            azure_ad_token_provider: false,
+            terminal: Arc::default(),
+        }
+    }
+
+    pub async fn resume(&mut self, result: Option<OcrHostResult>) -> Result<OcrCallStep, Error> {
+        if self.completed {
+            return Err(Error::InvalidRequest(
+                "OCR call cannot be resumed after completion".into(),
+            ));
+        }
+        match (self.pending_result.take(), result) {
+            (Some(sender), Some(result)) => sender
+                .send(result)
+                .map_err(|_| Error::InvalidRequest("OCR host operation was abandoned".into()))?,
+            (None, None) if self.execution.is_none() => self.start(),
+            (Some(sender), None) => {
+                self.pending_result = Some(sender);
+                return Err(Error::InvalidRequest(
+                    "OCR host operation result is required".into(),
+                ));
+            }
+            (None, Some(_)) => {
+                return Err(Error::InvalidRequest(
+                    "unexpected OCR host operation result".into(),
+                ));
+            }
+            (None, None) => {}
+        }
+
+        let execution = self.execution.as_mut().ok_or_else(|| {
+            Error::InvalidRequest("OCR call cannot be resumed after completion".into())
+        })?;
+        tokio::select! {
+            operation = self.operations_rx.recv() => {
+                let operation = operation.ok_or_else(|| Error::InvalidRequest("OCR operation channel closed".into()))?;
+                self.pending_result = Some(operation.result);
+                Ok(OcrCallStep::Host(operation.operation))
+            }
+            result = execution => {
+                self.execution = None;
+                self.completed = true;
+                result
+                    .map_err(|error| Error::Network(format!("OCR execution task failed: {error}")))?
+                    .map(OcrCallStep::Complete)
+            }
+        }
+    }
+
+    fn start(&mut self) {
+        let client = self.client.take().expect("admitted OCR call has a client");
+        let mut request = self
+            .request
+            .take()
+            .expect("admitted OCR call has a request");
+        let intercepts_requests = request.hooks.intercepts_requests();
+        if self.azure_ad_token_provider {
+            request.azure_ad_token_provider = Some(TokenProviderHandle::new(Arc::new(
+                OcrAzureAdTokenProvider {
+                    operations: self.operations_tx.clone(),
+                },
+            )));
+        }
+        request.hooks = Arc::new(ProtocolHooks {
+            operations: self.operations_tx.clone(),
+            intercepts_requests,
+            terminal: self.terminal.clone(),
+        });
+        self.execution = Some(tokio::spawn(async move {
+            perform_ocr_request(&client, request).await
+        }));
+    }
+
+    fn cancel(&mut self) {
+        self.pending_result = None;
+        if let Some(execution) = &self.execution {
+            execution.abort();
+        }
+    }
+
+    async fn stop(&mut self) {
+        self.cancel();
+        if let Some(execution) = self.execution.as_mut() {
+            let _ = execution.await;
+        }
+        self.execution = None;
+    }
+}
+
+impl Drop for OcrExecution {
+    fn drop(&mut self) {
+        if let Some(execution) = &self.execution {
+            execution.abort();
+        }
+    }
+}
+
+struct ProtocolHooks {
+    operations: mpsc::UnboundedSender<PendingOperation>,
+    intercepts_requests: bool,
+    terminal: Arc<std::sync::Mutex<Option<(CallLifecycleContext, CallLifecycleTiming)>>>,
+}
+
+#[derive(Debug)]
+struct OcrAzureAdTokenProvider {
+    operations: mpsc::UnboundedSender<PendingOperation>,
+}
+
+impl TokenProvider for OcrAzureAdTokenProvider {
+    fn acquire(&self) -> TokenFuture<'_> {
+        Box::pin(async move {
+            let (result, receiver) = oneshot::channel();
+            self.operations
+                .send(PendingOperation {
+                    operation: OcrHostOperation::AcquireAzureAdToken,
+                    result,
+                })
+                .map_err(|_| {
+                    AuthError::AzureTokenAcquisition("OCR host driver was abandoned".into())
+                })?;
+            match receiver.await.map_err(|_| {
+                AuthError::AzureTokenAcquisition(
+                    "OCR token provider operation was abandoned".into(),
+                )
+            })? {
+                OcrHostResult::AzureAdToken(result) => result,
+                _ => Err(AuthError::AzureTokenAcquisition(
+                    "invalid OCR token provider host result".into(),
+                )),
+            }
+        })
+    }
+}
+
+impl ProtocolHooks {
+    async fn invoke(&self, operation: OcrHostOperation) -> Result<OcrHostResult, Error> {
+        let (result, receiver) = oneshot::channel();
+        self.operations
+            .send(PendingOperation { operation, result })
+            .map_err(|_| Error::InvalidRequest("OCR host driver was abandoned".into()))?;
+        receiver
+            .await
+            .map_err(|_| Error::InvalidRequest("OCR host operation was abandoned".into()))
+    }
+}
+
+impl OcrHooks for ProtocolHooks {
+    fn intercepts_requests(&self) -> bool {
+        self.intercepts_requests
+    }
+
+    fn pre_call(&self, request: OcrPreCallRequest) -> OcrHookFuture<'_, OcrPreCallRequest> {
+        Box::pin(async move {
+            match self.invoke(OcrHostOperation::PreCall(request)).await? {
+                OcrHostResult::PreCall(result) => result,
+                _ => Err(Error::InvalidRequest(
+                    "invalid OCR pre-call host result".into(),
+                )),
+            }
+        })
+    }
+
+    fn during_call(
+        &self,
+        request: OcrDuringCallRequest,
+    ) -> OcrHookFuture<'_, OcrDuringCallRequest> {
+        Box::pin(async move {
+            match self.invoke(OcrHostOperation::DuringCall(request)).await? {
+                OcrHostResult::DuringCall(result) => result,
+                _ => Err(Error::InvalidRequest(
+                    "invalid OCR during-call host result".into(),
+                )),
+            }
+        })
+    }
+
+    fn post_call(&self, request: OcrPostCallRequest) -> OcrHookFuture<'_, OcrPostCallRequest> {
+        Box::pin(async move {
+            match self.invoke(OcrHostOperation::PostCall(request)).await? {
+                OcrHostResult::PostCall(result) => result,
+                _ => Err(Error::InvalidRequest(
+                    "invalid OCR post-call host result".into(),
+                )),
+            }
+        })
+    }
+
+    fn success<'a>(
+        &'a self,
+        context: &'a CallLifecycleContext,
+        _response: &'a LiteLLMOcrResponse,
+        timing: &'a CallLifecycleTiming,
+    ) -> OcrLogFuture<'a> {
+        Box::pin(async move {
+            *self
+                .terminal
+                .lock()
+                .unwrap_or_else(|error| error.into_inner()) =
+                Some((context.clone(), timing.clone()));
+        })
+    }
+
+    fn failure<'a>(
+        &'a self,
+        context: &'a CallLifecycleContext,
+        _error: &'a Error,
+        timing: &'a CallLifecycleTiming,
+    ) -> OcrLogFuture<'a> {
+        Box::pin(async move {
+            *self
+                .terminal
+                .lock()
+                .unwrap_or_else(|error| error.into_inner()) =
+                Some((context.clone(), timing.clone()));
+        })
+    }
+}
+
+pub type OcrHostFuture<'a> = Pin<Box<dyn Future<Output = OcrHostResult> + Send + 'a>>;
+
+pub trait OcrHost: Send + Sync {
+    fn invoke(&self, operation: OcrHostOperation) -> OcrHostFuture<'_>;
+}
+
+pub struct NoopOcrHost;
+
+impl OcrHost for NoopOcrHost {
+    fn invoke(&self, operation: OcrHostOperation) -> OcrHostFuture<'_> {
+        Box::pin(async move {
+            match operation {
+                OcrHostOperation::ProjectRequest => OcrHostResult::Request(Err(
+                    Error::InvalidRequest("OCR host has no request projection".into()),
+                )),
+                OcrHostOperation::Lifecycle(_)
+                | OcrHostOperation::ConstructResponse(_)
+                | OcrHostOperation::MapFailure(_)
+                | OcrHostOperation::Success { .. }
+                | OcrHostOperation::Failure { .. } => OcrHostResult::Lifecycle(Ok(())),
+                OcrHostOperation::AcquireAzureAdToken => {
+                    OcrHostResult::AzureAdToken(Err(AuthError::AzureTokenAcquisition(
+                        "OCR host has no Azure AD token provider".into(),
+                    )))
+                }
+                OcrHostOperation::PreCall(request) => OcrHostResult::PreCall(Ok(request)),
+                OcrHostOperation::DuringCall(request) => OcrHostResult::DuringCall(Ok(request)),
+                OcrHostOperation::PostCall(request) => OcrHostResult::PostCall(Ok(request)),
+            }
+        })
+    }
+}
+
+pub struct OcrHookHost {
+    hooks: Arc<dyn OcrHooks>,
+}
+
+impl OcrHookHost {
+    pub fn new(hooks: Arc<dyn OcrHooks>) -> Self {
+        Self { hooks }
+    }
+}
+
+impl OcrHost for OcrHookHost {
+    fn invoke(&self, operation: OcrHostOperation) -> OcrHostFuture<'_> {
+        Box::pin(async move {
+            match operation {
+                OcrHostOperation::ProjectRequest => OcrHostResult::Request(Err(
+                    Error::InvalidRequest("OCR hook host has no request projection".into()),
+                )),
+                OcrHostOperation::Success {
+                    context,
+                    response,
+                    timing,
+                } => {
+                    self.hooks.success(&context, &response, &timing).await;
+                    OcrHostResult::Lifecycle(Ok(()))
+                }
+                OcrHostOperation::Failure {
+                    context,
+                    error,
+                    timing,
+                } => {
+                    self.hooks.failure(&context, &error, &timing).await;
+                    OcrHostResult::Lifecycle(Ok(()))
+                }
+                OcrHostOperation::Lifecycle(_)
+                | OcrHostOperation::ConstructResponse(_)
+                | OcrHostOperation::MapFailure(_) => OcrHostResult::Lifecycle(Ok(())),
+                OcrHostOperation::AcquireAzureAdToken => {
+                    OcrHostResult::AzureAdToken(Err(AuthError::AzureTokenAcquisition(
+                        "OCR hook host has no Azure AD token provider".into(),
+                    )))
+                }
+                OcrHostOperation::PreCall(request) => {
+                    OcrHostResult::PreCall(self.hooks.pre_call(request).await)
+                }
+                OcrHostOperation::DuringCall(request) => {
+                    OcrHostResult::DuringCall(self.hooks.during_call(request).await)
+                }
+                OcrHostOperation::PostCall(request) => {
+                    OcrHostResult::PostCall(self.hooks.post_call(request).await)
+                }
+            }
+        })
+    }
+}

litellm-rust/crates/core/src/ocr/mod.rs

@@ -5,12 +5,18 @@ mod document;
 pub mod error;
 mod handler;
 pub mod hooks;
+mod lifecycle;
 mod prepare;
 mod registry;
 pub mod types;
 pub mod wire;
 
 pub use client::{OcrClient, ocr};
+pub use document::{encode_file_document, mime_type_for_name, upload_mime_type};
+pub use lifecycle::{
+    NativeOutcome, NativeResult, NoopOcrHost, OcrAdmission, OcrCall, OcrCallStep, OcrDecline,
+    OcrHookHost, OcrHost, OcrHostOperation, OcrHostResult,
+};
 pub use types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection, OcrDocument};
 
 #[cfg(test)]

litellm-rust/crates/core/src/ocr/prepare.rs

@@ -62,34 +62,48 @@ pub(crate) async fn transform_request_body<B>(
     request: &LiteLLMOcrRequest,
     url: &str,
     headers: &[(String, String)],
+    retains_document: bool,
     body: B,
     validate: impl FnOnce(&B) -> Result<(), OcrRequestError>,
 ) -> Result<reqwest::Request, OcrError>
 where
     B: Serialize + DeserializeOwned,
 {
-    let body = if request.hooks.has_guardrails() {
+    let (body, headers) = if request.hooks.intercepts_requests() {
+        let body = serde_json::to_value(body).map_err(|_| OcrRequestError::RequestField {
+            path: "body".into(),
+        })?;
+        let retained_fields = request
+            .optional_params
+            .keys()
+            .filter(|name| body.get(*name).is_some())
+            .cloned()
+            .chain(retains_document.then(|| "document".to_string()))
+            .collect();
         let changed = request
             .hooks
             .during_call(OcrDuringCallRequest {
                 model: request.model.clone(),
                 custom_llm_provider: request.adapter.provider().as_str().into(),
                 url: url.into(),
-                body: serde_json::to_value(body).map_err(|_| OcrRequestError::RequestField {
-                    path: "body".into(),
-                })?,
+                headers: headers.to_vec(),
+                body,
+                retained_fields,
             })
             .await?;
         let body = OcrWireBody::<B>::decode(changed.body)?;
         validate(&body.body)?;
-        body
+        (body, changed.headers)
     } else {
-        OcrWireBody {
-            body,
-            extra: Map::new(),
-        }
+        (
+            OcrWireBody {
+                body,
+                extra: Map::new(),
+            },
+            headers.to_vec(),
+        )
     };
-    build_http_request(client, request, url, headers, &body)
+    build_http_request(client, request, url, &headers, &body)
 }
 
 pub(crate) fn build_http_request<B: Serialize>(
@@ -113,9 +127,10 @@ pub(crate) fn build_http_request<B: Serialize>(
 pub(crate) async fn guardrail_document(
     request: &LiteLLMOcrRequest,
     url: &str,
-) -> Result<OcrDocument, OcrError> {
-    if !request.hooks.has_guardrails() {
-        return Ok(request.document.clone());
+    headers: &[(String, String)],
+) -> Result<(OcrDocument, Vec<(String, String)>), OcrError> {
+    if !request.hooks.intercepts_requests() {
+        return Ok((request.document.clone(), headers.to_vec()));
     }
     let changed = request
         .hooks
@@ -123,14 +138,17 @@ pub(crate) async fn guardrail_document(
             model: request.model.clone(),
             custom_llm_provider: request.adapter.provider().as_str().into(),
             url: url.into(),
+            headers: headers.to_vec(),
             body: serde_json::to_value(&request.document).map_err(|_| {
                 OcrRequestError::RequestField {
                     path: "document".into(),
                 }
             })?,
+            retained_fields: Vec::new(),
         })
         .await?;
-    super::wire::decode_request_value(changed.body, "guardrail.document").map_err(OcrError::from)
+    let document = super::wire::decode_request_value(changed.body, "guardrail.document")?;
+    Ok((document, changed.headers))
 }
 
 #[derive(Serialize)]

litellm-rust/crates/core/src/ocr/registry.rs

@@ -23,6 +23,7 @@ super::adapters::for_each_ocr_adapter!(define_adapter_types);
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub(crate) enum OcrProvider {
+    Cohere,
     Mistral,
     AzureAi,
     Reducto,
@@ -32,6 +33,7 @@ pub(crate) enum OcrProvider {
 impl OcrProvider {
     pub(crate) const fn as_str(self) -> &'static str {
         match self {
+            Self::Cohere => "cohere",
             Self::Mistral => "mistral",
             Self::AzureAi => "azure_ai",
             Self::Reducto => "reducto",
@@ -50,6 +52,7 @@ pub(crate) fn resolve_wire_adapter(
             custom_llm_provider: OcrProvider::Mistral.as_str(),
         });
     let typed_provider = match provider.custom_llm_provider {
+        "cohere" => OcrProvider::Cohere,
         "mistral" => OcrProvider::Mistral,
         "azure_ai" => OcrProvider::AzureAi,
         "reducto" => OcrProvider::Reducto,
@@ -57,10 +60,17 @@ pub(crate) fn resolve_wire_adapter(
         value => return Err(Error::InvalidProvider(value.to_string())),
     };
     let adapter = match typed_provider {
+        OcrProvider::Cohere => OcrAdapterKind::Cohere,
         OcrProvider::Mistral => OcrAdapterKind::Mistral,
         OcrProvider::AzureAi if is_document_intelligence_model(provider.model) => {
             OcrAdapterKind::AzureDocumentIntelligence
         }
+        OcrProvider::AzureAi
+            if provider.model.to_ascii_lowercase().contains("cohere")
+                && provider.model.to_ascii_lowercase().contains("parse") =>
+        {
+            OcrAdapterKind::AzureCohere
+        }
         OcrProvider::AzureAi => OcrAdapterKind::AzureMistral,
         OcrProvider::Reducto if provider.model.eq_ignore_ascii_case("parse-legacy") => {
             OcrAdapterKind::ReductoLegacy
@@ -68,12 +78,7 @@ pub(crate) fn resolve_wire_adapter(
         OcrProvider::Reducto if provider.model.eq_ignore_ascii_case("parse-v3") => {
             OcrAdapterKind::ReductoV3
         }
-        OcrProvider::Reducto => {
-            return Err(Error::InvalidRequest(format!(
-                "unsupported Reducto OCR model: {}",
-                provider.model
-            )));
-        }
+        OcrProvider::Reducto => OcrAdapterKind::ReductoV3,
         OcrProvider::VertexAi if provider.model.to_ascii_lowercase().contains("deepseek") => {
             OcrAdapterKind::VertexDeepSeek
         }
@@ -107,11 +112,10 @@ mod tests {
     }
 
     #[test]
-    fn unknown_reducto_models_are_rejected() {
-        assert!(matches!(
-            resolve_wire_adapter("reducto/future-parse-model", None),
-            Err(Error::InvalidRequest(_))
-        ));
+    fn unknown_reducto_models_use_the_current_protocol() {
+        let (model, adapter) = resolve_wire_adapter("reducto/future-parse-model", None).unwrap();
+        assert_eq!(model, "future-parse-model");
+        assert_eq!(adapter, OcrAdapterKind::ReductoV3);
     }
 
     #[test]

litellm-rust/crates/core/src/ocr/types.rs

@@ -8,7 +8,7 @@ use serde_json::{Map, Value};
 use super::hooks::{NoopOcrHooks, OcrHooks};
 use super::registry::{OcrAdapterKind, resolve_wire_adapter};
 use crate::Error;
-use crate::auth::InputSource;
+use crate::auth::{InputSource, TokenProviderHandle};
 use crate::constants::OCR_HTTP_TIMEOUT_SECS;
 
 #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
@@ -68,6 +68,7 @@ pub struct OcrConnection {
     pub extra_headers_source: InputSource,
     pub timeout: Duration,
     pub max_download_bytes: u64,
+    pub max_response_bytes: usize,
     pub poll_timeout: Duration,
 }
 
@@ -82,6 +83,7 @@ impl Default for OcrConnection {
             extra_headers_source: InputSource::Deployment,
             timeout: Duration::from_secs(OCR_HTTP_TIMEOUT_SECS),
             max_download_bytes: crate::constants::OCR_DOWNLOAD_MAX_BYTES,
+            max_response_bytes: crate::constants::OCR_RESPONSE_MAX_BYTES,
             poll_timeout: Duration::from_secs(crate::constants::OCR_POLL_TIMEOUT_SECS),
         }
     }
@@ -95,6 +97,7 @@ pub struct LiteLLMOcrRequest {
     pub litellm_call_id: Option<String>,
     pub optional_params: Map<String, Value>,
     pub input_sources: BTreeMap<String, InputSource>,
+    pub azure_ad_token_provider: Option<TokenProviderHandle>,
     pub(crate) adapter: OcrAdapterKind,
 }
 
@@ -115,6 +118,7 @@ impl LiteLLMOcrRequest {
             litellm_call_id: None,
             optional_params,
             input_sources: BTreeMap::new(),
+            azure_ad_token_provider: None,
             adapter: adapter_kind,
         })
     }
@@ -132,6 +136,10 @@ impl LiteLLMOcrRequest {
             .map(|format| format.unwrap_or_default())
     }
 
+    pub fn provider_name(&self) -> &'static str {
+        self.adapter.provider().as_str()
+    }
+
     pub fn with_host_hooks(
         self,
         hooks: Arc<dyn OcrHooks>,
@@ -169,6 +177,47 @@ mod tests {
     use super::*;
     use serde_json::json;
 
+    #[test]
+    fn document_variants_preserve_provider_fields_when_rewriting_sources() {
+        for (value, original, replacement, expected) in [
+            (
+                json!({
+                    "type":"document_url",
+                    "document_url":"https://example.com/input.pdf",
+                    "document_name":"input.pdf"
+                }),
+                "https://example.com/input.pdf",
+                "data:application/pdf;base64,AA==",
+                json!({
+                    "type":"document_url",
+                    "document_url":"data:application/pdf;base64,AA==",
+                    "document_name":"input.pdf"
+                }),
+            ),
+            (
+                json!({
+                    "type":"image_url",
+                    "image_url":"https://example.com/input.png",
+                    "detail":"high"
+                }),
+                "https://example.com/input.png",
+                "data:image/png;base64,AA==",
+                json!({
+                    "type":"image_url",
+                    "image_url":"data:image/png;base64,AA==",
+                    "detail":"high"
+                }),
+            ),
+        ] {
+            let document: OcrDocument = serde_json::from_value(value).unwrap();
+            assert_eq!(document.source(), original);
+            assert_eq!(
+                serde_json::to_value(document.with_source(replacement.into())).unwrap(),
+                expected
+            );
+        }
+    }
+
     #[test]
     fn response_serialization_flattens_extra_fields_and_omits_absent_native_response() {
         let response = LiteLLMOcrResponse {

litellm-rust/crates/core/src/ocr/wire.rs

@@ -3,7 +3,6 @@ use crate::ocr::error::OcrResponseError;
 use std::collections::BTreeMap;
 use std::time::Duration;
 
-use super::hooks::{OcrDuringCallRequest, OcrPreCallRequest};
 use super::types::{LiteLLMOcrRequest, OcrConnection, OcrDocument};
 use crate::Error;
 use crate::auth::InputSource;
@@ -13,10 +12,58 @@ use serde::{
 };
 use serde_json::{Map, Value};
 
+const COMMON_OPTION_FIELDS: &[&str] = &["req_format", "extra_body", "max_response_bytes"];
+const MISTRAL_OPTION_FIELDS: &[&str] = &[
+    "pages",
+    "include_image_base64",
+    "image_limit",
+    "image_min_size",
+    "bbox_annotation_format",
+    "document_annotation_format",
+    "document_annotation_prompt",
+    "extract_header",
+    "extract_footer",
+    "table_format",
+    "confidence_scores_granularity",
+    "include_blocks",
+    "id",
+];
+const DEEPSEEK_OPTION_FIELDS: &[&str] =
+    &["stream", "temperature", "max_tokens", "top_p", "n", "stop"];
+const DOCUMENT_INTELLIGENCE_OPTION_FIELDS: &[&str] = &["pages", "features"];
+const REDUCTO_V3_OPTION_FIELDS: &[&str] = &["formatting", "retrieval", "settings"];
+const REDUCTO_LEGACY_OPTION_FIELDS: &[&str] = &["enhance"];
+const AZURE_AUTH_OPTION_FIELDS: &[&str] = &[
+    "azure_ad_token",
+    "tenant_id",
+    "client_id",
+    "client_secret",
+    "azure_scope",
+    "azure_authority_host",
+    "azure_credential",
+    "azure_federated_token_file",
+    "enable_azure_ad_token_refresh",
+];
+const VERTEX_AUTH_OPTION_FIELDS: &[&str] = &[
+    "vertex_credentials",
+    "vertex_ai_credentials",
+    "vertex_project",
+    "vertex_ai_project",
+    "vertex_location",
+    "vertex_ai_location",
+];
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub struct OptionalParamSpec {
+    pub name: &'static str,
+    pub secret: bool,
+}
+
 #[derive(Debug)]
 pub struct DecodedOcrResponse<T> {
     pub data: T,
     pub native: Option<Value>,
+    pub text: String,
 }
 
 #[derive(Deserialize)]
@@ -39,11 +86,65 @@ pub fn is_supported_request(model: &str, custom_llm_provider: Option<&str>) -> b
     super::registry::resolve_wire_adapter(model, custom_llm_provider).is_ok()
 }
 
+pub fn consumed_optional_param_names(
+    model: &str,
+    custom_llm_provider: Option<&str>,
+) -> Result<Vec<&'static str>, Error> {
+    use super::registry::OcrAdapterKind;
+
+    let (_, adapter) = super::registry::resolve_wire_adapter(model, custom_llm_provider)?;
+    let provider_fields: &[&str] = match adapter {
+        OcrAdapterKind::Cohere | OcrAdapterKind::AzureCohere => &["output_format"],
+        OcrAdapterKind::Mistral | OcrAdapterKind::AzureMistral | OcrAdapterKind::VertexMistral => {
+            MISTRAL_OPTION_FIELDS
+        }
+        OcrAdapterKind::AzureDocumentIntelligence => DOCUMENT_INTELLIGENCE_OPTION_FIELDS,
+        OcrAdapterKind::ReductoV3 => REDUCTO_V3_OPTION_FIELDS,
+        OcrAdapterKind::ReductoLegacy => REDUCTO_LEGACY_OPTION_FIELDS,
+        OcrAdapterKind::VertexDeepSeek => DEEPSEEK_OPTION_FIELDS,
+    };
+    let auth_fields: &[&str] = match adapter {
+        OcrAdapterKind::AzureMistral
+        | OcrAdapterKind::AzureDocumentIntelligence
+        | OcrAdapterKind::AzureCohere => AZURE_AUTH_OPTION_FIELDS,
+        OcrAdapterKind::VertexMistral | OcrAdapterKind::VertexDeepSeek => VERTEX_AUTH_OPTION_FIELDS,
+        _ => &[],
+    };
+    Ok(COMMON_OPTION_FIELDS
+        .iter()
+        .chain(provider_fields)
+        .chain(auth_fields)
+        .copied()
+        .collect())
+}
+
+pub fn consumed_optional_params(
+    model: &str,
+    custom_llm_provider: Option<&str>,
+) -> Result<Vec<OptionalParamSpec>, Error> {
+    consumed_optional_param_names(model, custom_llm_provider).map(|names| {
+        names
+            .into_iter()
+            .map(|name| OptionalParamSpec {
+                name,
+                secret: matches!(
+                    name,
+                    "azure_ad_token"
+                        | "client_secret"
+                        | "azure_federated_token_file"
+                        | "vertex_credentials"
+                        | "vertex_ai_credentials"
+                ),
+            })
+            .collect()
+    })
+}
+
 pub fn decode_request(wire: OcrWireRequest) -> Result<LiteLLMOcrRequest, Error> {
     let api_key_source = source_for(&wire.input_sources, "api_key");
     let api_base_source = source_for(&wire.input_sources, "api_base");
     let extra_headers_source = source_for(&wire.input_sources, "extra_headers");
-    let document = decode_request_value(wire.document, "document")?;
+    let document = decode_document(wire.document)?;
     let headers = wire
         .extra_headers
         .unwrap_or_default()
@@ -66,11 +167,28 @@ pub fn decode_request(wire: OcrWireRequest) -> Result<LiteLLMOcrRequest, Error>
         })
         .transpose()?;
     let defaults = OcrConnection::default();
+    let max_response_bytes = wire
+        .optional_params
+        .get("max_response_bytes")
+        .map(|value| {
+            value
+                .as_u64()
+                .and_then(|value| usize::try_from(value).ok())
+                .filter(|value| *value > 0 && *value <= defaults.max_response_bytes)
+                .ok_or_else(|| OcrRequestError::RequestField {
+                    path: "max_response_bytes".into(),
+                })
+        })
+        .transpose()?
+        .unwrap_or(defaults.max_response_bytes);
     let request = LiteLLMOcrRequest::new(
         wire.model,
         document,
         wire.custom_llm_provider.as_deref(),
-        wire.optional_params,
+        wire.optional_params
+            .into_iter()
+            .filter(|(name, _)| name != "max_response_bytes")
+            .collect(),
     )?;
     let connection = OcrConnection {
         api_key: nonblank(wire.api_key),
@@ -81,6 +199,7 @@ pub fn decode_request(wire: OcrWireRequest) -> Result<LiteLLMOcrRequest, Error>
         extra_headers_source,
         timeout: timeout.unwrap_or(defaults.timeout),
         max_download_bytes: defaults.max_download_bytes,
+        max_response_bytes,
         poll_timeout: defaults.poll_timeout,
     };
     Ok(LiteLLMOcrRequest {
@@ -90,6 +209,16 @@ pub fn decode_request(wire: OcrWireRequest) -> Result<LiteLLMOcrRequest, Error>
     })
 }
 
+fn decode_document(value: Value) -> Result<OcrDocument, OcrRequestError> {
+    let kind = value.get("type").and_then(Value::as_str);
+    let missing_url = matches!(kind, Some("document_url")) && value.get("document_url").is_none()
+        || matches!(kind, Some("image_url")) && value.get("image_url").is_none();
+    if missing_url {
+        return Err(OcrRequestError::MissingDocumentUrl);
+    }
+    decode_request_value(value, "document")
+}
+
 fn source_for(sources: &BTreeMap<String, InputSource>, name: &str) -> InputSource {
     sources.get(name).copied().unwrap_or_default()
 }
@@ -134,38 +263,81 @@ pub fn decode_response<T: DeserializeOwned>(
     } else {
         None
     };
-    Ok(DecodedOcrResponse { data, native })
-}
-
-pub fn decode_pre_call_result(
-    original: OcrPreCallRequest,
-    value: Value,
-) -> Result<OcrPreCallRequest, OcrRequestError> {
-    #[derive(Deserialize)]
-    struct Changed {
-        document: OcrDocument,
-        #[serde(default)]
-        optional_params: Map<String, Value>,
-    }
-    let changed: Changed = decode_request_value(value, "guardrail")?;
-    Ok(OcrPreCallRequest {
-        document: changed.document,
-        optional_params: Value::Object(changed.optional_params),
-        ..original
+    Ok(DecodedOcrResponse {
+        data,
+        native,
+        text: String::from_utf8_lossy(bytes).into_owned(),
     })
 }
 
-pub fn decode_during_call_result(
-    original: OcrDuringCallRequest,
-    value: Value,
-) -> Result<OcrDuringCallRequest, OcrRequestError> {
-    #[derive(Deserialize)]
-    struct Changed {
-        body: Value,
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn option_projection_is_provider_specific_and_excludes_opaque_fields() {
+        let mistral = consumed_optional_param_names("mistral/model", None).unwrap();
+        assert!(mistral.contains(&"pages"));
+        assert!(mistral.contains(&"req_format"));
+        assert!(!mistral.contains(&"vertex_project"));
+        assert!(!mistral.contains(&"opaque_extension"));
+
+        let vertex = consumed_optional_param_names("vertex_ai/deepseek-ocr", None).unwrap();
+        assert!(vertex.contains(&"temperature"));
+        assert!(vertex.contains(&"vertex_credentials"));
+        assert!(!vertex.contains(&"pages"));
+    }
+
+    #[test]
+    fn optional_param_metadata_marks_only_credentials_as_secret() {
+        let azure = consumed_optional_params("model", Some("azure_ai")).unwrap();
+        assert!(
+            azure
+                .iter()
+                .any(|spec| spec.name == "client_secret" && spec.secret)
+        );
+        assert!(
+            azure
+                .iter()
+                .any(|spec| spec.name == "tenant_id" && !spec.secret)
+        );
+        let vertex = consumed_optional_params("deepseek-ocr", Some("vertex_ai")).unwrap();
+        assert!(
+            vertex
+                .iter()
+                .any(|spec| spec.name == "vertex_credentials" && spec.secret)
+        );
+        assert!(
+            vertex
+                .iter()
+                .any(|spec| spec.name == "vertex_project" && !spec.secret)
+        );
+    }
+
+    #[test]
+    fn activation_includes_migrated_providers() {
+        assert!(is_supported_request("model", Some("mistral")));
+        assert!(is_supported_request("pixtral-12b", Some("azure_ai")));
+        assert!(is_supported_request(
+            "documentintelligence/prebuilt-read",
+            Some("azure_ai")
+        ));
+        assert!(is_supported_request("parse-v3", Some("reducto")));
+        assert!(is_supported_request("parse-legacy", Some("reducto")));
+        assert!(is_supported_request("mistral-ocr", Some("vertex_ai")));
+        assert!(is_supported_request("deepseek-ocr", Some("vertex_ai")));
+    }
+
+    #[test]
+    fn missing_document_source_has_a_typed_public_error() {
+        for document in [
+            serde_json::json!({"type": "document_url"}),
+            serde_json::json!({"type": "image_url"}),
+        ] {
+            assert_eq!(
+                decode_document(document),
+                Err(OcrRequestError::MissingDocumentUrl)
+            );
+        }
     }
-    let changed: Changed = decode_request_value(value, "guardrail")?;
-    Ok(OcrDuringCallRequest {
-        body: changed.body,
-        ..original
-    })
 }

litellm-rust/crates/core/src/responses/websocket.rs

@@ -1,3 +1,21 @@
+use std::collections::HashMap;
+use std::io;
+use std::sync::{Arc, OnceLock};
+use std::time::Duration;
+
+use futures_util::{SinkExt, StreamExt};
+use rustls::{ClientConfig, RootCertStore};
+use tokio::net::TcpStream;
+use tokio::sync::Mutex;
+use tokio_tungstenite::tungstenite::Message;
+use tokio_tungstenite::tungstenite::client::IntoClientRequest;
+use tokio_tungstenite::tungstenite::error::TlsError;
+use tokio_tungstenite::tungstenite::handshake::client::Response;
+use tokio_tungstenite::tungstenite::http::{HeaderName, HeaderValue};
+use tokio_tungstenite::{
+    Connector, MaybeTlsStream, WebSocketStream, connect_async_tls_with_config,
+};
+
 use crate::Error;
 use crate::constants::{OPENAI_RESPONSES_DEFAULT_API_BASE, OPENAI_RESPONSES_PATH};
 use crate::responses::types::{ResponsesWsEvent, ResponsesWsEventType, ResponsesWsTransformResult};
@@ -125,6 +143,137 @@ pub fn is_terminal_event(event_type: &ResponsesWsEventType) -> bool {
     )
 }
 
+pub type ResponsesUpstreamWs = WebSocketStream<MaybeTlsStream<TcpStream>>;
+
+static TLS_CONFIG: OnceLock<Arc<ClientConfig>> = OnceLock::new();
+
+fn build_tls_config() -> Result<ClientConfig, Box<tokio_tungstenite::tungstenite::Error>> {
+    let native = rustls_native_certs::load_native_certs();
+    let mut store = RootCertStore::empty();
+    let (added, _ignored) = store.add_parsable_certificates(native.certs);
+    if added == 0 {
+        return Err(Box::new(tokio_tungstenite::tungstenite::Error::Io(
+            io::Error::other(format!(
+                "no usable native root certificates: {:?}",
+                native.errors
+            )),
+        )));
+    }
+    ClientConfig::builder_with_provider(Arc::new(rustls::crypto::ring::default_provider()))
+        .with_safe_default_protocol_versions()
+        .map(|builder| builder.with_root_certificates(store).with_no_client_auth())
+        .map_err(|error| {
+            Box::new(tokio_tungstenite::tungstenite::Error::Tls(
+                TlsError::Rustls(error),
+            ))
+        })
+}
+
+fn tls_config() -> Result<Arc<ClientConfig>, Box<tokio_tungstenite::tungstenite::Error>> {
+    if let Some(config) = TLS_CONFIG.get() {
+        return Ok(Arc::clone(config));
+    }
+    let built = Arc::new(build_tls_config()?);
+    Ok(Arc::clone(TLS_CONFIG.get_or_init(|| built)))
+}
+
+pub async fn connect_upstream<R>(
+    request: R,
+) -> Result<(ResponsesUpstreamWs, Response), Box<tokio_tungstenite::tungstenite::Error>>
+where
+    R: IntoClientRequest + Unpin,
+{
+    let request = request.into_client_request().map_err(Box::new)?;
+    let connector = match request.uri().scheme_str() {
+        Some("wss") => Some(Connector::Rustls(tls_config()?)),
+        _ => None,
+    };
+    connect_async_tls_with_config(request, None, false, connector)
+        .await
+        .map_err(Box::new)
+}
+
+#[derive(Clone)]
+pub struct ResponsesWebSocketConnection {
+    socket: Arc<Mutex<Option<ResponsesUpstreamWs>>>,
+}
+
+impl ResponsesWebSocketConnection {
+    pub async fn connect_url(
+        url: &str,
+        headers: &HashMap<String, String>,
+        timeout: Option<Duration>,
+    ) -> Result<Self, Error> {
+        let mut request = url
+            .into_client_request()
+            .map_err(|error| Error::Network(error.to_string()))?;
+        for (name, value) in headers {
+            let header_name = name
+                .parse::<HeaderName>()
+                .map_err(|error| Error::InvalidRequest(error.to_string()))?;
+            let header_value = HeaderValue::from_str(value)
+                .map_err(|error| Error::InvalidRequest(error.to_string()))?;
+            request.headers_mut().insert(header_name, header_value);
+        }
+        let connect = connect_upstream(request);
+        let result = match timeout {
+            Some(timeout) => tokio::time::timeout(timeout, connect)
+                .await
+                .map_err(|_| Error::Network("Responses WebSocket connection timed out".into()))?,
+            None => connect.await,
+        };
+        let (socket, _) = result.map_err(|error| match *error {
+            tokio_tungstenite::tungstenite::Error::Http(response) => Error::Http {
+                status: response.status().as_u16(),
+                body: String::new(),
+            },
+            other => Error::Network(other.to_string()),
+        })?;
+        Ok(Self {
+            socket: Arc::new(Mutex::new(Some(socket))),
+        })
+    }
+
+    pub async fn send_text(&self, text: String) -> Result<(), Error> {
+        let mut socket = self.socket.lock().await;
+        let Some(socket) = socket.as_mut() else {
+            return Err(Error::Network("Responses WebSocket is closed".into()));
+        };
+        socket
+            .send(Message::Text(text))
+            .await
+            .map_err(|error| Error::Network(error.to_string()))
+    }
+
+    pub async fn recv_text(&self) -> Result<Option<String>, Error> {
+        let mut socket = self.socket.lock().await;
+        let Some(socket) = socket.as_mut() else {
+            return Ok(None);
+        };
+        match socket.next().await {
+            Some(Ok(Message::Text(text))) => Ok(Some(text)),
+            Some(Ok(Message::Binary(bytes))) => String::from_utf8(bytes.to_vec())
+                .map(Some)
+                .map_err(|error| Error::InvalidResponse(error.to_string())),
+            Some(Ok(Message::Close(_))) | None => Ok(None),
+            Some(Ok(_)) => Ok(None),
+            Some(Err(error)) => Err(Error::Network(error.to_string())),
+        }
+    }
+
+    pub async fn close(&self) -> Result<(), Error> {
+        let mut socket = self.socket.lock().await;
+        if let Some(socket) = socket.as_mut() {
+            socket
+                .close(None)
+                .await
+                .map_err(|error| Error::Network(error.to_string()))?;
+        }
+        *socket = None;
+        Ok(())
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

litellm-rust/crates/core/tests/azure_ai_ocr.rs

@@ -70,7 +70,7 @@ async fn facade_acquires_supplied_entra_token_for_final_request() {
 struct ReplaceBodyDocument;
 
 impl OcrHooks for ReplaceBodyDocument {
-    fn has_guardrails(&self) -> bool {
+    fn intercepts_requests(&self) -> bool {
         true
     }
 

litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs

@@ -1,4 +1,5 @@
 use serde_json::{Value, json};
+use std::sync::{Arc, Mutex};
 
 use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request};
 use super::wire::{OcrWireRequest, decode_request};
@@ -124,6 +125,14 @@ async fn immediate_response_normalizes_pages_and_preserves_native() {
         json!({"width":816,"height":1056,"dpi":96})
     );
     assert_eq!(result.usage_info, Some(json!({"pages_processed":1})));
+    let serialized = result.clone().into_json();
+    assert_eq!(serialized["content"], "A\n\nB");
+    assert_eq!(serialized["tables"], json!([{"cells":[]}]));
+    assert_eq!(
+        serialized["keyValuePairs"],
+        json!([{"key":{"content":"A"}}])
+    );
+    assert!(serialized.get("key_value_pairs").is_none());
     assert_eq!(result.provider_native_response, Some(operation));
 }
 
@@ -169,6 +178,55 @@ async fn accepted_response_polls_to_success_with_only_credentials() {
     }
 }
 
+struct SubmissionBoundary {
+    request_count: Arc<Mutex<Vec<String>>>,
+}
+
+impl super::hooks::OcrHooks for SubmissionBoundary {
+    fn post_call(
+        &self,
+        request: super::hooks::OcrPostCallRequest,
+    ) -> super::hooks::OcrHookFuture<'_, super::hooks::OcrPostCallRequest> {
+        Box::pin(async move {
+            match self.request_count.lock().unwrap().len() {
+                1 => assert_eq!(request.original_response, json!(r#"{"submitted":true}"#)),
+                2 => assert!(
+                    request
+                        .original_response
+                        .as_str()
+                        .unwrap()
+                        .contains("succeeded")
+                ),
+                count => panic!("unexpected callback after {count} requests"),
+            }
+            Ok(request)
+        })
+    }
+}
+
+#[tokio::test]
+async fn accepted_response_runs_post_call_before_polling() {
+    let (base, seen, server) = mock_server(vec![
+        MockResponse {
+            status: 202,
+            headers: vec![("Operation-Location", "{base}/operation".into())],
+            body: json!({"submitted": true}),
+        },
+        MockResponse::json(json!({"status":"succeeded"})),
+    ])
+    .await;
+    let request = super::LiteLLMOcrRequest {
+        hooks: Arc::new(SubmissionBoundary {
+            request_count: seen.clone(),
+        }),
+        ..wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({}))
+    };
+
+    perform_ocr(request).await.unwrap();
+    server.await.unwrap();
+    assert_eq!(seen.lock().unwrap().len(), 2);
+}
+
 #[tokio::test]
 async fn polling_forwards_bearer_credentials() {
     let (base, seen, server) = mock_server(vec![
@@ -361,7 +419,7 @@ async fn pre_call_guardrail_receives_caller_pages_before_mapping() {
 
     struct RewritePages;
     impl OcrHooks for RewritePages {
-        fn has_guardrails(&self) -> bool {
+        fn intercepts_requests(&self) -> bool {
             true
         }
 

litellm-rust/crates/core/tests/deepseek_ocr.rs

@@ -34,6 +34,28 @@ fn request_mapping_matches_python(#[case] name: &str, #[case] value: Value) {
     assert!(result.get("ignored").is_none());
 }
 
+#[rstest]
+#[case(json!({"type":"image_url","image_url":"data:image/png;base64,AA=="}))]
+#[case(json!({"type":"document_url","document_url":"data:application/pdf;base64,AA=="}))]
+fn request_maps_both_document_types_to_image_content(#[case] document: Value) {
+    let source = document
+        .get("image_url")
+        .or_else(|| document.get("document_url"))
+        .unwrap()
+        .clone();
+    let request = transform_ocr_request(
+        "deepseek-ai/deepseek-ocr-maas",
+        serde_json::from_value(document).unwrap(),
+        &DeepSeekOcrParams::default(),
+    )
+    .unwrap();
+    let result = serde_json::to_value(request).unwrap();
+    assert_eq!(
+        result["messages"][0]["content"][0],
+        json!({"type":"image_url","image_url":source})
+    );
+}
+
 #[rstest]
 #[case(json!("# hello"), "# hello")]
 #[case(json!("{broken"), "{broken")]

litellm-rust/crates/core/tests/host_lifecycle.rs

@@ -0,0 +1,116 @@
+use crate::Error;
+use crate::call_lifecycle::host::{HostFailure, HostLifecycle, HostPhase};
+
+fn run(fail_at: Option<HostPhase>, asynchronous: bool) -> (Vec<HostPhase>, Vec<Error>) {
+    let mut lifecycle = HostLifecycle::new(asynchronous);
+    let mut events = Vec::new();
+    let mut failures = Vec::new();
+    while lifecycle.phase() != HostPhase::Complete {
+        let phase = lifecycle.phase();
+        events.push(phase);
+        let result = if Some(phase) == fail_at {
+            Err(HostFailure::Error(Error::InvalidRequest(
+                "selected failure".into(),
+            )))
+        } else {
+            Ok(())
+        };
+        if let Some(error) = lifecycle.accept(result) {
+            failures.push(error);
+        }
+    }
+    (events, failures)
+}
+
+#[test]
+fn public_outcome_is_finalized_before_a_single_terminal_dispatch() {
+    for asynchronous in [false, true] {
+        let (events, failures) = run(None, asynchronous);
+        assert!(failures.is_empty());
+        assert_eq!(
+            &events[events.len() - 2..],
+            &[HostPhase::Finalize, HostPhase::Success]
+        );
+        assert_eq!(
+            events
+                .iter()
+                .filter(|phase| **phase == HostPhase::Execute)
+                .count(),
+            1
+        );
+        assert_eq!(
+            events.contains(&HostPhase::DeploymentPostCall),
+            asynchronous
+        );
+    }
+}
+
+#[test]
+fn only_provider_and_response_construction_failures_use_provider_mapping() {
+    for phase in [
+        HostPhase::Setup,
+        HostPhase::DeploymentPreCall,
+        HostPhase::Prepare,
+        HostPhase::Execute,
+        HostPhase::ConstructResponse,
+        HostPhase::DeploymentPostCall,
+        HostPhase::Finalize,
+    ] {
+        let (events, failures) = run(Some(phase), true);
+        assert_eq!(failures.len(), 1);
+        assert!(!events.contains(&HostPhase::Success));
+        let mapped = matches!(phase, HostPhase::Execute | HostPhase::ConstructResponse);
+        assert_eq!(events.contains(&HostPhase::MapFailure), mapped);
+        assert_eq!(events.contains(&HostPhase::DeploymentFailure), mapped);
+        assert_eq!(
+            &events[events.len() - 2..],
+            &[HostPhase::Failure, HostPhase::AsyncFailure]
+        );
+        assert!(
+            events
+                .iter()
+                .filter(|phase| **phase == HostPhase::Execute)
+                .count()
+                <= 1
+        );
+    }
+}
+
+#[test]
+fn failure_handler_errors_do_not_replace_selected_failure_or_suppress_async_dispatch() {
+    let mut lifecycle = HostLifecycle::new(true);
+    while lifecycle.phase() != HostPhase::Execute {
+        lifecycle.accept(Ok(()));
+    }
+    let selected = Error::InvalidRequest("provider".into());
+    assert_eq!(
+        lifecycle.accept(Err(HostFailure::Error(selected.clone()))),
+        Some(selected)
+    );
+    lifecycle.accept(Ok(()));
+    for phase in [
+        HostPhase::DeploymentFailure,
+        HostPhase::Failure,
+        HostPhase::AsyncFailure,
+    ] {
+        assert_eq!(lifecycle.phase(), phase);
+        assert_eq!(
+            lifecycle.accept(Err(HostFailure::Error(Error::InvalidRequest(
+                "callback".into()
+            )))),
+            None
+        );
+    }
+    assert_eq!(lifecycle.phase(), HostPhase::Complete);
+}
+
+#[test]
+fn cancellation_skips_terminal_dispatch() {
+    let mut lifecycle = HostLifecycle::new(true);
+    let error = Error::InvalidRequest("cancelled".into());
+    assert_eq!(
+        lifecycle.accept(Err(HostFailure::Cancelled(error.clone()))),
+        Some(error)
+    );
+    assert_eq!(lifecycle.phase(), HostPhase::Complete);
+}

litellm-rust/crates/core/tests/ocr.rs

@@ -3,9 +3,16 @@ use std::sync::{Arc, Mutex};
 use serde_json::{Value, json};
 
 use super::OcrClient;
-use super::hooks::{OcrHookFuture, OcrHooks, OcrLogFuture, OcrPreCallRequest};
+use super::hooks::{
+    OcrDuringCallRequest, OcrHookFuture, OcrHooks, OcrLogFuture, OcrPostCallRequest,
+    OcrPreCallRequest,
+};
 use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request};
 use super::wire::{OcrWireRequest, decode_request};
+use super::{
+    NativeOutcome, NoopOcrHost, OcrAdmission, OcrCall, OcrCallStep, OcrDecline, OcrHost,
+    OcrHostOperation, OcrHostResult,
+};
 use crate::call_lifecycle::{CallLifecycleContext, CallLifecycleTiming};
 
 #[test]
@@ -51,7 +58,7 @@ async fn facade_executes_direct_mistral_once() {
     let result = perform_ocr(wire_request(
         "mistral/model",
         &base,
-        json!({"extract_header":true,"unknown":"ignored"}),
+        json!({"pages":"0,2-4","extract_header":true,"unknown":"ignored"}),
     ))
     .await
     .unwrap();
@@ -72,6 +79,7 @@ async fn facade_executes_direct_mistral_once() {
         json!({
             "model":"model",
             "document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"},
+            "pages":"0,2-4",
             "extract_header":true
         })
     );
@@ -124,7 +132,7 @@ struct RecordingHooks {
 }
 
 impl OcrHooks for RecordingHooks {
-    fn has_guardrails(&self) -> bool {
+    fn intercepts_requests(&self) -> bool {
         true
     }
 
@@ -148,6 +156,13 @@ impl OcrHooks for RecordingHooks {
         })
     }
 
+    fn post_call(&self, request: OcrPostCallRequest) -> OcrHookFuture<'_, OcrPostCallRequest> {
+        Box::pin(async move {
+            self.events.lock().unwrap().push("post");
+            Ok(request)
+        })
+    }
+
     fn success<'a>(
         &'a self,
         _context: &'a CallLifecycleContext,
@@ -171,6 +186,38 @@ impl OcrHooks for RecordingHooks {
     }
 }
 
+struct HeaderEditHooks;
+
+impl OcrHooks for HeaderEditHooks {
+    fn intercepts_requests(&self) -> bool {
+        true
+    }
+
+    fn during_call(
+        &self,
+        mut request: OcrDuringCallRequest,
+    ) -> OcrHookFuture<'_, OcrDuringCallRequest> {
+        request
+            .headers
+            .push(("x-core-callback".into(), "edited".into()));
+        Box::pin(async move { Ok(request) })
+    }
+}
+
+#[tokio::test]
+async fn lifecycle_sends_headers_returned_by_the_typed_during_call_operation() {
+    let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await;
+    let request = super::LiteLLMOcrRequest {
+        hooks: Arc::new(HeaderEditHooks),
+        ..wire_request("mistral/model", &base, json!({}))
+    };
+
+    perform_ocr(request).await.unwrap();
+    server.await.unwrap();
+
+    assert!(seen.lock().unwrap()[0].contains("x-core-callback: edited"));
+}
+
 #[tokio::test]
 async fn lifecycle_orders_hooks_and_emits_one_success() {
     let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await;
@@ -185,7 +232,10 @@ async fn lifecycle_orders_hooks_and_emits_one_success() {
     };
     perform_ocr(request).await.unwrap();
     server.await.unwrap();
-    assert_eq!(*events.lock().unwrap(), ["pre", "during", "success"]);
+    assert_eq!(
+        *events.lock().unwrap(),
+        ["pre", "during", "post", "success"]
+    );
     assert_eq!(seen.lock().unwrap().len(), 1);
 }
 
@@ -227,3 +277,562 @@ async fn upstream_failure_emits_one_terminal_failure() {
     assert_eq!(*events.lock().unwrap(), ["pre", "during", "failure"]);
     assert_eq!(seen.lock().unwrap().len(), 1);
 }
+
+struct AdmissionSpy {
+    effects: Arc<Mutex<usize>>,
+}
+
+impl OcrHooks for AdmissionSpy {
+    fn intercepts_requests(&self) -> bool {
+        *self.effects.lock().unwrap() += 1;
+        true
+    }
+
+    fn pre_call(&self, request: OcrPreCallRequest) -> OcrHookFuture<'_, OcrPreCallRequest> {
+        *self.effects.lock().unwrap() += 1;
+        Box::pin(async move { Ok(request) })
+    }
+}
+
+#[test]
+fn admission_declines_without_invoking_hooks_or_transport() {
+    for (admission, expected) in [
+        (
+            OcrAdmission {
+                provider_workflow: false,
+                host_operations: true,
+                asynchronous: false,
+            },
+            OcrDecline::ProviderWorkflow,
+        ),
+        (
+            OcrAdmission {
+                provider_workflow: true,
+                host_operations: false,
+                asynchronous: false,
+            },
+            OcrDecline::HostOperations,
+        ),
+    ] {
+        let outcome = OcrCall::admit(super::test_support::ocr_client(), admission);
+        assert!(matches!(outcome, NativeOutcome::Declined(reason) if reason == expected));
+    }
+}
+
+#[tokio::test]
+async fn fallible_host_phases_do_not_replay_or_reach_transport() {
+    for failure_phase in ["pre", "during"] {
+        let request = super::LiteLLMOcrRequest {
+            hooks: Arc::new(AdmissionSpy {
+                effects: Arc::new(Mutex::new(0)),
+            }),
+            ..wire_request("mistral/model", "http://127.0.0.1:1", json!({}))
+        };
+        let NativeOutcome::Completed(mut call) =
+            OcrCall::admit(super::test_support::ocr_client(), OcrAdmission::all())
+        else {
+            panic!("supported call declined")
+        };
+        let mut request = Some(request);
+        let mut result = None;
+        let mut phases = Vec::new();
+        let error = loop {
+            match call.resume(result.take()).await {
+                Ok(OcrCallStep::Host(operation)) => match operation {
+                    OcrHostOperation::Lifecycle(_)
+                    | OcrHostOperation::ConstructResponse(_)
+                    | OcrHostOperation::MapFailure(_)
+                    | OcrHostOperation::Success { .. }
+                    | OcrHostOperation::Failure { .. } => {
+                        result = Some(OcrHostResult::Lifecycle(Ok(())))
+                    }
+                    OcrHostOperation::ProjectRequest => {
+                        result = Some(OcrHostResult::Request(Ok((
+                            Box::new(request.take().unwrap()),
+                            false,
+                        ))))
+                    }
+                    OcrHostOperation::AcquireAzureAdToken => {
+                        panic!("test request has no token provider")
+                    }
+                    OcrHostOperation::PreCall(request) => {
+                        phases.push("pre");
+                        result = Some(OcrHostResult::PreCall(if failure_phase == "pre" {
+                            Err(crate::Error::InvalidRequest("pre failed".into()))
+                        } else {
+                            Ok(request)
+                        }));
+                    }
+                    OcrHostOperation::DuringCall(request) => {
+                        phases.push("during");
+                        result = Some(OcrHostResult::DuringCall(if failure_phase == "during" {
+                            Err(crate::Error::InvalidRequest("during failed".into()))
+                        } else {
+                            Ok(request)
+                        }));
+                    }
+                    OcrHostOperation::PostCall(_) => panic!("transport should not be reached"),
+                },
+                Err(error) => break error,
+                Ok(OcrCallStep::Complete(_)) => panic!("failed call completed"),
+            }
+        };
+        assert!(matches!(error, crate::Error::InvalidRequest(_)));
+        assert_eq!(
+            phases
+                .iter()
+                .filter(|phase| **phase == failure_phase)
+                .count(),
+            1
+        );
+    }
+}
+
+#[tokio::test]
+async fn invalid_provider_response_runs_post_call_before_normalization_failure() {
+    let (base, seen, server) =
+        mock_server(vec![MockResponse::json(json!({"pages":"invalid"}))]).await;
+    let mut request = Some(wire_request("mistral/model", &base, json!({})));
+    let NativeOutcome::Completed(mut call) =
+        OcrCall::admit(super::test_support::ocr_client(), OcrAdmission::all())
+    else {
+        panic!("supported call declined")
+    };
+    let host = NoopOcrHost;
+    let mut result = None;
+    let mut post_calls = Vec::new();
+    let error = loop {
+        match call.resume(result.take()).await {
+            Ok(OcrCallStep::Host(OcrHostOperation::ProjectRequest)) => {
+                result = Some(OcrHostResult::Request(Ok((
+                    Box::new(request.take().unwrap()),
+                    false,
+                ))));
+            }
+            Ok(OcrCallStep::Host(operation)) => {
+                if let OcrHostOperation::PostCall(request) = &operation {
+                    post_calls.push(request.original_response.clone());
+                }
+                result = Some(host.invoke(operation).await);
+            }
+            Err(error) => break error,
+            Ok(OcrCallStep::Complete(_)) => panic!("invalid provider response completed"),
+        }
+    };
+    server.await.unwrap();
+    assert!(matches!(error, crate::Error::InvalidResponse(_)));
+    assert_eq!(seen.lock().unwrap().len(), 1);
+    assert_eq!(post_calls, [json!(r#"{"pages":"invalid"}"#)]);
+}
+
+#[tokio::test]
+async fn direct_native_host_drives_the_same_state_machine() {
+    let (base, seen, server) = mock_server(vec![MockResponse::json(json!({
+        "pages":[{"index":0,"markdown":"native"}]
+    }))])
+    .await;
+    let request = super::LiteLLMOcrRequest {
+        hooks: Arc::new(AdmissionSpy {
+            effects: Arc::new(Mutex::new(0)),
+        }),
+        ..wire_request("mistral/model", &base, json!({}))
+    };
+    let NativeOutcome::Completed(mut call) = OcrCall::admit(
+        super::test_support::ocr_client(),
+        OcrAdmission {
+            asynchronous: true,
+            ..OcrAdmission::all()
+        },
+    ) else {
+        panic!("supported call declined")
+    };
+    let mut request = Some(request);
+    let host = NoopOcrHost;
+    let mut result = None;
+    let mut operations = Vec::new();
+    let response = loop {
+        match call.resume(result.take()).await.unwrap() {
+            OcrCallStep::Host(operation) => {
+                operations.push(match &operation {
+                    OcrHostOperation::ProjectRequest => "ProjectRequest".into(),
+                    OcrHostOperation::Lifecycle(phase) => format!("{phase:?}"),
+                    OcrHostOperation::PreCall(_) => "PreCall".into(),
+                    OcrHostOperation::DuringCall(_) => "DuringCall".into(),
+                    OcrHostOperation::PostCall(_) => "PostCall".into(),
+                    OcrHostOperation::ConstructResponse(_) => "ConstructResponse".into(),
+                    OcrHostOperation::Success { response, .. } => {
+                        assert_eq!(response.pages[0]["markdown"], "native");
+                        "Success".into()
+                    }
+                    _ => panic!("unexpected OCR operation"),
+                });
+                result = Some(match operation {
+                    OcrHostOperation::ProjectRequest => {
+                        OcrHostResult::Request(Ok((Box::new(request.take().unwrap()), false)))
+                    }
+                    operation => host.invoke(operation).await,
+                });
+            }
+            OcrCallStep::Complete(response) => break response,
+        }
+    };
+    server.await.unwrap();
+    assert_eq!(response.pages[0]["markdown"], "native");
+    assert_eq!(seen.lock().unwrap().len(), 1);
+    assert_eq!(
+        operations,
+        [
+            "Setup",
+            "DeploymentPreCall",
+            "Prepare",
+            "ProjectRequest",
+            "PreCall",
+            "DuringCall",
+            "PostCall",
+            "ConstructResponse",
+            "DeploymentPostCall",
+            "Finalize",
+            "Success",
+        ]
+    );
+    assert!(matches!(
+        call.resume(None).await,
+        Err(crate::Error::InvalidRequest(_))
+    ));
+}
+
+#[tokio::test]
+async fn public_finalization_failure_never_dispatches_success_or_replays_provider() {
+    use crate::call_lifecycle::host::{HostFailure, HostPhase};
+
+    let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await;
+    let mut request = Some(wire_request("mistral/model", &base, json!({})));
+    let NativeOutcome::Completed(mut call) = OcrCall::admit(
+        super::test_support::ocr_client(),
+        OcrAdmission {
+            asynchronous: true,
+            ..OcrAdmission::all()
+        },
+    ) else {
+        panic!("supported call declined")
+    };
+    let selected = crate::Error::InvalidRequest("public metadata failed".into());
+    let host = NoopOcrHost;
+    let mut result = None;
+    let mut failures = Vec::new();
+    let error = loop {
+        match call.resume(result.take()).await {
+            Ok(OcrCallStep::Host(operation)) => {
+                result = Some(match operation {
+                    OcrHostOperation::Lifecycle(HostPhase::Finalize) => {
+                        OcrHostResult::Lifecycle(Err(HostFailure::Error(selected.clone())))
+                    }
+                    OcrHostOperation::Failure { error, .. } => {
+                        assert_eq!(error, selected);
+                        failures.push("sync");
+                        OcrHostResult::Lifecycle(Err(HostFailure::Error(
+                            crate::Error::InvalidRequest("failure callback failed".into()),
+                        )))
+                    }
+                    OcrHostOperation::Lifecycle(HostPhase::AsyncFailure) => {
+                        failures.push("async");
+                        OcrHostResult::Lifecycle(Ok(()))
+                    }
+                    OcrHostOperation::Success { .. }
+                    | OcrHostOperation::MapFailure(_)
+                    | OcrHostOperation::Lifecycle(HostPhase::DeploymentFailure) => {
+                        panic!("finalization failure used provider/success dispatch")
+                    }
+                    OcrHostOperation::ProjectRequest => {
+                        OcrHostResult::Request(Ok((Box::new(request.take().unwrap()), false)))
+                    }
+                    operation => host.invoke(operation).await,
+                });
+            }
+            Ok(OcrCallStep::Complete(_)) => panic!("failed call completed successfully"),
+            Err(error) => break error,
+        }
+    };
+    server.await.unwrap();
+    assert_eq!(error, selected);
+    assert_eq!(failures, ["sync", "async"]);
+    assert_eq!(seen.lock().unwrap().len(), 1);
+}
+
+#[tokio::test]
+async fn cancellation_at_provider_hook_prevents_execution_and_further_resumption() {
+    use crate::call_lifecycle::host::HostFailure;
+
+    let request = super::LiteLLMOcrRequest {
+        hooks: Arc::new(AdmissionSpy {
+            effects: Arc::new(Mutex::new(0)),
+        }),
+        ..wire_request("mistral/model", "http://127.0.0.1:1", json!({}))
+    };
+    let NativeOutcome::Completed(mut call) =
+        OcrCall::admit(super::test_support::ocr_client(), OcrAdmission::all())
+    else {
+        panic!("supported call declined")
+    };
+    let mut request = Some(request);
+    let host = NoopOcrHost;
+    let mut result = None;
+    loop {
+        match call.resume(result.take()).await.unwrap() {
+            OcrCallStep::Host(OcrHostOperation::PreCall(_)) => break,
+            OcrCallStep::Host(OcrHostOperation::ProjectRequest) => {
+                result = Some(OcrHostResult::Request(Ok((
+                    Box::new(request.take().unwrap()),
+                    false,
+                ))))
+            }
+            OcrCallStep::Host(operation) => result = Some(host.invoke(operation).await),
+            OcrCallStep::Complete(_) => panic!("provider executed before pre-call result"),
+        }
+    }
+    let selected = crate::Error::InvalidRequest("cancelled".into());
+    assert!(matches!(
+        call.interrupt(HostFailure::Cancelled(selected.clone())).await,
+        Err(error) if error == selected
+    ));
+    assert!(
+        call.resume(Some(OcrHostResult::Lifecycle(Ok(()))))
+            .await
+            .is_err()
+    );
+}
+
+#[tokio::test]
+async fn missing_host_result_preserves_pending_operation() {
+    use crate::call_lifecycle::host::HostPhase;
+
+    let NativeOutcome::Completed(mut call) =
+        OcrCall::admit(super::test_support::ocr_client(), OcrAdmission::all())
+    else {
+        panic!("supported call declined")
+    };
+    assert!(matches!(
+        call.resume(None).await.unwrap(),
+        OcrCallStep::Host(OcrHostOperation::Lifecycle(HostPhase::Setup))
+    ));
+    assert!(call.resume(None).await.is_err());
+    assert!(matches!(
+        call.resume(Some(OcrHostResult::Lifecycle(Ok(()))))
+            .await
+            .unwrap(),
+        OcrCallStep::Host(OcrHostOperation::Lifecycle(HostPhase::Prepare))
+    ));
+}
+
+async fn read_bounded_response(
+    response: Vec<u8>,
+    limit: usize,
+) -> Result<bytes::Bytes, super::error::OcrError> {
+    use tokio::io::{AsyncReadExt, AsyncWriteExt};
+
+    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let address = listener.local_addr().unwrap();
+    let server = tokio::spawn(async move {
+        let (mut socket, _) = listener.accept().await.unwrap();
+        let mut request = [0; 4096];
+        assert!(socket.read(&mut request).await.unwrap() > 0);
+        socket.write_all(&response).await.unwrap();
+        std::future::pending::<()>().await;
+    });
+    let response = reqwest::Client::new()
+        .get(format!("http://{address}"))
+        .send()
+        .await
+        .unwrap();
+    let result = tokio::time::timeout(
+        std::time::Duration::from_secs(2),
+        super::client::read_response_bytes(response, limit),
+    )
+    .await;
+    server.abort();
+    let _ = server.await;
+    result.expect("bounded reads must finish without waiting for the rest of an oversized body")
+}
+
+#[tokio::test]
+async fn response_limit_accepts_exact_size_and_rejects_declared_and_chunked_overflow() {
+    use super::error::{OcrError, OcrResponseError};
+
+    for response in [
+        "HTTP/1.1 200 OK\r\nContent-Length: 8\r\n\r\nabcdefgh",
+        "HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n4\r\nabcd\r\n4\r\nefgh\r\n0\r\n\r\n",
+    ] {
+        assert_eq!(
+            read_bounded_response(response.as_bytes().to_vec(), 8)
+                .await
+                .unwrap(),
+            "abcdefgh"
+        );
+    }
+    for response in [
+        "HTTP/1.1 200 OK\r\nContent-Length: 9\r\n\r\n",
+        "HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n4\r\nabcd\r\n5\r\nefghi\r\n",
+    ] {
+        assert!(matches!(
+            read_bounded_response(response.as_bytes().to_vec(), 8).await,
+            Err(OcrError::Response(OcrResponseError::TooLarge { limit: 8 }))
+        ));
+    }
+}
+
+#[tokio::test]
+async fn oversized_error_retains_http_status_and_bounded_diagnostics_without_draining() {
+    let prefix = "x".repeat(4 * (crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS + 1));
+    for headers in ["Content-Length: 1000000", "Transfer-Encoding: chunked"] {
+        let body = if headers.starts_with("Transfer") {
+            format!("{:x}\r\n{prefix}\r\n", prefix.len())
+        } else {
+            prefix.clone()
+        };
+        let response = format!("HTTP/1.1 429 Too Many Requests\r\n{headers}\r\n\r\n{body}");
+        let error = read_bounded_response(response.into_bytes(), 4096)
+            .await
+            .unwrap_err();
+        match error {
+            super::error::OcrError::Transport(crate::error::TransportError::Http {
+                status,
+                body,
+            }) => {
+                assert_eq!(status, 429);
+                assert_eq!(
+                    body,
+                    format!(
+                        "{}... (truncated)",
+                        "x".repeat(crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS)
+                    )
+                );
+            }
+            error => panic!("unexpected error: {error}"),
+        }
+    }
+}
+
+#[test]
+fn response_limit_is_validated_and_not_forwarded_to_the_provider() {
+    let request = wire_request(
+        "mistral/model",
+        "http://localhost",
+        json!({"max_response_bytes": 123}),
+    );
+    assert_eq!(request.connection.max_response_bytes, 123);
+    assert!(!request.optional_params.contains_key("max_response_bytes"));
+    for value in [
+        json!(0),
+        json!(-1),
+        json!(true),
+        json!("123"),
+        json!(1.5),
+        json!(crate::constants::OCR_RESPONSE_MAX_BYTES + 1),
+        Value::Null,
+    ] {
+        let wire = serde_json::from_value(json!({
+            "model": "mistral/model", "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"},
+            "optional_params": {"max_response_bytes": value}
+        })).unwrap();
+        let Err(error) = decode_request(wire) else {
+            panic!("invalid response limit accepted")
+        };
+        assert!(error.to_string().contains("max_response_bytes"));
+    }
+}
+
+#[derive(Debug)]
+struct PendingToken {
+    entered: Arc<tokio::sync::Notify>,
+    dropped: Arc<std::sync::atomic::AtomicBool>,
+}
+
+struct TokenFutureDrop(Arc<std::sync::atomic::AtomicBool>);
+
+impl Drop for TokenFutureDrop {
+    fn drop(&mut self) {
+        self.0.store(true, std::sync::atomic::Ordering::SeqCst);
+    }
+}
+
+impl crate::auth::TokenProvider for PendingToken {
+    fn acquire(&self) -> crate::auth::TokenFuture<'_> {
+        Box::pin(async move {
+            let _guard = TokenFutureDrop(self.dropped.clone());
+            self.entered.notify_one();
+            std::future::pending().await
+        })
+    }
+}
+
+#[tokio::test]
+async fn cancellation_waits_for_provider_capture_drop_even_when_acknowledgement_is_cancelled() {
+    use crate::call_lifecycle::host::HostFailure;
+    use std::future::Future;
+    use std::sync::atomic::{AtomicBool, Ordering};
+    use std::task::Poll;
+
+    for interrupt_acknowledgement in [false, true] {
+        let entered = Arc::new(tokio::sync::Notify::new());
+        let dropped = Arc::new(AtomicBool::new(false));
+        let request = wire_request("azure_ai/mistral-ocr", "https://example.invalid", json!({}));
+        let request = super::LiteLLMOcrRequest {
+            connection: super::OcrConnection {
+                extra_headers: vec![("authorization".into(), "Bearer test-key".into())],
+                ..request.connection
+            },
+            azure_ad_token_provider: Some(crate::auth::TokenProviderHandle::new(Arc::new(
+                PendingToken {
+                    entered: entered.clone(),
+                    dropped: dropped.clone(),
+                },
+            ))),
+            ..request
+        };
+        let NativeOutcome::Completed(mut call) =
+            OcrCall::admit(super::test_support::ocr_client(), OcrAdmission::all())
+        else {
+            panic!("supported call declined")
+        };
+        let mut request = Some(request);
+        let mut result = None;
+        tokio::time::timeout(std::time::Duration::from_secs(2), async {
+            loop {
+                tokio::select! {
+                    _ = entered.notified() => break,
+                    step = call.resume(result.take()) => {
+                        result = Some(match step.unwrap() {
+                            OcrCallStep::Host(OcrHostOperation::ProjectRequest) => OcrHostResult::Request(Ok((Box::new(request.take().unwrap()), false))),
+                            OcrCallStep::Host(operation) => NoopOcrHost.invoke(operation).await,
+                            OcrCallStep::Complete(_) => panic!("pending provider completed"),
+                        });
+                    }
+                }
+            }
+        }).await.unwrap();
+        assert!(!dropped.load(Ordering::SeqCst));
+        let selected = crate::Error::InvalidRequest("cancelled".into());
+        if interrupt_acknowledgement {
+            let mut acknowledgement =
+                Box::pin(call.interrupt(HostFailure::Cancelled(selected.clone())));
+            std::future::poll_fn(|cx| {
+                assert!(acknowledgement.as_mut().poll(cx).is_pending());
+                Poll::Ready(())
+            })
+            .await;
+            drop(acknowledgement);
+            assert!(!dropped.load(Ordering::SeqCst));
+        }
+        let result = tokio::time::timeout(
+            std::time::Duration::from_secs(2),
+            call.interrupt(HostFailure::Cancelled(selected.clone())),
+        )
+        .await
+        .unwrap();
+        assert!(matches!(result, Err(error) if error == selected));
+        assert!(
+            dropped.load(Ordering::SeqCst),
+            "cancellation returned while provider captures were still alive"
+        );
+    }
+}

litellm-rust/crates/core/tests/reducto_ocr.rs

@@ -3,7 +3,7 @@ use std::sync::Arc;
 use rstest::rstest;
 use serde_json::{Value, json};
 
-use super::hooks::{OcrDuringCallRequest, OcrHookFuture, OcrHooks};
+use super::hooks::{OcrDuringCallRequest, OcrHookFuture, OcrHooks, OcrPostCallRequest};
 use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request};
 
 fn request_body(request: &str) -> Value {
@@ -100,6 +100,42 @@ async fn data_uri_upload_preserves_multipart_headers(#[case] model: &str) {
     assert!(requests[1].starts_with("POST /parse "));
 }
 
+struct ParseBoundary {
+    request_count: Arc<std::sync::Mutex<Vec<String>>>,
+}
+
+impl OcrHooks for ParseBoundary {
+    fn post_call(&self, request: OcrPostCallRequest) -> OcrHookFuture<'_, OcrPostCallRequest> {
+        Box::pin(async move {
+            assert_eq!(self.request_count.lock().unwrap().len(), 2);
+            assert_eq!(
+                request.original_response,
+                json!(r#"{"result":{"chunks":[]}}"#)
+            );
+            Ok(request)
+        })
+    }
+}
+
+#[tokio::test]
+async fn post_call_stays_after_reducto_upload_and_parse() {
+    let (base, seen, server) = mock_server(vec![
+        MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})),
+        MockResponse::json(json!({"result":{"chunks":[]}})),
+    ])
+    .await;
+    let request = super::LiteLLMOcrRequest {
+        hooks: Arc::new(ParseBoundary {
+            request_count: seen.clone(),
+        }),
+        ..wire_request("reducto/parse-v3", &base, json!({}))
+    };
+
+    perform_ocr(request).await.unwrap();
+    server.await.unwrap();
+    assert_eq!(seen.lock().unwrap().len(), 2);
+}
+
 #[rstest]
 #[case(json!({"file_id":""}))]
 #[case(json!({}))]
@@ -148,9 +184,16 @@ async fn rejects_invalid_document_sources_before_network(#[case] source: &str) {
 fn response_normalization_groups_blocks_and_distinguishes_null_result() {
     use crate::ocr::codecs::reducto::{ReductoResponse, transform_ocr_response};
 
-    let raw = json!({"usage":{"num_pages":"2","credits":"3"},"result":{"chunks":[
-        {"blocks":[{"content":"B","bbox":{"page":2},"kind":"table"}]},
-        {"blocks":[{"content":"A","bbox":{"page":1},"kind":"text"},{"content":"C","bbox":{"page":1}}]}
+    let raw = json!({"usage":{"num_pages":"2","credits":"3"},"result":{"type":"full","chunks":[
+        {"blocks":[{
+            "type":"Table",
+            "content":"B",
+            "bbox":{"left":0.1,"top":0.2,"width":0.8,"height":0.3,"page":2,"original_page":4},
+            "confidence":"high",
+            "granular_confidence":{"parse_confidence":0.95,"extract_confidence":null},
+            "image_url":null
+        }]},
+        {"blocks":[{"content":"A","bbox":{"page":1},"type":"Text"},{"content":"C","bbox":{"page":1}}]}
     ]}});
     let response: ReductoResponse = serde_json::from_value(raw).unwrap();
     let normalized = transform_ocr_response("parse-v3", response)
@@ -158,7 +201,17 @@ fn response_normalization_groups_blocks_and_distinguishes_null_result() {
         .into_json();
     assert_eq!(normalized["pages"][0]["markdown"], "A\n\nC");
     assert_eq!(normalized["pages"][1]["markdown"], "B");
-    assert_eq!(normalized["pages"][1]["blocks"][0]["kind"], "table");
+    assert_eq!(normalized["pages"][1]["blocks"][0]["type"], "Table");
+    assert_eq!(
+        normalized["pages"][1]["blocks"][0]["bbox"],
+        json!({"left":0.1,"top":0.2,"width":0.8,"height":0.3,"page":2,"original_page":4})
+    );
+    assert_eq!(normalized["pages"][1]["blocks"][0]["confidence"], "high");
+    assert_eq!(
+        normalized["pages"][1]["blocks"][0]["granular_confidence"]["parse_confidence"],
+        0.95
+    );
+    assert!(normalized["pages"][1]["blocks"][0]["image_url"].is_null());
     assert_eq!(normalized["usage_info"]["pages_processed"], 2);
     assert_eq!(normalized["usage_info"]["credits"], 3.0);
 
@@ -195,7 +248,7 @@ async fn facade_omits_native_response_by_default_and_preserves_auth_priority() {
 struct RewriteDocument;
 
 impl OcrHooks for RewriteDocument {
-    fn has_guardrails(&self) -> bool {
+    fn intercepts_requests(&self) -> bool {
         true
     }
 

litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs

@@ -49,7 +49,7 @@ async fn facade_executes_vertex_deepseek_at_the_openai_endpoint() {
     assert!(body.get("extra_body").is_none());
     assert_eq!(
         body["messages"][0]["content"][0],
-        json!({"type":"document_url","document_url":"gs://bucket/document.pdf"})
+        json!({"type":"image_url","image_url":"gs://bucket/document.pdf"})
     );
 }
 

litellm-rust/crates/core/tests/workspace_crate_allowlist.rs

@@ -1,115 +0,0 @@
-//! Enforcement: the litellm-rust workspace has exactly six crates.
-//!
-//! `core` (the Rust SDK), `token-counter` (standalone input token counting),
-//! `config` (the config-loading boundary),
-//! `ai-gateway` (the HTTP/WebSocket host),
-//! `python-interop` (domain-neutral PyO3 primitives), and `python-bridge` (the
-//! PyO3 cdylib). Adding or removing a crate must be a
-//! deliberate act: this test fails until the allowlist here is updated, forcing
-//! whoever changes the crate set to justify the new crate per the rule that a
-//! crate is a layer needing independent compilation / its own deps / a separate
-//! artifact — and to keep `litellm-rust/AGENTS.md` in sync.
-//!
-//! Std-only (no toml crate): we scan the workspace manifest's `members = [...]`
-//! block and the `crates/` directory directly.
-
-use std::collections::BTreeSet;
-use std::fs;
-use std::path::{Path, PathBuf};
-
-/// The one true crate set. Update BOTH this and `litellm-rust/AGENTS.md` when the
-/// workspace legitimately gains or loses a crate.
-const EXPECTED_MEMBERS: &[&str] = &[
-    "crates/core",
-    "crates/token-counter",
-    "crates/config",
-    "crates/ai-gateway",
-    "crates/python-interop",
-    "crates/python-bridge",
-];
-
-/// The crate subdirectory names that must exist under `crates/`.
-const EXPECTED_CRATE_DIRS: &[&str] = &[
-    "core",
-    "token-counter",
-    "config",
-    "ai-gateway",
-    "python-interop",
-    "python-bridge",
-];
-
-const MISMATCH: &str = "litellm-rust crate set changed — update this allowlist AND litellm-rust/AGENTS.md, and justify the crate per the rule (crate = layer needing independent compilation / its own deps / a separate artifact).";
-
-/// Absolute path to the workspace root (`litellm-rust/`).
-fn workspace_root() -> PathBuf {
-    // CARGO_MANIFEST_DIR is `.../litellm-rust/crates/core`; the workspace root is
-    // two levels up.
-    Path::new(concat!(env!("CARGO_MANIFEST_DIR"), "/../.."))
-        .canonicalize()
-        .expect("workspace root should resolve")
-}
-
-/// Parse the `members = [ ... ]` array out of the workspace `[workspace]` table.
-///
-/// Minimal hand-rolled scan: find `members`, then collect every double-quoted
-/// string up to the closing `]`. Good enough for our fixed manifest shape and
-/// keeps this test dependency-free.
-fn parse_members(manifest: &str) -> BTreeSet<String> {
-    let after_members = manifest
-        .split_once("members")
-        .map(|(_, rest)| rest)
-        .expect("workspace manifest should declare members");
-    let open = after_members.find('[').expect("members should be an array");
-    let close = after_members[open..]
-        .find(']')
-        .map(|offset| open + offset)
-        .expect("members array should be closed");
-    let body = &after_members[open + 1..close];
-
-    let mut members = BTreeSet::new();
-    let mut rest = body;
-    while let Some(start) = rest.find('"') {
-        let after_quote = &rest[start + 1..];
-        let end = after_quote
-            .find('"')
-            .expect("opening quote should be matched");
-        members.insert(after_quote[..end].to_string());
-        rest = &after_quote[end + 1..];
-    }
-    members
-}
-
-/// The crate subdirectory names under `crates/`.
-///
-/// A directory counts as a crate only when it holds a `Cargo.toml`; non-crate
-/// directories (e.g. docs like `CODING_STANDARDS/`) are ignored so they can live
-/// under `crates/` without tripping the crate-proliferation guard.
-fn crate_dirs(root: &Path) -> BTreeSet<String> {
-    fs::read_dir(root.join("crates"))
-        .expect("crates/ directory should exist")
-        .filter_map(Result::ok)
-        .filter(|entry| entry.file_type().map(|ty| ty.is_dir()).unwrap_or(false))
-        .filter(|entry| entry.path().join("Cargo.toml").is_file())
-        .map(|entry| entry.file_name().to_string_lossy().into_owned())
-        .collect()
-}
-
-#[test]
-fn workspace_members_match_allowlist() {
-    let root = workspace_root();
-    let manifest = fs::read_to_string(root.join("Cargo.toml"))
-        .expect("workspace Cargo.toml should be readable");
-
-    let actual = parse_members(&manifest);
-    let expected: BTreeSet<String> = EXPECTED_MEMBERS.iter().map(|s| s.to_string()).collect();
-    assert_eq!(actual, expected, "{MISMATCH}");
-}
-
-#[test]
-fn crates_directory_matches_allowlist() {
-    let root = workspace_root();
-
-    let actual = crate_dirs(&root);
-    let expected: BTreeSet<String> = EXPECTED_CRATE_DIRS.iter().map(|s| s.to_string()).collect();
-    assert_eq!(actual, expected, "{MISMATCH}");
-}

litellm-rust/crates/python-bridge/AGENTS.md

@@ -1,3 +1,42 @@
-litellm-python-bridge is the PyO3 cdylib that exposes LiteLLM Rust APIs to the Python SDK. Keep API registration, domain dependency wiring, request assembly, and Python exception mapping here. Put domain-neutral Python/Serde conversion and GIL primitives in litellm-python-interop.
-
-Keep it thin: no business logic, no transforms, no I/O orchestration — just marshal in/out and call the core entrypoint.
+- Target invariants, not completion claims; these supersede older conflicting bridge guidance
+- Keep this crate the product-specific PyO3 consumer of `litellm-python-interop`
+  - Own registration, input projection, retained Python state, callback invocation, public response/error construction and host scheduling
+  - Keep value-oriented execution, sync waiting, nested-runtime checks, signal polling and panic containment in `execution.rs`; native async work uses `pyo3-async-runtimes`, Serde output uses `Pythonized<T>`
+  - Core owns typed native state, admission, lifecycle sequencing, provider preparation/I/O, normalization and terminal-outcome/dispatch decisions
+  - Python, Rust SDK and gateway use one lifecycle-bearing core route entrypoint; provider helpers stay private, never bridge-accessible transport drivers
+  - Built-in provider/config/secret/auth/document preparation stays in Rust; caller-authored callbacks and focused Python-file reads run only at core-selected points
+- Target GIL-enabled CPython explicitly with `#[pymodule(gil_used = true)]`; detach Rust-only work
+  - Free-threading requires separate runtime/concurrency validation; omitting the attribute does not opt out on PyO3 0.28+
+- Preserve public argument binding and Python object provenance
+  - Retain complete boundary arguments, opaque unknown values, aliases, omitted/default distinctions and deliberate copies; preserve the established deployment-hook kwargs view
+  - Retain independently captured body/header roots; in-place mutation and logging-envelope field replacement have different effects
+  - Project only consumed fields at reference read points; no eager whole-graph serialization or equality-based alias reconstruction
+  - Preserve provider-specific upload/submission/poll observation and encoding boundaries; signed/build-captured bytes must not be silently reserialized
+- Only core's typed, effect-free admission may return `Declined`; conversion errors and all post-admission failures are terminal
+  - Admission cannot invoke hooks, acquire credentials, consume files/iterators, prepare requests or perform I/O
+  - Disabled/unavailable native execution or an admission decline may select legacy once; callback exceptions never authorize fallback or replay
+- Use one ordinary inline `async def` driver in `litellm/rust_bridge/lifecycle.py`, with the native handle in `src/lifecycle.rs`
+  - Contract: `start`, `resume_value`, `resume_error`, idempotent `close`; explicitly tagged `Await`/`Complete` preserve awaitable final values
+  - Validate Created/Running/Suspended/Closed protocol states; core alone chooses lifecycle phases and result/error policy
+  - Defer effectful setup/context reads/timestamps until start; unstarted-handle destruction releases inputs independently of Python `finally`
+  - Catch only the selected await's errors; start/resume errors propagate, `GeneratorExit` closes without further awaits
+  - Inline hooks preserve caller task/thread/loop and context writes; `into_future` creates a separate task and cannot satisfy this contract
+  - Delivery follows the binding, not callable type; keep direct, awaited, worker, background and deferred behavior distinct
+- Finalize fallible public response/error construction, replacements and metadata under core control before terminal dispatch
+  - Success/failure handler entry receives the exact selected public response/exception; logging projections/redaction/snapshots retain their own copy contracts
+  - Ordinary failure-callback errors cannot suppress later eligible sync/async callbacks or replace the mapped provider error; control-flow exceptions have phase-specific policy
+  - Dispatch errors never replay provider work/accepted dispatch or trigger the opposite outcome; proxy acceptance/rejection releases core-owned deferred success at most once
+- Make ownership safe across suspension, re-entry, cancellation and GC
+  - Keep native provider state typed in core; do not shuttle it through opaque Python transport/response classes
+  - Prefer one retained `Py<PyBaseException>` via `PyErr::into_value(py)`; reconstruct transient `PyErr`s, preserving identity, traceback, cause and context
+  - Traverse every owned Python edge, including duplicate references; traversal cannot call Python
+  - Take state out and mark Running under a short borrow, release borrows/locks before Python invocation, publish terminal state before finalizer-capable drops
+  - Close/GC/deferred release are idempotent and re-entry-safe, including during Rust unwinding; release only owned references, never clear caller containers or mask the selected error
+  - Cancellation signaling is not termination; retain captures until work actually finishes and use a Rust-selected awaited acknowledgement where required, never synchronous close/GC
+- Verify behavior through a fresh, provenance-checked installed extension and positive native execution evidence before replacing the custom coroutine
+  - Cover admitted provider workflows, binding/read-point/identity behavior, failure continuation, finalization, no replay, deferred gates, re-entry, GC and cancellation termination
+  - Measure real conversion/copy costs before optimizing; preserve input contracts and capture lifetimes with `PyBackedBytes`, and lookup timing when interning names
+  - Ship accurate `_native.pyi` declarations and typing markers; distinguish Future-returning bindings from coroutine-returning bindings
+- References: [ownership](https://pyo3.rs/v0.29.2/types.html), [GC](https://pyo3.rs/v0.29.2/class/protocols.html#garbage-collector-integration), [exception transfer](https://docs.rs/pyo3/0.29.2/pyo3/struct.PyErr.html#method.into_value), [re-entry](https://pyo3.rs/v0.29.2/class/call.html)
+  - [GIL policy](https://pyo3.rs/v0.29.2/free-threading.html), [experimental async limits](https://pyo3.rs/v0.29.2/async-await.html), [task conversion](https://docs.rs/pyo3-async-runtimes/0.29.0/pyo3_async_runtimes/fn.into_future_with_locals.html), [native cancellation/delivery](https://docs.rs/pyo3-async-runtimes/0.29.0/pyo3_async_runtimes/tokio/fn.future_into_py.html)
+  - [performance](https://pyo3.rs/v0.29.2/performance.html), [PyBackedBytes](https://docs.rs/pyo3/0.29.2/pyo3/pybacked/struct.PyBackedBytes.html), [typing](https://pyo3.rs/v0.29.2/python-typing-hints.html)

litellm-rust/crates/python-bridge/Cargo.toml

@@ -17,7 +17,6 @@ panic-test = []
 trace-parity = [
     "dep:tracing",
     "litellm-core/observability",
-    "litellm-ai-gateway/trace-parity",
 ]
 
 [dependencies]
@@ -25,7 +24,6 @@ futures-util.workspace = true
 tracing = { workspace = true, optional = true }
 litellm-core = { workspace = true, features = ["bedrock-auth"] }
 litellm-token-counter.workspace = true
-litellm-ai-gateway = { workspace = true, default-features = false }
 litellm-python-interop.workspace = true
 pyo3.workspace = true
 pyo3-async-runtimes.workspace = true
@@ -35,6 +33,7 @@ tokio = { workspace = true, features = ["sync"] }
 
 [dev-dependencies]
 criterion.workspace = true
+rstest.workspace = true
 tokio-tungstenite.workspace = true
 tracing.workspace = true
 

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

@@ -0,0 +1,194 @@
+use litellm_core::auth::{ResolvedCredential, SecretValue};
+use pyo3::exceptions::{PyException, PyRuntimeError, PyTypeError};
+use pyo3::gc::{PyTraverseError, PyVisit};
+use pyo3::prelude::*;
+use pyo3::types::PyString;
+
+#[derive(Clone, Copy)]
+pub(crate) struct TokenProviderContract {
+    callable_error: &'static str,
+    token_type_error: &'static str,
+    callback_error: &'static str,
+}
+
+pub(crate) const AZURE_AD_TOKEN_PROVIDER: TokenProviderContract = TokenProviderContract {
+    callable_error: "Azure AD token provider must be callable",
+    token_type_error: "Azure AD token must be a string, got {}",
+    callback_error: "Failed to get Azure AD token: {}",
+};
+
+pub(crate) struct PythonTokenProvider {
+    callback: Py<PyAny>,
+    contract: TokenProviderContract,
+}
+
+impl PythonTokenProvider {
+    pub(crate) fn select(
+        provider: Bound<'_, PyAny>,
+        contract: TokenProviderContract,
+    ) -> Option<Self> {
+        (provider.is_callable() && provider.is_truthy().unwrap_or(false)).then(|| Self {
+            callback: provider.unbind(),
+            contract,
+        })
+    }
+
+    pub(crate) fn acquire(&self, py: Python<'_>) -> PyResult<ResolvedCredential> {
+        let provider = self.callback.bind(py);
+        if !provider.is_callable() {
+            return Err(PyTypeError::new_err(self.contract.callable_error));
+        }
+        let token = (|| {
+            let token = provider.call0()?;
+            if !token.is_instance_of::<PyString>() {
+                let message = PyString::new(py, self.contract.token_type_error)
+                    .call_method1("format", (token.get_type(),))?;
+                return Err(PyTypeError::new_err(message.unbind()));
+            }
+            Ok(token)
+        })()
+        .map_err(|error| {
+            if error.is_instance_of::<PyTypeError>(py) || !error.is_instance_of::<PyException>(py) {
+                return error;
+            }
+            match PyString::new(py, self.contract.callback_error)
+                .call_method1("format", (error.value(py),))
+            {
+                Ok(message) => {
+                    let wrapped = PyRuntimeError::new_err(message.unbind());
+                    wrapped.set_context(py, Some(error.clone_ref(py)));
+                    wrapped.set_cause(py, Some(error));
+                    wrapped
+                }
+                Err(format_error) => {
+                    format_error.set_context(py, Some(error));
+                    format_error
+                }
+            }
+        })?;
+        Ok(ResolvedCredential::AccessToken {
+            token: SecretValue::new(token.extract::<String>()?),
+            expires_on: None,
+        })
+    }
+
+    pub(crate) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
+        visit.call(&self.callback)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use pyo3::exceptions::PyRuntimeError;
+    use pyo3::types::PyDict;
+
+    use super::*;
+
+    #[test]
+    fn token_callback_preserves_exception_identity_and_explicit_chaining() {
+        Python::initialize();
+        Python::attach(|py| {
+            let locals = PyDict::new(py);
+            py.run(
+                pyo3::ffi::c_str!(
+                    r#"
+class ProviderError(Exception):
+    def __format__(self, specification):
+        return 'unavailable'
+ordinary = ProviderError('must use __format__')
+type_error = TypeError('signature')
+abort = KeyboardInterrupt('cancelled')
+def provider(error):
+    def acquire():
+        raise error
+    return acquire
+"#
+                ),
+                Some(&locals),
+                Some(&locals),
+            )
+            .unwrap();
+            for name in ["ordinary", "type_error", "abort"] {
+                let original = locals.get_item(name).unwrap().unwrap();
+                let callback = locals
+                    .get_item("provider")
+                    .unwrap()
+                    .unwrap()
+                    .call1((&original,))
+                    .unwrap();
+                let provider =
+                    PythonTokenProvider::select(callback, AZURE_AD_TOKEN_PROVIDER).unwrap();
+                let error = provider.acquire(py).unwrap_err();
+                if name == "ordinary" {
+                    assert!(error.is_instance_of::<PyRuntimeError>(py));
+                    assert!(error.cause(py).unwrap().value(py).is(&original));
+                    assert!(
+                        error
+                            .value(py)
+                            .getattr("__context__")
+                            .unwrap()
+                            .is(&original)
+                    );
+                    assert_eq!(
+                        error.value(py).str().unwrap().to_str().unwrap(),
+                        "Failed to get Azure AD token: unavailable"
+                    );
+                } else {
+                    assert!(error.value(py).is(&original));
+                }
+            }
+        });
+    }
+
+    #[test]
+    fn invalid_token_type_formatting_preserves_python_failure_semantics() {
+        Python::initialize();
+        Python::attach(|py| {
+            let locals = PyDict::new(py);
+            py.run(
+                pyo3::ffi::c_str!(
+                    r#"
+failure = ValueError('formatting failed')
+class TokenType(type):
+    def __format__(cls, specification):
+        raise failure
+class Token(metaclass=TokenType):
+    pass
+def provider():
+    return Token()
+"#
+                ),
+                Some(&locals),
+                Some(&locals),
+            )
+            .unwrap();
+            let provider = PythonTokenProvider::select(
+                locals.get_item("provider").unwrap().unwrap(),
+                AZURE_AD_TOKEN_PROVIDER,
+            )
+            .unwrap();
+            let error = provider.acquire(py).unwrap_err();
+            assert!(error.is_instance_of::<PyRuntimeError>(py));
+            assert!(
+                error
+                    .cause(py)
+                    .unwrap()
+                    .value(py)
+                    .is(locals.get_item("failure").unwrap().unwrap())
+            );
+        });
+    }
+
+    #[test]
+    fn token_string_extraction_errors_are_not_wrapped_as_callback_failures() {
+        Python::initialize();
+        Python::attach(|py| {
+            let callback = py
+                .eval(pyo3::ffi::c_str!("lambda: '\\ud800'"), None, None)
+                .unwrap();
+            let provider = PythonTokenProvider::select(callback, AZURE_AD_TOKEN_PROVIDER).unwrap();
+            let error = provider.acquire(py).unwrap_err();
+            assert!(error.is_instance_of::<pyo3::exceptions::PyUnicodeEncodeError>(py));
+        });
+    }
+}

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

@@ -22,7 +22,8 @@ pub(crate) fn core_error_to_pyerr(err: Error) -> PyErr {
         Error::InvalidProvider(_)
         | Error::InvalidRequest(_)
         | Error::InvalidType { .. }
-        | Error::MissingField(_) => PyValueError::new_err(err.to_string()),
+        | Error::MissingField(_)
+        | Error::MissingDocumentUrl => PyValueError::new_err(err.to_string()),
         other => PyRuntimeError::new_err(other.to_string()),
     }
 }
@@ -41,6 +42,7 @@ pub(crate) fn chat_completions_error_to_pyerr(err: Error) -> PyErr {
         | Error::InvalidRequest(_)
         | Error::InvalidType { .. }
         | Error::MissingField(_)
+        | Error::MissingDocumentUrl
         | Error::MissingApiKey { .. }
         | Error::MissingAzureAiCredentials
         | Error::MissingAzureDocumentIntelligenceCredentials
@@ -49,9 +51,7 @@ pub(crate) fn chat_completions_error_to_pyerr(err: Error) -> PyErr {
         // Nothing reached the provider, so serving it on Python cannot double
         // bill and is the only way the caller gets an answer at all.
         | Error::Connect(_) => RustBridgeDeclined::new_err(err.to_string()),
-        Error::Http { status, body } => {
-            RustUpstreamError::new_err((status, format!("{status}: {body}")))
-        }
+        Error::Http { status, body } => RustUpstreamError::new_err((status, body)),
         Error::Network(message) | Error::InvalidResponse(message) => {
             RustUpstreamError::new_err((0u16, message))
         }
@@ -63,41 +63,3 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
     module.add("RustBridgeDeclined", py.get_type::<RustBridgeDeclined>())?;
     module.add("RustUpstreamError", py.get_type::<RustUpstreamError>())
 }
-
-pub(crate) fn ocr_error_to_pyerr(err: Error) -> PyErr {
-    match err {
-        Error::MissingField("document_url" | "image_url") => {
-            PyValueError::new_err("Document URL is required")
-        }
-        Error::Http { status, body } => RustUpstreamError::new_err((status, body)),
-        other => core_error_to_pyerr(other),
-    }
-}
-
-#[cfg(test)]
-mod ocr_error_tests {
-    use super::*;
-
-    #[test]
-    fn ocr_errors_preserve_python_validation_and_provider_details() {
-        Python::initialize();
-        Python::attach(|py| {
-            for field in ["document_url", "image_url"] {
-                let mapped = ocr_error_to_pyerr(Error::MissingField(field));
-                assert!(mapped.is_instance_of::<PyValueError>(py));
-                assert_eq!(mapped.value(py).to_string(), "Document URL is required");
-            }
-            let mapped = ocr_error_to_pyerr(Error::Http {
-                status: 429,
-                body: r#"{"message":"rate limited"}"#.to_string(),
-            });
-            assert!(mapped.is_instance_of::<RustUpstreamError>(py));
-            let args: (u16, String) = mapped
-                .value(py)
-                .getattr("args")
-                .and_then(|args| args.extract())
-                .expect("OCR failures retain status and unprefixed provider message");
-            assert_eq!(args, (429, r#"{"message":"rate limited"}"#.to_string()));
-        });
-    }
-}

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

@@ -1,5 +1,7 @@
 use std::future::Future;
 use std::panic::AssertUnwindSafe;
+use std::pin::Pin;
+use std::task::{Context, Poll, Waker};
 use std::time::Duration;
 
 use futures_util::FutureExt;
@@ -28,6 +30,27 @@ where
     )
 }
 
+pub(crate) fn run_sync_value<T, F>(py: Python<'_>, future: F) -> PyResult<T>
+where
+    T: Send + 'static,
+    F: Future<Output = PyResult<T>> + Send + 'static,
+{
+    run_sync_value_on(py, pyo3_async_runtimes::tokio::get_runtime(), future)
+}
+
+fn run_sync_value_on<T, F>(py: Python<'_>, runtime: &Runtime, future: F) -> PyResult<T>
+where
+    T: Send + 'static,
+    F: Future<Output = PyResult<T>> + Send + 'static,
+{
+    if Handle::try_current().is_ok() {
+        return Err(PyRuntimeError::new_err(
+            "synchronous native routes cannot run from a Tokio context; use the async route",
+        ));
+    }
+    release_gil(py, move || runtime.block_on(wait_for_sync_result(future)))?
+}
+
 fn run_sync_on<T, E, F>(
     py: Python<'_>,
     runtime: &Runtime,
@@ -67,6 +90,32 @@ where
     })
 }
 
+pub(crate) fn run_async_value<T, F>(py: Python<'_>, future: F) -> PyResult<Bound<'_, PyAny>>
+where
+    T: for<'py> IntoPyObject<'py> + Send + 'static,
+    F: Future<Output = PyResult<T>> + Send + 'static,
+{
+    pyo3_async_runtimes::tokio::future_into_py(py, async move { catch_future_panic(future).await? })
+}
+
+pub(crate) fn poll_async_value<T, F>(py: Python<'_>, future: Pin<&mut F>) -> PyResult<Poll<T>>
+where
+    T: Send,
+    F: Future<Output = PyResult<T>> + Send,
+{
+    let result = release_gil(py, || {
+        let _runtime = pyo3_async_runtimes::tokio::get_runtime().enter();
+        std::panic::catch_unwind(AssertUnwindSafe(|| {
+            future.poll(&mut Context::from_waker(Waker::noop()))
+        }))
+        .map_err(panic_to_pyerr)
+    })?;
+    match result {
+        Poll::Ready(result) => result.map(Poll::Ready),
+        Poll::Pending => Ok(Poll::Pending),
+    }
+}
+
 fn map_core_result<T, E>(result: Result<T, E>, map_error: fn(E) -> PyErr) -> PyResult<T> {
     match result {
         Ok(value) => Ok(value),
@@ -119,11 +168,30 @@ mod tests {
     use litellm_core::error::Error;
     use pyo3::panic::PanicException;
     use pyo3::types::{PyDict, PyModule};
+    use rstest::{fixture, rstest};
     use serde::Serializer;
     use tokio::runtime::Builder;
 
     use super::*;
 
+    struct InitializedPython;
+
+    impl InitializedPython {
+        fn attach<F, R>(&self, f: F) -> R
+        where
+            F: for<'py> FnOnce(Python<'py>) -> R,
+        {
+            Python::attach(f)
+        }
+    }
+
+    #[fixture]
+    #[once]
+    fn initialized_python() -> InitializedPython {
+        Python::initialize();
+        InitializedPython
+    }
+
     fn runtime_error(error: Error) -> PyErr {
         PyRuntimeError::new_err(error.to_string())
     }
@@ -194,10 +262,84 @@ mod tests {
             .expect("result should convert")
     }
 
-    #[test]
-    fn sync_runner_polls_future_on_the_caller_thread() {
-        Python::initialize();
-        Python::attach(|py| {
+    #[rstest]
+    fn inline_poll_releases_gil_and_enters_runtime(
+        #[from(initialized_python)] python: &InitializedPython,
+    ) {
+        python.attach(|py| {
+            let (sender, receiver) = mpsc::sync_channel(1);
+            let worker = thread::spawn(move || Python::attach(|_| sender.send(()).unwrap()));
+            let mut future = Box::pin(async move {
+                receiver.recv_timeout(Duration::from_secs(2)).unwrap();
+                Ok(Handle::try_current().is_ok())
+            });
+            assert_eq!(
+                poll_async_value(py, future.as_mut()).unwrap(),
+                Poll::Ready(true)
+            );
+            worker.join().unwrap();
+        });
+    }
+
+    #[rstest]
+    fn inline_poll_contains_panics_and_preserves_python_errors(
+        #[from(initialized_python)] python: &InitializedPython,
+    ) {
+        python.attach(|py| {
+            let mut panicking = Box::pin(poll_fn(|_| -> Poll<PyResult<()>> {
+                panic!("inline native panic")
+            }));
+            let error = poll_async_value(py, panicking.as_mut()).unwrap_err();
+            assert!(error.is_instance_of::<PanicException>(py));
+            let original = PyRuntimeError::new_err("inline failure");
+            let identity = original.value(py).clone().unbind();
+            let mut failing = Box::pin(async move { Err::<(), _>(original) });
+            let error = poll_async_value(py, failing.as_mut()).unwrap_err();
+            assert!(error.value(py).is(identity.bind(py)));
+        });
+    }
+
+    #[pyfunction]
+    fn pending_after_inline_poll(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
+        let starts = Arc::new(AtomicUsize::new(0));
+        let observed = Arc::clone(&starts);
+        let mut future = Box::pin(async move {
+            starts.fetch_add(1, Ordering::SeqCst);
+            tokio::time::sleep(Duration::from_millis(5)).await;
+            Ok(starts.load(Ordering::SeqCst))
+        });
+        assert!(poll_async_value(py, future.as_mut())?.is_pending());
+        assert_eq!(observed.load(Ordering::SeqCst), 1);
+        run_async_value(py, future)
+    }
+
+    #[rstest]
+    fn inline_pending_future_resumes_on_tokio_without_restarting(
+        #[from(initialized_python)] python: &InitializedPython,
+    ) {
+        python.attach(|py| {
+            let locals = PyDict::new(py);
+            locals
+                .set_item(
+                    "pending",
+                    wrap_pyfunction!(pending_after_inline_poll, py).unwrap(),
+                )
+                .unwrap();
+            py.run(
+                pyo3::ffi::c_str!(
+                    "import asyncio\nasync def exercise():\n    assert await asyncio.wait_for(pending(), 2) == 1\nasyncio.run(exercise())"
+                ),
+                Some(&locals),
+                Some(&locals),
+            ).unwrap();
+        });
+    }
+
+    #[rstest]
+    fn sync_runner_polls_future_on_the_caller_thread(
+        #[from(initialized_python)] python: &InitializedPython,
+    ) {
+        python.attach(|py| {
             let caller_thread = std::thread::current().id();
             let result = run_sync(
                 py,
@@ -209,10 +351,11 @@ mod tests {
         });
     }
 
-    #[test]
-    fn sync_runner_releases_gil_while_waiting() {
-        Python::initialize();
-        Python::attach(|py| {
+    #[rstest]
+    fn sync_runner_releases_gil_while_waiting(
+        #[from(initialized_python)] python: &InitializedPython,
+    ) {
+        python.attach(|py| {
             let result = run_sync(
                 py,
                 async {
@@ -230,16 +373,17 @@ mod tests {
         });
     }
 
-    #[test]
-    fn sync_runner_rejects_calls_from_a_tokio_context() {
-        Python::initialize();
+    #[rstest]
+    fn sync_runner_rejects_calls_from_a_tokio_context(
+        #[from(initialized_python)] python: &InitializedPython,
+    ) {
         let runtime = Builder::new_current_thread()
             .enable_all()
             .build()
             .expect("runtime should build");
 
         let error = runtime.block_on(async {
-            Python::attach(|py| {
+            python.attach(|py| {
                 run_sync::<bool, Error, _>(py, async { Ok(true) }, runtime_error)
                     .expect_err("sync route should reject a nested Tokio runtime")
             })
@@ -251,14 +395,15 @@ mod tests {
         );
     }
 
-    #[test]
-    fn sync_runner_can_drive_a_current_thread_runtime() {
-        Python::initialize();
+    #[rstest]
+    fn sync_runner_can_drive_a_current_thread_runtime(
+        #[from(initialized_python)] python: &InitializedPython,
+    ) {
         let runtime = Builder::new_current_thread()
             .enable_all()
             .build()
             .expect("runtime should build");
-        Python::attach(|py| {
+        python.attach(|py| {
             let result = run_sync_on(
                 py,
                 &runtime,
@@ -272,10 +417,9 @@ mod tests {
         });
     }
 
-    #[test]
-    fn sync_runner_maps_a_panicked_future() {
-        Python::initialize();
-        Python::attach(|py| {
+    #[rstest]
+    fn sync_runner_maps_a_panicked_future(#[from(initialized_python)] python: &InitializedPython) {
+        python.attach(|py| {
             let error = run_sync::<bool, Error, _>(
                 py,
                 poll_fn(|_| -> Poll<Result<bool, Error>> { panic!("route future panicked") }),
@@ -288,10 +432,11 @@ mod tests {
         });
     }
 
-    #[test]
-    fn sync_runner_maps_a_panicked_error_mapper() {
-        Python::initialize();
-        Python::attach(|py| {
+    #[rstest]
+    fn sync_runner_maps_a_panicked_error_mapper(
+        #[from(initialized_python)] python: &InitializedPython,
+    ) {
+        python.attach(|py| {
             let error = run_sync::<bool, Error, _>(
                 py,
                 async { Err(Error::InvalidRequest("invalid".to_string())) },
@@ -304,10 +449,11 @@ mod tests {
         });
     }
 
-    #[test]
-    fn sync_runner_surfaces_serializer_panics() {
-        Python::initialize();
-        Python::attach(|py| {
+    #[rstest]
+    fn sync_runner_surfaces_serializer_panics(
+        #[from(initialized_python)] python: &InitializedPython,
+    ) {
+        python.attach(|py| {
             let error = run_sync(py, async { Ok(PanickingOutput) }, runtime_error)
                 .expect_err("serializer panic should become a Python exception");
 
@@ -316,9 +462,10 @@ mod tests {
         });
     }
 
-    #[test]
-    fn sync_runner_supports_concurrent_callers_on_the_shared_runtime() {
-        Python::initialize();
+    #[rstest]
+    fn sync_runner_supports_concurrent_callers_on_the_shared_runtime(
+        #[from(initialized_python)] _python: &InitializedPython,
+    ) {
         let barrier = Arc::new(tokio::sync::Barrier::new(2));
         let callers: Vec<_> = (0..2)
             .map(|_| {
@@ -349,10 +496,11 @@ mod tests {
         assert_eq!(results, vec![true, true]);
     }
 
-    #[test]
-    fn async_runner_surfaces_serializer_panics() {
-        Python::initialize();
-        Python::attach(|py| {
+    #[rstest]
+    fn async_runner_surfaces_serializer_panics(
+        #[from(initialized_python)] python: &InitializedPython,
+    ) {
+        python.attach(|py| {
             let module = PyModule::new(py, "runtime").expect("module should be created");
             module
                 .add_function(
@@ -386,11 +534,12 @@ asyncio.run(exercise())
         });
     }
 
-    #[test]
-    fn async_result_delivery_does_not_stall_tokio_workers() {
-        Python::initialize();
+    #[rstest]
+    fn async_result_delivery_does_not_stall_tokio_workers(
+        #[from(initialized_python)] python: &InitializedPython,
+    ) {
         ASYNC_PROBE_COMPLETED.store(0, Ordering::SeqCst);
-        Python::attach(|py| {
+        python.attach(|py| {
             let module = PyModule::new(py, "runtime").expect("module should be created");
             for function in [
                 wrap_pyfunction!(async_runtime_probe, &module).expect("function should wrap"),

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

@@ -1,14 +1,16 @@
+mod auth;
 mod constants;
 mod diagnostics;
 mod errors;
 mod execution;
 #[cfg(feature = "trace-parity")]
 mod function_trace;
+mod lifecycle;
 mod marshal;
 mod routes;
 mod token_counter;
 
-use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection as RustResponsesWebSocketConnection;
+use litellm_core::responses::websocket::ResponsesWebSocketConnection as RustResponsesWebSocketConnection;
 use pyo3::prelude::*;
 use pyo3::types::PyAny;
 use serde_json::Value;
@@ -64,7 +66,7 @@ impl ResponsesWebSocketConnection {
     }
 }
 
-#[pymodule(gil_used = false)]
+#[pymodule(gil_used = true)]
 mod _native {
     use pyo3::prelude::*;
 
@@ -152,7 +154,6 @@ mod tests {
                         "amessages",
                         "chat_completions",
                         "achat_completions",
-                        "gateway_messages",
                     ]
                 );
             }

litellm-rust/crates/python-bridge/src/lifecycle/bindings.rs

@@ -0,0 +1,391 @@
+use pyo3::exceptions::PyBaseException;
+use pyo3::gc::{PyTraverseError, PyVisit};
+use pyo3::prelude::*;
+use pyo3::types::{PyDict, PyTuple};
+
+#[derive(FromPyObject)]
+pub(crate) struct PythonLogger(Py<PyAny>);
+
+impl PythonLogger {
+    pub(crate) fn object<'py>(&self, py: Python<'py>) -> &Bound<'py, PyAny> {
+        self.0.bind(py)
+    }
+
+    pub(crate) fn clone_ref(&self, py: Python<'_>) -> Self {
+        Self(self.0.clone_ref(py))
+    }
+
+    pub(crate) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
+        visit.call(&self.0)
+    }
+
+    pub(crate) fn callbacks_needed(&self, py: Python<'_>, phase: &str) -> PyResult<bool> {
+        if !self
+            .object(py)
+            .getattr("_native_callback_fast_path")
+            .is_ok_and(|value| value.is_truthy().unwrap_or(false))
+        {
+            return Ok(true);
+        }
+        py.import("litellm.rust_bridge.lifecycle")?
+            .getattr("callbacks_needed")?
+            .call1((self.object(py), phase))?
+            .extract()
+    }
+
+    pub(super) fn success_bookkeeping(
+        &self,
+        py: Python<'_>,
+        response: &Option<Py<PyAny>>,
+        start: &Py<PyAny>,
+        end: &Option<Py<PyAny>>,
+        asynchronous: bool,
+    ) -> PyResult<()> {
+        py.import("litellm.rust_bridge.lifecycle")?
+            .getattr("success_bookkeeping")?
+            .call1((self.object(py), response, start, end, asynchronous))?;
+        Ok(())
+    }
+
+    pub(super) fn defers_async_logging(&self, py: Python<'_>) -> bool {
+        self.object(py)
+            .getattr("_defer_async_logging")
+            .is_ok_and(|value| value.is_truthy().unwrap_or(false))
+    }
+
+    pub(super) fn defer_success(
+        &self,
+        py: Python<'_>,
+        pending: Py<super::PendingLogging>,
+    ) -> PyResult<()> {
+        self.object(py).setattr("_native_pending_logging", pending)
+    }
+
+    pub(super) fn sync_success_for_async_call(
+        &self,
+        py: Python<'_>,
+        response: &Option<Py<PyAny>>,
+        start: &Py<PyAny>,
+        end: &Option<Py<PyAny>>,
+    ) -> PyResult<()> {
+        if !self.callbacks_needed(py, "sync_success_async")? {
+            return Ok(());
+        }
+        self.object(py).call_method1(
+            "handle_sync_success_callbacks_for_async_calls",
+            (response, start, end),
+        )?;
+        Ok(())
+    }
+
+    pub(super) fn failure(
+        &self,
+        py: Python<'_>,
+        error: &Py<PyBaseException>,
+        start: &Py<PyAny>,
+        end: &Option<Py<PyAny>>,
+        asynchronous: bool,
+    ) -> PyResult<Option<Py<PyAny>>> {
+        if !self.callbacks_needed(
+            py,
+            if asynchronous {
+                "async_failure"
+            } else {
+                "sync_failure"
+            },
+        )? {
+            py.import("litellm.rust_bridge.lifecycle")?
+                .getattr("failure_bookkeeping")?
+                .call1((self.object(py), error, start, end, asynchronous))?;
+            return Ok(None);
+        }
+        let trace = py
+            .import("traceback")?
+            .getattr("format_exception")?
+            .call1((error,))?;
+        let trace = pyo3::types::PyString::new(py, "").call_method1("join", (trace,))?;
+        let value = self.object(py).call_method1(
+            if asynchronous {
+                "async_failure_handler"
+            } else {
+                "failure_handler"
+            },
+            (error, trace, start, end),
+        )?;
+        Ok(asynchronous.then(|| value.unbind()))
+    }
+
+    pub(super) fn restore_context(&self, py: Python<'_>) -> PyResult<()> {
+        py.import("litellm.utils")?
+            .getattr("_restore_correlation_context_if_supported")?
+            .call1((self.object(py),))?;
+        Ok(())
+    }
+
+    pub(super) fn submit_success(
+        &self,
+        py: Python<'_>,
+        response: &Option<Py<PyAny>>,
+        start: &Py<PyAny>,
+        end: &Option<Py<PyAny>>,
+    ) -> PyResult<()> {
+        if !self.callbacks_needed(py, "sync_success")? {
+            return self.success_bookkeeping(py, response, start, end, false);
+        }
+        let context = py.import("contextvars")?.call_method0("copy_context")?;
+        py.import("litellm.litellm_core_utils.litellm_logging")?
+            .getattr("executor")?
+            .call_method1(
+                "submit",
+                (
+                    context.getattr("run")?,
+                    self.object(py).getattr("success_handler")?,
+                    response,
+                    start,
+                    end,
+                ),
+            )?;
+        Ok(())
+    }
+
+    pub(super) fn enqueue_success(
+        &self,
+        py: Python<'_>,
+        response: &Option<Py<PyAny>>,
+        start: &Py<PyAny>,
+        end: &Option<Py<PyAny>>,
+    ) -> PyResult<()> {
+        if !self.callbacks_needed(py, "async_success")? {
+            return self.success_bookkeeping(py, response, start, end, true);
+        }
+        let context = py.import("contextvars")?.call_method0("copy_context")?;
+        let worker = py
+            .import("litellm.litellm_core_utils.logging_worker")?
+            .getattr("GLOBAL_LOGGING_WORKER")?
+            .getattr("ensure_initialized_and_enqueue")?;
+        let coroutine = self
+            .object(py)
+            .call_method1("async_success_handler", (response, start, end))?;
+        let enqueue = context.call_method1("run", (worker, &coroutine));
+        if enqueue.is_err()
+            && let Err(error) = coroutine.call_method0("close")
+        {
+            error.write_unraisable(py, Some(&coroutine));
+        }
+        enqueue.map(|_| ())
+    }
+}
+
+pub(super) struct SetupResult<'py>(Bound<'py, PyAny>);
+
+impl SetupResult<'_> {
+    pub(super) fn logger(&self) -> PyResult<PythonLogger> {
+        self.0.getattr("logger")?.extract()
+    }
+
+    pub(super) fn kwargs(&self) -> PyResult<Py<PyDict>> {
+        Ok(self.0.getattr("kwargs")?.extract()?)
+    }
+}
+
+pub(super) fn setup<'py>(
+    py: Python<'py>,
+    call_type: &str,
+    args: &Py<PyTuple>,
+    kwargs: &Py<PyDict>,
+    start: &Py<PyAny>,
+    asynchronous: bool,
+) -> PyResult<SetupResult<'py>> {
+    py.import("litellm.rust_bridge.lifecycle")?
+        .getattr("setup")?
+        .call1((call_type, args, kwargs, start, asynchronous))
+        .map(SetupResult)
+}
+
+pub(super) fn finalize(
+    py: Python<'_>,
+    response: &Option<Py<PyAny>>,
+    logger: &PythonLogger,
+    kwargs: &Py<PyDict>,
+    start: &Py<PyAny>,
+    end: &Option<Py<PyAny>>,
+) -> PyResult<()> {
+    py.import("litellm.rust_bridge.lifecycle")?
+        .getattr("finalize")?
+        .call1((response, logger.object(py), kwargs, start, end))?;
+    Ok(())
+}
+
+pub(super) fn is_internal_call(py: Python<'_>) -> PyResult<bool> {
+    py.import("litellm._internal_context")?
+        .getattr("is_internal_call")?
+        .call_method0("get")?
+        .extract()
+}
+
+pub(super) struct DeploymentHooks;
+
+impl DeploymentHooks {
+    pub(super) fn needed(py: Python<'_>) -> PyResult<bool> {
+        py.import("litellm.rust_bridge.lifecycle")?
+            .getattr("deployment_callbacks_needed")?
+            .call0()?
+            .extract()
+    }
+
+    pub(super) fn before_call(
+        py: Python<'_>,
+        kwargs: &Py<PyDict>,
+        call_type: &str,
+    ) -> PyResult<Py<PyAny>> {
+        py.import("litellm.utils")?
+            .getattr("async_pre_call_deployment_hook")?
+            .call1((kwargs, call_type))
+            .map(Bound::unbind)
+    }
+
+    pub(super) fn after_success(
+        py: Python<'_>,
+        kwargs: &Py<PyDict>,
+        response: &Option<Py<PyAny>>,
+        call_type: &str,
+    ) -> PyResult<Py<PyAny>> {
+        py.import("litellm.utils")?
+            .getattr("async_post_call_success_deployment_hook")?
+            .call1((kwargs, response, call_type))
+            .map(Bound::unbind)
+    }
+
+    pub(super) fn after_failure(
+        py: Python<'_>,
+        kwargs: &Py<PyDict>,
+        error: &Py<PyBaseException>,
+        call_type: &str,
+    ) -> PyResult<Py<PyAny>> {
+        py.import("litellm.utils")?
+            .getattr("async_post_call_failure_deployment_hook")?
+            .call1((kwargs, error, call_type))
+            .map(Bound::unbind)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use pyo3::exceptions::PyTypeError;
+
+    #[test]
+    fn setup_fields_are_checked_in_order_without_eager_logger_method_reads() {
+        Python::initialize();
+        Python::attach(|py| {
+            let locals = PyDict::new(py);
+            py.run(
+                pyo3::ffi::c_str!(
+                    r#"
+reads = []
+class Logger:
+    def __getattribute__(self, name):
+        reads.append(name)
+        raise AssertionError('logger methods must remain lazy')
+logger = Logger()
+class Setup:
+    @property
+    def logger(self):
+        reads.append('logger')
+        return logger
+    @property
+    def kwargs(self):
+        reads.append('kwargs')
+        return []
+result = Setup()
+"#
+                ),
+                Some(&locals),
+                Some(&locals),
+            )
+            .unwrap();
+            let result = SetupResult(locals.get_item("result").unwrap().unwrap());
+            let logger = result.logger().unwrap();
+            assert!(
+                logger
+                    .object(py)
+                    .is(locals.get_item("logger").unwrap().unwrap())
+            );
+            assert_eq!(
+                locals
+                    .get_item("reads")
+                    .unwrap()
+                    .unwrap()
+                    .extract::<Vec<String>>()
+                    .unwrap(),
+                ["logger"]
+            );
+            assert!(
+                result
+                    .kwargs()
+                    .unwrap_err()
+                    .is_instance_of::<PyTypeError>(py)
+            );
+            assert_eq!(
+                locals
+                    .get_item("reads")
+                    .unwrap()
+                    .unwrap()
+                    .extract::<Vec<String>>()
+                    .unwrap(),
+                ["logger", "kwargs"]
+            );
+        });
+    }
+
+    #[test]
+    fn logger_resolves_each_callback_at_invocation_and_preserves_arguments() {
+        Python::initialize();
+        Python::attach(|py| {
+            let locals = PyDict::new(py);
+            py.run(
+                pyo3::ffi::c_str!(
+                    r#"
+calls = []
+response, start, end = object(), object(), object()
+class Logger:
+    @property
+    def handle_sync_success_callbacks_for_async_calls(self):
+        generation = len(calls)
+        def callback(*args):
+            assert args == (response, start, end)
+            calls.append(generation)
+        return callback
+logger = Logger()
+"#
+                ),
+                Some(&locals),
+                Some(&locals),
+            )
+            .unwrap();
+            let logger: PythonLogger = locals
+                .get_item("logger")
+                .unwrap()
+                .unwrap()
+                .extract()
+                .unwrap();
+            let response = Some(locals.get_item("response").unwrap().unwrap().unbind());
+            let start = locals.get_item("start").unwrap().unwrap().unbind();
+            let end = Some(locals.get_item("end").unwrap().unwrap().unbind());
+            for _ in 0..2 {
+                logger
+                    .sync_success_for_async_call(py, &response, &start, &end)
+                    .unwrap();
+            }
+            assert_eq!(
+                locals
+                    .get_item("calls")
+                    .unwrap()
+                    .unwrap()
+                    .extract::<Vec<usize>>()
+                    .unwrap(),
+                [0, 1]
+            );
+        });
+    }
+}

litellm-rust/crates/python-bridge/src/lifecycle/handle.rs

@@ -0,0 +1,139 @@
+use std::panic::{AssertUnwindSafe, catch_unwind};
+
+use litellm_python_interop::panic_to_pyerr;
+use pyo3::exceptions::{PyBaseException, PyRuntimeError};
+use pyo3::gc::{PyTraverseError, PyVisit};
+use pyo3::prelude::*;
+
+pub(super) enum ExecutionStep {
+    Return(Py<PyAny>),
+    Await(Py<PyAny>),
+}
+
+pub(super) trait ExecutionBody: Send + Sync {
+    fn resume(&mut self, result: Option<PyResult<Py<PyAny>>>) -> PyResult<ExecutionStep>;
+    fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError>;
+}
+
+enum ExecutionState {
+    Created(Box<dyn ExecutionBody>),
+    Running,
+    Suspended(Box<dyn ExecutionBody>),
+    Closed,
+}
+
+#[pyclass]
+pub(super) struct Execution {
+    state: ExecutionState,
+}
+
+impl Execution {
+    pub(super) fn new(body: impl ExecutionBody + 'static) -> Self {
+        Self {
+            state: ExecutionState::Created(Box::new(body)),
+        }
+    }
+
+    fn advance(
+        slf: &Bound<'_, Self>,
+        py: Python<'_>,
+        result: Option<PyResult<Py<PyAny>>>,
+    ) -> PyResult<Py<PyAny>> {
+        let mut body = {
+            let mut execution = slf.borrow_mut();
+            match (&execution.state, result.is_some()) {
+                (ExecutionState::Created(_), false) | (ExecutionState::Suspended(_), true) => {}
+                (ExecutionState::Running, _) => {
+                    return Err(PyRuntimeError::new_err("execution is already running"));
+                }
+                (ExecutionState::Closed, _) => {
+                    return Err(PyRuntimeError::new_err("execution is closed"));
+                }
+                _ => {
+                    return Err(PyRuntimeError::new_err(
+                        "execution requires start before resume and can only start once",
+                    ));
+                }
+            }
+            match std::mem::replace(&mut execution.state, ExecutionState::Running) {
+                ExecutionState::Created(body) | ExecutionState::Suspended(body) => body,
+                _ => unreachable!(),
+            }
+        };
+        let outcome = catch_unwind(AssertUnwindSafe(|| {
+            let step = body.resume(result)?;
+            let (tag, value, suspended) = match step {
+                ExecutionStep::Await(value) => ("Await", value, true),
+                ExecutionStep::Return(value) => ("Complete", value, false),
+            };
+            let step = py
+                .import("litellm.rust_bridge.lifecycle")?
+                .getattr(tag)?
+                .call1((value,))?
+                .unbind();
+            Ok((step, suspended))
+        }))
+        .map_err(panic_to_pyerr)
+        .and_then(|result| result);
+        match outcome {
+            Ok((step, true)) if matches!(slf.borrow().state, ExecutionState::Running) => {
+                slf.borrow_mut().state = ExecutionState::Suspended(body);
+                Ok(step)
+            }
+            outcome => {
+                slf.borrow_mut().state = ExecutionState::Closed;
+                drop(body);
+                outcome.and_then(|(step, suspended)| {
+                    if suspended {
+                        Err(PyRuntimeError::new_err(
+                            "execution was closed while running",
+                        ))
+                    } else {
+                        Ok(step)
+                    }
+                })
+            }
+        }
+    }
+}
+
+#[pymethods]
+impl Execution {
+    fn start(slf: &Bound<'_, Self>, py: Python<'_>) -> PyResult<Py<PyAny>> {
+        Self::advance(slf, py, None)
+    }
+
+    fn resume_value(
+        slf: &Bound<'_, Self>,
+        py: Python<'_>,
+        value: Py<PyAny>,
+    ) -> PyResult<Py<PyAny>> {
+        Self::advance(slf, py, Some(Ok(value)))
+    }
+
+    fn resume_error(
+        slf: &Bound<'_, Self>,
+        py: Python<'_>,
+        error: Bound<'_, PyBaseException>,
+    ) -> PyResult<Py<PyAny>> {
+        Self::advance(slf, py, Some(Err(PyErr::from_value(error.into_any()))))
+    }
+
+    fn close(slf: &Bound<'_, Self>) {
+        let state = std::mem::replace(&mut slf.borrow_mut().state, ExecutionState::Closed);
+        drop(state);
+    }
+
+    fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
+        match &self.state {
+            ExecutionState::Created(body) | ExecutionState::Suspended(body) => {
+                body.traverse(&visit)
+            }
+            _ => Ok(()),
+        }
+    }
+
+    fn __clear__(slf: &Bound<'_, Self>) {
+        Self::close(slf);
+    }
+}

litellm-rust/crates/python-bridge/src/lifecycle/preparation.rs

@@ -0,0 +1,314 @@
+use litellm_core::auth::{credential_default_fields, credential_index};
+use pyo3::prelude::*;
+use pyo3::types::{PyDict, PyList};
+
+struct CredentialEntry<'py>(Bound<'py, PyAny>);
+
+impl<'py> CredentialEntry<'py> {
+    fn name(&self) -> PyResult<String> {
+        self.0.getattr("credential_name")?.extract()
+    }
+
+    fn values(&self) -> PyResult<Bound<'py, PyDict>> {
+        Ok(self.0.getattr("credential_values")?.cast_into::<PyDict>()?)
+    }
+}
+
+pub(super) fn prepare<'py>(
+    py: Python<'py>,
+    kwargs: &Bound<'py, PyDict>,
+    logger: &super::PythonLogger,
+) -> PyResult<Bound<'py, PyDict>> {
+    let arguments = kwargs.copy()?;
+    arguments.set_item("litellm_logging_obj", logger.object(py))?;
+    let litellm = py.import("litellm")?;
+    inherit_credentials(py, &litellm, &arguments)?;
+    py.import("litellm.rust_bridge.lifecycle")?
+        .getattr("check_limits")?
+        .call1((&arguments,))?;
+    Ok(arguments)
+}
+
+fn inherit_credentials(
+    py: Python<'_>,
+    litellm: &Bound<'_, PyModule>,
+    arguments: &Bound<'_, PyDict>,
+) -> PyResult<()> {
+    let Some(requested) = arguments
+        .get_item("litellm_credential_name")?
+        .filter(|value| !value.is_none())
+    else {
+        return Ok(());
+    };
+    if !requested.is_truthy()? {
+        return Ok(());
+    }
+    let requested: String = requested.extract()?;
+    let credentials = litellm.getattr("credential_list")?.cast_into::<PyList>()?;
+    let names = credentials
+        .iter()
+        .map(|credential| CredentialEntry(credential).name())
+        .collect::<PyResult<Vec<_>>>()?;
+    let Some(index) = credential_index(&requested, &names) else {
+        py.import("litellm._logging")?.getattr("verbose_logger")?.call_method1(
+            "warning",
+            ("litellm_credential_name=%s matched none of the %d loaded credentials; the request runs without it", requested, names.len()),
+        )?;
+        return Ok(());
+    };
+    let selected = CredentialEntry(credentials.get_item(index)?);
+    let values = selected.values()?;
+    let supplied: Vec<String> = arguments.keys().extract()?;
+    let fields: Vec<String> = values.keys().extract()?;
+    for name in credential_default_fields(&supplied, &fields) {
+        if let Some(value) = values.get_item(name)? {
+            arguments.set_item(name, value)?;
+        }
+    }
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn eval<'py>(py: Python<'py>, source: &std::ffi::CStr) -> Bound<'py, PyDict> {
+        let locals = PyDict::new(py);
+        py.run(source, Some(&locals), Some(&locals)).unwrap();
+        locals
+    }
+
+    fn inherit(py: Python<'_>, locals: &Bound<'_, PyDict>) -> PyResult<()> {
+        let litellm = PyModule::new(py, "credential_host")?;
+        litellm.setattr(
+            "credential_list",
+            locals.get_item("credentials").unwrap().unwrap(),
+        )?;
+        inherit_credentials(
+            py,
+            &litellm,
+            &locals
+                .get_item("arguments")
+                .unwrap()
+                .unwrap()
+                .cast_into::<PyDict>()?,
+        )
+    }
+
+    #[test]
+    fn duplicate_names_select_the_first_entry_without_reading_other_values() {
+        Python::initialize();
+        Python::attach(|py| {
+            let locals = eval(
+                py,
+                c"
+accesses = []
+class Credential:
+    def __init__(self, name, values):
+        self._name = name
+        self._values = values
+    @property
+    def credential_name(self):
+        accesses.append(('name', self._name))
+        return self._name
+    @property
+    def credential_values(self):
+        accesses.append(('values', self._name))
+        return self._values
+credentials = [
+    Credential('ocr-test', {'api_key': 'first'}),
+    Credential('other', {'api_key': 'unused'}),
+    Credential('ocr-test', {'api_key': 'later'}),
+]
+arguments = {'litellm_credential_name': 'ocr-test'}
+",
+            );
+            inherit(py, &locals).unwrap();
+            let arguments = locals
+                .get_item("arguments")
+                .unwrap()
+                .unwrap()
+                .cast_into::<PyDict>()
+                .unwrap();
+            assert_eq!(
+                arguments
+                    .get_item("api_key")
+                    .unwrap()
+                    .unwrap()
+                    .extract::<String>()
+                    .unwrap(),
+                "first"
+            );
+            let accesses: Vec<(String, String)> = locals
+                .get_item("accesses")
+                .unwrap()
+                .unwrap()
+                .extract()
+                .unwrap();
+            assert_eq!(
+                accesses,
+                [
+                    ("name".into(), "ocr-test".into()),
+                    ("name".into(), "other".into()),
+                    ("name".into(), "ocr-test".into()),
+                    ("values".into(), "ocr-test".into()),
+                ]
+            );
+        });
+    }
+
+    #[test]
+    fn later_invalid_name_still_fails_after_an_earlier_match() {
+        Python::initialize();
+        Python::attach(|py| {
+            let locals = eval(
+                py,
+                c"
+failure = LookupError('later name')
+class Good:
+    credential_name = 'ocr-test'
+    credential_values = {'api_key': 'first'}
+class Bad:
+    @property
+    def credential_name(self):
+        raise failure
+credentials = [Good(), Bad()]
+arguments = {'litellm_credential_name': 'ocr-test'}
+",
+            );
+            let error = inherit(py, &locals).unwrap_err();
+            assert!(
+                error
+                    .value(py)
+                    .is(locals.get_item("failure").unwrap().unwrap())
+            );
+        });
+    }
+
+    #[test]
+    fn selected_values_must_be_a_dictionary_and_property_errors_keep_identity() {
+        Python::initialize();
+        Python::attach(|py| {
+            let locals = eval(
+                py,
+                c"
+class Listed:
+    credential_name = 'ocr-test'
+    credential_values = ['not-a-dict']
+credentials = [Listed()]
+arguments = {'litellm_credential_name': 'ocr-test'}
+",
+            );
+            assert!(
+                inherit(py, &locals)
+                    .unwrap_err()
+                    .is_instance_of::<pyo3::exceptions::PyTypeError>(py)
+            );
+
+            let locals = eval(
+                py,
+                c"
+failure = RuntimeError('values failed')
+class Broken:
+    credential_name = 'ocr-test'
+    @property
+    def credential_values(self):
+        raise failure
+credentials = [Broken()]
+arguments = {'litellm_credential_name': 'ocr-test'}
+",
+            );
+            let error = inherit(py, &locals).unwrap_err();
+            assert!(
+                error
+                    .value(py)
+                    .is(locals.get_item("failure").unwrap().unwrap())
+            );
+        });
+    }
+
+    #[test]
+    fn explicit_none_is_not_overwritten_and_inherited_objects_keep_identity() {
+        Python::initialize();
+        Python::attach(|py| {
+            let locals = eval(
+                py,
+                c"
+opaque = object()
+class Credential:
+    credential_name = 'ocr-test'
+    credential_values = {'api_key': 'credential-key', 'opaque': opaque}
+credentials = [Credential()]
+arguments = {'litellm_credential_name': 'ocr-test', 'api_key': None}
+",
+            );
+            inherit(py, &locals).unwrap();
+            let arguments = locals
+                .get_item("arguments")
+                .unwrap()
+                .unwrap()
+                .cast_into::<PyDict>()
+                .unwrap();
+            assert!(arguments.get_item("api_key").unwrap().unwrap().is_none());
+            assert!(
+                arguments
+                    .get_item("opaque")
+                    .unwrap()
+                    .unwrap()
+                    .is(locals.get_item("opaque").unwrap().unwrap())
+            );
+        });
+    }
+
+    #[test]
+    fn selection_rereads_the_list_after_name_properties_run() {
+        Python::initialize();
+        Python::attach(|py| {
+            let locals = eval(
+                py,
+                c"
+class First:
+    @property
+    def credential_name(self):
+        credentials[0] = Second()
+        return 'ocr-test'
+    credential_values = {'api_key': 'first'}
+class Second:
+    credential_name = 'ocr-test'
+    credential_values = {'api_key': 'replaced'}
+credentials = [First()]
+arguments = {'litellm_credential_name': 'ocr-test'}
+",
+            );
+            inherit(py, &locals).unwrap();
+            let arguments = locals
+                .get_item("arguments")
+                .unwrap()
+                .unwrap()
+                .cast_into::<PyDict>()
+                .unwrap();
+            assert_eq!(
+                arguments
+                    .get_item("api_key")
+                    .unwrap()
+                    .unwrap()
+                    .extract::<String>()
+                    .unwrap(),
+                "replaced"
+            );
+        });
+    }
+
+    #[test]
+    fn falsy_credential_names_return_before_loading_credentials() {
+        Python::initialize();
+        Python::attach(|py| {
+            let litellm = PyModule::new(py, "credential_host").unwrap();
+            for name in [py.None(), py.eval(c"''", None, None).unwrap().unbind()] {
+                let arguments = PyDict::new(py);
+                arguments.set_item("litellm_credential_name", name).unwrap();
+                inherit_credentials(py, &litellm, &arguments).unwrap();
+            }
+        });
+    }
+}

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

@@ -1,10 +1,14 @@
-use std::collections::HashMap;
+use std::collections::{BTreeMap, HashMap};
 use std::time::Duration;
 
 use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
+use pyo3::types::PyDict;
 use serde_json::{Map, Value};
 
+use litellm_core::auth::InputSource;
+use litellm_python_interop::from_py_preserving_errors as from_py;
+
 pub(crate) struct RouteOptions {
     pub(crate) model: String,
     pub(crate) api_key: Option<String>,
@@ -36,18 +40,18 @@ impl RouteOptions {
     }
 }
 
-pub(crate) fn required_value(
-    name: &'static str,
-    value: Value,
-    expected: fn(&Value) -> bool,
-    expected_name: &'static str,
-) -> PyResult<Value> {
-    if expected(&value) {
-        return Ok(value);
+pub(crate) fn required_array(name: &'static str, value: Value) -> PyResult<Vec<Value>> {
+    match value {
+        Value::Array(values) => Ok(values),
+        _ => Err(PyValueError::new_err(format!("{name} must be a list"))),
+    }
+}
+
+pub(crate) fn required_object(name: &'static str, value: Value) -> PyResult<Map<String, Value>> {
+    match value {
+        Value::Object(values) => Ok(values),
+        _ => Err(PyValueError::new_err(format!("{name} must be a dict"))),
     }
-    Err(PyValueError::new_err(format!(
-        "{name} must be a {expected_name}"
-    )))
 }
 
 pub(crate) fn object_or_empty(
@@ -55,7 +59,7 @@ pub(crate) fn object_or_empty(
     value: Option<Value>,
 ) -> PyResult<Map<String, Value>> {
     match value {
-        Some(value) => object(name, value),
+        Some(value) => required_object(name, value),
         None => Ok(Map::new()),
     }
 }
@@ -64,14 +68,7 @@ fn optional_object(
     name: &'static str,
     value: Option<Value>,
 ) -> PyResult<Option<Map<String, Value>>> {
-    value.map(|value| object(name, value)).transpose()
-}
-
-fn object(name: &'static str, value: Value) -> PyResult<Map<String, Value>> {
-    match value {
-        Value::Object(map) => Ok(map),
-        _ => Err(PyValueError::new_err(format!("{name} must be a dict"))),
-    }
+    value.map(|value| required_object(name, value)).transpose()
 }
 
 pub(crate) fn optional_timeout(timeout_seconds: Option<f64>) -> Option<Duration> {
@@ -84,6 +81,72 @@ pub(crate) fn optional_timeout(timeout_seconds: Option<f64>) -> Option<Duration>
     })
 }
 
+pub(crate) fn python_timeout_seconds(py: Python<'_>, timeout: Py<PyAny>) -> PyResult<Option<f64>> {
+    py.import("litellm.rust_bridge.timeouts")?
+        .getattr("timeout_to_seconds")?
+        .call1((timeout,))?
+        .extract()
+}
+
+pub(crate) fn project_optional_fields(
+    kwargs: &Bound<'_, PyDict>,
+    names: &[&str],
+) -> PyResult<Map<String, Value>> {
+    names
+        .iter()
+        .filter_map(|name| match kwargs.get_item(name) {
+            Ok(Some(value)) => Some(from_py(&value).map(|value| ((*name).to_string(), value))),
+            Ok(None) => None,
+            Err(error) => Some(Err(error)),
+        })
+        .collect()
+}
+
+struct RequestFieldSources<'py> {
+    body: Option<Bound<'py, PyAny>>,
+    credentials: Option<Bound<'py, PyAny>>,
+}
+
+impl<'py> RequestFieldSources<'py> {
+    fn extract(proxy_request: &Bound<'py, PyAny>) -> PyResult<Self> {
+        let proxy_request = proxy_request.cast::<PyDict>()?;
+
+        let body = proxy_request
+            .get_item("body_fields")?
+            .or(proxy_request.get_item("body")?);
+
+        let credentials = proxy_request.get_item("credential_fields")?;
+
+        Ok(Self { body, credentials })
+    }
+
+    fn contains(&self, name: &str) -> bool {
+        self.body
+            .as_ref()
+            .is_some_and(|fields| fields.contains(name).unwrap_or(false))
+            || self
+                .credentials
+                .as_ref()
+                .is_some_and(|fields| fields.contains(name).unwrap_or(false))
+    }
+}
+
+pub(crate) fn request_input_sources<'a>(
+    kwargs: &Bound<'_, PyDict>,
+    names: impl Iterator<Item = &'a str>,
+) -> PyResult<BTreeMap<String, InputSource>> {
+    let Some(proxy_request) = kwargs.get_item("proxy_server_request")? else {
+        return Ok(BTreeMap::new());
+    };
+
+    let sources = RequestFieldSources::extract(&proxy_request)?;
+
+    Ok(names
+        .filter(|name| sources.contains(name))
+        .map(|name| (name.to_string(), InputSource::Request))
+        .collect())
+}
+
 pub(crate) fn marshal_headers(headers: Option<Value>) -> PyResult<HashMap<String, String>> {
     let value = match headers {
         Some(headers) => headers,
@@ -102,3 +165,199 @@ pub(crate) fn marshal_headers(headers: Option<Value>) -> PyResult<HashMap<String
         })
         .collect()
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use pyo3::exceptions::PyTypeError;
+    use serde_json::json;
+
+    fn eval<'py>(py: Python<'py>, source: &std::ffi::CStr) -> Bound<'py, PyDict> {
+        let locals = PyDict::new(py);
+        py.run(source, Some(&locals), Some(&locals)).unwrap();
+        locals
+    }
+
+    fn sources(
+        py: Python<'_>,
+        proxy: &Bound<'_, PyAny>,
+        names: &[&str],
+    ) -> PyResult<BTreeMap<String, InputSource>> {
+        let kwargs = PyDict::new(py);
+        kwargs.set_item("proxy_server_request", proxy)?;
+        request_input_sources(&kwargs, names.iter().copied())
+    }
+
+    #[test]
+    fn required_shapes_preserve_nested_values_and_existing_errors() {
+        let nested = json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}]);
+        assert_eq!(
+            Value::Array(required_array("messages", nested.clone()).unwrap()),
+            nested
+        );
+
+        let body = json!({"model": "claude", "metadata": {"user": "1"}});
+        assert_eq!(
+            Value::Object(required_object("body", body.clone()).unwrap()),
+            body
+        );
+
+        assert_eq!(
+            required_array("messages", json!({"role": "user"}))
+                .unwrap_err()
+                .to_string(),
+            "ValueError: messages must be a list"
+        );
+        assert_eq!(
+            required_object("body", json!([])).unwrap_err().to_string(),
+            "ValueError: body must be a dict"
+        );
+    }
+
+    #[test]
+    fn optional_parameters_treat_missing_as_empty() {
+        assert_eq!(
+            object_or_empty("optional_params", None).unwrap(),
+            Map::new()
+        );
+        assert_eq!(
+            object_or_empty("optional_params", Some(json!({"temperature": 0.2}))).unwrap(),
+            required_object("optional_params", json!({"temperature": 0.2})).unwrap()
+        );
+    }
+
+    #[test]
+    fn missing_none_and_empty_proxy_metadata_are_distinct() {
+        Python::initialize();
+        Python::attach(|py| {
+            let kwargs = PyDict::new(py);
+            assert!(
+                request_input_sources(&kwargs, ["api_key"].into_iter())
+                    .unwrap()
+                    .is_empty()
+            );
+
+            kwargs.set_item("proxy_server_request", py.None()).unwrap();
+            assert!(
+                request_input_sources(&kwargs, ["api_key"].into_iter())
+                    .unwrap_err()
+                    .is_instance_of::<PyTypeError>(py)
+            );
+
+            kwargs
+                .set_item("proxy_server_request", PyDict::new(py))
+                .unwrap();
+            assert!(
+                request_input_sources(&kwargs, ["api_key"].into_iter())
+                    .unwrap()
+                    .is_empty()
+            );
+        });
+    }
+
+    #[test]
+    fn body_fields_win_over_body_and_explicit_none_does_not_fall_back() {
+        Python::initialize();
+        Python::attach(|py| {
+            let locals = eval(
+                py,
+                c"
+proxy = {'body_fields': ['api_key'], 'body': ['api_base']}
+none_fields = {'body_fields': None, 'body': ['api_key']}
+body_only = {'body': ['api_base']}
+",
+            );
+            let named = sources(
+                py,
+                &locals.get_item("proxy").unwrap().unwrap(),
+                &["api_key", "api_base"],
+            )
+            .unwrap();
+            assert_eq!(named.get("api_key").copied(), Some(InputSource::Request));
+            assert!(!named.contains_key("api_base"));
+
+            assert!(
+                sources(
+                    py,
+                    &locals.get_item("none_fields").unwrap().unwrap(),
+                    &["api_key"],
+                )
+                .unwrap()
+                .is_empty()
+            );
+
+            let body_only = sources(
+                py,
+                &locals.get_item("body_only").unwrap().unwrap(),
+                &["api_base"],
+            )
+            .unwrap();
+            assert_eq!(
+                body_only.get("api_base").copied(),
+                Some(InputSource::Request)
+            );
+        });
+    }
+
+    #[test]
+    fn body_and_credential_membership_can_mark_request_fields() {
+        Python::initialize();
+        Python::attach(|py| {
+            let locals = eval(
+                py,
+                c"
+class Raising:
+    def __contains__(self, item):
+        raise RuntimeError('credential membership')
+proxy = {
+    'body_fields': ['api_key'],
+    'credential_fields': Raising(),
+}
+credentials_only = {'credential_fields': ['extra_headers']}
+erroring = {'body_fields': Raising()}
+extra = {'body_fields': ['api_key', 'unused']}
+",
+            );
+            let skipped = sources(
+                py,
+                &locals.get_item("proxy").unwrap().unwrap(),
+                &["api_key"],
+            )
+            .unwrap();
+            assert_eq!(skipped.get("api_key").copied(), Some(InputSource::Request));
+
+            let credentials = sources(
+                py,
+                &locals.get_item("credentials_only").unwrap().unwrap(),
+                &["extra_headers"],
+            )
+            .unwrap();
+            assert_eq!(
+                credentials.get("extra_headers").copied(),
+                Some(InputSource::Request)
+            );
+
+            assert!(
+                sources(
+                    py,
+                    &locals.get_item("erroring").unwrap().unwrap(),
+                    &["api_key"],
+                )
+                .unwrap()
+                .is_empty()
+            );
+
+            let requested = sources(
+                py,
+                &locals.get_item("extra").unwrap().unwrap(),
+                &["api_key"],
+            )
+            .unwrap();
+            assert_eq!(requested.len(), 1);
+            assert_eq!(
+                requested.get("api_key").copied(),
+                Some(InputSource::Request)
+            );
+        });
+    }
+}

litellm-rust/crates/python-bridge/src/routes/audio_transcription/mod.rs

@@ -0,0 +1,12 @@
+mod value;
+
+use pyo3::prelude::*;
+
+pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
+    value::register(module)
+}
+
+#[cfg(feature = "trace-parity")]
+pub(super) fn register_trace(module: &Bound<'_, PyModule>) -> PyResult<()> {
+    value::register_trace(module)
+}

litellm-rust/crates/python-bridge/src/routes/chat_completions/mod.rs

@@ -0,0 +1,12 @@
+mod value;
+
+use pyo3::prelude::*;
+
+pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
+    value::register(module)
+}
+
+#[cfg(feature = "trace-parity")]
+pub(super) fn register_trace(module: &Bound<'_, PyModule>) -> PyResult<()> {
+    value::register_trace(module)
+}

litellm-rust/crates/python-bridge/src/routes/chat_completions/value.rs

@@ -9,12 +9,12 @@ use pyo3::prelude::*;
 use serde_json::Value;
 
 use crate::errors::chat_completions_error_to_pyerr;
-use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty, required_value};
+use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty, required_array};
 
 fn prepare_chat_completions(
     inputs: ChatCompletionsInputs,
 ) -> PyResult<impl Future<Output = Result<ChatCompletionsResponse, Error>> + Send + 'static> {
-    let messages = required_value("messages", inputs.messages, Value::is_array, "list")?;
+    let messages = required_array("messages", inputs.messages)?;
     let optional_params = object_or_empty("optional_params", inputs.optional_params)?;
     let options = RouteOptions::from_python(RouteOptionsInputs {
         model: inputs.model,
@@ -36,7 +36,7 @@ fn prepare_chat_completions(
         } = options;
         run_chat_completions(ChatCompletionsRequest {
             model: &model,
-            messages,
+            messages: Value::Array(messages),
             optional_params,
             api_key: api_key.as_deref(),
             api_base: api_base.as_deref(),

litellm-rust/crates/python-bridge/src/routes/definition.rs

@@ -389,6 +389,82 @@ mod tests {
         });
     }
 
+    #[test]
+    fn missing_and_explicit_none_optional_params_share_the_next_error() {
+        Python::initialize();
+        Python::attach(|py| {
+            let module = PyModule::new(py, "routes").expect("module should be created");
+            crate::routes::register(&module).expect("routes should register");
+            let messages = PyList::empty(py);
+            let headers = PyList::empty(py);
+            let omitted = PyDict::new(py);
+            omitted
+                .set_item("extra_headers", &headers)
+                .expect("kwargs should accept extra_headers");
+            let explicit = PyDict::new(py);
+            explicit
+                .set_item("optional_params", py.None())
+                .expect("kwargs should accept optional_params");
+            explicit
+                .set_item("extra_headers", &headers)
+                .expect("kwargs should accept extra_headers");
+
+            let omitted_error = module
+                .getattr("chat_completions")
+                .and_then(|function| function.call(("model", &messages), Some(&omitted)))
+                .expect_err("omitted optional_params should reach header validation");
+            let explicit_error = module
+                .getattr("chat_completions")
+                .and_then(|function| function.call(("model", &messages), Some(&explicit)))
+                .expect_err("None optional_params should reach header validation");
+            assert_eq!(
+                omitted_error.to_string(),
+                "ValueError: extra_headers must be a dict"
+            );
+            assert_eq!(explicit_error.to_string(), omitted_error.to_string());
+        });
+    }
+
+    #[test]
+    fn chat_completions_decline_keeps_existing_reasons() {
+        Python::initialize();
+        Python::attach(|py| {
+            let module = PyModule::new(py, "routes").expect("module should be created");
+            crate::routes::register(&module).expect("routes should register");
+            let decline = module
+                .getattr("chat_completions_decline")
+                .expect("decline helper should be registered");
+            let empty = PyList::empty(py);
+            let unreadable = py
+                .eval(c"'nope'", None, None)
+                .expect("string messages should convert");
+
+            let unknown: Option<String> = decline
+                .call1(("unknown-model", &empty))
+                .and_then(|value| value.extract())
+                .expect("unknown providers should decline");
+            assert_eq!(
+                unknown.as_deref(),
+                Some("provider is not on the rust chat completions path")
+            );
+
+            let empty_reason: Option<String> = decline
+                .call1(("anthropic/claude-sonnet-4-5", &empty))
+                .and_then(|value| value.extract())
+                .expect("empty lists should decline");
+            assert_eq!(empty_reason.as_deref(), Some("empty message list"));
+
+            let unreadable_reason: Option<String> = decline
+                .call1(("anthropic/claude-sonnet-4-5", unreadable))
+                .and_then(|value| value.extract())
+                .expect("non-list messages should decline");
+            assert_eq!(
+                unreadable_reason.as_deref(),
+                Some("unreadable message list")
+            );
+        });
+    }
+
     #[test]
     fn generated_routes_execute_sync_and_async_contracts() {
         Python::initialize();

litellm-rust/crates/python-bridge/src/routes/gateway_messages.rs

@@ -1,29 +0,0 @@
-use pyo3::prelude::*;
-use serde_json::Value;
-
-use crate::errors::core_error_to_pyerr;
-
-#[pyfunction]
-fn gateway_messages<'py>(
-    py: Python<'py>,
-    model_alias: String,
-    provider_model: String,
-    api_base: String,
-    #[pyo3(from_py_with = litellm_python_interop::from_py)] body: Value,
-) -> PyResult<Bound<'py, PyAny>> {
-    let future = litellm_ai_gateway::trace_parity::messages_request(
-        model_alias,
-        provider_model,
-        api_base,
-        body,
-    );
-    crate::execution::run_async(
-        py,
-        crate::function_trace::capture(future),
-        core_error_to_pyerr,
-    )
-}
-
-pub(super) fn register_trace(module: &Bound<'_, PyModule>) -> PyResult<()> {
-    super::definition::add_function(module, wrap_pyfunction!(gateway_messages, module)?)
-}

litellm-rust/crates/python-bridge/src/routes/messages/mod.rs

@@ -0,0 +1,12 @@
+mod value;
+
+use pyo3::prelude::*;
+
+pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
+    value::register(module)
+}
+
+#[cfg(feature = "trace-parity")]
+pub(super) fn register_trace(module: &Bound<'_, PyModule>) -> PyResult<()> {
+    value::register_trace(module)
+}

litellm-rust/crates/python-bridge/src/routes/messages/value.rs

@@ -6,12 +6,12 @@ use serde_json::Value;
 use std::future::Future;
 
 use crate::errors::core_error_to_pyerr;
-use crate::marshal::{RouteOptions, RouteOptionsInputs, required_value};
+use crate::marshal::{RouteOptions, RouteOptionsInputs, required_object};
 
 fn prepare_messages(
     inputs: MessagesInputs,
 ) -> PyResult<impl Future<Output = Result<AnthropicMessagesResponse, Error>> + Send + 'static> {
-    let body = required_value("body", inputs.body, Value::is_object, "dict")?;
+    let body = required_object("body", inputs.body)?;
     let options = RouteOptions::from_python(RouteOptionsInputs {
         model: inputs.model,
         api_key: inputs.api_key,
@@ -32,7 +32,7 @@ fn prepare_messages(
         } = options;
         run_messages(MessagesRequest {
             model: &model,
-            body,
+            body: Value::Object(body),
             api_key: api_key.as_deref(),
             api_base: api_base.as_deref(),
             custom_llm_provider: custom_llm_provider.as_deref(),

litellm-rust/crates/python-bridge/src/routes/mod.rs

@@ -3,9 +3,6 @@ use pyo3::prelude::*;
 #[macro_use]
 mod definition;
 
-#[cfg(feature = "trace-parity")]
-mod gateway_messages;
-
 mod audio_transcription;
 mod chat_completions;
 mod messages;
@@ -16,6 +13,7 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
     audio_transcription::register(module)?;
     messages::register(module)?;
     chat_completions::register(module)?;
+
     #[cfg(feature = "trace-parity")]
     {
         let trace = PyModule::new(module.py(), "_trace")?;
@@ -23,7 +21,6 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
         audio_transcription::register_trace(&trace)?;
         messages::register_trace(&trace)?;
         chat_completions::register_trace(&trace)?;
-        gateway_messages::register_trace(&trace)?;
         module.add_submodule(&trace)?;
     }
     Ok(())

litellm-rust/crates/python-bridge/src/routes/ocr/callbacks.rs

@@ -0,0 +1,161 @@
+use pyo3::exceptions::PyBaseException;
+use pyo3::prelude::*;
+use pyo3::types::PyDict;
+use serde_json::Value;
+
+use litellm_core::ocr::LiteLLMOcrResponse;
+use litellm_core::ocr::hooks::OcrPreCallRequest;
+use litellm_python_interop::to_py_preserving_errors as to_py;
+
+use crate::lifecycle::PythonLogger;
+
+pub(super) struct OcrLoggingFields {
+    model: String,
+    custom_llm_provider: String,
+    optional_params: Value,
+}
+
+impl From<&OcrPreCallRequest> for OcrLoggingFields {
+    fn from(request: &OcrPreCallRequest) -> Self {
+        Self {
+            model: request.model.clone(),
+            custom_llm_provider: request.custom_llm_provider.clone(),
+            optional_params: request.optional_params.clone(),
+        }
+    }
+}
+
+impl PythonLogger {
+    pub(super) fn update_ocr(
+        &self,
+        py: Python<'_>,
+        kwargs: &Py<PyDict>,
+        pre_call: &OcrLoggingFields,
+        secret_fields: &[&str],
+        url: &str,
+    ) -> PyResult<()> {
+        let update = PyDict::new(py);
+        update.set_item("kwargs", redact(py, kwargs.bind(py), secret_fields)?)?;
+        update.set_item("model", &pre_call.model)?;
+        update.set_item(
+            "optional_params",
+            redact(
+                py,
+                &to_py(py, &pre_call.optional_params)?
+                    .into_bound(py)
+                    .cast_into::<PyDict>()?,
+                secret_fields,
+            )?,
+        )?;
+        let params = PyDict::new(py);
+        params.set_item(
+            "litellm_call_id",
+            kwargs.bind(py).get_item("litellm_call_id")?,
+        )?;
+        params.set_item("api_base", url)?;
+        for name in ["logger_fn", "litellm_request_debug"] {
+            if let Some(value) = kwargs.bind(py).get_item(name)? {
+                params.set_item(name, value)?;
+            }
+        }
+        update.set_item("litellm_params", params)?;
+        update.set_item("custom_llm_provider", &pre_call.custom_llm_provider)?;
+        self.object(py)
+            .call_method("update_from_kwargs", (), Some(&update))?;
+        Ok(())
+    }
+
+    pub(crate) fn pre_ocr(
+        &self,
+        py: Python<'_>,
+        api_key: &Option<Py<PyAny>>,
+        body: &Bound<'_, PyDict>,
+        headers: &Bound<'_, PyDict>,
+        url: &str,
+    ) -> PyResult<()> {
+        let additional = PyDict::new(py);
+        additional.set_item("complete_input_dict", body)?;
+        additional.set_item("headers", headers)?;
+        additional.set_item("api_base", url)?;
+        let kwargs = PyDict::new(py);
+        kwargs.set_item("input", "OCR document processing")?;
+        kwargs.set_item("api_key", api_key)?;
+        kwargs.set_item("additional_args", &additional)?;
+        if self.callbacks_needed(py, "input")? {
+            self.object(py).call_method("pre_call", (), Some(&kwargs))?;
+        } else {
+            self.object(py)
+                .call_method("_pre_call", (), Some(&kwargs))?;
+            self.object(py).call_method0("record_api_call_start_time")?;
+        }
+        Ok(())
+    }
+
+    pub(crate) fn post_ocr(
+        &self,
+        py: Python<'_>,
+        original_response: &Value,
+        body: Option<&Py<PyDict>>,
+        headers: Option<&Py<PyDict>>,
+    ) -> PyResult<()> {
+        let additional = PyDict::new(py);
+        additional.set_item("complete_input_dict", body)?;
+        additional.set_item("headers", headers)?;
+        if self.callbacks_needed(py, "input")? {
+            let kwargs = PyDict::new(py);
+            kwargs.set_item("original_response", to_py(py, original_response)?)?;
+            kwargs.set_item("additional_args", &additional)?;
+            self.object(py)
+                .call_method("post_call", (), Some(&kwargs))?;
+        } else {
+            let response = py
+                .import("json")?
+                .call_method1("dumps", (to_py(py, original_response)?,))?;
+            self.object(py).call_method1(
+                "record_post_call",
+                (response, py.None(), py.None(), additional),
+            )?;
+        }
+        Ok(())
+    }
+}
+
+fn redact(
+    py: Python<'_>,
+    params: &Bound<'_, PyDict>,
+    secret_fields: &[&str],
+) -> PyResult<Py<PyDict>> {
+    let redacted = PyDict::new(py);
+    for (name, value) in params {
+        let name = name.extract::<String>()?;
+        if name == "proxy_server_request" {
+            continue;
+        }
+        if secret_fields.contains(&name.as_str()) {
+            redacted.set_item(name, "****")?;
+        } else {
+            redacted.set_item(name, value)?;
+        }
+    }
+    Ok(redacted.unbind())
+}
+
+pub(super) fn response(py: Python<'_>, response: &LiteLLMOcrResponse) -> PyResult<Py<PyAny>> {
+    py.import("litellm.rust_bridge.ocr")?
+        .getattr("_response")?
+        .call1((to_py(py, response)?,))
+        .map(Bound::unbind)
+}
+
+pub(super) fn map_failure(
+    py: Python<'_>,
+    error: &Py<PyBaseException>,
+    request: &Bound<'_, PyAny>,
+    provider: &str,
+) -> PyResult<Py<PyBaseException>> {
+    Ok(py
+        .import("litellm.rust_bridge.ocr_lifecycle")?
+        .getattr("map_failure")?
+        .call1((error, request, provider))?
+        .extract()?)
+}

litellm-rust/crates/python-bridge/src/routes/ocr/document.rs

@@ -0,0 +1,264 @@
+use std::io::Read;
+use std::path::PathBuf;
+
+use pyo3::exceptions::{PyFileNotFoundError, PyTypeError, PyValueError};
+use pyo3::prelude::*;
+use pyo3::pybacked::PyBackedBytes;
+#[cfg(test)]
+use pyo3::types::PyDict;
+use pyo3::types::{PyBytes, PyString};
+
+use litellm_core::constants::OCR_INLINE_MAX_BYTES;
+use litellm_core::ocr::{OcrDocument, encode_file_document, mime_type_for_name, upload_mime_type};
+use litellm_python_interop::to_py_preserving_errors;
+
+enum FileBytes {
+    Python(PyBackedBytes),
+    Native(Vec<u8>),
+}
+
+impl AsRef<[u8]> for FileBytes {
+    fn as_ref(&self) -> &[u8] {
+        match self {
+            Self::Python(bytes) => bytes,
+            Self::Native(bytes) => bytes,
+        }
+    }
+}
+
+fn read_file_input(
+    py: Python<'_>,
+    file: &Bound<'_, PyAny>,
+) -> PyResult<(FileBytes, Option<String>)> {
+    if file.is_instance_of::<PyString>() {
+        return Err(PyValueError::new_err(
+            "OCR file input does not accept bare str values. Pass bytes, a pathlib.Path, or a file-like object.",
+        ));
+    }
+    if file.is_instance(&py.import("os")?.getattr("PathLike")?)? {
+        let path: PathBuf = file.extract()?;
+        let name = path
+            .file_name()
+            .map(|value| value.to_string_lossy().into_owned());
+        let bytes = py
+            .detach(|| {
+                let mut bytes = Vec::new();
+                std::fs::File::open(&path)?
+                    .take(OCR_INLINE_MAX_BYTES as u64 + 1)
+                    .read_to_end(&mut bytes)?;
+                Ok::<_, std::io::Error>(bytes)
+            })
+            .map_err(|error| {
+                if error.kind() == std::io::ErrorKind::NotFound {
+                    PyFileNotFoundError::new_err(format!("File not found: {}", path.display()))
+                } else {
+                    error.into()
+                }
+            })?;
+        return Ok((FileBytes::Native(bytes), name));
+    }
+    if file.is_instance_of::<PyBytes>() {
+        return Ok((FileBytes::Python(file.extract()?), None));
+    }
+    let reader = file
+        .getattr_opt("read")?
+        .filter(|value| value.is_callable());
+    let Some(reader) = reader else {
+        return Err(PyValueError::new_err(format!(
+            "Unsupported file input type: {}. Expected pathlib.Path, bytes, or a file-like object.",
+            file.get_type(),
+        )));
+    };
+    let name = file
+        .getattr_opt("name")?
+        .filter(|value| !value.is_none())
+        .map(|value| value.extract::<String>())
+        .transpose()?;
+    let value = reader.call0()?;
+    let bytes = if value.is_instance_of::<PyString>() {
+        FileBytes::Native(value.extract::<String>()?.into_bytes())
+    } else if value.is_instance_of::<PyBytes>() {
+        FileBytes::Python(value.extract()?)
+    } else {
+        return Err(PyTypeError::new_err(format!(
+            "OCR file read must return bytes or str, got {}",
+            value.get_type(),
+        )));
+    };
+    Ok((bytes, name))
+}
+
+pub(super) struct FileDocumentInput {
+    bytes: FileBytes,
+    name: Option<String>,
+    mime_type: Option<String>,
+}
+
+impl FromPyObject<'_, '_> for FileDocumentInput {
+    type Error = PyErr;
+
+    fn extract(document: Borrowed<'_, '_, PyAny>) -> PyResult<Self> {
+        let py = document.py();
+        let mime_type = match document.get_item("mime_type") {
+            Ok(value) => Some(value.extract::<String>()?),
+            Err(error) if error.is_instance_of::<pyo3::exceptions::PyKeyError>(py) => None,
+            Err(error) => return Err(error),
+        };
+        let file = document.get_item("file").map_err(|error| {
+            if error.is_instance_of::<pyo3::exceptions::PyKeyError>(py) {
+                PyValueError::new_err("document with type='file' must include a 'file' field containing a pathlib.Path, file-like object, or bytes")
+            } else {
+                error
+            }
+        })?;
+        if file.is_none() {
+            return Err(PyValueError::new_err(
+                "document with type='file' must include a 'file' field containing a pathlib.Path, file-like object, or bytes",
+            ));
+        }
+        let (bytes, name) = read_file_input(py, &file)?;
+        Ok(Self {
+            bytes,
+            name,
+            mime_type,
+        })
+    }
+}
+
+pub(super) fn file_document(py: Python<'_>, document: FileDocumentInput) -> PyResult<OcrDocument> {
+    py.detach(|| {
+        encode_file_document(
+            document.bytes.as_ref(),
+            document.name.as_deref(),
+            document.mime_type.as_deref(),
+        )
+    })
+    .map_err(|error| PyValueError::new_err(error.to_string()))
+}
+
+#[pyfunction]
+fn _ocr_file_document(py: Python<'_>, document: Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
+    to_py_preserving_errors(py, &file_document(py, document.extract()?)?)
+}
+
+#[pyfunction]
+fn _ocr_mime_type(file_name: &str) -> String {
+    mime_type_for_name(file_name).into()
+}
+
+#[pyfunction]
+#[pyo3(signature = (file_content, file_name=None, content_type=None))]
+fn _ocr_upload_document(
+    py: Python<'_>,
+    file_content: &Bound<'_, PyBytes>,
+    file_name: Option<&str>,
+    content_type: Option<&str>,
+) -> PyResult<Py<PyAny>> {
+    let bytes: PyBackedBytes = file_content.extract()?;
+    let document = py
+        .detach(|| {
+            encode_file_document(
+                &bytes,
+                None,
+                Some(upload_mime_type(file_name, content_type)),
+            )
+        })
+        .map_err(|error| PyValueError::new_err(error.to_string()))?;
+    to_py_preserving_errors(py, &document)
+}
+
+pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
+    module.add("_OCR_MAX_FILE_BYTES", OCR_INLINE_MAX_BYTES)?;
+    module.add_function(wrap_pyfunction!(_ocr_upload_document, module)?)?;
+    module.add_function(wrap_pyfunction!(_ocr_file_document, module)?)?;
+    module.add_function(wrap_pyfunction!(_ocr_mime_type, module)?)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn extraction_validates_required_file_and_optional_mime_type() {
+        Python::initialize();
+        Python::attach(|py| {
+            for expression in [c"{}", c"{'file': None}"] {
+                let document = py.eval(expression, None, None).unwrap();
+                let error = document.extract::<FileDocumentInput>().err().unwrap();
+                assert!(error.is_instance_of::<PyValueError>(py));
+                assert!(error.to_string().contains("must include a 'file' field"));
+            }
+            for expression in [
+                c"{'file': b'abc', 'mime_type': None}",
+                c"{'file': b'abc', 'mime_type': 7}",
+            ] {
+                let document = py.eval(expression, None, None).unwrap();
+                let error = document.extract::<FileDocumentInput>().err().unwrap();
+                assert!(error.is_instance_of::<PyTypeError>(py));
+            }
+            let document = py.eval(c"{'file': b'abc'}", None, None).unwrap();
+            let input: FileDocumentInput = document.extract().unwrap();
+            assert_eq!(input.bytes.as_ref(), b"abc");
+            assert_eq!(input.name, None);
+            assert_eq!(input.mime_type, None);
+        });
+    }
+
+    #[test]
+    fn extraction_validates_mime_type_before_consuming_file() {
+        Python::initialize();
+        Python::attach(|py| {
+            let locals = PyDict::new(py);
+            py.run(
+                c"class Reader:
+    def __init__(self):
+        self.reads = 0
+    def read(self):
+        self.reads += 1
+        return b'abc'
+reader = Reader()
+document = {'file': reader, 'mime_type': 7}",
+                Some(&locals),
+                Some(&locals),
+            )
+            .unwrap();
+            let document = locals.get_item("document").unwrap().unwrap();
+            let error = document.extract::<FileDocumentInput>().err().unwrap();
+            assert!(error.is_instance_of::<PyTypeError>(py));
+            let reads: usize = locals
+                .get_item("reader")
+                .unwrap()
+                .unwrap()
+                .getattr("reads")
+                .unwrap()
+                .extract()
+                .unwrap();
+            assert_eq!(reads, 0);
+        });
+    }
+
+    #[test]
+    fn extraction_preserves_reader_key_error_identity() {
+        Python::initialize();
+        Python::attach(|py| {
+            let locals = PyDict::new(py);
+            py.run(
+                c"failure = KeyError('reader failed')
+class Reader:
+    def read(self):
+        raise failure
+document = {'file': Reader()}",
+                Some(&locals),
+                Some(&locals),
+            )
+            .unwrap();
+            let document = locals.get_item("document").unwrap().unwrap();
+            let error = document.extract::<FileDocumentInput>().err().unwrap();
+            assert!(
+                error
+                    .value(py)
+                    .is(locals.get_item("failure").unwrap().unwrap())
+            );
+        });
+    }
+}

litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs

@@ -0,0 +1,72 @@
+use litellm_core::error::Error;
+use pyo3::prelude::*;
+
+use crate::errors::{RustUpstreamError, core_error_to_pyerr};
+
+pub(super) fn to_pyerr(error: Error) -> PyErr {
+    let status = error.http_status_code();
+    let mapped = match error {
+        Error::Http { status, body } => RustUpstreamError::new_err((status, body)),
+        other => core_error_to_pyerr(other),
+    };
+    attach_status(mapped, status)
+}
+
+fn attach_status(error: PyErr, status: Option<u16>) -> PyErr {
+    if let Some(status) = status {
+        Python::attach(|py| {
+            let value = error.value(py);
+            value.setattr("status_code", status).ok();
+            value.setattr("message", value.to_string()).ok();
+        });
+    }
+    error
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use pyo3::exceptions::PyValueError;
+
+    #[test]
+    fn preserves_python_validation_and_provider_details() {
+        Python::initialize();
+        Python::attach(|py| {
+            let mapped = to_pyerr(Error::MissingDocumentUrl);
+            assert!(mapped.is_instance_of::<pyo3::exceptions::PyValueError>(py));
+            assert_eq!(mapped.value(py).to_string(), "Document URL is required");
+            assert_eq!(
+                mapped
+                    .value(py)
+                    .getattr("status_code")
+                    .unwrap()
+                    .extract::<u16>()
+                    .unwrap(),
+                500
+            );
+            let mapped = to_pyerr(Error::Http {
+                status: 429,
+                body: r#"{"message":"rate limited"}"#.to_string(),
+            });
+            assert!(mapped.is_instance_of::<RustUpstreamError>(py));
+            let args: (u16, String) = mapped
+                .value(py)
+                .getattr("args")
+                .and_then(|args| args.extract())
+                .expect("OCR failures retain status and unprefixed provider message");
+            assert_eq!(args, (429, r#"{"message":"rate limited"}"#.to_string()));
+
+            let mapped = to_pyerr(Error::InvalidRequest("invalid format".into()));
+            assert!(mapped.is_instance_of::<PyValueError>(py));
+            assert_eq!(
+                mapped
+                    .value(py)
+                    .getattr("status_code")
+                    .unwrap()
+                    .extract::<u16>()
+                    .unwrap(),
+                400
+            );
+        });
+    }
+}

litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs

@@ -0,0 +1,311 @@
+use pyo3::prelude::*;
+use pyo3::types::{PyDict, PyTuple};
+
+use litellm_core::auth::ResolvedCredential;
+use litellm_core::ocr::hooks::{OcrDuringCallRequest, OcrPostCallRequest, OcrPreCallRequest};
+use litellm_core::ocr::{OcrAdmission, OcrCall, OcrClient, OcrHostOperation, OcrHostResult};
+use litellm_python_interop::{
+    from_py_preserving_errors as from_py, to_py_preserving_errors as to_py,
+};
+
+use super::callbacks;
+use super::errors::to_pyerr as ocr_error_to_pyerr;
+use super::project::{ProjectedOcrFields, admitted_call, project_request};
+use crate::lifecycle::{
+    OperationClass, PythonCallState, PythonRoute, missing_state, now, run_call,
+};
+
+struct PythonOcrHost {
+    state: PythonCallState,
+    data: OcrHostData,
+}
+
+enum OcrHostData {
+    Unprojected { request: Py<PyAny> },
+    Projected(Box<ProjectedOcrHost>),
+    Released,
+}
+
+struct ProjectedOcrHost {
+    fields: ProjectedOcrFields,
+    pre_call: Option<callbacks::OcrLoggingFields>,
+    retained_fields: Option<Py<PyDict>>,
+    body: Option<Py<PyDict>>,
+    headers: Option<Py<PyDict>>,
+}
+
+impl PythonOcrHost {
+    fn projected(&self) -> PyResult<&ProjectedOcrHost> {
+        match &self.data {
+            OcrHostData::Projected(projected) => Ok(projected),
+            _ => Err(missing_state()),
+        }
+    }
+
+    fn projected_mut(&mut self) -> PyResult<&mut ProjectedOcrHost> {
+        match &mut self.data {
+            OcrHostData::Projected(projected) => Ok(projected),
+            _ => Err(missing_state()),
+        }
+    }
+
+    fn pre_call(
+        &mut self,
+        py: Python<'_>,
+        request: OcrPreCallRequest,
+    ) -> PyResult<OcrPreCallRequest> {
+        let kwargs = self.state.kwargs.bind(py);
+        let retained_fields = PyDict::new(py);
+        for name in request
+            .optional_params
+            .as_object()
+            .ok_or_else(missing_state)?
+            .keys()
+        {
+            if let Some(value) = kwargs.get_item(name)? {
+                retained_fields.set_item(name, value)?;
+            }
+        }
+        retained_fields.set_item("document", &self.projected()?.fields.document)?;
+        let projected = self.projected_mut()?;
+        projected.retained_fields = Some(retained_fields.unbind());
+        projected.pre_call = Some((&request).into());
+        Ok(request)
+    }
+
+    fn acquire_azure_ad_token(&self, py: Python<'_>) -> PyResult<ResolvedCredential> {
+        let provider = self
+            .projected()?
+            .fields
+            .azure_ad_token_provider
+            .as_ref()
+            .ok_or_else(missing_state)?;
+        provider.acquire(py)
+    }
+
+    fn python_pre_call(
+        &mut self,
+        py: Python<'_>,
+        mut request: OcrDuringCallRequest,
+    ) -> PyResult<OcrDuringCallRequest> {
+        let projected = self.projected()?;
+        let pre_call = projected.pre_call.as_ref().ok_or_else(missing_state)?;
+        self.state.logger()?.update_ocr(
+            py,
+            &self.state.kwargs,
+            pre_call,
+            &projected.fields.secret_fields,
+            &request.url,
+        )?;
+        if !self.state.logger()?.callbacks_needed(py, "payload")? {
+            self.state
+                .logger()?
+                .object(py)
+                .call_method0("record_api_call_start_time")?;
+            return Ok(request);
+        }
+        if let Some(body) = request.body.as_object_mut() {
+            for name in &request.retained_fields {
+                body.remove(name);
+            }
+        }
+        let body = to_py(py, &request.body)?
+            .into_bound(py)
+            .cast_into::<PyDict>()?;
+        if let Some(retained) = &self.projected()?.retained_fields {
+            for name in &request.retained_fields {
+                if let Some(value) = retained.bind(py).get_item(name)? {
+                    body.set_item(name, value)?;
+                }
+            }
+        }
+        let headers = PyDict::new(py);
+        for (name, value) in &request.headers {
+            headers.set_item(name, value)?;
+        }
+        let api_key = self.projected()?.fields.api_key.clone_ref(py);
+        let projected = self.projected_mut()?;
+        projected.body = Some(body.clone().unbind());
+        projected.headers = Some(headers.clone().unbind());
+        self.state
+            .logger()?
+            .pre_ocr(py, &Some(api_key), &body, &headers, &request.url)?;
+        let headers = headers
+            .iter()
+            .map(|(name, value)| Ok((name.extract::<String>()?, value.extract::<String>()?)))
+            .collect::<PyResult<Vec<_>>>()?;
+        request.body = from_py(&body)?;
+        request.headers = headers;
+        Ok(request)
+    }
+
+    fn python_post_call(
+        &mut self,
+        py: Python<'_>,
+        request: OcrPostCallRequest,
+    ) -> PyResult<OcrPostCallRequest> {
+        let logger = self.state.logger()?;
+        if logger.callbacks_needed(py, "payload")? {
+            let projected = self.projected()?;
+            logger.post_ocr(
+                py,
+                &request.original_response,
+                projected.body.as_ref(),
+                projected.headers.as_ref(),
+            )?;
+        }
+        Ok(request)
+    }
+}
+
+impl PythonRoute for PythonOcrHost {
+    type Call = OcrCall;
+
+    fn state(&self) -> &PythonCallState {
+        &self.state
+    }
+
+    fn state_mut(&mut self) -> &mut PythonCallState {
+        &mut self.state
+    }
+
+    fn classify(operation: &OcrHostOperation) -> OperationClass {
+        operation
+            .phase()
+            .map_or(OperationClass::Route, OperationClass::Phase)
+    }
+
+    fn lifecycle_result() -> OcrHostResult {
+        OcrHostResult::Lifecycle(Ok(()))
+    }
+
+    fn map_error(error: litellm_core::Error) -> PyErr {
+        ocr_error_to_pyerr(error)
+    }
+
+    fn invoke(&mut self, py: Python<'_>, operation: OcrHostOperation) -> PyResult<OcrHostResult> {
+        Ok(match operation {
+            OcrHostOperation::ProjectRequest => {
+                let OcrHostData::Unprojected { request } = &self.data else {
+                    return Err(missing_state());
+                };
+                let projected = project_request(py, request.bind(py), self.state.kwargs.bind(py))?;
+                let has_token_provider = projected.fields.azure_ad_token_provider.is_some();
+                let request = projected.request;
+                self.data = OcrHostData::Projected(Box::new(ProjectedOcrHost {
+                    fields: projected.fields,
+                    pre_call: None,
+                    retained_fields: None,
+                    body: None,
+                    headers: None,
+                }));
+                OcrHostResult::Request(Ok((Box::new(request), has_token_provider)))
+            }
+            OcrHostOperation::AcquireAzureAdToken => {
+                OcrHostResult::AzureAdToken(Ok(self.acquire_azure_ad_token(py)?))
+            }
+            OcrHostOperation::PreCall(request) => {
+                OcrHostResult::PreCall(Ok(self.pre_call(py, request)?))
+            }
+            OcrHostOperation::DuringCall(request) => {
+                OcrHostResult::DuringCall(Ok(self.python_pre_call(py, request)?))
+            }
+            OcrHostOperation::PostCall(request) => {
+                OcrHostResult::PostCall(Ok(self.python_post_call(py, request)?))
+            }
+            OcrHostOperation::ConstructResponse(response) => {
+                self.state.end = Some(now(py)?);
+                self.state.response = Some(callbacks::response(py, response.as_ref())?);
+                OcrHostResult::Lifecycle(Ok(()))
+            }
+            OcrHostOperation::MapFailure(error) => {
+                if self.state.error.is_none() {
+                    self.state.retain_error(py, ocr_error_to_pyerr(error));
+                }
+                if self.state.end.is_none() {
+                    self.state.end = Some(now(py)?);
+                }
+                let error = self.state.error.as_ref().ok_or_else(missing_state)?;
+                let (request, provider) = match &self.data {
+                    OcrHostData::Unprojected { request } => (request.bind(py), ""),
+                    OcrHostData::Projected(projected) => (
+                        projected.fields.boundary_request.bind(py),
+                        projected.fields.provider,
+                    ),
+                    OcrHostData::Released => return Err(missing_state()),
+                };
+                let mapped = callbacks::map_failure(py, error, request, provider)?;
+                self.state
+                    .retain_error(py, PyErr::from_value(mapped.into_bound(py).into_any()));
+                OcrHostResult::Lifecycle(Ok(()))
+            }
+            OcrHostOperation::Lifecycle(_)
+            | OcrHostOperation::Success { .. }
+            | OcrHostOperation::Failure { .. } => return Err(missing_state()),
+        })
+    }
+
+    fn cleanup(&mut self) {
+        self.data = OcrHostData::Released;
+    }
+    fn traverse(&self, visit: &pyo3::gc::PyVisit<'_>) -> Result<(), pyo3::gc::PyTraverseError> {
+        match &self.data {
+            OcrHostData::Unprojected { request } => visit.call(request),
+            OcrHostData::Projected(projected) => {
+                visit.call(&projected.fields.boundary_request)?;
+                visit.call(&projected.fields.document)?;
+                visit.call(&projected.fields.api_key)?;
+                if let Some(provider) = &projected.fields.azure_ad_token_provider {
+                    provider.traverse(visit)?;
+                }
+                visit.call(&projected.retained_fields)?;
+                visit.call(&projected.body)?;
+                visit.call(&projected.headers)
+            }
+            OcrHostData::Released => Ok(()),
+        }
+    }
+}
+
+pub(super) struct BridgeOcrHooks;
+
+impl litellm_core::ocr::hooks::OcrHooks for BridgeOcrHooks {
+    fn intercepts_requests(&self) -> bool {
+        true
+    }
+}
+
+#[pyfunction]
+fn _ocr_lifecycle(
+    py: Python<'_>,
+    request: Bound<'_, PyAny>,
+    args: Bound<'_, PyTuple>,
+    kwargs: Bound<'_, PyDict>,
+    asynchronous: bool,
+) -> PyResult<Py<PyAny>> {
+    let client = OcrClient::shared().map_err(ocr_error_to_pyerr)?;
+    let call = admitted_call(OcrCall::admit(
+        client,
+        OcrAdmission {
+            asynchronous,
+            ..OcrAdmission::all()
+        },
+    ))?;
+    let host = PythonOcrHost {
+        state: PythonCallState::new(
+            py,
+            args.unbind(),
+            kwargs.copy()?.unbind(),
+            asynchronous,
+            if asynchronous { "aocr" } else { "ocr" },
+        )?,
+        data: OcrHostData::Unprojected {
+            request: request.unbind(),
+        },
+    };
+    run_call(py, call, host)
+}
+
+pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
+    module.add_function(wrap_pyfunction!(_ocr_lifecycle, module)?)
+}

litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs

@@ -0,0 +1,19 @@
+mod callbacks;
+mod document;
+mod errors;
+mod lifecycle;
+mod project;
+mod value;
+
+use pyo3::prelude::*;
+
+pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
+    value::register(module)?;
+    document::register(module)?;
+    lifecycle::register(module)
+}
+
+#[cfg(feature = "trace-parity")]
+pub(super) fn register_trace(module: &Bound<'_, PyModule>) -> PyResult<()> {
+    value::register_trace(module)
+}

litellm-rust/crates/python-bridge/src/routes/ocr/project.rs

@@ -0,0 +1,579 @@
+use std::sync::Arc;
+
+use litellm_core::ocr::wire::{OcrWireRequest, consumed_optional_params, decode_request};
+use litellm_core::ocr::{LiteLLMOcrRequest, NativeOutcome, OcrCall};
+use litellm_python_interop::{
+    from_py_preserving_errors as from_py, to_py_preserving_errors as to_py,
+};
+use pyo3::prelude::*;
+use pyo3::types::PyDict;
+use serde_json::{Map, Value};
+
+use super::errors::to_pyerr as ocr_error_to_pyerr;
+use super::lifecycle::BridgeOcrHooks;
+use crate::auth::{AZURE_AD_TOKEN_PROVIDER, PythonTokenProvider};
+use crate::errors::RustBridgeDeclined;
+use crate::marshal::{project_optional_fields, python_timeout_seconds, request_input_sources};
+
+pub(super) struct ProjectedOcrFields {
+    pub boundary_request: Py<PyAny>,
+    pub document: Py<PyAny>,
+    pub api_key: Py<PyAny>,
+    pub azure_ad_token_provider: Option<PythonTokenProvider>,
+    pub provider: &'static str,
+    pub secret_fields: Vec<&'static str>,
+}
+
+pub(super) struct ProjectedOcrCall {
+    pub request: LiteLLMOcrRequest,
+    pub fields: ProjectedOcrFields,
+}
+
+struct OcrArguments<'a, 'py> {
+    request: &'a Bound<'py, PyAny>,
+    kwargs: &'a Bound<'py, PyDict>,
+}
+
+impl<'py> OcrArguments<'_, 'py> {
+    fn lookup(&self, name: &str) -> PyResult<Bound<'py, PyAny>> {
+        match self.kwargs.get_item(name)? {
+            Some(value) => Ok(value),
+            None => self.request.getattr(name),
+        }
+    }
+
+    fn model(&self) -> PyResult<String> {
+        self.lookup("model")?.extract()
+    }
+
+    fn custom_llm_provider(&self) -> PyResult<Option<String>> {
+        self.lookup("custom_llm_provider")?.extract()
+    }
+
+    fn document(&self) -> PyResult<Bound<'py, PyAny>> {
+        self.lookup("document")
+    }
+
+    fn api_key(&self) -> PyResult<Bound<'py, PyAny>> {
+        self.lookup("api_key")
+    }
+
+    fn api_base(&self) -> PyResult<Option<String>> {
+        self.lookup("api_base")?.extract()
+    }
+
+    fn extra_headers(&self) -> PyResult<Option<Map<String, Value>>> {
+        self.lookup("extra_headers")?
+            .extract::<Option<Py<PyAny>>>()?
+            .map(|value| from_py(value.bind(self.request.py())))
+            .transpose()
+    }
+
+    fn timeout_seconds(&self) -> PyResult<Option<f64>> {
+        Ok(self
+            .lookup("timeout")?
+            .extract::<Option<Py<PyAny>>>()?
+            .map(|value| python_timeout_seconds(self.request.py(), value))
+            .transpose()?
+            .flatten())
+    }
+}
+
+enum ProjectedDocument {
+    File { wire: Value, retained: Py<PyAny> },
+    Other { wire: Value, retained: Py<PyAny> },
+}
+
+impl ProjectedDocument {
+    fn project(py: Python<'_>, document: &Bound<'_, PyAny>) -> PyResult<Self> {
+        let kind: String = document.get_item("type")?.extract()?;
+        if kind != "file" {
+            return Ok(Self::Other {
+                wire: from_py(document)?,
+                retained: document.clone().unbind(),
+            });
+        }
+        let input = document.extract()?;
+        let encoded = super::document::file_document(py, input)?;
+        let wire = serde_json::to_value(encoded)
+            .map_err(|error| pyo3::exceptions::PyValueError::new_err(error.to_string()))?;
+        Ok(Self::File {
+            retained: to_py(py, &wire)?,
+            wire,
+        })
+    }
+
+    fn into_parts(self) -> (Value, Py<PyAny>) {
+        match self {
+            Self::File { wire, retained } | Self::Other { wire, retained } => (wire, retained),
+        }
+    }
+}
+
+pub(super) fn project_request(
+    py: Python<'_>,
+    request: &Bound<'_, PyAny>,
+    kwargs: &Bound<'_, PyDict>,
+) -> PyResult<ProjectedOcrCall> {
+    let boundary_request = request.clone().unbind();
+    let arguments = OcrArguments { request, kwargs };
+    let model = arguments.model()?;
+    let custom_llm_provider = arguments.custom_llm_provider()?;
+    let (wire_document, retained_document) =
+        ProjectedDocument::project(py, &arguments.document()?)?.into_parts();
+    let api_key = arguments.api_key()?;
+    let specs = consumed_optional_params(&model, custom_llm_provider.as_deref())
+        .map_err(ocr_error_to_pyerr)?;
+    let names = specs.iter().map(|spec| spec.name).collect::<Vec<_>>();
+    let optional_params = project_optional_fields(kwargs, &names)?;
+    let input_sources = request_input_sources(
+        kwargs,
+        names
+            .iter()
+            .copied()
+            .chain(["api_key", "api_base", "extra_headers"]),
+    )?;
+    let azure_ad_token_provider = kwargs
+        .get_item("azure_ad_token_provider")?
+        .and_then(|provider| PythonTokenProvider::select(provider, AZURE_AD_TOKEN_PROVIDER));
+    let wire = OcrWireRequest {
+        model,
+        document: wire_document,
+        api_key: api_key.extract()?,
+        api_base: arguments.api_base()?,
+        custom_llm_provider,
+        extra_headers: arguments.extra_headers()?,
+        optional_params,
+        input_sources,
+        timeout_seconds: arguments.timeout_seconds()?,
+    };
+    let request = decode_request(wire).map_err(ocr_error_to_pyerr)?;
+    let provider = request.provider_name();
+    Ok(ProjectedOcrCall {
+        request: request.with_host_hooks(Arc::new(BridgeOcrHooks), None),
+        fields: ProjectedOcrFields {
+            boundary_request,
+            document: retained_document,
+            api_key: api_key.unbind(),
+            azure_ad_token_provider,
+            provider,
+            secret_fields: specs
+                .into_iter()
+                .filter(|spec| spec.secret)
+                .map(|spec| spec.name)
+                .collect(),
+        },
+    })
+}
+
+pub(super) fn admitted_call(outcome: NativeOutcome<OcrCall>) -> PyResult<OcrCall> {
+    match outcome {
+        NativeOutcome::Completed(call) => Ok(call),
+        NativeOutcome::Declined(reason) => Err(RustBridgeDeclined::new_err(format!(
+            "native OCR admission declined: {reason:?}"
+        ))),
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use litellm_core::Error;
+    use litellm_core::ocr::OcrDecline;
+    use pyo3::exceptions::{PyKeyError, PyTypeError, PyValueError};
+
+    use super::*;
+
+    fn eval<'py>(py: Python<'py>, source: &std::ffi::CStr) -> Bound<'py, PyDict> {
+        let locals = PyDict::new(py);
+        py.run(source, Some(&locals), Some(&locals)).unwrap();
+        locals
+    }
+
+    fn arguments<'a, 'py>(
+        request: &'a Bound<'py, PyAny>,
+        kwargs: &'a Bound<'py, PyDict>,
+    ) -> OcrArguments<'a, 'py> {
+        OcrArguments { request, kwargs }
+    }
+
+    fn project_document(
+        py: Python<'_>,
+        document: &Bound<'_, PyAny>,
+    ) -> PyResult<(Value, Py<PyAny>)> {
+        ProjectedDocument::project(py, document).map(ProjectedDocument::into_parts)
+    }
+
+    fn stub_timeout_conversion(py: Python<'_>) {
+        eval(
+            py,
+            c"
+import sys
+import types
+timeouts = types.ModuleType('litellm.rust_bridge.timeouts')
+timeouts.timeout_to_seconds = lambda timeout: None if timeout is None else float(timeout)
+sys.modules.setdefault('litellm', types.ModuleType('litellm'))
+sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bridge'))
+sys.modules['litellm.rust_bridge.timeouts'] = timeouts
+",
+        );
+    }
+
+    #[test]
+    fn typed_initial_decline_uses_bridge_decline_contract() {
+        Python::initialize();
+        Python::attach(|py| {
+            let Err(error) = admitted_call(NativeOutcome::Declined(OcrDecline::HostOperations))
+            else {
+                panic!("unsupported host operations should decline admission");
+            };
+            assert!(error.is_instance_of::<RustBridgeDeclined>(py));
+        });
+    }
+
+    #[test]
+    fn post_admission_error_does_not_use_bridge_decline_contract() {
+        Python::initialize();
+        Python::attach(|py| {
+            let error = ocr_error_to_pyerr(Error::InvalidRequest("callback result".into()));
+            assert!(error.is_instance_of::<PyValueError>(py));
+            assert!(!error.is_instance_of::<RustBridgeDeclined>(py));
+        });
+    }
+
+    #[test]
+    fn kwargs_override_request_attributes_including_explicit_none() {
+        Python::initialize();
+        Python::attach(|py| {
+            let locals = eval(
+                py,
+                c"
+class Request:
+    def __init__(self):
+        self.accesses = []
+    def __getattribute__(self, name):
+        if name != 'accesses':
+            object.__getattribute__(self, 'accesses').append(name)
+        return object.__getattribute__(self, name)
+request = Request()
+request.model = 'from-request'
+request.custom_llm_provider = 'mistral'
+kwargs = {'model': 'from-kwargs', 'custom_llm_provider': None}
+",
+            );
+            let request = locals.get_item("request").unwrap().unwrap();
+            let kwargs = locals
+                .get_item("kwargs")
+                .unwrap()
+                .unwrap()
+                .cast_into::<PyDict>()
+                .unwrap();
+            let arguments = arguments(&request, &kwargs);
+            assert_eq!(arguments.model().unwrap(), "from-kwargs");
+            assert_eq!(arguments.custom_llm_provider().unwrap(), None);
+            let accesses: Vec<String> = request.getattr("accesses").unwrap().extract().unwrap();
+            assert_eq!(accesses, Vec::<String>::new());
+        });
+    }
+
+    #[test]
+    fn missing_kwargs_read_the_request_property_once() {
+        Python::initialize();
+        Python::attach(|py| {
+            let locals = eval(
+                py,
+                c"
+class Request:
+    def __init__(self):
+        self.reads = 0
+    @property
+    def model(self):
+        self.reads += 1
+        return 'mistral-ocr-latest'
+request = Request()
+kwargs = {}
+",
+            );
+            let request = locals.get_item("request").unwrap().unwrap();
+            let kwargs = locals
+                .get_item("kwargs")
+                .unwrap()
+                .unwrap()
+                .cast_into::<PyDict>()
+                .unwrap();
+            assert_eq!(
+                arguments(&request, &kwargs).model().unwrap(),
+                "mistral-ocr-latest"
+            );
+            assert_eq!(
+                request.getattr("reads").unwrap().extract::<i32>().unwrap(),
+                1
+            );
+        });
+    }
+
+    #[test]
+    fn request_property_exceptions_keep_their_identity() {
+        Python::initialize();
+        Python::attach(|py| {
+            let locals = eval(
+                py,
+                c"
+failure = LookupError('model failed')
+class Request:
+    @property
+    def model(self):
+        raise failure
+request = Request()
+kwargs = {}
+",
+            );
+            let request = locals.get_item("request").unwrap().unwrap();
+            let kwargs = locals
+                .get_item("kwargs")
+                .unwrap()
+                .unwrap()
+                .cast_into::<PyDict>()
+                .unwrap();
+            let error = arguments(&request, &kwargs).model().unwrap_err();
+            assert!(
+                error
+                    .value(py)
+                    .is(locals.get_item("failure").unwrap().unwrap())
+            );
+        });
+    }
+
+    #[test]
+    fn unused_raising_property_is_never_inspected() {
+        Python::initialize();
+        Python::attach(|py| {
+            let locals = eval(
+                py,
+                c"
+class Request:
+    @property
+    def unused(self):
+        raise RuntimeError('unused')
+    model = 'mistral-ocr-latest'
+    custom_llm_provider = None
+request = Request()
+kwargs = {}
+",
+            );
+            let request = locals.get_item("request").unwrap().unwrap();
+            let kwargs = locals
+                .get_item("kwargs")
+                .unwrap()
+                .unwrap()
+                .cast_into::<PyDict>()
+                .unwrap();
+            let arguments = arguments(&request, &kwargs);
+            assert_eq!(arguments.model().unwrap(), "mistral-ocr-latest");
+            assert_eq!(arguments.custom_llm_provider().unwrap(), None);
+        });
+    }
+
+    #[test]
+    fn document_reader_mutations_are_visible_to_later_field_reads() {
+        Python::initialize();
+        Python::attach(|py| {
+            stub_timeout_conversion(py);
+            let locals = eval(
+                py,
+                c"
+class Request:
+    api_base = 'original'
+    timeout = 1
+    @property
+    def document(self):
+        return document
+class Reader:
+    def read(self):
+        Request.api_base = 'mutated'
+        Request.timeout = 9
+        return b'abc'
+document = {'type': 'file', 'file': Reader()}
+request = Request()
+kwargs = {}
+",
+            );
+            let request = locals.get_item("request").unwrap().unwrap();
+            let kwargs = locals
+                .get_item("kwargs")
+                .unwrap()
+                .unwrap()
+                .cast_into::<PyDict>()
+                .unwrap();
+            let arguments = arguments(&request, &kwargs);
+            let document = arguments.document().unwrap();
+            project_document(py, &document).unwrap();
+            assert_eq!(arguments.api_base().unwrap().as_deref(), Some("mutated"));
+            assert_eq!(arguments.timeout_seconds().unwrap(), Some(9.0));
+        });
+    }
+
+    #[test]
+    fn captured_api_key_keeps_the_original_python_object() {
+        Python::initialize();
+        Python::attach(|py| {
+            let locals = eval(
+                py,
+                c"
+key = object()
+class Request:
+    api_key = None
+request = Request()
+kwargs = {'api_key': key}
+",
+            );
+            let request = locals.get_item("request").unwrap().unwrap();
+            let kwargs = locals
+                .get_item("kwargs")
+                .unwrap()
+                .unwrap()
+                .cast_into::<PyDict>()
+                .unwrap();
+            let captured = arguments(&request, &kwargs).api_key().unwrap();
+            assert!(
+                captured
+                    .unbind()
+                    .bind(py)
+                    .is(locals.get_item("key").unwrap().unwrap())
+            );
+        });
+    }
+
+    #[test]
+    fn file_documents_are_encoded_and_other_documents_keep_the_python_object() {
+        Python::initialize();
+        Python::attach(|py| {
+            let file = py
+                .eval(
+                    c"{'type': 'file', 'file': b'%PDF-1.4', 'mime_type': 'application/pdf'}",
+                    None,
+                    None,
+                )
+                .unwrap();
+            assert_eq!(
+                project_document(py, &file).unwrap().0,
+                serde_json::json!({
+                    "type": "document_url",
+                    "document_url": "data:application/pdf;base64,JVBERi0xLjQ=",
+                })
+            );
+
+            let original = py
+                .eval(
+                    c"{'type': 'document_url', 'document_url': 'https://example.com/a.pdf'}",
+                    None,
+                    None,
+                )
+                .unwrap();
+            let (wire, retained) = project_document(py, &original).unwrap();
+            assert_eq!(
+                wire,
+                serde_json::json!({
+                    "type": "document_url",
+                    "document_url": "https://example.com/a.pdf",
+                })
+            );
+            assert!(retained.bind(py).is(&original));
+        });
+    }
+
+    #[test]
+    fn unknown_document_types_reach_existing_downstream_validation() {
+        Python::initialize();
+        Python::attach(|py| {
+            let document = py
+                .eval(c"{'type': 'mystery', 'mystery': 'x'}", None, None)
+                .unwrap();
+            let wire_document = project_document(py, &document).unwrap().0;
+            assert_eq!(
+                wire_document,
+                serde_json::json!({"type": "mystery", "mystery": "x"})
+            );
+            let error = match decode_request(OcrWireRequest {
+                model: "mistral/mistral-ocr-latest".into(),
+                document: wire_document,
+                api_key: None,
+                api_base: None,
+                custom_llm_provider: None,
+                extra_headers: None,
+                optional_params: Map::new(),
+                input_sources: Default::default(),
+                timeout_seconds: None,
+            }) {
+                Ok(_) => panic!("unknown discriminators belong to core validation"),
+                Err(error) => error,
+            };
+            assert!(error.to_string().contains("document"));
+        });
+    }
+
+    #[test]
+    fn document_discriminator_errors_keep_their_existing_exceptions() {
+        Python::initialize();
+        Python::attach(|py| {
+            let missing = py.eval(c"{}", None, None).unwrap();
+            assert!(
+                project_document(py, &missing)
+                    .unwrap_err()
+                    .is_instance_of::<PyKeyError>(py)
+            );
+
+            let non_string = py.eval(c"{'type': 1}", None, None).unwrap();
+            assert!(
+                project_document(py, &non_string)
+                    .unwrap_err()
+                    .is_instance_of::<PyTypeError>(py)
+            );
+
+            let locals = eval(
+                py,
+                c"
+failure = RuntimeError('type lookup failed')
+class Document:
+    def __getitem__(self, key):
+        raise failure
+document = Document()
+",
+            );
+            let error =
+                project_document(py, &locals.get_item("document").unwrap().unwrap()).unwrap_err();
+            assert!(
+                error
+                    .value(py)
+                    .is(locals.get_item("failure").unwrap().unwrap())
+            );
+        });
+    }
+
+    #[test]
+    fn document_classification_happens_once() {
+        Python::initialize();
+        Python::attach(|py| {
+            let locals = eval(
+                py,
+                c"
+class Document(dict):
+    def __init__(self):
+        super().__init__({'file': b'abc'})
+        self.reads = []
+    def __getitem__(self, key):
+        self.reads.append(key)
+        if key == 'type':
+            return 'file' if self.reads.count('type') == 1 else 'document_url'
+        return super().__getitem__(key)
+document = Document()
+",
+            );
+            let document = locals.get_item("document").unwrap().unwrap();
+            let (wire, retained) = project_document(py, &document).unwrap();
+            assert_eq!(wire["type"], "document_url");
+            assert!(!retained.bind(py).is(&document));
+            let reads: Vec<String> = document.getattr("reads").unwrap().extract().unwrap();
+            assert_eq!(reads, ["type", "mime_type", "file"]);
+        });
+    }
+}

litellm-rust/crates/python-bridge/src/routes/ocr/value.rs

@@ -1,12 +1,11 @@
 use litellm_core::Error;
 use std::future::Future;
 
-use litellm_ai_gateway::io::ocr::{OcrRequest, ocr as run_ocr};
-use litellm_core::ocr::wire::{OcrWireRequest, decode_request, is_supported_request};
+use litellm_core::ocr::wire::{OcrWireRequest, decode_request};
 use pyo3::prelude::*;
 use serde_json::Value;
 
-use crate::errors::ocr_error_to_pyerr;
+use super::errors::to_pyerr as ocr_error_to_pyerr;
 use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty};
 
 fn prepare_ocr(
@@ -38,37 +37,20 @@ fn prepare_ocr(
             extra_headers,
             timeout,
         } = options;
-        if is_supported_request(&model, custom_llm_provider.as_deref()) {
-            let request = decode_request(OcrWireRequest {
-                model,
-                document,
-                api_key,
-                api_base,
-                custom_llm_provider,
-                extra_headers,
-                optional_params,
-                input_sources,
-                timeout_seconds: timeout.map(|value| value.as_secs_f64()),
-            })?;
-            return litellm_core::ocr::ocr(request)
-                .await
-                .map(|response| response.into_json());
-        }
-        run_ocr(OcrRequest {
-            model: &model,
+        let request = decode_request(OcrWireRequest {
+            model,
             document,
-            api_key: api_key.as_deref(),
-            api_base: api_base.as_deref(),
-            custom_llm_provider: custom_llm_provider.as_deref(),
+            api_key,
+            api_base,
+            custom_llm_provider,
             extra_headers,
             optional_params,
-            timeout,
-            callbacks: Vec::new(),
-            guardrails: Vec::new(),
-            request_metadata: Default::default(),
-            litellm_call_id: None,
-        })
-        .await
+            input_sources,
+            timeout_seconds: timeout.map(|value| value.as_secs_f64()),
+        })?;
+        litellm_core::ocr::ocr(request)
+            .await
+            .map(|response| response.into_json())
     })
 }
 
@@ -96,22 +78,3 @@ bridge_route! {
     prepare = prepare_ocr,
     errors = ocr_error_to_pyerr,
 }
-
-#[cfg(test)]
-mod tests {
-    use litellm_core::ocr::wire::is_supported_request;
-
-    #[test]
-    fn native_activation_includes_migrated_providers() {
-        assert!(is_supported_request("model", Some("mistral")));
-        assert!(is_supported_request("pixtral-12b", Some("azure_ai")));
-        assert!(is_supported_request(
-            "documentintelligence/prebuilt-read",
-            Some("azure_ai")
-        ));
-        assert!(is_supported_request("parse-v3", Some("reducto")));
-        assert!(is_supported_request("parse-legacy", Some("reducto")));
-        assert!(is_supported_request("mistral-ocr", Some("vertex_ai")));
-        assert!(is_supported_request("deepseek-ocr", Some("vertex_ai")));
-    }
-}

litellm-rust/crates/python-bridge/tests/lifecycle.py

@@ -0,0 +1,186 @@
+import asyncio
+import gc
+import threading
+import weakref
+from contextvars import ContextVar
+
+
+async def exercise():
+    caller = asyncio.current_task()
+    thread = threading.get_ident()
+    loop = asyncio.get_running_loop()
+    marker = ContextVar("driver", default="before")
+    entered = asyncio.Event()
+    released = asyncio.Event()
+    result = object()
+
+    class CustomAwaitable:
+        def __await__(self):
+            return operation().__await__()
+
+    async def operation():
+        assert asyncio.current_task() is caller
+        assert threading.get_ident() == thread
+        assert asyncio.get_running_loop() is loop
+        marker.set("inside")
+        entered.set()
+        await released.wait()
+        assert asyncio.current_task() is caller
+        assert marker.get() == "inside"
+        return result
+
+    async def release():
+        await entered.wait()
+        released.set()
+
+    releaser = asyncio.create_task(release())
+    execution = await_execution(CustomAwaitable())
+    try:
+        execution.resume_value(None)
+    except RuntimeError:
+        pass
+    else:
+        raise AssertionError("resumed an unstarted execution")
+    wrapped = drive(execution)
+    try:
+        wrapped.send(1)
+    except TypeError:
+        pass
+    else:
+        raise AssertionError("accepted initial value")
+    assert await wrapped is result
+    assert marker.get() == "inside"
+    await releaser
+    execution.close()
+    execution.close()
+    try:
+        await wrapped
+    except RuntimeError:
+        pass
+    else:
+        raise AssertionError("accepted coroutine reuse")
+
+    final_awaitable = CustomAwaitable()
+    assert await drive(calling_execution(lambda: final_awaitable)) is final_awaitable
+
+    cause = KeyError("cause")
+    failure = ValueError("original")
+
+    async def failing():
+        await asyncio.sleep(0)
+        raise failure from cause
+
+    try:
+        await drive(await_execution(failing()))
+    except ValueError as error:
+        assert error is failure
+        assert error.__cause__ is cause
+        names = []
+        traceback = error.__traceback__
+        while traceback:
+            names.append(traceback.tb_frame.f_code.co_name)
+            traceback = traceback.tb_next
+        assert "failing" in names
+    else:
+        raise AssertionError("lost original exception")
+
+    for suppress in (False, True):
+        pending = asyncio.Event()
+        cleanup_entered = asyncio.Event()
+        cleanup_release = asyncio.Event()
+        cleaned = []
+
+        async def cancel_operation():
+            try:
+                pending.set()
+                await asyncio.Event().wait()
+            except asyncio.CancelledError:
+                if suppress:
+                    return result
+                raise
+            finally:
+                cleanup_entered.set()
+                try:
+                    await cleanup_release.wait()
+                except asyncio.CancelledError:
+                    await cleanup_release.wait()
+                cleaned.append(asyncio.current_task())
+
+        task = asyncio.create_task(drive(await_execution(cancel_operation())))
+        await pending.wait()
+        task.cancel()
+        await cleanup_entered.wait()
+        assert not task.done()
+        task.cancel()
+        await asyncio.sleep(0)
+        cleanup_release.set()
+        if suppress:
+            assert await task is result
+        else:
+            try:
+                await task
+            except asyncio.CancelledError:
+                pass
+            else:
+                raise AssertionError("lost cancellation")
+        assert cleaned == [task]
+
+    observed = []
+
+    def reenter():
+        try:
+            active.start()
+        except RuntimeError as error:
+            observed.append(str(error))
+        return result
+
+    active = calling_execution(reenter)
+    assert await drive(active) is result
+    assert observed == ["execution is already running"]
+
+    class Finalizer:
+        def __call__(self):
+            return result
+
+        def __del__(self):
+            self.owner.close()
+            observed.append("released")
+
+    def cycle(started):
+        callback = Finalizer()
+        execution = calling_execution(callback)
+        callback.owner = execution
+        if started:
+            assert execution.start().value is result
+        return weakref.ref(callback)
+
+    for started in (False, True):
+        reference = cycle(started)
+        gc.collect()
+        assert reference() is None
+    assert observed[-2:] == ["released", "released"]
+
+    class Awaitable:
+        def __await__(self):
+            try:
+                yield self
+            finally:
+                observed.append("unwound")
+
+    def abandoned(started):
+        awaitable = Awaitable()
+        coroutine = drive(await_execution(awaitable))
+        awaitable.owner = coroutine
+        if started:
+            assert coroutine.send(None) is awaitable
+        coroutine.close()
+        return weakref.ref(awaitable)
+
+    for started in (False, True):
+        reference = abandoned(started)
+        gc.collect()
+        assert reference() is None
+    assert observed[-1] == "unwound"
+
+
+asyncio.run(asyncio.wait_for(exercise(), 10))