refactor(ocr): complete native lifecycle and preserve Azure auth

[yujonglee-berri]

TLDR

Problem this solves:

• OCR lifecycle ownership was split across Python and Rust
• Generic Azure credentials could override Document Intelligence credentials
• Azure polling could drop accepted request authentication

How it solves it:

• Rust owns supported OCR execution by default
• LITELLM_RUST=0 keeps the Python path available
• OCR adapters own service-specific connection resolution
• Azure polling preserves subscription-key and Entra authentication

User Flow

Before: OCR behavior depends on the selected runtime and Azure Document Intelligence can fail with valid credentials

1. The developer calls litellm.ocr() or litellm.aocr() with a supported document
2. Default execution enters the native OCR lifecycle, while LITELLM_RUST=0 selects Python
3. A Python Document Intelligence call can send the generic Azure AI key and return 401
4. An accepted Entra-authenticated analysis can lose authentication while polling and fail with Unknown operation status: None

After: both runtime choices preserve the public SDK behavior and Azure authentication

1. The developer calls litellm.ocr() or litellm.aocr() with the same supported document
2. Default execution enters the native OCR lifecycle, while LITELLM_RUST=0 selects Python
3. A Python Document Intelligence call uses its dedicated key and returns parsed pages
4. An accepted Entra-authenticated analysis keeps its bearer token while polling and returns parsed pages

Design

This PR replaces the former #40734 / #40676 stack and consolidates #40437 and #40499

Core selects setup, request preparation, provider work, public response construction, deployment processing, metadata finalization, and terminal dispatch. Python performs selected host operations while retaining caller-owned objects, exception identity, and scheduling behavior

The native execution handle exposes start, resume_value, resume_error, and idempotent close. An ordinary Python coroutine awaits selected operations in the caller task. Retained Python references participate in garbage collection, and cancellation signals provider work with an awaited termination acknowledgement where required

Asynchronous native operations are polled once while detached from Python and inside the Tokio runtime. Ready operations complete directly; pending operations retain the same future and use the existing cancellation-aware scheduler. Logging retains only its needed fields, avoiding an extra document copy

Bridge-created OCR loggers check live registrations at each phase. Requests without optional consumers skip deployment hook coroutines, logging executor scheduling, deferred objects, callback context copies, traceback formatting, and logging-only payload conversions. Credentials, limits, cost accounting, timing, response metadata, and context cleanup remain active

Provider responses use a 64 MiB server-side byte limit across OCR, upload, and polling requests. max_response_bytes may lower that limit for one call but cannot raise it. Oversized success responses fail without draining the body; upstream errors retain their HTTP status and a bounded diagnostic prefix

Provider workflows cover Mistral, Cohere, Azure Mistral, Azure Cohere, Azure Document Intelligence, Vertex Mistral, Vertex DeepSeek, and Reducto

Relevant issues

Supersedes #40676, #40437, and #40499

Linear ticket

Pre-Submission checklist

• I have added meaningful tests
• Focused Python and installed-extension tests pass locally
• My PR passes all required CI/CD checks
• This PR consolidates OCR lifecycle ownership into one change
• Greptile confidence is at least 4/5 on the current tip

Screenshots / Proof of Fix

Shared setup: run /tmp/ocr_auth_matrix.py from the repository virtual environment with live provider credentials

Before (9132a5343b)

Default Rust OCR

1. Run .venv/bin/python /tmp/ocr_auth_matrix.py --rust --jobs 4 --delay 0.25
2. Observe 25 passes, including Azure Document Intelligence sync and async calls; remaining failures are missing Entra credentials or provider rate limits

Python fallback Azure Document Intelligence

1. Run LITELLM_RUST=0 .venv/bin/python /tmp/ocr_auth_matrix.py --jobs 2 --delay 0.75
2. Observe environment-key sync and async calls return 401, and one polling call fails with Unknown operation status: None

After (af77246920)

Default Rust OCR

1. Run .venv/bin/python /tmp/ocr_auth_matrix.py --rust --case azure-di-env-key --case azure-di-explicit-key --jobs 1 --delay 10
2. Observe all four sync and async calls return one parsed page after quota-limited calls are retried

Python fallback Azure Document Intelligence

1. Run LITELLM_RUST=0 .venv/bin/python /tmp/ocr_auth_matrix.py --case azure-di-env-key --case azure-di-explicit-key --jobs 1 --delay 1
2. Observe all four sync and async calls return one parsed page

Type

Refactoring
Bug Fix

Caveats (if any)

Low

• Azure free-tier calls require spacing to avoid polling limits
• Local Azure default-credential probes lack Entra credentials
• Bugbot is unavailable for this PR and is not counted as a passing review

Validation

The current tip passes 201 focused Python OCR, native configuration, and Azure transformation tests. The isolated compiled-wheel suite passes all 183 Rust-extension Python tests. The response-cap regression, WebSocket crypto-provider integration, and all 63 Python bridge library tests pass. Workspace Clippy passes for all targets with warnings denied. make check passes Python formatting, Ruff, type discipline, basedpyright, test-quality budgets, import checks, and generated API checks. Hosted CI is green with 88 passing checks, one intentional skip, and passing patch coverage. Greptile reports 5/5 confidence with no blocking findings, and Veria reports no open security concerns. No lint or coverage thresholds were changed

Live matrix runs exercised sync and async requests through default Rust execution and the LITELLM_RUST=0 Python fallback. Rust passed Azure Document Intelligence environment and explicit key cases, plus the previously rate-limited Azure Cohere explicit-key cases. Python passed all four Azure Document Intelligence cases after the port. Subsequent Azure Cohere calls reached provider rate limits after one successful Python request

Final Attestation

• Behavioral tests cover lifecycle ordering, identity, mutations, failures, authentication, and cleanup

[greptile-apps]

Greptile Summary

This PR moves supported OCR execution into the native lifecycle while retaining a Python fallback and preserving Azure Document Intelligence authentication across submission and polling.

• Makes Rust responsible for OCR lifecycle sequencing, provider execution, response construction, and terminal dispatch.
• Preserves dedicated Azure Document Intelligence credentials over generic Azure AI configuration.
• Carries subscription-key and Entra authentication through asynchronous polling.
• Adds response-size limits and lifecycle, callback, provider, bridge, and authentication regression coverage.
• The latest change isolates the Azure auth-precedence test by supplying its Document Intelligence endpoint explicitly.

Confidence Score: 5/5

The PR appears safe to merge, with no outstanding blocking findings at the current tip.

The latest test change fully isolates endpoint selection while retaining a decisive assertion that the dedicated Document Intelligence key is used instead of the generic Azure AI key. The previously reported streaming helpers now have complete parameter and return typing. yujonglee-berri accepted the earlier disconnected-completion behavior as intended for a subsequent change, and that thread was resolved. The final-response callback, core callback ownership, and public helper typing threads were manually resolved without explanation; current code also addresses their reported behavior.

Important Files Changed

Filename	Overview
tests/test_litellm/ocr/test_legacy.py	Explicitly supplies the Document Intelligence endpoint so the auth-precedence regression test is independent of developer environment configuration.
litellm/llms/azure_ai/ocr/document_intelligence/transformation.py	Resolves Document Intelligence connection settings independently and builds service-specific API-key or Entra authentication.
litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/polling.rs	Preserves authentication during polling and dispatches post-call handling for the final provider response.
litellm-rust/crates/core/src/ocr/lifecycle.rs	Implements native OCR lifecycle sequencing while delegating selected host operations through the bridge.
litellm/rust_bridge/ocr_lifecycle.py	Drives native lifecycle suspension and resumption while preserving Python task and exception behavior.

_{Reviews (14): Last reviewed commit: "test(ocr): isolate Azure auth precedence..." | Re-trigger Greptile}

[codecov]

Codecov Report

❌ Patch coverage is 92.34568% with 31 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
litellm/rust_bridge/lifecycle.py	83.50%	16 Missing ⚠️
litellm/ocr/legacy.py	91.02%	14 Missing ⚠️
litellm/rust_bridge/ocr_lifecycle.py	96.29%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[yujonglee-berri]

@greptileai Please re-review the wrapper handoff fixes and new replacement, failure, deferral, context, and identity coverage.

[yujonglee-berri]

@greptileai please review the current tip after the immutable invocation-kwargs fix and replacement-hook regression coverage.

[yujonglee-berri]

@greptileai please review f8043c370b, including completion retention through final metadata and embedding processing.

[yujonglee-berri]

@greptileai Please review 505016e for typed OCR completion handoff, wrapper-owned cleanup, cancellation handling, and deferred callback behavior

[yujonglee-berri]

@greptileai Please review 1eeaaad: explicit lazy executor import, corrected submit protocol, real-thread context coverage, and Rust fixture isolation

[codspeed]

Merging this PR will not alter performance

✅ 31 untouched benchmarks

---

_{Comparing litellm_ocr_completion_boundary (af77246) with litellm_internal_staging (0c98afa)}

[veria-ai]

PR overview

All previously flagged issues have been addressed. No open security concerns remain on this pull request.

Security review

No open security issues remain on this pull request.

Fixed/addressed: 1 · PR risk: 0/10

[yujonglee-berri]

@greptileai Please review current tip 02e150c after Azure auth consolidation, response-cap enforcement, and successful local Clippy and regression tests

[yujonglee-berri]

@greptileai Please review current tip fa6bdf2 after aligning compiled-wheel OCR tests with explicit GIL module and public validation contracts

[yujonglee-berri]

@greptileai Please review current tip af77246 after isolating Azure Document Intelligence auth precedence coverage from developer environment configuration