done

litellm-rust/README.md

@@ -51,10 +51,14 @@ function per top-level route, mirroring the core entrypoints.
 
 ## Checks
 
-### Native OCR Boundary
+### Private Native OCR Proof
 
-`LITELLM_RUST=1` selects native OCR before Python provider preparation. With Rust
-disabled, the existing Python execution and authentication paths are unchanged
+Public `litellm.ocr` and `litellm.aocr` always use the existing Python lifecycle,
+including when `litellm.rust(True)` or `LITELLM_RUST=1` enables other Rust paths.
+Native OCR remains a private proof until full lifecycle parity is established.
+Only tests requesting the private `native_ocr` fixture replace those public
+functions with test-only route selection: Rust enabled calls the native bridge,
+and Rust disabled calls the captured production Python functions
 
 The bridge retains the complete call argument dictionary as a Python object,
 including opaque callback and metadata objects. It creates callback-visible
@@ -65,7 +69,7 @@ chat messages using the existing Rust transform. Rust performs provider preparat
 encoding, HTTP and response normalization. Python continues to dispatch existing
 logging operations and construct the public response object
 
-This is an opt-in implementation scaffold, not full OCR parity. Azure Mistral and
+This is a private implementation scaffold, not full OCR parity. Azure Mistral and
 Vertex Mistral accept inline data URIs with supplied keys/tokens, native environment
 keys or auth headers. Azure also accepts a supplied `azure_ad_token`. Vertex
 DeepSeek uses its existing chat request and OCR response transforms. Cloud
@@ -74,9 +78,9 @@ HTTP document URL conversion fails only for configs requiring data URIs. Azure
 Document Intelligence selects its own config but fails at the polling capability
 check before sending a billable analyze request. Cohere transforms, file inputs,
 streaming, native response format and compression remain unsupported.
-An enabled but missing native extension also fails;
-neither case falls back to Python execution. Transport failures currently use a
-generic error rather than the SDK's timeout-specific exception
+Direct private bridge calls with a missing native extension also fail;
+neither case falls back to Python execution within the private route. Transport
+failures currently use a generic error rather than the SDK's timeout-specific exception
 
 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
@@ -100,7 +104,7 @@ make lint-rust-python-fixtures
 `lint-rust-python-fixtures` runs pinned Ruff lint and formatting checks without
 syncing the project environment
 
-Run the native OCR acceptance gate from the repository root:
+Run the private native OCR proof gate from the repository root:
 
 ```bash
 make test-rust-ocr
@@ -126,7 +130,30 @@ invocation context and ownership against Python behavior, including existing
 LiteLLM components. Short synthetic pre-call contracts use Rust-owned table-driven
 cases with inline Python callbacks; larger component scenarios share Python
 fixtures. These generic proofs complement, rather than replace, native OCR
-acceptance tests
+private proof tests
+
+The standard-library-only tests in
+`crates/python-interop/tests/callback_patterns.rs` define small inline Python
+callbacks, with Rust controlling invocation, ownership and assertions. They
+compare Python-reference and Rust-retained calls using both direct and awaited
+invocation. They model the behavior
+groups in the callback use-case inventory: live versus serialized queues,
+mutation before an error, ignored returns, identity-based redaction, block-state
+stash, background writes after return, parallel live data versus snapshots, and
+shallow/deep copies with uncopyable-value fallback. Copy controls deliberately
+produce different observations; event gates establish ordering without sleeps.
+Existing synthetic lifecycle cases also cover streams, context and cancellation
+
+Run this matrix without LiteLLM, vendor SDKs, credentials or services:
+
+```bash
+cargo test --manifest-path litellm-rust/Cargo.toml -p litellm-python-interop --test callback_patterns
+```
+
+These are behavioral models, not tests of vendor authentication, delivery or
+production dispatcher policy. The optional component and integration fixtures
+exercise existing LiteLLM implementations with fake transports and credentials
+as supplementary coverage; run them with `make test-rust-python`
 
 The callback lifecycle scenarios use
 `#[serial(python_interpreter)]` to isolate CPython GC and interpreter-wide

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

@@ -0,0 +1,600 @@
+use crate::Error;
+
+use super::{OcrRequest, prepare};
+
+#[derive(Debug)]
+pub enum NativeOutcome<T> {
+    Completed(T),
+    Declined(Decline),
+}
+
+#[derive(Debug, PartialEq, Eq)]
+pub struct Decline(&'static str);
+
+impl Decline {
+    pub fn reason(&self) -> &'static str {
+        self.0
+    }
+}
+
+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
+pub enum CredentialMethod {
+    #[default]
+    Configured,
+    Acquisition,
+}
+
+#[derive(Default)]
+pub struct Options {
+    pub asynchronous: bool,
+    pub internal_call: bool,
+    pub call_id: Option<String>,
+    pub trace_id: Option<String>,
+    pub credential_method: CredentialMethod,
+}
+
+#[derive(Debug, PartialEq, Eq)]
+pub struct Identity {
+    pub requested_model: String,
+    pub call_id: String,
+    pub trace_id: Option<String>,
+    pub generated_call_id: bool,
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum Outcome {
+    Success,
+    Failure,
+    Abort,
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum Operation {
+    Setup,
+    DeploymentPre,
+    Prepare,
+    Send,
+    DeploymentSuccess,
+    DeploymentFailure,
+    SyncSuccess,
+    AsyncSuccess,
+    SyncSuccessIfNeeded,
+    SyncFailure,
+    AsyncFailure,
+    Restore,
+    Complete(Outcome),
+}
+
+#[derive(Clone, Copy, Debug, Default)]
+pub struct Observations {
+    pub logger_available: bool,
+    pub has_fallbacks: bool,
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum ErrorDisposition {
+    Preserve,
+    Replace,
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub struct Transition {
+    pub operation: Operation,
+    pub error: ErrorDisposition,
+}
+
+#[derive(Debug)]
+pub struct Lifecycle {
+    operation: Operation,
+    outcome: Outcome,
+    asynchronous: bool,
+    internal_call: bool,
+    identity: Identity,
+}
+
+impl Lifecycle {
+    pub fn new(admission: &OcrRequest, options: Options) -> Result<NativeOutcome<Self>, Error> {
+        match prepare::admission_capabilities(admission) {
+            Err(Error::Unsupported(reason)) => return Ok(NativeOutcome::Declined(Decline(reason))),
+            Err(error) => return Err(error),
+            Ok(()) => {}
+        }
+        if options.credential_method == CredentialMethod::Acquisition {
+            return Ok(NativeOutcome::Declined(Decline(
+                "OCR credential acquisition",
+            )));
+        }
+        let generated_call_id = options.call_id.is_none();
+        let call_id = options.call_id.unwrap_or_else(|| {
+            let id = (rand::random::<u128>() & !(0xf000_u128 << 64 | 0xc000_u128 << 48))
+                | (0x4000_u128 << 64 | 0x8000_u128 << 48);
+            let hex = format!("{id:032x}");
+            format!(
+                "{}-{}-{}-{}-{}",
+                &hex[..8],
+                &hex[8..12],
+                &hex[12..16],
+                &hex[16..20],
+                &hex[20..]
+            )
+        });
+        Ok(NativeOutcome::Completed(Self {
+            operation: Operation::Setup,
+            outcome: Outcome::Success,
+            asynchronous: options.asynchronous,
+            internal_call: options.internal_call,
+            identity: Identity {
+                requested_model: admission.model.clone(),
+                call_id,
+                trace_id: options.trace_id,
+                generated_call_id,
+            },
+        }))
+    }
+
+    pub fn operation(&self) -> Operation {
+        self.operation
+    }
+
+    pub fn identity(&self) -> &Identity {
+        &self.identity
+    }
+
+    pub fn advance(
+        &mut self,
+        outcome: Outcome,
+        observations: Observations,
+    ) -> Result<Transition, Error> {
+        use Operation::*;
+
+        if matches!(self.operation, Complete(_)) {
+            return Err(Error::InvalidRequest(
+                "OCR lifecycle is already complete".into(),
+            ));
+        }
+        let failure = if observations.logger_available && !(self.asynchronous && self.internal_call)
+        {
+            SyncFailure
+        } else {
+            Restore
+        };
+        let error = if outcome != Outcome::Success && self.operation != DeploymentFailure {
+            self.outcome = outcome;
+            ErrorDisposition::Replace
+        } else {
+            ErrorDisposition::Preserve
+        };
+        self.operation = match (self.operation, outcome) {
+            (Restore, _) => Complete(self.outcome),
+            (DeploymentFailure, _) => failure,
+            (_, Outcome::Abort) => Restore,
+            (SyncFailure | AsyncFailure, Outcome::Failure) => Restore,
+            (Prepare | Send, Outcome::Failure) if self.asynchronous => DeploymentFailure,
+            (_, Outcome::Failure) => failure,
+            (Setup, Outcome::Success) if self.asynchronous => DeploymentPre,
+            (Setup | DeploymentPre, Outcome::Success) => Prepare,
+            (Prepare, Outcome::Success) => Send,
+            (Send, Outcome::Success) if self.asynchronous => DeploymentSuccess,
+            (Send, Outcome::Success) => SyncSuccess,
+            (DeploymentSuccess, Outcome::Success) => {
+                if self.internal_call || observations.has_fallbacks {
+                    SyncSuccessIfNeeded
+                } else {
+                    AsyncSuccess
+                }
+            }
+            (AsyncSuccess, Outcome::Success) => SyncSuccessIfNeeded,
+            (SyncFailure, Outcome::Success) if self.asynchronous => AsyncFailure,
+            (SyncSuccess | SyncSuccessIfNeeded | SyncFailure | AsyncFailure, Outcome::Success) => {
+                Restore
+            }
+            (Complete(_), _) => unreachable!(),
+        };
+        Ok(Transition {
+            operation: self.operation,
+            error,
+        })
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::rc::Rc;
+
+    use super::*;
+    use crate::ocr::types::OcrDocument;
+
+    fn request() -> OcrRequest {
+        OcrRequest {
+            model: "mistral/requested-model".into(),
+            custom_llm_provider: None,
+            api_key: Some("test-key".into()),
+            api_base: Some("https://example.test".into()),
+            extra_headers: vec![],
+            timeout_seconds: 2.0,
+            request_format: None,
+            document: OcrDocument::DocumentUrl {
+                document_url: "https://example.test/doc.pdf".into(),
+            },
+            azure_ad_token: None,
+            vertex_project: None,
+            vertex_location: None,
+            stream: false,
+        }
+    }
+
+    fn machine(asynchronous: bool) -> Lifecycle {
+        let NativeOutcome::Completed(machine) = Lifecycle::new(
+            &request(),
+            Options {
+                asynchronous,
+                ..Options::default()
+            },
+        )
+        .unwrap() else {
+            panic!("expected admission")
+        };
+        machine
+    }
+
+    fn observed() -> Observations {
+        Observations {
+            logger_available: true,
+            has_fallbacks: false,
+        }
+    }
+
+    fn reach(machine: &mut Lifecycle, operation: Operation) {
+        for _ in 0..12 {
+            if machine.operation() == operation {
+                return;
+            }
+            machine.advance(Outcome::Success, observed()).unwrap();
+        }
+        panic!("operation not reached: {operation:?}")
+    }
+
+    #[test]
+    fn success_sequences_and_completion_are_core_selected() {
+        use Operation::*;
+        for (asynchronous, expected) in [
+            (false, vec![Setup, Prepare, Send, SyncSuccess, Restore]),
+            (
+                true,
+                vec![
+                    Setup,
+                    DeploymentPre,
+                    Prepare,
+                    Send,
+                    DeploymentSuccess,
+                    AsyncSuccess,
+                    SyncSuccessIfNeeded,
+                    Restore,
+                ],
+            ),
+        ] {
+            let mut machine = machine(asynchronous);
+            for operation in expected {
+                assert_eq!(machine.operation(), operation);
+                assert_eq!(
+                    machine.advance(Outcome::Success, observed()).unwrap().error,
+                    ErrorDisposition::Preserve
+                );
+            }
+            assert_eq!(machine.operation(), Complete(Outcome::Success));
+            assert!(machine.advance(Outcome::Success, observed()).is_err());
+        }
+    }
+
+    #[test]
+    fn ordinary_errors_and_cancellation_at_every_execution_stage() {
+        use Operation::*;
+        for asynchronous in [false, true] {
+            let stages = if asynchronous {
+                vec![
+                    Setup,
+                    DeploymentPre,
+                    Prepare,
+                    Send,
+                    DeploymentSuccess,
+                    AsyncSuccess,
+                    SyncSuccessIfNeeded,
+                ]
+            } else {
+                vec![Setup, Prepare, Send, SyncSuccess]
+            };
+            for stage in stages {
+                for outcome in [Outcome::Failure, Outcome::Abort] {
+                    let mut machine = machine(asynchronous);
+                    reach(&mut machine, stage);
+                    let transition = machine.advance(outcome, observed()).unwrap();
+                    assert_eq!(transition.error, ErrorDisposition::Replace);
+                    let expected = if outcome == Outcome::Abort {
+                        Restore
+                    } else if asynchronous && matches!(stage, Prepare | Send) {
+                        DeploymentFailure
+                    } else {
+                        SyncFailure
+                    };
+                    assert_eq!(transition.operation, expected, "{stage:?}, {outcome:?}");
+                    if expected == DeploymentFailure {
+                        assert_eq!(
+                            machine
+                                .advance(Outcome::Success, observed())
+                                .unwrap()
+                                .operation,
+                            SyncFailure
+                        );
+                    }
+                    if outcome == Outcome::Failure {
+                        assert_eq!(
+                            machine
+                                .advance(Outcome::Success, observed())
+                                .unwrap()
+                                .operation,
+                            if asynchronous { AsyncFailure } else { Restore }
+                        );
+                        if asynchronous {
+                            assert_eq!(
+                                machine
+                                    .advance(Outcome::Success, observed())
+                                    .unwrap()
+                                    .operation,
+                                Restore
+                            );
+                        }
+                    }
+                    assert_eq!(
+                        machine
+                            .advance(Outcome::Success, observed())
+                            .unwrap()
+                            .operation,
+                        Complete(outcome)
+                    );
+                }
+            }
+        }
+    }
+
+    #[test]
+    fn deployment_observer_preserves_opaque_original_error_even_on_abort() {
+        for observer_outcome in [Outcome::Success, Outcome::Failure, Outcome::Abort] {
+            let mut machine = machine(true);
+            reach(&mut machine, Operation::Send);
+            let original = Rc::new("original opaque error");
+            let mut retained = Rc::clone(&original);
+            assert_eq!(
+                machine
+                    .advance(Outcome::Failure, observed())
+                    .unwrap()
+                    .operation,
+                Operation::DeploymentFailure
+            );
+            let transition = machine.advance(observer_outcome, observed()).unwrap();
+            if transition.error == ErrorDisposition::Replace {
+                retained = Rc::new("observer error");
+            }
+            assert!(Rc::ptr_eq(&original, &retained));
+            assert_eq!(transition.operation, Operation::SyncFailure);
+            reach(&mut machine, Operation::Restore);
+            assert_eq!(
+                machine
+                    .advance(Outcome::Success, observed())
+                    .unwrap()
+                    .operation,
+                Operation::Complete(Outcome::Failure)
+            );
+        }
+    }
+
+    #[test]
+    fn failure_handlers_and_restore_propagate_their_own_errors_without_recursion() {
+        for stage in [
+            Operation::SyncFailure,
+            Operation::AsyncFailure,
+            Operation::Restore,
+        ] {
+            for outcome in [Outcome::Failure, Outcome::Abort] {
+                let mut machine = machine(true);
+                machine.advance(Outcome::Failure, observed()).unwrap();
+                reach(&mut machine, stage);
+                let transition = machine.advance(outcome, observed()).unwrap();
+                assert_eq!(transition.error, ErrorDisposition::Replace);
+                if stage != Operation::Restore {
+                    assert_eq!(transition.operation, Operation::Restore);
+                    machine.advance(Outcome::Success, observed()).unwrap();
+                }
+                assert_eq!(machine.operation(), Operation::Complete(outcome));
+            }
+        }
+    }
+
+    #[test]
+    fn logger_availability_internal_calls_and_fallbacks_control_logging_only() {
+        let mut failed_setup = machine(true);
+        assert_eq!(
+            failed_setup
+                .advance(Outcome::Failure, Observations::default())
+                .unwrap()
+                .operation,
+            Operation::Restore
+        );
+        for internal_call in [false, true] {
+            for has_fallbacks in [false, true] {
+                let NativeOutcome::Completed(mut machine) = Lifecycle::new(
+                    &request(),
+                    Options {
+                        asynchronous: true,
+                        internal_call,
+                        ..Options::default()
+                    },
+                )
+                .unwrap() else {
+                    panic!("expected admission")
+                };
+                reach(&mut machine, Operation::DeploymentSuccess);
+                let next = machine
+                    .advance(
+                        Outcome::Success,
+                        Observations {
+                            has_fallbacks,
+                            ..observed()
+                        },
+                    )
+                    .unwrap();
+                assert_eq!(
+                    next.operation,
+                    if internal_call || has_fallbacks {
+                        Operation::SyncSuccessIfNeeded
+                    } else {
+                        Operation::AsyncSuccess
+                    }
+                );
+            }
+        }
+        let NativeOutcome::Completed(mut internal) = Lifecycle::new(
+            &request(),
+            Options {
+                asynchronous: true,
+                internal_call: true,
+                ..Options::default()
+            },
+        )
+        .unwrap() else {
+            panic!("expected admission")
+        };
+        reach(&mut internal, Operation::Prepare);
+        assert_eq!(
+            internal
+                .advance(Outcome::Failure, observed())
+                .unwrap()
+                .operation,
+            Operation::DeploymentFailure
+        );
+        assert_eq!(
+            internal
+                .advance(Outcome::Abort, observed())
+                .unwrap()
+                .operation,
+            Operation::Restore
+        );
+    }
+
+    #[test]
+    fn decline_is_admission_only_and_file_is_an_inert_descriptor() {
+        for request in [
+            OcrRequest {
+                document: OcrDocument::File,
+                ..request()
+            },
+            OcrRequest {
+                document: OcrDocument::Unsupported,
+                ..request()
+            },
+            OcrRequest {
+                stream: true,
+                ..request()
+            },
+            OcrRequest {
+                request_format: Some("native".into()),
+                ..request()
+            },
+            OcrRequest {
+                model: "openai/model".into(),
+                ..request()
+            },
+        ] {
+            assert!(matches!(
+                Lifecycle::new(&request, Options::default()),
+                Ok(NativeOutcome::Declined(_))
+            ));
+        }
+        assert!(matches!(
+            Lifecycle::new(
+                &request(),
+                Options {
+                    credential_method: CredentialMethod::Acquisition,
+                    ..Options::default()
+                }
+            ),
+            Ok(NativeOutcome::Declined(_))
+        ));
+        assert!(matches!(
+            Lifecycle::new(
+                &OcrRequest {
+                    timeout_seconds: f64::NAN,
+                    ..request()
+                },
+                Options::default()
+            ),
+            Err(Error::InvalidRequest(_))
+        ));
+        let mut machine = machine(true);
+        reach(&mut machine, Operation::Prepare);
+        assert!(matches!(
+            prepare::prepare(OcrRequest {
+                document: OcrDocument::File,
+                ..request()
+            }),
+            Err(Error::Unsupported(_))
+        ));
+        assert_eq!(
+            machine
+                .advance(Outcome::Failure, observed())
+                .unwrap()
+                .operation,
+            Operation::DeploymentFailure
+        );
+        reach(&mut machine, Operation::Restore);
+        assert_eq!(
+            machine
+                .advance(Outcome::Success, observed())
+                .unwrap()
+                .operation,
+            Operation::Complete(Outcome::Failure)
+        );
+    }
+
+    #[test]
+    fn identity_keeps_supplied_provenance_across_hook_replacement_and_sdk_attempts() {
+        for _ in 0..2 {
+            let NativeOutcome::Completed(mut machine) = Lifecycle::new(
+                &request(),
+                Options {
+                    asynchronous: true,
+                    call_id: Some("logical-call".into()),
+                    trace_id: Some("trace".into()),
+                    ..Options::default()
+                },
+            )
+            .unwrap() else {
+                panic!("expected admission")
+            };
+            reach(&mut machine, Operation::Prepare);
+            let prepared = prepare::prepare(OcrRequest {
+                model: "mistral/replacement".into(),
+                ..request()
+            })
+            .unwrap();
+            assert_eq!(prepared.model, "replacement");
+            assert_eq!(
+                machine.identity(),
+                &Identity {
+                    requested_model: "mistral/requested-model".into(),
+                    call_id: "logical-call".into(),
+                    trace_id: Some("trace".into()),
+                    generated_call_id: false,
+                }
+            );
+            reach(&mut machine, Operation::Restore);
+            assert_eq!(machine.identity().call_id, "logical-call");
+        }
+        let first = machine(false);
+        let second = machine(false);
+        assert!(first.identity().generated_call_id);
+        assert_eq!(first.identity().call_id.len(), 36);
+        assert_eq!(&first.identity().call_id[14..15], "4");
+        assert_ne!(first.identity().call_id, second.identity().call_id);
+        assert_eq!(first.identity().trace_id, None);
+    }
+}

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

@@ -1,3 +1,4 @@
+pub mod lifecycle;
 pub mod prepare;
 pub mod transformation;
 pub mod types;

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

@@ -6,13 +6,15 @@ use crate::Error;
 use crate::providers::azure_ai::ocr::transformation as azure_ai;
 use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG;
 use crate::providers::vertex_ai::ocr::transformation as vertex_ai;
-use crate::routing_utils::provider::get_custom_llm_provider;
+use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider};
 
 use super::transformation::{OcrProviderConfig, OcrResponseHandling};
 use super::types::OcrRequest;
 pub use super::types::PreparedOcr;
 
-pub fn prepare(request: OcrRequest) -> Result<PreparedOcr, Error> {
+fn request_config(
+    request: &OcrRequest,
+) -> Result<(CustomLlmProvider<'_>, &'static dyn OcrProviderConfig), Error> {
     match request.request_format.as_deref() {
         None | Some("litellm") => {}
         Some("native") => return Err(Error::Unsupported("native OCR request format")),
@@ -33,10 +35,58 @@ pub fn prepare(request: OcrRequest) -> Result<PreparedOcr, Error> {
         .filter(|timeout| !timeout.is_zero())
         .ok_or_else(|| Error::InvalidRequest("timeout must be positive and finite".into()))?;
 
+    Ok((provider, config))
+}
+
+pub(super) fn admission_capabilities(request: &OcrRequest) -> Result<(), Error> {
+    check_admission_capabilities(request, &|key| std::env::var(key).ok())
+}
+
+fn check_admission_capabilities(
+    request: &OcrRequest,
+    env_lookup: &dyn Fn(&str) -> Option<String>,
+) -> Result<(), Error> {
+    let (provider, config) = request_config(request)?;
+    validate_capabilities(config)?;
+    request
+        .document
+        .validate(config.requires_data_uri_document())?;
+    if request.stream {
+        return Err(Error::Unsupported("OCR streaming response handling"));
+    }
+    if let Some(operation) = config.credential_acquisition_operation() {
+        let supplied = request
+            .api_key
+            .as_deref()
+            .is_some_and(|key| !key.trim().is_empty())
+            || crate::http_utils::has_header(&request.extra_headers, "authorization");
+        let configured = match provider.custom_llm_provider {
+            "azure_ai" => {
+                crate::http_utils::has_header(&request.extra_headers, "api-key")
+                    || request
+                        .azure_ad_token
+                        .as_deref()
+                        .is_some_and(|key| !key.trim().is_empty())
+                    || env_lookup("AZURE_AI_API_KEY").is_some_and(|key| !key.trim().is_empty())
+            }
+            "vertex_ai" => ["VERTEX_AI_API_KEY", "VERTEXAI_API_KEY"]
+                .into_iter()
+                .any(|name| env_lookup(name).is_some_and(|key| !key.trim().is_empty())),
+            _ => false,
+        };
+        if !supplied && !configured {
+            return Err(Error::Unsupported(operation));
+        }
+    }
+    Ok(())
+}
+
+pub fn prepare(request: OcrRequest) -> Result<PreparedOcr, Error> {
+    let (provider, config) = request_config(&request)?;
     let env_lookup = |key: &str| std::env::var(key).ok();
     let headers = config
         .validate_credentials(
-            request.extra_headers,
+            request.extra_headers.clone(),
             request.api_key.as_deref(),
             request.azure_ad_token.as_deref(),
             &env_lookup,
@@ -55,11 +105,11 @@ pub fn prepare(request: OcrRequest) -> Result<PreparedOcr, Error> {
         return Err(Error::Unsupported("OCR streaming response handling"));
     }
     let url_params = [
-        ("vertex_project", request.vertex_project),
-        ("vertex_location", request.vertex_location),
+        ("vertex_project", request.vertex_project.as_ref()),
+        ("vertex_location", request.vertex_location.as_ref()),
     ]
     .into_iter()
-    .filter_map(|(name, value)| value.map(|value| (name.into(), Value::String(value))))
+    .filter_map(|(name, value)| value.map(|value| (name.into(), Value::String(value.clone()))))
     .collect();
     let url = config.complete_url(
         request.api_base.as_deref(),
@@ -72,7 +122,7 @@ pub fn prepare(request: OcrRequest) -> Result<PreparedOcr, Error> {
     if !matches!(parsed_url.scheme(), "http" | "https") || parsed_url.host_str().is_none() {
         return Err(Error::InvalidRequest("invalid OCR API URL".into()));
     }
-    let document = serde_json::to_value(request.document)
+    let document = serde_json::to_value(&request.document)
         .map_err(|_| Error::InvalidRequest("could not project OCR document".into()))?;
     let template = config.transform_ocr_request(provider.model, document, Map::new())?;
     if template.files.is_some() {
@@ -113,3 +163,133 @@ pub(super) fn provider_config(
         _ => Err(Error::Unsupported("OCR provider")),
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::ocr::types::OcrDocument;
+
+    fn request() -> OcrRequest {
+        OcrRequest {
+            model: "mistral/mistral-ocr-latest".into(),
+            custom_llm_provider: None,
+            api_key: Some("test-key".into()),
+            api_base: Some("not a URL".into()),
+            extra_headers: vec![],
+            timeout_seconds: 2.0,
+            request_format: None,
+            document: OcrDocument::DocumentUrl {
+                document_url: "https://example.test/document.pdf".into(),
+            },
+            azure_ad_token: None,
+            vertex_project: None,
+            vertex_location: None,
+            stream: false,
+        }
+    }
+
+    #[rstest::rstest]
+    #[case("openai/model", Some("native"), "native OCR request format")]
+    #[case("openai/model", None, "OCR provider")]
+    #[case(
+        "azure_ai/doc-intelligence/prebuilt-read",
+        None,
+        "Azure Document Intelligence OCR polling"
+    )]
+    #[case(
+        "vertex_ai/mistral-ocr-latest",
+        None,
+        "OCR HTTP document URL to data URI conversion"
+    )]
+    #[case("mistral/mistral-ocr-latest", None, "OCR streaming response handling")]
+    fn admission_preserves_unsupported_error_order(
+        #[case] model: &str,
+        #[case] format: Option<&str>,
+        #[case] expected: &str,
+    ) {
+        let request = OcrRequest {
+            model: model.into(),
+            request_format: format.map(str::to_owned),
+            stream: true,
+            ..request()
+        };
+        assert!(
+            matches!(check_admission_capabilities(&request, &|_| None), Err(Error::Unsupported(message)) if message == expected)
+        );
+        assert!(
+            matches!(prepare(request), Err(Error::Unsupported(message)) if message == expected)
+        );
+    }
+
+    #[test]
+    fn admission_leaves_url_preparation_and_revalidation_until_prepare() {
+        let request = request();
+        assert!(check_admission_capabilities(&request, &|_| None).is_ok());
+        assert!(matches!(prepare(request), Err(Error::InvalidRequest(_))));
+
+        let mut request = self::request();
+        assert!(check_admission_capabilities(&request, &|_| None).is_ok());
+        request.timeout_seconds = 0.0;
+        request.stream = true;
+        assert!(matches!(
+            check_admission_capabilities(&request, &|_| None),
+            Err(Error::InvalidRequest(_))
+        ));
+        assert!(matches!(prepare(request), Err(Error::InvalidRequest(_))));
+    }
+
+    #[test]
+    fn admission_does_not_validate_credentials_or_resolve_headers() {
+        let request = OcrRequest {
+            api_key: None,
+            ..request()
+        };
+        assert!(
+            check_admission_capabilities(&request, &|_| panic!(
+                "Mistral admission needs no key lookup"
+            ))
+            .is_ok()
+        );
+        assert!(
+            MISTRAL_OCR_CONFIG
+                .validate_credentials(vec![], None, None, &|_| None)
+                .is_err()
+        );
+    }
+
+    #[rstest::rstest]
+    #[case("azure_ai/model", "AZURE_AI_API_KEY")]
+    #[case("vertex_ai/model", "VERTEX_AI_API_KEY")]
+    #[case("vertex_ai/model", "VERTEXAI_API_KEY")]
+    fn admission_checks_credential_method_using_only_inert_configuration(
+        #[case] model: &str,
+        #[case] env_name: &str,
+    ) {
+        let request = OcrRequest {
+            model: model.into(),
+            api_key: None,
+            document: OcrDocument::ImageUrl {
+                image_url: "data:image/png;base64,AA==".into(),
+            },
+            ..request()
+        };
+        assert!(matches!(
+            check_admission_capabilities(&request, &|_| None),
+            Err(Error::Unsupported(_))
+        ));
+        assert!(matches!(
+            check_admission_capabilities(&request, &|_| Some(" ".into())),
+            Err(Error::Unsupported(_))
+        ));
+        assert!(
+            check_admission_capabilities(&request, &|name| (name == env_name)
+                .then(|| "configured".into()))
+            .is_ok()
+        );
+        let request = OcrRequest {
+            extra_headers: vec![("AUTHORIZATION".into(), "inert".into())],
+            ..request
+        };
+        assert!(check_admission_capabilities(&request, &|_| None).is_ok());
+    }
+}

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

@@ -3,14 +3,19 @@
 //! callbacks; no Python preparation, auth, encoding, or provider transforms run.
 
 use litellm_core::error::Error;
+use litellm_core::ocr::lifecycle::{
+    ErrorDisposition, Lifecycle, NativeOutcome, Observations, Operation, Options, Outcome,
+};
 use litellm_core::ocr::types::{OcrDocumentProjection, OcrRequest, PreparedOcr};
 use litellm_core::routing_utils::provider::get_custom_llm_provider;
-use litellm_python_interop::{Pythonized, from_py, to_py};
+use litellm_python_interop::{
+    InvocationMode, InvocationOutcome, PreparedCall, Pythonized, from_py, to_py,
+};
 use pyo3::exceptions::{PyNotImplementedError, PyRuntimeError, PyValueError};
 use pyo3::prelude::*;
 use pyo3::pyclass::{PyTraverseError, PyVisit};
 use pyo3::sync::PyOnceLock;
-use pyo3::types::PyDict;
+use pyo3::types::{PyDict, PyTuple};
 use serde_json::Value;
 
 use crate::errors::core_error_to_pyerr;
@@ -50,7 +55,7 @@ impl OcrState {
 
 fn ocr_error_to_pyerr(py: Python<'_>, error: Error, model: &str, provider: &str) -> PyErr {
     let status = match error {
-        Error::Unsupported(message) => return PyNotImplementedError::new_err(message),
+        Error::Unsupported(message) => return PyRuntimeError::new_err(message),
         Error::Auth(_) => 401,
         Error::Http { status, .. } => status,
         Error::Network(_) | Error::Connect(_) => {
@@ -120,24 +125,11 @@ fn header_pairs(headers: &Bound<'_, PyDict>) -> PyResult<Vec<(String, String)>>
         .collect()
 }
 
-#[pyfunction]
-#[pyo3(signature = (arguments, asynchronous=false))]
-fn prepare(py: Python<'_>, arguments: Py<PyDict>, asynchronous: bool) -> PyResult<Py<OcrState>> {
-    let bag = arguments.bind(py);
+fn decode_request(py: Python<'_>, bag: &Bound<'_, PyDict>) -> PyResult<OcrRequest> {
     let document = bag
         .get_item("document")?
         .ok_or_else(|| PyValueError::new_err("OCR requires document"))?
         .cast_into::<PyDict>()?;
-    if scalar(&document, "type")?.as_deref() == Some("file") {
-        return Err(ocr_error_to_pyerr(
-            py,
-            Error::Unsupported(
-                "Native OCR does not support file documents; pass a document_url or image_url dict",
-            ),
-            "",
-            "",
-        ));
-    }
     let timeout_seconds = match bag.get_item("timeout")?.filter(|value| !value.is_none()) {
         None => py
             .import("litellm.constants")?
@@ -163,9 +155,9 @@ fn prepare(py: Python<'_>, arguments: Py<PyDict>, asynchronous: bool) -> PyResul
             document_input.set_item(name, value)?;
         }
     }
-    let prepared = litellm_core::ocr::prepare::prepare(OcrRequest {
-        model: model.clone(),
-        custom_llm_provider: custom_llm_provider.clone(),
+    Ok(OcrRequest {
+        model,
+        custom_llm_provider,
         api_key: scalar(bag, "api_key")?,
         api_base: scalar(bag, "api_base")?,
         extra_headers,
@@ -182,20 +174,179 @@ fn prepare(py: Python<'_>, arguments: Py<PyDict>, asynchronous: bool) -> PyResul
             .transpose()?
             .unwrap_or(false),
     })
-    .map_err(|error| {
-        let resolved = get_custom_llm_provider(&model, custom_llm_provider.as_deref());
-        ocr_error_to_pyerr(
-            py,
-            error,
-            resolved
-                .as_ref()
-                .map_or(model.as_str(), |value| value.model),
-            resolved
-                .as_ref()
-                .map_or("", |value| value.custom_llm_provider),
-        )
-    })?;
+}
 
+fn request_error_to_pyerr(
+    py: Python<'_>,
+    error: Error,
+    model: &str,
+    custom_llm_provider: Option<&str>,
+) -> PyErr {
+    let resolved = get_custom_llm_provider(model, custom_llm_provider);
+    ocr_error_to_pyerr(
+        py,
+        error,
+        resolved.as_ref().map_or(model, |value| value.model),
+        resolved
+            .as_ref()
+            .map_or("", |value| value.custom_llm_provider),
+    )
+}
+
+#[pyclass]
+struct OcrLifecycle {
+    machine: Lifecycle,
+    asynchronous: bool,
+}
+
+#[pymethods]
+impl OcrLifecycle {
+    #[new]
+    fn new(
+        py: Python<'_>,
+        arguments: &Bound<'_, PyDict>,
+        asynchronous: bool,
+        internal_call: bool,
+    ) -> PyResult<Self> {
+        let request = decode_request(py, arguments)?;
+        let logger = arguments
+            .get_item("litellm_logging_obj")?
+            .filter(|value| !value.is_none());
+        let identity = |name: &str| -> PyResult<Option<String>> {
+            if let Some(logger) = &logger {
+                match logger.getattr(name) {
+                    Ok(value) => {
+                        if let Ok(value) = value.extract::<String>() {
+                            return Ok(Some(value));
+                        }
+                    }
+                    Err(error)
+                        if !error.is_instance_of::<pyo3::exceptions::PyAttributeError>(py) =>
+                    {
+                        return Err(error);
+                    }
+                    _ => {}
+                }
+            }
+            scalar(arguments, name)
+        };
+        let machine = Lifecycle::new(
+            &request,
+            Options {
+                asynchronous,
+                internal_call,
+                call_id: identity("litellm_call_id")?,
+                trace_id: identity("litellm_trace_id")?,
+                ..Options::default()
+            },
+        )
+        .map_err(|error| {
+            request_error_to_pyerr(
+                py,
+                error,
+                &request.model,
+                request.custom_llm_provider.as_deref(),
+            )
+        })?;
+        match machine {
+            NativeOutcome::Completed(machine) => Ok(Self {
+                machine,
+                asynchronous,
+            }),
+            NativeOutcome::Declined(decline) => {
+                Err(PyNotImplementedError::new_err(decline.reason()))
+            }
+        }
+    }
+
+    fn identity(&self) -> (String, Option<String>) {
+        let identity = self.machine.identity();
+        (identity.call_id.clone(), identity.trace_id.clone())
+    }
+
+    fn advance(
+        &mut self,
+        outcome: u8,
+        logger_available: bool,
+        has_fallbacks: bool,
+    ) -> PyResult<bool> {
+        let outcome = match outcome {
+            0 => Outcome::Success,
+            1 => Outcome::Failure,
+            _ => Outcome::Abort,
+        };
+        self.machine
+            .advance(
+                outcome,
+                Observations {
+                    logger_available,
+                    has_fallbacks,
+                },
+            )
+            .map(|transition| transition.error == ErrorDisposition::Replace)
+            .map_err(core_error_to_pyerr)
+    }
+
+    fn complete(&self) -> Option<bool> {
+        match self.machine.operation() {
+            Operation::Complete(outcome) => Some(outcome == Outcome::Success),
+            _ => None,
+        }
+    }
+}
+
+#[pyfunction]
+fn invoke(
+    py: Python<'_>,
+    machine: Py<OcrLifecycle>,
+    host: Py<PyAny>,
+) -> PyResult<(bool, Py<PyAny>)> {
+    let (operation, asynchronous) = {
+        let machine = machine.borrow(py);
+        (machine.machine.operation(), machine.asynchronous)
+    };
+    let (method, mode) = match operation {
+        Operation::Setup => ("setup", InvocationMode::Direct),
+        Operation::DeploymentPre => ("deployment_pre", InvocationMode::Await),
+        Operation::Prepare => ("prepare", InvocationMode::Direct),
+        Operation::Send if asynchronous => ("send", InvocationMode::Await),
+        Operation::Send => ("send_sync", InvocationMode::Direct),
+        Operation::DeploymentSuccess => ("deployment_success", InvocationMode::Await),
+        Operation::DeploymentFailure => ("deployment_failure", InvocationMode::Await),
+        Operation::SyncSuccess => ("sync_success", InvocationMode::Direct),
+        Operation::AsyncSuccess => ("async_success", InvocationMode::Direct),
+        Operation::SyncSuccessIfNeeded => ("sync_success_if_needed", InvocationMode::Direct),
+        Operation::SyncFailure => ("sync_failure", InvocationMode::Direct),
+        Operation::AsyncFailure => ("async_failure", InvocationMode::Await),
+        Operation::Restore => ("restore", InvocationMode::Direct),
+        Operation::Complete(_) => return Err(PyRuntimeError::new_err("OCR lifecycle is complete")),
+    };
+    let call = PreparedCall::new(
+        mode,
+        host.getattr(py, method)?,
+        PyTuple::empty(py).unbind(),
+        None,
+    );
+    match call.invoke(py)? {
+        InvocationOutcome::Returned(value) => Ok((false, value)),
+        InvocationOutcome::Awaitable(value) => Ok((true, value)),
+    }
+}
+
+#[pyfunction]
+#[pyo3(signature = (arguments, asynchronous=false))]
+fn prepare(py: Python<'_>, arguments: Py<PyDict>, asynchronous: bool) -> PyResult<Py<OcrState>> {
+    let bag = arguments.bind(py);
+    let request = decode_request(py, bag)?;
+    let model = request.model.clone();
+    let custom_llm_provider = request.custom_llm_provider.clone();
+    let prepared = litellm_core::ocr::prepare::prepare(request).map_err(|error| {
+        request_error_to_pyerr(py, error, &model, custom_llm_provider.as_deref())
+    })?;
+    let document = bag
+        .get_item("document")?
+        .ok_or_else(|| PyValueError::new_err("OCR requires document"))?
+        .cast_into::<PyDict>()?;
     let body = to_py(py, &prepared.body)?
         .into_bound(py)
         .cast_into::<PyDict>()?;
@@ -355,64 +506,128 @@ fn driver(py: Python<'_>) -> PyResult<&Bound<'_, PyModule>> {
     let module = PyModule::from_code(
         py,
         c"from datetime import datetime
-from litellm.rust_bridge.ocr import invoke_terminal
+from litellm import utils
+from litellm.types.utils import CallTypes
+from litellm.rust_bridge.ocr import initialize_logging, invoke_terminal
+
+class Host:
+    def __init__(self, arguments, asynchronous):
+        self.machine = _Lifecycle(arguments, asynchronous, utils.is_internal_call.get())
+        self.arguments = arguments
+        self.current = arguments
+        self.asynchronous = asynchronous
+        self.logger = arguments.get('litellm_logging_obj')
+        self.state = None
+        self.response = None
+        self.error = None
+        self.start = datetime.now()
+        self.end = None
+
+    def setup(self):
+        call_id, trace_id = self.machine.identity()
+        self.arguments['litellm_call_id'] = call_id
+        self.arguments['litellm_trace_id'] = trace_id
+        self.logger = initialize_logging(self.arguments, self.asynchronous)
+        self.arguments['litellm_logging_obj'] = self.logger
+
+    async def deployment_pre(self):
+        modified = await utils.async_pre_call_deployment_hook(self.current, 'aocr')
+        if modified is not None:
+            self.current = modified
+        self.current['litellm_logging_obj'] = self.logger
+        call_id, trace_id = self.machine.identity()
+        self.current['litellm_call_id'] = call_id
+        self.current['litellm_trace_id'] = trace_id
+
+    def prepare(self):
+        self.state = _prepare(self.current, self.asynchronous)
+
+    def send_sync(self):
+        self.response = _send_sync(self.state)
+        self.end = datetime.now()
+
+    async def send(self):
+        self.response = _finish(await _send(self.state))
+        self.end = datetime.now()
+
+    async def deployment_success(self):
+        self.response = await utils.async_post_call_success_deployment_hook(self.current, self.response, CallTypes.aocr)
+
+    async def deployment_failure(self):
+        await utils.async_post_call_failure_deployment_hook(self.current, self.error, 'aocr')
+
+    def terminal(self, action, value):
+        return invoke_terminal(action, (self.arguments, self.current, self.state), self.logger, value, self.start, self.end)
+
+    def sync_success(self):
+        return self.terminal('sync_success', self.response)
+
+    def async_success(self):
+        return self.terminal('async_success', self.response)
+
+    def sync_success_if_needed(self):
+        return self.terminal('sync_success_if_needed', self.response)
+
+    def sync_failure(self):
+        return self.terminal('sync_failure', self.error)
+
+    def async_failure(self):
+        return self.terminal('async_failure', self.error)
+
+    def restore(self):
+        utils._restore_correlation_context_if_supported(self.logger)
+
+    def advance(self, outcome, error=None):
+        if error is not None and self.end is None:
+            self.end = datetime.now()
+        if self.logger is None:
+            self.logger = self.arguments.get('litellm_logging_obj')
+        replace = self.machine.advance(outcome, self.logger is not None, self.current.get('fallbacks') is not None)
+        if replace:
+            self.error = error
+
+    def result(self):
+        if self.machine.complete():
+            return self.response
+        raise self.error
 
 def drive_sync(arguments):
-    start = datetime.now()
-    state = None
-    try:
-        state = _prepare(arguments, False)
-        response = _send_sync(state)
-    except Exception as error:
-        logger = arguments.get('litellm_logging_obj')
-        if state is not None or (logger is not None and not isinstance(error, NotImplementedError)):
-            end = datetime.now()
-            for action in _sync_failure:
-                invoke_terminal(action, (arguments, state), logger, error, start, end)
-        raise
-    end = datetime.now()
-    for action in _sync_success:
-        invoke_terminal(action, (arguments, state), arguments['litellm_logging_obj'], response, start, end)
-    return response
+    host = Host(arguments, False)
+    while host.machine.complete() is None:
+        try:
+            _invoke(host.machine, host)
+        except Exception as error:
+            host.advance(1, error)
+        except BaseException as error:
+            host.advance(2, error)
+        else:
+            host.advance(0)
+    return host.result()
 
 async def drive(arguments):
-    start = datetime.now()
-    state = None
-    try:
-        state = _prepare(arguments, True)
-        response = _finish(await _send(state))
-    except Exception as error:
-        logger = arguments.get('litellm_logging_obj')
-        if state is not None or (logger is not None and not isinstance(error, NotImplementedError)):
-            end = datetime.now()
-            for action in _async_failure:
-                pending = invoke_terminal(action, (arguments, state), logger, error, start, end)
-                if pending is not None:
-                    await pending
-        raise
-    end = datetime.now()
-    for action in _async_success:
-        invoke_terminal(action, (arguments, state), arguments['litellm_logging_obj'], response, start, end)
-    return response
+    host = Host(arguments, True)
+    while host.machine.complete() is None:
+        try:
+            awaiting, value = _invoke(host.machine, host)
+            if awaiting:
+                await value
+        except Exception as error:
+            host.advance(1, error)
+        except BaseException as error:
+            host.advance(2, error)
+        else:
+            host.advance(0)
+    return host.result()
 ",
         c"ocr_driver.py",
         c"_ocr_driver",
     )?;
+    module.add("_Lifecycle", py.get_type::<OcrLifecycle>())?;
+    module.add("_invoke", wrap_pyfunction!(invoke, &module)?)?;
     module.add("_prepare", wrap_pyfunction!(prepare, &module)?)?;
     module.add("_send", wrap_pyfunction!(send, &module)?)?;
     module.add("_send_sync", wrap_pyfunction!(send_sync, &module)?)?;
     module.add("_finish", wrap_pyfunction!(finish, &module)?)?;
-    for (name, asynchronous, success) in [
-        ("_sync_success", false, true),
-        ("_async_success", true, true),
-        ("_sync_failure", false, false),
-        ("_async_failure", true, false),
-    ] {
-        module.add(
-            name,
-            litellm_core::ocr::terminal_callbacks(asynchronous, success).to_vec(),
-        )?;
-    }
     Ok(DRIVER.get_or_init(py, || module.unbind()).bind(py))
 }
 
@@ -519,6 +734,88 @@ mod tests {
         });
     }
 
+    #[test]
+    fn callback_decline_is_terminal_and_identity_is_reused() {
+        Python::initialize();
+        Python::attach(|py| {
+            let module = PyModule::new(py, "ocr_test").unwrap();
+            module
+                .add_function(wrap_pyfunction!(ocr, &module).unwrap())
+                .unwrap();
+            module
+                .add_function(wrap_pyfunction!(aocr, &module).unwrap())
+                .unwrap();
+            let globals = PyDict::new(py);
+            globals.set_item("native", module).unwrap();
+            py.run(
+                c"
+import asyncio
+import contextvars
+import threading
+from datetime import datetime
+
+marker = contextvars.ContextVar('terminal_marker')
+
+class Logger:
+    litellm_call_id = 'supplied-call'
+    litellm_trace_id = 'supplied-trace'
+
+    def update_from_kwargs(self, **values):
+        assert values['kwargs']['litellm_call_id'] == self.litellm_call_id
+        assert values['kwargs']['litellm_trace_id'] == self.litellm_trace_id
+        assert threading.get_ident() == self.thread
+        marker.set('update')
+        raise self.original
+
+    def failure_handler(self, error, trace, start, end):
+        assert error is self.original
+        assert marker.get() == 'update'
+        assert start <= end <= datetime.now()
+        self.end = end
+        self.calls.append('failure')
+
+    async def async_failure_handler(self, error, trace, start, end):
+        await asyncio.sleep(0)
+        assert asyncio.current_task() is self.task
+        assert marker.get() == 'update'
+        assert error is self.original
+        assert end is self.end
+        self.calls.append('async_failure')
+
+    def _restore_correlation_context(self):
+        self.calls.append('restore')
+
+async def exercise():
+    for asynchronous in (False, True):
+        logger = Logger()
+        logger.thread = threading.get_ident()
+        logger.task = asyncio.current_task()
+        logger.calls = []
+        logger.original = NotImplementedError('callback declined, not admission')
+        arguments = dict(model='mistral/mistral-ocr-latest', api_key='test-key', timeout=1.0,
+                         document={'type': 'document_url', 'document_url': 'https://example.test/doc.pdf'},
+                         litellm_logging_obj=logger)
+        try:
+            if asynchronous:
+                await native.aocr(arguments)
+            else:
+                native.ocr(arguments)
+        except NotImplementedError as error:
+            assert error is logger.original
+        else:
+            raise AssertionError('callback exception was lost')
+        assert logger.calls == (['failure', 'async_failure', 'restore'] if asynchronous else ['failure', 'restore'])
+        assert arguments['litellm_call_id'] == 'supplied-call'
+        assert arguments['litellm_trace_id'] == 'supplied-trace'
+
+asyncio.run(exercise())
+",
+                Some(&globals),
+                Some(&globals),
+            ).unwrap();
+        });
+    }
+
     #[test]
     fn native_send_owns_state_without_the_python_driver() {
         Python::initialize();
@@ -733,7 +1030,8 @@ class Logger:
         assert asyncio.current_task() is caller
         assert threading.get_ident() == caller_thread
         assert marker.get() == 'caller'
-        assert values['kwargs'] is arguments
+        assert values['kwargs'] is not arguments
+        assert values['kwargs']['opaque'] is arguments['opaque']
         self.calls.append('update')
         marker.set('updated')
 

litellm-rust/crates/python-interop/tests/callback_controls.rs

@@ -0,0 +1,227 @@
+use std::ffi::CStr;
+
+use litellm_python_interop::{InvocationMode, InvocationOutcome, PreparedCall};
+use pyo3::prelude::*;
+use pyo3::types::{PyDict, PyTuple};
+use rstest::rstest;
+use serial_test::serial;
+
+#[path = "support/mod.rs"]
+mod support;
+
+use support::python::{InitializedPython, initialized_python, item, scope};
+
+enum ControlCall {
+    Reference(Py<PyAny>),
+    Prepared(PreparedCall),
+}
+
+impl ControlCall {
+    fn new(
+        py: Python<'_>,
+        callback: Bound<'_, PyAny>,
+        args: Bound<'_, PyTuple>,
+        kwargs: Option<Bound<'_, PyDict>>,
+        retained: bool,
+        mode: InvocationMode,
+    ) -> PyResult<Self> {
+        if retained {
+            return Ok(Self::Prepared(PreparedCall::new(
+                mode,
+                callback.unbind(),
+                args.unbind(),
+                kwargs.map(Bound::unbind),
+            )));
+        }
+        let factory = py
+            .import("callback_lifecycle")?
+            .getattr("ReferenceFactory")?
+            .call0()?;
+        Ok(Self::Reference(
+            factory
+                .call_method1(
+                    "prepare",
+                    (callback, args, kwargs, mode == InvocationMode::Await),
+                )?
+                .unbind(),
+        ))
+    }
+
+    fn invoke(&self, py: Python<'_>, mode: InvocationMode) -> PyResult<Py<PyAny>> {
+        match self {
+            Self::Reference(owner) => owner.call_method0(py, "invoke"),
+            Self::Prepared(call) => match call.invoke(py)? {
+                InvocationOutcome::Returned(value) => {
+                    assert_eq!(mode, InvocationMode::Direct);
+                    Ok(value)
+                }
+                InvocationOutcome::Awaitable(value) => {
+                    assert_eq!(mode, InvocationMode::Await);
+                    Ok(value)
+                }
+            },
+        }
+    }
+}
+
+#[rstest]
+#[case::original_graph(c"(positional, keywords)", (true, true, true), (1, 1, 1))]
+#[case::rebuilt_envelopes(c"(tuple(value for value in positional), dict(keywords))", (true, true, true), (1, 1, 1))]
+#[case::shallow_payload(c"((copy.copy(positional[0]),), keywords)", (false, true, true), (0, 1, 1))]
+#[case::copied_graph_keeps_cross_argument_alias(c"copy.deepcopy((positional, keywords))", (false, false, true), (0, 0, 0))]
+#[case::separate_copies_break_cross_argument_alias(c"(copy.deepcopy(positional), copy.deepcopy(keywords))", (false, false, false), (0, 0, 0))]
+#[serial(python_interpreter)]
+fn argument_copy_boundaries_determine_identity_and_mutation_visibility(
+    initialized_python: &InitializedPython,
+    #[case] transform: &CStr,
+    #[case] expected_identity: (bool, bool, bool),
+    #[case] live_stages: (u8, u8, u8),
+    #[values(false, true)] retained: bool,
+    #[values(InvocationMode::Direct, InvocationMode::Await)] mode: InvocationMode,
+) -> PyResult<()> {
+    initialized_python.attach(|py| {
+        let globals = scope(
+            py,
+            c"
+import copy
+from types import SimpleNamespace
+
+nested = SimpleNamespace(stage=0)
+original = SimpleNamespace(nested=nested, stage=0)
+positional, keywords = (original,), {'alias': nested}
+
+def observe(value, *, alias):
+    return value, alias, (value.stage, value.nested.stage, alias.stage)
+
+async def observe_async(value, *, alias):
+    return observe(value, alias=alias)
+",
+        )?;
+        let transformed = py.eval(transform, Some(&globals), None)?;
+        let call = ControlCall::new(
+            py,
+            item(
+                &globals,
+                if mode == InvocationMode::Await {
+                    "observe_async"
+                } else {
+                    "observe"
+                },
+            ),
+            transformed.get_item(0)?.cast_into::<PyTuple>()?,
+            Some(transformed.get_item(1)?.cast_into::<PyDict>()?),
+            retained,
+            mode,
+        )?;
+        let original = item(&globals, "original");
+        let nested = item(&globals, "nested");
+        original.setattr("stage", 1)?;
+        nested.setattr("stage", 1)?;
+        let pending = call.invoke(py, mode)?;
+        original.setattr("stage", 2)?;
+        nested.setattr("stage", 2)?;
+        let result = if mode == InvocationMode::Await {
+            let lifecycle = py.import("callback_lifecycle")?;
+            lifecycle.call_method1(
+                "run_checked",
+                (lifecycle.call_method0("ReferenceFactory")?, pending),
+            )?
+        } else {
+            pending.into_bound(py)
+        };
+        drop(call);
+        let value = result.get_item(0)?;
+        let alias = result.get_item(1)?;
+        assert_eq!(
+            (
+                value.is(&original),
+                value.getattr("nested")?.is(&nested),
+                value.getattr("nested")?.is(&alias)
+            ),
+            expected_identity
+        );
+        let stage = if mode == InvocationMode::Await { 2 } else { 1 };
+        assert_eq!(
+            result.get_item(2)?.extract::<(u8, u8, u8)>()?,
+            (
+                live_stages.0 * stage,
+                live_stages.1 * stage,
+                live_stages.2 * stage
+            )
+        );
+        Ok(())
+    })
+}
+
+#[rstest]
+#[case::original_result(None, (true, true))]
+#[case::passthrough_result(Some(c"lambda value: value"), (true, true))]
+#[case::shallow_result(Some(c"copy.copy"), (false, true))]
+#[case::deep_result(Some(c"copy.deepcopy"), (false, false))]
+#[serial(python_interpreter)]
+fn result_copy_boundaries_determine_root_and_nested_identity(
+    initialized_python: &InitializedPython,
+    #[case] transform: Option<&CStr>,
+    #[case] expected_identity: (bool, bool),
+    #[values(false, true)] retained: bool,
+    #[values(InvocationMode::Direct, InvocationMode::Await)] mode: InvocationMode,
+) -> PyResult<()> {
+    initialized_python.attach(|py| {
+        let globals = scope(
+            py,
+            c"
+import copy
+from types import SimpleNamespace
+
+original = SimpleNamespace(nested=SimpleNamespace(stage=0), stage=0)
+
+def callback():
+    return original
+
+async def callback_async():
+    return original
+",
+        )?;
+        let call = ControlCall::new(
+            py,
+            item(
+                &globals,
+                if mode == InvocationMode::Await {
+                    "callback_async"
+                } else {
+                    "callback"
+                },
+            ),
+            PyTuple::empty(py),
+            None,
+            retained,
+            mode,
+        )?;
+        let pending = call.invoke(py, mode)?;
+        let settled = if mode == InvocationMode::Await {
+            let lifecycle = py.import("callback_lifecycle")?;
+            lifecycle.call_method1(
+                "run_checked",
+                (lifecycle.call_method0("ReferenceFactory")?, pending),
+            )?
+        } else {
+            pending.into_bound(py)
+        };
+        let result = if let Some(transform) = transform {
+            py.eval(transform, Some(&globals), None)?
+                .call1((settled,))?
+        } else {
+            settled
+        };
+        drop(call);
+        let original = item(&globals, "original");
+        assert_eq!(
+            (
+                result.is(&original),
+                result.getattr("nested")?.is(original.getattr("nested")?)
+            ),
+            expected_identity
+        );
+        Ok(())
+    })
+}

litellm-rust/crates/python-interop/tests/callback_lifecycle.rs

@@ -1,8 +1,9 @@
 use std::process::Command;
 use std::time::{Duration, Instant};
 
+use litellm_python_interop::{InvocationMode, InvocationOutcome, PreparedCall};
 use pyo3::prelude::*;
-use pyo3::types::PyDict;
+use pyo3::types::{PyDict, PyTuple};
 use rstest::{fixture, rstest};
 use serial_test::{parallel, serial};
 
@@ -12,7 +13,7 @@ mod callback_owner;
 #[path = "support/mod.rs"]
 mod support;
 
-use support::python::{InitializedPython, initialized_python, run_fixture};
+use support::python::{InitializedPython, initialized_python, item, run_fixture};
 
 #[test]
 fn cold_awaited_adapter_initialization_allows_reentry() -> PyResult<()> {
@@ -100,8 +101,6 @@ fn scenario_scope(initialized_python: &InitializedPython) -> Py<PyDict> {
 #[case::stream_lifecycle("stream_lifecycle")]
 #[case::sync_stream_lifecycle("sync_stream_lifecycle")]
 #[case::repeated_ownership("repeated_ownership")]
-#[case::retained_field_replacement("retained_field_replacement")]
-#[case::queued_graph_ownership("queued_graph_ownership")]
 #[case::detached_work_after_error("detached_work_after_error")]
 #[serial(python_interpreter)]
 fn lifecycle_contract(
@@ -176,8 +175,8 @@ fn component_contract(
 #[case::real_logging_queue_copy_control("real_logging_queue_copy_control")]
 #[case::real_crowdstrike_translator_identity("real_crowdstrike_translator_identity")]
 #[case::real_rubrik_block_lifecycle("real_rubrik_block_lifecycle")]
-#[case::real_parallel_guardrail_snapshots("real_parallel_guardrail_snapshots")]
-#[case::real_purview_sync_background("real_purview_sync_background")]
+#[case::real_parallel_guardrail_sharing_and_exception_order("real_parallel_guardrail_snapshots")]
+#[case::real_purview_sync_background_and_active_loop("real_purview_sync_background")]
 #[ignore = "requires the repository Python environment and LiteLLM on PYTHONPATH"]
 #[serial(python_interpreter)]
 fn integration_contract(
@@ -220,20 +219,6 @@ fn run_scenario_fixture(
 }
 
 #[rstest]
-#[case::original_arguments("argument_identity", "identity")]
-#[case::envelope_arguments("argument_identity", "envelope")]
-#[case::shallow_arguments("argument_identity", "shallow_payload")]
-#[case::copied_graph("argument_identity", "deep_graph")]
-#[case::independent_copies("argument_identity", "deep_separate")]
-#[case::original_read_timing("mutation_timing", "identity")]
-#[case::envelope_read_timing("mutation_timing", "envelope")]
-#[case::shallow_read_timing("mutation_timing", "shallow_payload")]
-#[case::deep_read_timing("mutation_timing", "deep_graph")]
-#[case::independent_read_timing("mutation_timing", "deep_separate")]
-#[case::original_result("result_identity", "identity")]
-#[case::passthrough_result("result_identity", "result_passthrough")]
-#[case::shallow_result("result_identity", "result_shallow")]
-#[case::deep_result("result_identity", "result_deep")]
 #[case::retained_lifetime("deferred_lifetime", "identity")]
 #[case::prepared_ownership("deferred_lifetime", "missing_handoff")]
 #[case::externally_owned_retained("borrowed_lifetime", "identity")]
@@ -305,6 +290,153 @@ fn run_control_fixture(
     })
 }
 
+fn invoke_direct_callback(
+    py: Python<'_>,
+    globals: &Bound<'_, PyDict>,
+    callback: &str,
+    argument: &str,
+    retained: bool,
+) -> PyResult<Py<PyAny>> {
+    let args = PyTuple::new(py, [item(globals, argument)])?;
+    if !retained {
+        let factory = item(globals, "ReferenceFactory").call0()?;
+        let owner = factory.call_method1("prepare", (item(globals, callback), args))?;
+        let result = owner.call_method0("invoke");
+        owner.call_method0("close")?;
+        assert_eq!(factory.getattr("live")?.extract::<usize>()?, 0);
+        return result.map(Bound::unbind);
+    }
+    let call = PreparedCall::new(
+        InvocationMode::Direct,
+        item(globals, callback).unbind(),
+        args.unbind(),
+        None,
+    );
+    match call.invoke(py)? {
+        InvocationOutcome::Returned(value) => Ok(value),
+        InvocationOutcome::Awaitable(_) => panic!("direct callback produced an awaitable outcome"),
+    }
+}
+
+#[rstest]
+#[serial(python_interpreter)]
+fn retained_field_survives_replacement_and_observes_original_mutations(
+    scenario_scope: Py<PyDict>,
+    #[values(false, true)] retained: bool,
+) -> PyResult<()> {
+    Python::attach(|py| {
+        let globals = scenario_scope.bind(py);
+        py.run(
+            c"
+original = {'messages': [{'content': 'original'}]}
+replacement = {'messages': [{'content': 'replacement'}]}
+event = {'payload': original, 'alias': original}
+saved = []
+
+def retain(value):
+    saved.append(value['payload'])
+
+def replace(value):
+    value['payload'] = replacement
+    value['alias']['messages'][0]['content'] = 'mutated original'
+",
+            Some(globals),
+            None,
+        )?;
+        for callback in ["retain", "replace"] {
+            assert!(invoke_direct_callback(py, globals, callback, "event", retained)?.is_none(py));
+        }
+        let original = item(globals, "original");
+        let replacement = item(globals, "replacement");
+        let event = item(globals, "event");
+        let saved = item(globals, "saved").get_item(0)?;
+        assert!(saved.is(&original));
+        assert!(event.get_item("alias")?.is(&original));
+        assert!(event.get_item("payload")?.is(&replacement));
+        assert_eq!(
+            saved
+                .get_item("messages")?
+                .get_item(0)?
+                .get_item("content")?
+                .extract::<String>()?,
+            "mutated original"
+        );
+        replacement
+            .get_item("messages")?
+            .get_item(0)?
+            .set_item("content", "mutated replacement")?;
+        assert_eq!(
+            event
+                .get_item("payload")?
+                .get_item("messages")?
+                .get_item(0)?
+                .get_item("content")?
+                .extract::<String>()?,
+            "mutated replacement"
+        );
+        assert_eq!(
+            original
+                .get_item("messages")?
+                .get_item(0)?
+                .get_item("content")?
+                .extract::<String>()?,
+            "mutated original"
+        );
+        Ok(())
+    })
+}
+
+#[rstest]
+#[serial(python_interpreter)]
+fn queued_graph_outlives_invocation_and_stays_live_until_serialized(
+    scenario_scope: Py<PyDict>,
+    #[values(false, true)] retained: bool,
+) -> PyResult<()> {
+    Python::attach(|py| {
+        let globals = scenario_scope.bind(py);
+        py.run(
+            c"
+queue = asyncio.Queue()
+sentinel = Value()
+reference = weakref.ref(sentinel)
+payload = {'sentinel': sentinel, 'nested': {'status': 'queued'}}
+snapshot = json.dumps(payload['nested'])
+
+def enqueue(value):
+    queue.put_nowait(value)
+",
+            Some(globals),
+            None,
+        )?;
+        assert!(invoke_direct_callback(py, globals, "enqueue", "payload", retained)?.is_none(py));
+        py.run(c"del sentinel, payload\ngc.collect()", Some(globals), None)?;
+        let reference = item(globals, "reference");
+        assert!(!reference.call0()?.is_none());
+        let queue = item(globals, "queue");
+        let queued = queue.call_method0("get_nowait")?;
+        queued
+            .get_item("nested")?
+            .set_item("status", "changed before flush")?;
+        let json = item(globals, "json");
+        let flushed = json.call_method1(
+            "loads",
+            (json.call_method1("dumps", (queued.get_item("nested")?,))?,),
+        )?;
+        assert_eq!(
+            flushed.get_item("status")?.extract::<String>()?,
+            "changed before flush"
+        );
+        let snapshot = json.call_method1("loads", (item(globals, "snapshot"),))?;
+        assert_eq!(snapshot.get_item("status")?.extract::<String>()?, "queued");
+        assert!(queued.get_item("sentinel")?.is(reference.call0()?));
+        queue.call_method0("task_done")?;
+        drop(queued);
+        py.run(c"gc.collect()", Some(globals), None)?;
+        assert!(reference.call0()?.is_none());
+        Ok(())
+    })
+}
+
 #[rstest]
 #[parallel(python_interpreter)]
 fn detached_release(initialized_python: &InitializedPython) -> PyResult<()> {

litellm-rust/crates/python-interop/tests/callback_patterns.rs

@@ -0,0 +1,720 @@
+use std::ffi::CStr;
+
+use litellm_python_interop::{InvocationMode, InvocationOutcome, PreparedCall};
+use pyo3::prelude::*;
+use pyo3::types::{PyDict, PyTuple};
+use rstest::rstest;
+use serial_test::serial;
+
+#[path = "support/mod.rs"]
+mod support;
+
+use support::python::{InitializedPython, initialized_python, item, run_fixture, scope};
+
+#[derive(Clone, Copy, Debug)]
+enum Backend {
+    Python,
+    PreparedCall,
+}
+
+#[pyfunction]
+fn invoke_prepared<'py>(
+    py: Python<'py>,
+    callback: Py<PyAny>,
+    args: Py<PyTuple>,
+    awaited: bool,
+) -> PyResult<Bound<'py, PyAny>> {
+    let mode = if awaited {
+        InvocationMode::Await
+    } else {
+        InvocationMode::Direct
+    };
+    let call = PreparedCall::new(mode, callback, args, None);
+    let outcome = call.invoke(py);
+    drop(call);
+    match outcome? {
+        InvocationOutcome::Returned(value) => {
+            assert_eq!(mode, InvocationMode::Direct);
+            Ok(value.into_bound(py))
+        }
+        InvocationOutcome::Awaitable(value) => {
+            assert_eq!(mode, InvocationMode::Await);
+            Ok(value.into_bound(py))
+        }
+    }
+}
+
+fn pattern_scope<'py>(
+    py: Python<'py>,
+    backend: Backend,
+    mode: InvocationMode,
+) -> PyResult<Bound<'py, PyDict>> {
+    let globals = scope(
+        py,
+        c"
+import asyncio
+import copy
+import gc
+import json
+import threading
+import weakref
+from unittest import TestCase
+from callback_lifecycle import ReferenceFactory, run_checked
+
+class Value:
+    pass
+
+def invoke(callback, *args):
+    if retained:
+        return invoke_prepared(callback, args, awaited)
+    if not awaited:
+        return callback(*args)
+
+    async def await_callback():
+        return await callback(*args)
+
+    return await_callback()
+
+async def async_call(callback, *args):
+    async def async_callback(*values):
+        return callback(*values)
+
+    result = invoke(async_callback if awaited else callback, *args)
+    return await result if awaited else result
+
+def call(callback, *args):
+    if not awaited:
+        return invoke(callback, *args)
+    return run_async(async_call(callback, *args))
+
+def run_async(scenario):
+    return run_checked(ReferenceFactory(), scenario)
+",
+    )?;
+    globals.set_item("invoke_prepared", wrap_pyfunction!(invoke_prepared, py)?)?;
+    globals.set_item("retained", matches!(backend, Backend::PreparedCall))?;
+    globals.set_item("awaited", mode == InvocationMode::Await)?;
+    Ok(globals)
+}
+
+#[track_caller]
+fn run_pattern(
+    python: &InitializedPython,
+    backend: Backend,
+    mode: InvocationMode,
+    parameters: impl for<'py> FnOnce(&Bound<'py, PyDict>) -> PyResult<()>,
+    scenario: &CStr,
+) -> PyResult<()> {
+    let location = std::panic::Location::caller();
+    python.attach(|py| {
+        let globals = pattern_scope(py, backend, mode)?;
+        parameters(&globals)?;
+        run_fixture(
+            py,
+            &globals,
+            scenario.to_str().unwrap(),
+            &format!("{}:{}", location.file(), location.line()),
+        )?;
+        let result = py.run(
+            c"
+if asyncio.iscoroutinefunction(test_pattern):
+    run_async(test_pattern())
+else:
+    test_pattern()
+",
+            Some(&globals),
+            None,
+        );
+        if let Err(error) = &result {
+            error.print(py);
+        }
+        result
+    })
+}
+
+#[rstest]
+#[case::flush_after_edit_sees_mutation(false, false, false)]
+#[case::flush_before_edit_keeps_original(true, false, true)]
+#[case::copied_editor_leaves_queue_unchanged(false, true, true)]
+#[serial(python_interpreter)]
+fn reference_queue_observes_edits_until_serialization_and_owns_arguments(
+    initialized_python: &InitializedPython,
+    #[case] flush_before_edit: bool,
+    #[case] copy_before_edit: bool,
+    #[case] flushed_has_tools: bool,
+    #[values(Backend::Python, Backend::PreparedCall)] backend: Backend,
+    #[values(InvocationMode::Direct, InvocationMode::Await)] mode: InvocationMode,
+) -> PyResult<()> {
+    run_pattern(
+        initialized_python,
+        backend,
+        mode,
+        |globals| {
+            globals.set_item("flush_before_edit", flush_before_edit)?;
+            globals.set_item("copy_before_edit", copy_before_edit)?;
+            globals.set_item("flushed_has_tools", flushed_has_tools)
+        },
+        c"
+def test_pattern():
+    queue, snapshots, observations = [], [], []
+    payload = {'model_parameters': {'tools': ['lookup'], 'stream': True}}
+    response = Value()
+    response_ref = weakref.ref(response)
+    event = {'payload': payload}
+
+    def enqueue_live(kwargs, result):
+        queue.append((kwargs['payload'], kwargs, result))
+
+    def enqueue_snapshot(kwargs, result):
+        snapshots.append(json.dumps(kwargs['payload']))
+
+    def pop_tools(kwargs, result):
+        kwargs['payload']['model_parameters'].pop('tools')
+
+    def observe(kwargs, result):
+        observations.append('tools' in kwargs['payload']['model_parameters'])
+
+    call(enqueue_live, event, response)
+    call(enqueue_snapshot, event, response)
+    early_flush = json.dumps(queue[0][0]) if flush_before_edit else None
+    edited = copy.deepcopy(event) if copy_before_edit else event
+    call(pop_tools, edited, response)
+    call(observe, event, response)
+
+    assert queue[0][0] is payload
+    assert queue[0][1] is event
+    assert queue[0][2] is response
+    assert observations == [copy_before_edit]
+    assert json.loads(snapshots[0]) == {
+        'model_parameters': {'tools': ['lookup'], 'stream': True}
+    }
+
+    del response, event, payload, edited
+    gc.collect()
+    assert response_ref() is not None, 'queued response must outlive invocation'
+    flushed = json.loads(early_flush if early_flush is not None else json.dumps(queue[0][0]))
+    assert ('tools' in flushed['model_parameters']) == flushed_has_tools
+    assert flushed['model_parameters']['stream'] is True
+    queue.clear()
+    gc.collect()
+    assert response_ref() is None, 'clearing the queue must release its response'
+",
+    )
+}
+
+#[rstest]
+#[serial(python_interpreter)]
+fn shallow_queue_keeps_nested_aliases_but_not_replaced_fields(
+    initialized_python: &InitializedPython,
+    #[values(Backend::Python, Backend::PreparedCall)] backend: Backend,
+    #[values(InvocationMode::Direct, InvocationMode::Await)] mode: InvocationMode,
+) -> PyResult<()> {
+    initialized_python.attach(|py| {
+        let globals = pattern_scope(py, backend, mode)?;
+        py.run(
+            c"
+queue = []
+messages = [{'content': 'original'}]
+response = Value()
+response.messages = messages
+event = {'messages': messages, 'status': 'queued'}
+replacement = [{'content': 'replacement'}]
+
+def enqueue(data, result):
+    queue.append({**data})
+    queue.append(result.__dict__)
+",
+            Some(&globals),
+            None,
+        )?;
+        let event = item(&globals, "event");
+        let response = item(&globals, "response");
+        let messages = item(&globals, "messages");
+        item(&globals, "call").call1((item(&globals, "enqueue"), &event, &response))?;
+
+        messages.get_item(0)?.set_item("content", "edited")?;
+        event.set_item("messages", item(&globals, "replacement"))?;
+        event.set_item("status", "sent")?;
+
+        let queue = item(&globals, "queue");
+        let shallow_entry = queue.get_item(0)?;
+        let attributes_entry = queue.get_item(1)?;
+        assert!(!shallow_entry.is(&event));
+        assert!(attributes_entry.is(response.getattr("__dict__")?));
+        assert!(shallow_entry.get_item("messages")?.is(&messages));
+        assert!(attributes_entry.get_item("messages")?.is(&messages));
+        let serialized = item(&globals, "json").call_method1("dumps", (&shallow_entry,))?;
+        assert_eq!(
+            serde_json::from_str::<serde_json::Value>(serialized.extract::<&str>()?).unwrap(),
+            serde_json::json!({"messages": [{"content": "edited"}], "status": "queued"})
+        );
+        assert!(
+            event
+                .get_item("messages")?
+                .is(item(&globals, "replacement"))
+        );
+        Ok(())
+    })
+}
+
+#[rstest]
+#[case::ignored_return(false)]
+#[case::caught_error_does_not_roll_back(true)]
+#[serial(python_interpreter)]
+fn sequential_loggers_observe_prior_mutations_and_opaque_state(
+    initialized_python: &InitializedPython,
+    #[case] raises: bool,
+    #[values(Backend::Python, Backend::PreparedCall)] backend: Backend,
+    #[values(InvocationMode::Direct, InvocationMode::Await)] mode: InvocationMode,
+) -> PyResult<()> {
+    run_pattern(
+        initialized_python,
+        backend,
+        mode,
+        |globals| globals.set_item("raises", raises),
+        c"
+def test_pattern():
+    saved, seen, errors, order = [], [], [], []
+    state = {'optional_params': {'tools': ['lookup']}, 'cache_hit': None}
+    lock = threading.Lock()
+    failure = RuntimeError('after mutation')
+
+    def retain(kwargs):
+        order.append('retain')
+        saved.append(kwargs)
+
+    def mutate(kwargs):
+        order.append('mutate')
+        kwargs['optional_params'].pop('tools')
+        kwargs['cache_hit'] = False
+        kwargs['flush_lock'] = lock
+        if raises:
+            raise failure
+        return {'replacement': True}
+
+    def observe(kwargs):
+        order.append('observe')
+        seen.append((kwargs, kwargs['cache_hit'], kwargs['flush_lock']))
+
+    for callback in (retain, mutate, observe):
+        try:
+            call(callback, state)
+        except RuntimeError as error:
+            errors.append(error)
+
+    assert order == ['retain', 'mutate', 'observe']
+    assert errors == ([failure] if raises else [])
+    if raises:
+        assert errors[0] is failure
+    assert saved[0] is seen[0][0] is state
+    assert seen[0][1] is False
+    assert seen[0][2] is lock
+    assert state['optional_params'] == {}
+    assert 'replacement' not in state
+    with TestCase().assertRaises(TypeError):
+        json.dumps(state)
+",
+    )
+}
+
+#[rstest]
+#[case::replacement_list(false, false)]
+#[case::in_place_redaction(true, false)]
+#[case::equal_but_independently_copied_subset_does_not_match(false, true)]
+#[serial(python_interpreter)]
+fn redaction_matches_message_identity_and_preserves_unscanned_messages(
+    initialized_python: &InitializedPython,
+    #[case] redact_in_place: bool,
+    #[case] copy_subset: bool,
+    #[values(Backend::Python, Backend::PreparedCall)] backend: Backend,
+    #[values(InvocationMode::Direct, InvocationMode::Await)] mode: InvocationMode,
+) -> PyResult<()> {
+    run_pattern(
+        initialized_python,
+        backend,
+        mode,
+        |globals| {
+            globals.set_item("redact_in_place", redact_in_place)?;
+            globals.set_item("copy_subset", copy_subset)
+        },
+        c"
+def test_pattern():
+    first, second = {'content': 'private'}, {'content': 'keep'}
+    messages = [first, second]
+    subset = [first]
+    state = {'messages': messages}
+
+    def redact(full, selected):
+        if not {id(message) for message in selected} <= {id(message) for message in full}:
+            return None
+        if redact_in_place:
+            selected[0]['content'] = 'masked'
+            return full
+        replacements = {id(selected[0]): {'content': 'masked'}}
+        return [replacements.get(id(message), message) for message in full]
+
+    scanned = copy.deepcopy(subset) if copy_subset else subset
+    assert scanned == subset
+    result = call(redact, messages, scanned)
+    if result is not None and result is not messages:
+        state['messages'] = result
+
+    assert messages[0] is subset[0] is first
+    assert state['messages'][1] is second
+    if copy_subset:
+        assert scanned[0] is not first
+        assert result is None
+        assert state['messages'] is messages
+        assert first['content'] == 'private'
+    elif redact_in_place:
+        assert result is state['messages'] is messages
+        assert first['content'] == 'masked'
+    else:
+        assert state['messages'] is result
+        assert result is not messages
+        assert result[0] is not first
+        assert result[0] == {'content': 'masked'}
+        assert first['content'] == 'private'
+",
+    )
+}
+
+#[rstest]
+#[case::adopt_replacement(false)]
+#[case::merge_into_live_request(true)]
+#[serial(python_interpreter)]
+fn dispatcher_applies_returned_state_only_at_its_replacement_or_merge_boundary(
+    initialized_python: &InitializedPython,
+    #[case] merge: bool,
+    #[values(Backend::Python, Backend::PreparedCall)] backend: Backend,
+    #[values(InvocationMode::Direct, InvocationMode::Await)] mode: InvocationMode,
+) -> PyResult<()> {
+    initialized_python.attach(|py| {
+        let globals = pattern_scope(py, backend, mode)?;
+        py.run(
+            c"
+original = {'messages': [{'content': 'original'}], 'request_id': 'retained'}
+replacement = {'messages': [{'content': 'redacted'}], 'verdict': 'allow'}
+seen = []
+
+def rewrite(data):
+    data['checkpoint'] = True
+    return replacement
+
+def observe(data):
+    seen.append(data)
+",
+            Some(&globals),
+            None,
+        )?;
+        let original = item(&globals, "original");
+        let old_messages = original.get_item("messages")?;
+        let replacement = item(&globals, "replacement");
+        let call = item(&globals, "call");
+        let returned = call.call1((item(&globals, "rewrite"), &original))?;
+        assert!(returned.is(&replacement));
+        let current = if merge {
+            original.call_method1("update", (&returned,))?;
+            original.clone()
+        } else {
+            returned
+        };
+        call.call1((item(&globals, "observe"), &current))?;
+
+        assert!(item(&globals, "seen").get_item(0)?.is(&current));
+        assert!(
+            current
+                .get_item("messages")?
+                .is(replacement.get_item("messages")?)
+        );
+        assert_eq!(
+            old_messages
+                .get_item(0)?
+                .get_item("content")?
+                .extract::<String>()?,
+            "original"
+        );
+        assert!(original.get_item("checkpoint")?.extract::<bool>()?);
+        assert_eq!(current.is(&original), merge);
+        if merge {
+            assert_eq!(
+                current.get_item("request_id")?.extract::<String>()?,
+                "retained"
+            );
+            assert!(current.get_item("checkpoint")?.extract::<bool>()?);
+        } else {
+            assert!(original.get_item("messages")?.is(&old_messages));
+            assert!(!current.contains("request_id")?);
+            assert!(!current.contains("checkpoint")?);
+        }
+        Ok(())
+    })
+}
+
+#[rstest]
+#[case::failure_hook_only(false)]
+#[case::success_hook_checks_block_flag(true)]
+#[serial(python_interpreter)]
+fn block_exception_preserves_stashed_context_for_later_hooks(
+    initialized_python: &InitializedPython,
+    #[case] dispatch_success: bool,
+    #[values(Backend::Python, Backend::PreparedCall)] backend: Backend,
+    #[values(InvocationMode::Direct, InvocationMode::Await)] mode: InvocationMode,
+) -> PyResult<()> {
+    run_pattern(
+        initialized_python,
+        backend,
+        mode,
+        |globals| globals.set_item("dispatch_success", dispatch_success),
+        c"
+def test_pattern():
+    logger = Value()
+    logger.details = {}
+    request, failures, successes, success_calls = {}, [], [], []
+    block = RuntimeError('blocked')
+
+    def pre_call(data):
+        logger.details['blocked'] = True
+        data['logging_object'] = logger
+        raise block
+
+    def success(details):
+        success_calls.append(details)
+        if not details.get('blocked'):
+            successes.append(details)
+
+    def failure(data, error):
+        failures.append((data.pop('logging_object'), error))
+
+    with TestCase().assertRaises(RuntimeError) as caught:
+        call(pre_call, request)
+    assert caught.exception is block
+    assert request['logging_object'] is logger
+    assert logger.details == {'blocked': True}
+
+    if dispatch_success:
+        call(success, logger.details)
+        assert successes == []
+        assert success_calls[0] is logger.details
+        unblocked_details = {}
+        call(success, unblocked_details)
+        assert successes == [unblocked_details]
+        assert successes[0] is unblocked_details
+    assert len(success_calls) == (2 if dispatch_success else 0)
+    call(failure, request, caught.exception)
+
+    assert len(failures) == 1
+    assert failures[0][0] is logger
+    assert failures[0][1] is block
+    assert request == {}
+    assert logger.details == {'blocked': True}
+",
+    )
+}
+
+#[rstest]
+#[case::shared_inputs_lose_an_increment(false, false)]
+#[case::reverse_completion_changes_writer_order(false, true)]
+#[case::independent_snapshots_discard_edits(true, false)]
+#[serial(python_interpreter)]
+fn parallel_callbacks_share_state_without_isolation_or_automatic_return_merging(
+    initialized_python: &InitializedPython,
+    #[case] snapshots: bool,
+    #[case] reverse: bool,
+    #[values(Backend::Python, Backend::PreparedCall)] backend: Backend,
+    #[values(InvocationMode::Direct, InvocationMode::Await)] mode: InvocationMode,
+) -> PyResult<()> {
+    run_pattern(
+        initialized_python,
+        backend,
+        mode,
+        |globals| {
+            globals.set_item("snapshots", snapshots)?;
+            globals.set_item("reverse", reverse)
+        },
+        c"
+async def test_pattern():
+    state = {'count': 0, 'metadata': {'writers': []}}
+    entered = [asyncio.Event(), asyncio.Event()]
+    release = [asyncio.Event(), asyncio.Event()]
+    completed = [asyncio.Event(), asyncio.Event()]
+
+    async def increment(data, index):
+        previous = data['count']
+        entered[index].set()
+        await release[index].wait()
+        data['count'] = previous + 1
+        data['metadata']['writers'].append(index)
+        completed[index].set()
+        return {'ignored_replacement': True}
+
+    inputs = [copy.deepcopy(state) for _ in range(2)] if snapshots else [state, state]
+    tasks = [asyncio.create_task(invoke(increment, data, index)) for index, data in enumerate(inputs)]
+    await asyncio.gather(*(event.wait() for event in entered))
+    assert state == {'count': 0, 'metadata': {'writers': []}}
+    assert all(not task.done() for task in tasks)
+
+    order = [1, 0] if reverse else [0, 1]
+    for position, index in enumerate(order):
+        release[index].set()
+        await completed[index].wait()
+        assert inputs[index]['count'] == 1
+        if position == 0:
+            assert not tasks[1 - index].done()
+            assert inputs[index]['metadata']['writers'] == [index]
+
+    assert await asyncio.gather(*tasks) == [{'ignored_replacement': True}] * 2
+    if snapshots:
+        assert state == {'count': 0, 'metadata': {'writers': []}}
+        assert inputs[0] is not inputs[1]
+        for index, data in enumerate(inputs):
+            assert data is not state
+            assert data == {'count': 1, 'metadata': {'writers': [index]}}
+    else:
+        assert inputs[0] is inputs[1] is state
+        assert state == {'count': 1, 'metadata': {'writers': order}}
+",
+    )
+}
+
+#[rstest]
+#[case::after_return(false)]
+#[case::after_block_exception(true)]
+#[serial(python_interpreter)]
+fn background_task_keeps_arguments_and_updates_live_state_after_callback_finishes(
+    initialized_python: &InitializedPython,
+    #[case] raises: bool,
+    #[values(Backend::Python, Backend::PreparedCall)] backend: Backend,
+    #[values(InvocationMode::Direct, InvocationMode::Await)] mode: InvocationMode,
+) -> PyResult<()> {
+    run_pattern(
+        initialized_python,
+        backend,
+        mode,
+        |globals| globals.set_item("raises", raises),
+        c"
+async def test_pattern():
+    entered, release = asyncio.Event(), asyncio.Event()
+    tasks, observed_loops = [], []
+    state = Value()
+    state.metadata = {'audit': 'pending'}
+    reference = weakref.ref(state)
+    event_loop = asyncio.get_running_loop()
+
+    async def audit(data):
+        observed_loops.append(asyncio.get_running_loop())
+        entered.set()
+        await release.wait()
+        data.metadata['audit'] = 'complete'
+
+    def schedule(data):
+        observed_loops.append(asyncio.get_running_loop())
+        tasks.append(asyncio.create_task(audit(data)))
+        if raises:
+            raise RuntimeError('blocked after scheduling')
+
+    if raises:
+        with TestCase().assertRaisesRegex(RuntimeError, 'blocked after scheduling'):
+            await async_call(schedule, state)
+    else:
+        assert await async_call(schedule, state) is None
+
+    metadata = state.metadata
+    snapshot = copy.deepcopy(metadata)
+    del state
+    await entered.wait()
+    gc.collect()
+    assert reference() is not None
+    assert metadata == {'audit': 'pending'}
+    assert len(tasks) == 1
+    assert not tasks[0].done(), 'callback completion must not wait for background work'
+    assert tasks[0].get_loop() is event_loop
+    assert len(observed_loops) == 2
+    assert all(loop is event_loop for loop in observed_loops)
+
+    release.set()
+    assert await tasks[0] is None
+    assert metadata == {'audit': 'complete'}
+    assert snapshot == {'audit': 'pending'}
+    tasks.clear()
+    gc.collect()
+    assert reference() is None
+",
+    )
+}
+
+#[rstest]
+#[case::no_redaction_returns_original("passthrough")]
+#[case::shallow_redaction_copies_only_response("shallow_redaction")]
+#[case::per_key_copy_failure_shares_only_uncopyable_value("per_key_fallback")]
+#[case::whole_copy_success_isolates_nested_values("deep_copy")]
+#[case::whole_copy_failure_returns_original("whole_object_fallback")]
+#[serial(python_interpreter)]
+fn explicit_copy_boundaries_preserve_copy_depth_and_failure_fallback(
+    initialized_python: &InitializedPython,
+    #[case] copy_policy: &str,
+    #[values(Backend::Python, Backend::PreparedCall)] backend: Backend,
+    #[values(InvocationMode::Direct, InvocationMode::Await)] mode: InvocationMode,
+) -> PyResult<()> {
+    run_pattern(
+        initialized_python,
+        backend,
+        mode,
+        |globals| globals.set_item("copy_policy", copy_policy),
+        c"
+def test_pattern():
+    lock = threading.Lock()
+    state = {'messages': [{'content': 'private'}], 'response': {'content': 'private'}}
+    if copy_policy in ('per_key_fallback', 'whole_object_fallback'):
+        state['opaque'] = {'lock': lock, 'edits': []}
+
+    def passthrough(data):
+        return data
+
+    def shallow_redaction(data):
+        result = {**data, 'response': copy.deepcopy(data['response'])}
+        result['response']['content'] = 'masked'
+        return result
+
+    def per_key_fallback(data):
+        result = {}
+        for key, value in data.items():
+            try:
+                result[key] = copy.deepcopy(value)
+            except TypeError:
+                result[key] = value
+        return result
+
+    def whole_object_fallback(data):
+        try:
+            return copy.deepcopy(data)
+        except TypeError:
+            return data
+
+    callbacks = {
+        'passthrough': passthrough,
+        'shallow_redaction': shallow_redaction,
+        'per_key_fallback': per_key_fallback,
+        'deep_copy': whole_object_fallback,
+        'whole_object_fallback': whole_object_fallback,
+    }
+    result = call(callbacks[copy_policy], state)
+    returns_original = copy_policy in ('passthrough', 'whole_object_fallback')
+    shares_messages = returns_original or copy_policy == 'shallow_redaction'
+
+    assert (result is state) == returns_original
+    assert (result['response'] is state['response']) == returns_original
+    assert (result['messages'] is state['messages']) == shares_messages
+    result['messages'][0]['content'] = 'later edit'
+    assert state['messages'][0]['content'] == ('later edit' if shares_messages else 'private')
+    assert state['response']['content'] == 'private'
+    assert result['response']['content'] == ('masked' if copy_policy == 'shallow_redaction' else 'private')
+
+    if 'opaque' in state:
+        assert result['opaque'] is state['opaque']
+        assert result['opaque']['lock'] is lock
+        result['opaque']['edits'].append('shared despite copy')
+        assert state['opaque']['edits'] == ['shared despite copy']
+",
+    )
+}

litellm-rust/crates/python-interop/tests/fixtures/callback_controls.py

@@ -1,5 +1,4 @@
 import asyncio
-import copy
 import gc
 import inspect
 import weakref
@@ -31,46 +30,10 @@ class LiveCallFactory(CallFactory, Protocol):
     def live(self) -> int: ...
 
 
-Arguments = tuple[tuple[object, ...], dict[str, object] | None]
-ArgumentTransform = Callable[[tuple[object, ...], dict[str, object] | None], Arguments]
-
-
-def reconstruct_envelope(positional: tuple[object, ...], keywords: dict[str, object] | None) -> Arguments:
-    return tuple(value for value in positional), None if keywords is None else dict(keywords)
-
-
-def shallow_selected_payload(positional: tuple[object, ...], keywords: dict[str, object] | None) -> Arguments:
-    return (copy.copy(positional[0]), *positional[1:]), keywords
-
-
-def deepcopy_graph(positional: tuple[object, ...], keywords: dict[str, object] | None) -> Arguments:
-    return copy.deepcopy((positional, keywords))
-
-
-def deepcopy_separate(positional: tuple[object, ...], keywords: dict[str, object] | None) -> Arguments:
-    return copy.deepcopy(positional), copy.deepcopy(keywords)
-
-
 def unchanged_result(value: object) -> object:
     return value
 
 
-@dataclass(frozen=True, slots=True)
-class ArgumentTransformFactory:
-    inner: CallFactory
-    transform: ArgumentTransform
-
-    def prepare(
-        self,
-        callable: Callable[..., object],
-        positional: tuple[object, ...],
-        keywords: dict[str, object] | None = None,
-        awaited: bool = False,
-    ) -> PreparedInvocation:
-        args, kwargs = self.transform(positional, keywords)
-        return self.inner.prepare(callable, args, kwargs, awaited=awaited)
-
-
 @dataclass(frozen=True, slots=True)
 class ResultTransformInvocation:
     inner: PreparedInvocation
@@ -209,22 +172,7 @@ def control_factory(control: str, inner: CallFactory) -> CallFactory:
         return CheckedWeakFactory()
     if control == "missing_handoff":
         return MissingHandoffFactory(inner)
-    if control in ("result_passthrough", "result_shallow", "result_deep"):
-        return ResultTransformFactory(
-            inner,
-            {"result_passthrough": unchanged_result, "result_shallow": copy.copy, "result_deep": copy.deepcopy}[
-                control
-            ],
-        )
-    return ArgumentTransformFactory(
-        inner,
-        {
-            "envelope": reconstruct_envelope,
-            "shallow_payload": shallow_selected_payload,
-            "deep_graph": deepcopy_graph,
-            "deep_separate": deepcopy_separate,
-        }[control],
-    )
+    return {"result_passthrough": ResultTransformFactory(inner, unchanged_result)}[control]
 
 
 @dataclass
@@ -232,82 +180,6 @@ class ControlNode:
     stage: int = 0
 
 
-@dataclass
-class ControlPayload:
-    nested: ControlNode
-    stage: int = 0
-
-
-@dataclass(frozen=True, slots=True)
-class IdentityObservation:
-    root: bool
-    nested: bool
-    cross_argument: bool
-
-
-async def argument_identity(owners: CallFactory, awaited: bool) -> IdentityObservation:
-    nested = ControlNode()
-    original = ControlPayload(nested)
-
-    def observe(value: ControlPayload, *, alias: ControlNode) -> IdentityObservation:
-        return IdentityObservation(value is original, value.nested is nested, value.nested is alias)
-
-    async def observe_async(value: ControlPayload, *, alias: ControlNode) -> IdentityObservation:
-        return observe(value, alias=alias)
-
-    owner = owners.prepare(observe_async if awaited else observe, (original,), {"alias": nested}, awaited=awaited)
-    try:
-        pending = owner.invoke()
-        return await settle(pending, awaited)
-    finally:
-        owner.close()
-
-
-@dataclass(frozen=True, slots=True)
-class TimingObservation:
-    root: int
-    nested: int
-    alias: int
-
-
-async def mutation_timing(owners: CallFactory, awaited: bool) -> TimingObservation:
-    nested = ControlNode()
-    original = ControlPayload(nested)
-
-    def observe(value: ControlPayload, *, alias: ControlNode) -> TimingObservation:
-        return TimingObservation(value.stage, value.nested.stage, alias.stage)
-
-    async def observe_async(value: ControlPayload, *, alias: ControlNode) -> TimingObservation:
-        return observe(value, alias=alias)
-
-    owner = owners.prepare(observe_async if awaited else observe, (original,), {"alias": nested}, awaited=awaited)
-    try:
-        original.stage = nested.stage = 1
-        pending = owner.invoke()
-        original.stage = nested.stage = 2
-        return await settle(pending, awaited)
-    finally:
-        owner.close()
-
-
-async def result_identity(owners: CallFactory, awaited: bool) -> IdentityObservation:
-    original = ControlPayload(ControlNode())
-
-    def callback() -> ControlPayload:
-        return original
-
-    async def callback_async() -> ControlPayload:
-        return original
-
-    owner = owners.prepare(callback_async if awaited else callback, (), awaited=awaited)
-    try:
-        pending = owner.invoke()
-        result = await settle(pending, awaited)
-        return IdentityObservation(result is original, result.nested is original.nested, True)
-    finally:
-        owner.close()
-
-
 @dataclass
 class LifetimeCallback:
     awaited: bool
@@ -444,19 +316,6 @@ async def direct_coroutine(owners: CallFactory, awaited: bool) -> bool:
 
 
 def expected_control(witness: str, control: str, awaited: bool) -> object:
-    if witness == "argument_identity":
-        return IdentityObservation(
-            control in ("identity", "envelope"),
-            control in ("identity", "envelope", "shallow_payload"),
-            control != "deep_separate",
-        )
-    if witness == "mutation_timing":
-        stage = 2 if awaited else 1
-        if control in ("deep_graph", "deep_separate"):
-            return TimingObservation(0, 0, 0)
-        return TimingObservation(0 if control == "shallow_payload" else stage, stage, stage)
-    if witness == "result_identity":
-        return IdentityObservation(control in ("identity", "result_passthrough"), control != "result_deep", True)
     if witness in ("deferred_lifetime", "pending_handoff"):
         if control == "weak" or (witness == "pending_handoff" and control == "missing_handoff"):
             return LifetimeObservation(
@@ -480,9 +339,6 @@ def run_control(witness: str, control: str, retained: bool, awaited: bool, facto
 
 
 WITNESSES: dict[str, Callable[[CallFactory, bool], object]] = {
-    "argument_identity": argument_identity,
-    "mutation_timing": mutation_timing,
-    "result_identity": result_identity,
     "deferred_lifetime": deferred_lifetime,
     "borrowed_lifetime": borrowed_lifetime,
     "pending_handoff": pending_handoff,

litellm-rust/crates/python-interop/tests/fixtures/callback_integrations.py

@@ -10,6 +10,7 @@ from typing import Literal
 from unittest import TestCase
 
 import httpx
+from fastapi import HTTPException
 
 import litellm
 from litellm.integrations.custom_batch_logger import CustomBatchLogger
@@ -35,6 +36,12 @@ async def integration_invoke(owners, callback, *args, **kwargs):
     return await pending
 
 
+async def integration_checkpoint():
+    ready = asyncio.Event()
+    asyncio.get_running_loop().call_soon(ready.set)
+    await ready.wait()
+
+
 def integration_response(url, body, status=200, headers=None):
     return httpx.Response(status, json=body, headers=headers, request=httpx.Request("POST", url))
 
@@ -404,11 +411,18 @@ async def real_parallel_guardrail_snapshots(owners):
         for safe_memory_mode in (False, True):
             litellm.safe_memory_mode = safe_memory_mode
             await integration_parallel_snapshot_case(owners)
+            for ordinary_first in (False, True):
+                for reverse_completion in (False, True):
+                    await integration_parallel_snapshot_case(
+                        owners, ordinary_first=ordinary_first, reverse_completion=reverse_completion
+                    )
+            with TestCase().assertRaisesRegex(AssertionError, "copied live graph lost caller-visible mutation"):
+                await integration_parallel_snapshot_case(owners, copy_live=True)
     finally:
         litellm.safe_memory_mode = original_mode
 
 
-async def integration_parallel_snapshot_case(owners):
+async def integration_parallel_snapshot_case(owners, *, ordinary_first=None, reverse_completion=False, copy_live=False):
     from litellm.caching.dual_cache import DualCache
     from litellm.litellm_core_utils.core_helpers import independent_snapshot
     from litellm.proxy.utils import ProxyLogging
@@ -430,6 +444,21 @@ async def integration_parallel_snapshot_case(owners):
     assert raw is not live and raw["messages"][0] is not live["messages"][0]
     assert raw["uncopyable"] is sentinel and sentinel.attempts == 1
     live["messages"][0]["content"] = "masked"
+    dispatched = independent_snapshot(live) if copy_live else live
+    ordinary = RuntimeError("ordinary guardrail failure")
+    blocking = HTTPException(status_code=400, detail={"error": "blocked by policy"})
+    passthrough = ModifyResponseException("synthetic response", "fixture-model", dispatched)
+    errors = (
+        {}
+        if ordinary_first is None
+        else {
+            "fixture-live-writer": passthrough,
+            "fixture-raw-a": ordinary if ordinary_first else blocking,
+            "fixture-raw-b": blocking if ordinary_first else ordinary,
+        }
+    )
+    finished = []
+    permits, completed = {}, {}
 
     class Inspect(CustomGuardrail):
         async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
@@ -447,7 +476,12 @@ async def integration_parallel_snapshot_case(owners):
                 assert data["uncopyable"] is sentinel
                 data["uncopyable"].observed.append(self.guardrail_name)
             else:
-                assert data is live and data["messages"][0]["content"] == "shared mutation"
+                assert data is dispatched and data["messages"][0]["content"] == "shared mutation"
+            await permits[self.guardrail_name].wait()
+            finished.append(self.guardrail_name)
+            completed[self.guardrail_name].set()
+            if self.guardrail_name in errors:
+                raise errors[self.guardrail_name]
             return {"discarded": self.guardrail_name}
 
     guardrails = tuple(
@@ -462,30 +496,63 @@ async def integration_parallel_snapshot_case(owners):
     )
     proxy = ProxyLogging.__new__(ProxyLogging)
     proxy.call_details = {"user_api_key_cache": DualCache()}
+    permits.update((guardrail.guardrail_name, asyncio.Event()) for guardrail in guardrails)
+    completed.update((guardrail.guardrail_name, asyncio.Event()) for guardrail in guardrails)
     task = asyncio.create_task(
         integration_invoke(
-            owners, proxy._run_parallel_pre_call_guardrails, guardrails, live, raw, UserAPIKeyAuth(), "acompletion"
+            owners,
+            proxy._run_parallel_pre_call_guardrails,
+            guardrails,
+            dispatched,
+            raw,
+            UserAPIKeyAuth(),
+            "acompletion",
         )
     )
     try:
         await arrived.wait()
         assert not task.done()
-        assert observations["fixture-live-writer"] is observations["fixture-live-reader"] is live
+        assert observations["fixture-live-writer"] is observations["fixture-live-reader"] is dispatched
         first, second = observations["fixture-raw-a"], observations["fixture-raw-b"]
         assert first is not second and first is not raw and second is not raw
         assert first["messages"][0] is not second["messages"][0]
         assert first["messages"][0] is not raw["messages"][0]
         assert first["uncopyable"] is second["uncopyable"] is raw["uncopyable"] is sentinel
-        assert sentinel.attempts == 3 and not sentinel.observed
+        assert sentinel.attempts == 3 + copy_live and not sentinel.observed
         release.set()
-        assert await task is None
+        order = tuple(completed)[:: -1 if reverse_completion else 1]
+        for index, name in enumerate(order):
+            assert not task.done()
+            permits[name].set()
+            await completed[name].wait()
+            await integration_checkpoint()
+            assert finished == list(order[: index + 1])
+        if errors:
+            expected = ordinary if ordinary_first else blocking
+            with TestCase().assertRaises(type(expected)) as caught:
+                await task
+            assert caught.exception is expected
+            assert passthrough.request_data is dispatched
+            assert blocking.status_code == 400
+            assert blocking.detail == {
+                "error": "blocked by policy",
+                "guardrail_name": "fixture-raw-b" if ordinary_first else "fixture-raw-a",
+                "guardrail_mode": GuardrailEventHooks.pre_call,
+            }
+        else:
+            assert await task is None
         assert raw["messages"] == [{"role": "user", "content": "original"}]
-        assert live["messages"][0]["content"] == "shared mutation"
+        assert dispatched["messages"][0]["content"] == "shared mutation"
         assert set(sentinel.observed) == {"fixture-raw-a", "fixture-raw-b"} and len(sentinel.observed) == 2
-        assert "discarded" not in live
+        assert "discarded" not in dispatched
         assert first["messages"][0]["content"] == "fixture-raw-a"
         assert second["messages"][0]["content"] == "fixture-raw-b"
-        assert all(guardrail._pre_call_hook_already_ran(live) for guardrail in guardrails if guardrail.scan_raw_request)
+        assert all(
+            guardrail._pre_call_hook_already_ran(dispatched) is (guardrail.guardrail_name not in errors)
+            for guardrail in guardrails
+            if guardrail.scan_raw_request
+        )
+        assert live["messages"][0]["content"] == "shared mutation", "copied live graph lost caller-visible mutation"
     finally:
         release.set()
         mutated.set()
@@ -495,7 +562,13 @@ async def integration_parallel_snapshot_case(owners):
 
 
 async def real_purview_sync_background(owners):
+    for active_loop in (False, True):
+        await integration_purview_logging_case(owners, active_loop=active_loop)
+
+
+async def integration_purview_logging_case(owners, *, active_loop):
     entered, release = threading.Event(), threading.Event()
+    async_entered, async_release = asyncio.Event(), asyncio.Event()
     calls, workers = [], []
     main_thread = threading.get_ident()
 
@@ -505,7 +578,11 @@ async def real_purview_sync_background(owners):
             calls.append((url, kwargs))
             if url.endswith("/token"):
                 entered.set()
-                assert release.wait(5), "background audit was not released"
+                if active_loop:
+                    async_entered.set()
+                    await async_release.wait()
+                else:
+                    assert release.wait(5), "background audit was not released"
                 return integration_response(url, {"access_token": "fixture-token", "expires_in": 3600})
             assert kwargs["headers"]["Authorization"] == "Bearer fixture-token"
             if url.endswith("/compute"):
@@ -535,19 +612,39 @@ async def real_purview_sync_background(owners):
     }
     result = ModelResponse(model="fixture-model", choices=[{"message": {"role": "assistant", "content": "before"}}])
     owner = owners.prepare(purview.logging_hook, (kwargs, result, "completion"), awaited=False)
+    task = None
     try:
-        returned = await asyncio.to_thread(owner.invoke)
+        tasks_before = asyncio.all_tasks()
+        threads_before = set(threading.enumerate())
+        returned = owner.invoke() if active_loop else await asyncio.to_thread(owner.invoke)
         owner.close()
         assert returned[0] is kwargs and returned[1] is result
-        assert await asyncio.to_thread(entered.wait, 5)
-        assert len(calls) == 1 and workers[0].ident != main_thread and workers[0].daemon
+        if active_loop:
+            await integration_checkpoint()
+            assert not calls and not workers and not entered.is_set()
+            assert asyncio.all_tasks() == tasks_before and set(threading.enumerate()) == threads_before
+            task = asyncio.create_task(
+                integration_invoke(owners, purview.async_logging_hook, kwargs, result, "completion")
+            )
+            await async_entered.wait()
+            assert not task.done() and workers[0].ident == main_thread
+        else:
+            assert await asyncio.to_thread(entered.wait, 5)
+            assert workers[0].ident != main_thread and workers[0].daemon
+        assert len(calls) == 1
         assert workers[0].is_alive()
         kwargs["messages"][0]["content"] = "too late for prompt extraction"
         kwargs["litellm_call_id"] = "after-http"
         result.choices[0].message.content = "response mutated while audit waits"
         release.set()
-        await asyncio.to_thread(workers[0].join, 5)
-        assert not workers[0].is_alive() and all(worker is workers[0] for worker in workers)
+        async_release.set()
+        if active_loop:
+            audited = await task
+            assert audited[0] is kwargs and audited[1] is result
+        else:
+            await asyncio.to_thread(workers[0].join, 5)
+            assert not workers[0].is_alive()
+        assert all(worker is workers[0] for worker in workers)
         assert len(calls) == 4
         entries = [call[1]["json"]["contentToProcess"] for call in calls[2:]]
         assert [entry["activityMetadata"]["activity"] for entry in entries] == ["uploadText", "downloadText"]
@@ -561,5 +658,10 @@ async def real_purview_sync_background(owners):
     finally:
         owner.close()
         release.set()
-        if workers:
+        async_release.set()
+        if task is not None:
+            if not task.done():
+                task.cancel()
+            await asyncio.gather(task, return_exceptions=True)
+        if workers and not active_loop:
             await asyncio.to_thread(workers[0].join, 5)

litellm-rust/crates/python-interop/tests/fixtures/callback_lifecycle.py

@@ -563,58 +563,6 @@ async def repeated_ownership(owners):
         assert all(ref() is None for ref in refs)
 
 
-async def retained_field_replacement(owners):
-    original = {"messages": [{"content": "original"}]}
-    replacement = {"messages": [{"content": "replacement"}]}
-    event = {"payload": original, "alias": original}
-    saved = []
-
-    def retain(value):
-        saved.append(value["payload"])
-
-    def replace(value):
-        value["payload"] = replacement
-        value["alias"]["messages"][0]["content"] = "mutated original"
-
-    for callback in (retain, replace):
-        owner = owners.prepare(callback, (event,))
-        try:
-            assert owner.invoke() is None
-        finally:
-            owner.close()
-    assert saved[0] is original is event["alias"]
-    assert event["payload"] is replacement
-    assert saved[0]["messages"][0]["content"] == "mutated original"
-    replacement["messages"][0]["content"] = "mutated replacement"
-    assert event["payload"]["messages"][0]["content"] == "mutated replacement"
-    assert original["messages"][0]["content"] == "mutated original"
-
-
-async def queued_graph_ownership(owners):
-    queue = asyncio.Queue()
-    sentinel = Value()
-    reference = weakref.ref(sentinel)
-    payload = {"sentinel": sentinel, "nested": {"status": "queued"}}
-    snapshot = json.dumps(payload["nested"])
-    enqueue = owners.prepare(queue.put_nowait, (payload,))
-    try:
-        enqueue.invoke()
-    finally:
-        enqueue.close()
-    del sentinel, payload
-    gc.collect()
-    assert owners.live == 0 and reference() is not None
-    queued = queue.get_nowait()
-    queued["nested"]["status"] = "changed before flush"
-    assert json.loads(json.dumps(queued["nested"])) == {"status": "changed before flush"}
-    assert json.loads(snapshot) == {"status": "queued"}
-    assert queued["sentinel"] is reference()
-    queue.task_done()
-    del queued
-    gc.collect()
-    assert reference() is None
-
-
 async def detached_work_after_error(owners):
     for raises in (False, True):
         entered, release = asyncio.Event(), asyncio.Event()
@@ -665,16 +613,18 @@ def run_checked(owners, scenario):
 
     async def run():
         asyncio.get_running_loop().set_exception_handler(lambda loop, context: background_failures.append(context))
-        await asyncio.wait_for(scenario, timeout=15)
+        result = await asyncio.wait_for(scenario, timeout=15)
         gc.collect()
         assert owners.live == baseline
         pending = asyncio.all_tasks() - {asyncio.current_task()}
         assert not pending, f"undrained tasks: {pending}"
+        return result
 
-    asyncio.run(run())
+    result = asyncio.run(run())
     gc.collect()
     assert not background_failures, f"unhandled background failures: {background_failures}"
     assert owners.live == baseline
+    return result
 
 
 def run_scenario(name, retained, factory):

litellm-rust/crates/python-interop/tests/prepared_call.rs

@@ -1,6 +1,7 @@
 use litellm_python_interop::{InvocationMode, InvocationOutcome, PreparedCall};
+use pyo3::exceptions::{PyAssertionError, PyKeyboardInterrupt, PyValueError};
 use pyo3::prelude::*;
-use pyo3::types::{PyDict, PyTuple};
+use pyo3::types::{PyDict, PyList, PyTuple};
 use rstest::rstest;
 
 #[path = "support/mod.rs"]
@@ -17,19 +18,23 @@ fn retains_aliases_mutations_and_original_result(
         let globals = scope(
             py,
             c"
-shared = {'value': 'before'}
-payload = {'nested': shared}
-saved = []
 def callback(data, *, alias):
-    assert data['nested'] is alias
+    observed.append(data['nested'] is alias)
     saved.append(data)
     alias['value'] = 'during'
     return data
 ",
         )?;
-        let payload = item(&globals, "payload");
+        let shared = PyDict::new(py);
+        shared.set_item("value", "before")?;
+        let payload = PyDict::new(py);
+        payload.set_item("nested", &shared)?;
+        let saved = PyList::empty(py);
+        let observed = PyList::empty(py);
+        globals.set_item("saved", &saved)?;
+        globals.set_item("observed", &observed)?;
         let keywords = PyDict::new(py);
-        keywords.set_item("alias", item(&globals, "shared"))?;
+        keywords.set_item("alias", &shared)?;
         let invocation = PreparedCall::new(
             InvocationMode::Direct,
             item(&globals, "callback").unbind(),
@@ -39,16 +44,19 @@ def callback(data, *, alias):
         let result = invoke_direct(&invocation, py)?;
         assert!(result.bind(py).is(&payload));
         drop(invocation);
-        py.run(
-            c"
-assert saved[0] is payload
-assert shared['value'] == 'during'
-shared['value'] = 'after'
-assert saved[0]['nested']['value'] == 'after'
-",
-            Some(&globals),
-            None,
-        )
+        assert_eq!(observed.extract::<Vec<bool>>()?, [true]);
+        assert!(saved.get_item(0)?.is(&payload));
+        assert_eq!(item(&shared, "value").extract::<String>()?, "during");
+        shared.set_item("value", "after")?;
+        assert_eq!(
+            saved
+                .get_item(0)?
+                .get_item("nested")?
+                .get_item("value")?
+                .extract::<String>()?,
+            "after"
+        );
+        Ok(())
     })
 }
 
@@ -61,33 +69,38 @@ fn preserves_exception_identity_cause_traceback_and_prior_mutation(
         let globals = scope(
             py,
             c"
-payload = {}
-error = KeyboardInterrupt('original')
-cause = ValueError('cause')
 def callback(data):
     data['changed'] = True
     raise error from cause
 ",
         )?;
+        let payload = PyDict::new(py);
+        let original = PyKeyboardInterrupt::new_err("original");
+        let cause = PyValueError::new_err("cause");
+        globals.set_item("error", original.value(py))?;
+        globals.set_item("cause", cause.value(py))?;
         let invocation = PreparedCall::new(
             InvocationMode::Direct,
             item(&globals, "callback").unbind(),
-            PyTuple::new(py, [item(&globals, "payload")])?.unbind(),
+            PyTuple::new(py, [&payload])?.unbind(),
             None,
         );
         let error = invoke_direct(&invocation, py).unwrap_err();
-        assert!(error.value(py).is(item(&globals, "error")));
+        assert!(error.value(py).is(original.value(py)));
         drop(invocation);
-        py.run(
-            c"
-import traceback
-assert payload['changed'] is True
-assert error.__cause__ is cause
-assert traceback.extract_tb(error.__traceback__)[-1].name == 'callback'
-",
-            Some(&globals),
-            None,
-        )
+        assert!(item(&payload, "changed").extract::<bool>()?);
+        assert!(error.value(py).getattr("__cause__")?.is(cause.value(py)));
+        let frames = py
+            .import("traceback")?
+            .call_method1("extract_tb", (error.traceback(py),))?;
+        assert_eq!(
+            frames
+                .get_item(frames.len()? - 1)?
+                .getattr("name")?
+                .extract::<String>()?,
+            "callback"
+        );
+        Ok(())
     })
 }
 
@@ -98,15 +111,16 @@ fn returns_coroutine_without_executing_it(initialized_python: &InitializedPython
         let globals = scope(
             py,
             c"
-import inspect
-started = []
 async def work():
     started.append(True)
-coroutine = work()
 def callback():
     return coroutine
 ",
         )?;
+        let started = PyList::empty(py);
+        globals.set_item("started", &started)?;
+        let coroutine = item(&globals, "work").call0()?;
+        globals.set_item("coroutine", &coroutine)?;
         let invocation = PreparedCall::new(
             InvocationMode::Direct,
             item(&globals, "callback").unbind(),
@@ -114,16 +128,13 @@ def callback():
             None,
         );
         let result = invoke_direct(&invocation, py)?;
-        assert!(result.bind(py).is(item(&globals, "coroutine")));
-        py.run(
-            c"
-assert started == []
-assert inspect.getcoroutinestate(coroutine) == inspect.CORO_CREATED
-coroutine.close()
-",
-            Some(&globals),
-            None,
-        )
+        let inspect = py.import("inspect")?;
+        let state = inspect.call_method1("getcoroutinestate", (&coroutine,));
+        coroutine.call_method0("close")?;
+        assert!(result.bind(py).is(&coroutine));
+        assert!(started.is_empty());
+        assert!(state?.eq(inspect.getattr("CORO_CREATED")?)?);
+        Ok(())
     })
 }
 
@@ -149,25 +160,31 @@ fn preserves_current_context_thread_and_reentry(
         let globals = scope(
             py,
             c"
-import contextvars
 import threading
-context = contextvars.ContextVar('prepared_call_context')
-token = context.set('caller')
-thread = threading.get_ident()
-payload = {}
 def inner(data):
-    assert context.get() == 'outer'
-    assert threading.get_ident() == thread
+    observed.append((context.get(), threading.get_ident()))
     data['inner'] = True
     context.set('inner')
     return data
 def outer():
-    assert context.get() == 'caller'
-    assert threading.get_ident() == thread
+    observed.append((context.get(), threading.get_ident()))
     context.set('outer')
     return reenter(inner, payload)
 ",
         )?;
+        let context = py
+            .import("contextvars")?
+            .getattr("ContextVar")?
+            .call1(("prepared_call_context",))?;
+        let thread = py
+            .import("threading")?
+            .call_method0("get_ident")?
+            .extract::<u64>()?;
+        let payload = PyDict::new(py);
+        let observed = PyList::empty(py);
+        globals.set_item("context", &context)?;
+        globals.set_item("payload", &payload)?;
+        globals.set_item("observed", &observed)?;
         globals.set_item("reenter", wrap_pyfunction!(reenter, py)?)?;
         let invocation = PreparedCall::new(
             InvocationMode::Direct,
@@ -175,19 +192,18 @@ def outer():
             PyTuple::empty(py).unbind(),
             None,
         );
-        let result = invoke_direct(&invocation, py)?;
-        assert!(result.bind(py).is(item(&globals, "payload")));
-        py.run(
-            c"
-try:
-    assert payload['inner'] is True
-    assert context.get() == 'inner'
-finally:
-    context.reset(token)
-",
-            Some(&globals),
-            None,
-        )
+        let token = context.call_method1("set", ("caller",))?;
+        let result = invoke_direct(&invocation, py);
+        let final_context = context.call_method0("get");
+        context.call_method1("reset", (token,))?;
+        assert!(result?.bind(py).is(&payload));
+        assert!(item(&payload, "inner").extract::<bool>()?);
+        assert_eq!(final_context?.extract::<String>()?, "inner");
+        assert_eq!(
+            observed.extract::<Vec<(String, u64)>>()?,
+            [("caller".to_owned(), thread), ("outer".to_owned(), thread)]
+        );
+        Ok(())
     })
 }
 
@@ -200,57 +216,56 @@ fn owns_arguments_until_release_and_preserves_callback_retention(
         let globals = scope(
             py,
             c"
-import gc
-import weakref
-saved = []
 class Value:
     pass
 class Callback:
     def __call__(self, value, *, other):
         saved.append(value)
-        assert other is other_ref()
-value = Value()
-other = Value()
-callback = Callback()
-value_ref = weakref.ref(value)
-other_ref = weakref.ref(other)
-callback_ref = weakref.ref(callback)
+        observed.append(other is other_ref())
 ",
         )?;
+        globals.set_item("saved", PyList::empty(py))?;
+        globals.set_item("observed", PyList::empty(py))?;
+        let value = item(&globals, "Value").call0()?;
+        let other = item(&globals, "Value").call0()?;
+        let callback = item(&globals, "Callback").call0()?;
+        let weakref = py.import("weakref")?;
+        for (name, object) in [
+            ("value_ref", &value),
+            ("other_ref", &other),
+            ("callback_ref", &callback),
+        ] {
+            globals.set_item(name, weakref.call_method1("ref", (object,))?)?;
+        }
         let keywords = PyDict::new(py);
-        keywords.set_item("other", item(&globals, "other"))?;
+        keywords.set_item("other", other)?;
         let invocation = PreparedCall::new(
             InvocationMode::Direct,
-            item(&globals, "callback").unbind(),
-            PyTuple::new(py, [item(&globals, "value")])?.unbind(),
+            callback.unbind(),
+            PyTuple::new(py, [value])?.unbind(),
             Some(keywords.unbind()),
         );
-        py.run(c"del value, other, callback", Some(&globals), None)?;
         Ok::<_, PyErr>((invocation, globals.unbind()))
     })?;
     Python::attach(|py| {
         let globals = globals.bind(py);
-        py.run(
-            c"assert all(ref() is not None for ref in (value_ref, other_ref, callback_ref))",
-            Some(globals),
-            None,
-        )?;
+        for name in ["value_ref", "other_ref", "callback_ref"] {
+            assert!(!item(globals, name).call0()?.is_none(), "{name}");
+        }
         assert!(invoke_direct(&invocation, py)?.is_none(py));
         drop(invocation);
-        py.run(
-            c"
-gc.collect()
-assert callback_ref() is None
-assert other_ref() is None
-assert value_ref() is saved[0]
-saved[0].still_usable = True
-saved.clear()
-gc.collect()
-assert value_ref() is None
-",
-            Some(globals),
-            None,
-        )
+        let gc = py.import("gc")?;
+        gc.call_method0("collect")?;
+        assert_eq!(item(globals, "observed").extract::<Vec<bool>>()?, [true]);
+        assert!(item(globals, "callback_ref").call0()?.is_none());
+        assert!(item(globals, "other_ref").call0()?.is_none());
+        let saved = item(globals, "saved");
+        assert!(item(globals, "value_ref").call0()?.is(saved.get_item(0)?));
+        saved.get_item(0)?.setattr("still_usable", true)?;
+        saved.call_method0("clear")?;
+        gc.call_method0("collect")?;
+        assert!(item(globals, "value_ref").call0()?.is_none());
+        Ok(())
     })
 }
 
@@ -292,31 +307,51 @@ fn checked_runner_rejects_unhandled_background_failures(
 async def fail():
     raise RuntimeError('background task regression')
 
-for cyclic in (False, True):
-    for handled in (False, True):
-        async def scenario(cyclic=cyclic, handled=handled):
-            task = asyncio.create_task(fail())
-            if cyclic:
-                task.cycle = task
-            await checkpoint()
-            assert task.done()
-            if handled:
-                with TestCase().assertRaisesRegex(RuntimeError, 'background task regression'):
-                    task.result()
-            del task
-
-        owners = ReferenceFactory()
-        if handled:
-            run_checked(owners, scenario())
-        else:
-            with TestCase().assertRaisesRegex(
-                AssertionError, r'unhandled background failures: .*background task regression'
-            ):
-                run_checked(owners, scenario())
+async def scenario(cyclic, handled, observed):
+    task = asyncio.create_task(fail())
+    if cyclic:
+        task.cycle = task
+    await checkpoint()
+    observed['done'] = task.done()
+    if handled:
+        try:
+            task.result()
+        except RuntimeError as error:
+            observed['error'] = str(error)
+    del task
 ",
             Some(&globals),
             None,
-        )
+        )?;
+        for cyclic in [false, true] {
+            for handled in [false, true] {
+                let observed = PyDict::new(py);
+                let owners = item(&globals, "ReferenceFactory").call0()?;
+                let scenario = item(&globals, "scenario").call1((cyclic, handled, &observed))?;
+                let result = item(&globals, "run_checked").call1((owners, scenario));
+                assert!(
+                    item(&observed, "done").extract::<bool>()?,
+                    "cyclic={cyclic}, handled={handled}"
+                );
+                if handled {
+                    result?;
+                    assert_eq!(
+                        item(&observed, "error").extract::<String>()?,
+                        "background task regression"
+                    );
+                } else {
+                    let error = result.unwrap_err();
+                    assert!(error.is_instance_of::<PyAssertionError>(py), "{error}");
+                    let message = error.value(py).str()?.to_str()?.to_owned();
+                    assert!(
+                        message.starts_with("unhandled background failures: "),
+                        "{message}"
+                    );
+                    assert!(message.contains("background task regression"), "{message}");
+                }
+            }
+        }
+        Ok(())
     })
 }
 
@@ -330,11 +365,7 @@ fn real_ocr_logging_preserves_execution_roots_and_continues_after_error(
         let globals = scope(
             py,
             c"
-from datetime import datetime
 from litellm.integrations.custom_logger import CustomLogger
-from litellm.litellm_core_utils.litellm_logging import Logging
-
-order = []
 
 class Retain(CustomLogger):
     def log_pre_api_call(self, model, messages, kwargs):
@@ -365,60 +396,107 @@ class Observe(CustomLogger):
             'document' in self.view['complete_input_dict'],
         )
 
-first = Retain()
-last = Observe()
-document = {'value': 'original'}
-headers = {'X-Trace': 'original'}
-body = {'document': document, 'alias': document}
-view = {'headers': headers, 'complete_input_dict': body, 'api_base': 'https://example.invalid/ocr'}
-logger = Logging(
-    model='test', messages=[], stream=False, call_type='ocr',
-    start_time=datetime.now(), litellm_call_id='retained-test', function_id='retained-test',
-    dynamic_input_callbacks=[first, MutateThenFail(), last],
-)
 ",
         )?;
-        let headers = item(&globals, "headers").unbind();
-        let body = item(&globals, "body").unbind();
-        let view = item(&globals, "view").cast_into::<PyDict>()?.unbind();
-        let invocation = prepare_pre_call(py, &item(&globals, "logger"), view.bind(py))?;
+        let order = PyList::empty(py);
+        globals.set_item("order", &order)?;
+        let first = item(&globals, "Retain").call0()?;
+        let last = item(&globals, "Observe").call0()?;
+        let document = PyDict::new(py);
+        document.set_item("value", "original")?;
+        let headers = PyDict::new(py);
+        headers.set_item("X-Trace", "original")?;
+        let body = PyDict::new(py);
+        body.set_item("document", &document)?;
+        body.set_item("alias", &document)?;
+        let view = PyDict::new(py);
+        view.set_item("headers", &headers)?;
+        view.set_item("complete_input_dict", &body)?;
+        view.set_item("api_base", "https://example.invalid/ocr")?;
+        let keywords = PyDict::new(py);
+        keywords.set_item("model", "test")?;
+        keywords.set_item("messages", PyList::empty(py))?;
+        keywords.set_item("stream", false)?;
+        keywords.set_item("call_type", "ocr")?;
+        keywords.set_item(
+            "start_time",
+            py.import("datetime")?
+                .getattr("datetime")?
+                .call_method0("now")?,
+        )?;
+        keywords.set_item("litellm_call_id", "retained-test")?;
+        keywords.set_item("function_id", "retained-test")?;
+        keywords.set_item(
+            "dynamic_input_callbacks",
+            PyList::new(
+                py,
+                [&first, &item(&globals, "MutateThenFail").call0()?, &last],
+            )?,
+        )?;
+        let logger = py
+            .import("litellm.litellm_core_utils.litellm_logging")?
+            .getattr("Logging")?
+            .call((), Some(&keywords))?;
+        let invocation = prepare_pre_call(py, &logger, &view)?;
         assert!(invoke_direct(&invocation, py)?.is_none(py));
         drop(invocation);
-        py.run(c"del headers, body, view", Some(&globals), None)?;
-        assert!(
-            headers
-                .bind(py)
-                .is(item(&globals, "first").getattr("headers")?)
-        );
-        assert!(body.bind(py).is(item(&globals, "first").getattr("body")?));
-        assert!(view.bind(py).is(item(&globals, "last").getattr("view")?));
+        assert!(headers.is(first.getattr("headers")?));
+        assert!(body.is(first.getattr("body")?));
+        assert!(view.is(last.getattr("view")?));
+        assert_eq!(item(&headers, "X-Trace").extract::<String>()?, "mutated");
         assert_eq!(
-            headers.bind(py).get_item("X-Trace")?.extract::<String>()?,
+            order.extract::<Vec<String>>()?,
+            ["retain", "mutate_then_fail", "observe"]
+        );
+        assert_eq!(
+            first.getattr("snapshot")?.extract::<(String, String)>()?,
+            ("original".to_owned(), "original".to_owned())
+        );
+        assert_eq!(
+            last.getattr("snapshot")?
+                .extract::<(Vec<(String, String)>, bool, bool)>()?,
+            (
+                vec![("X-Trace".to_owned(), "replacement".to_owned())],
+                true,
+                false
+            )
+        );
+        assert!(first.getattr("view")?.is(last.getattr("view")?));
+        assert!(first.getattr("body")?.get_item("document")?.is(&document));
+        assert!(first.getattr("body")?.get_item("alias")?.is(&document));
+        assert_eq!(item(&document, "value").extract::<String>()?, "mutated");
+        assert_eq!(
+            last.getattr("view")?
+                .get_item("headers")?
+                .get_item("X-Trace")?
+                .extract::<String>()?,
+            "replacement"
+        );
+        let replacement = PyDict::new(py);
+        replacement.set_item("replacement", true)?;
+        assert!(
+            last.getattr("view")?
+                .get_item("complete_input_dict")?
+                .eq(replacement)?
+        );
+        document.set_item("value", "after invocation")?;
+        assert_eq!(
+            first
+                .getattr("body")?
+                .get_item("document")?
+                .get_item("value")?
+                .extract::<String>()?,
+            "after invocation"
+        );
+        drop((headers, body, view));
+        assert_eq!(
+            first
+                .getattr("headers")?
+                .get_item("X-Trace")?
+                .extract::<String>()?,
             "mutated"
         );
-        py.run(
-            c"
-assert order == ['retain', 'mutate_then_fail', 'observe']
-assert first.snapshot == ('original', 'original')
-assert last.snapshot == ((('X-Trace', 'replacement'),), True, False)
-assert first.view is last.view
-assert first.body['document'] is document
-assert first.body['alias'] is document
-assert document['value'] == 'mutated'
-assert last.view['headers']['X-Trace'] == 'replacement'
-assert last.view['complete_input_dict'] == {'replacement': True}
-document['value'] = 'after invocation'
-assert first.body['document']['value'] == 'after invocation'
-",
-            Some(&globals),
-            None,
-        )?;
-        drop((headers, body, view));
-        py.run(
-            c"assert first.headers['X-Trace'] == 'mutated'",
-            Some(&globals),
-            None,
-        )
+        Ok(())
     })
 }
 

litellm/ocr/main.py

@@ -28,8 +28,6 @@ from litellm.llms.base_llm.ocr.transformation import (
     parse_ocr_request_format,
 )
 from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
-from litellm.rust_bridge import ocr as rust_ocr_bridge
-from litellm.rust_bridge.configuration import rust_enabled
 from litellm.types.router import GenericLiteLLMParams
 from litellm.utils import ProviderConfigManager, client
 
@@ -242,19 +240,6 @@ async def aocr(
         )
         ```
     """
-    if rust_enabled():
-        return await rust_ocr_bridge.aocr(
-            {
-                **kwargs,
-                "model": model,
-                "document": document,
-                "api_key": api_key,
-                "api_base": api_base,
-                "timeout": timeout,
-                "custom_llm_provider": custom_llm_provider,
-                "extra_headers": extra_headers,
-            }
-        )
     return await _legacy_aocr(model, document, api_key, api_base, timeout, custom_llm_provider, extra_headers, **kwargs)
 
 
@@ -525,20 +510,6 @@ def ocr(
             print(f"Page {page.index}: {page.markdown}")
         ```
     """
-    if rust_enabled():
-        arguments: Final[dict[str, object]] = {
-            **kwargs,
-            "model": model,
-            "document": document,
-            "api_key": api_key,
-            "api_base": api_base,
-            "timeout": timeout,
-            "custom_llm_provider": custom_llm_provider,
-            "extra_headers": extra_headers,
-        }
-        if kwargs.get("aocr") is True:
-            return rust_ocr_bridge.aocr(arguments)
-        return rust_ocr_bridge.ocr(arguments)
     return _legacy_ocr(model, document, api_key, api_base, timeout, custom_llm_provider, extra_headers, **kwargs)
 
 

litellm/rust_bridge/ocr.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import inspect
 import traceback
 from collections.abc import Awaitable
 from contextvars import copy_context
@@ -58,14 +57,67 @@ async def aocr(arguments: dict[str, object]) -> OCRResponse:
 
 def initialize_logging(arguments: dict[str, object], asynchronous: bool) -> object:
     import litellm
-    from litellm.litellm_core_utils.litellm_logging import Logging
+    from litellm import utils
+    from litellm.integrations.custom_logger import CustomLogger
+    from litellm.litellm_core_utils import litellm_logging
+    from litellm.litellm_core_utils.coroutine_checker import coroutine_checker
+    from litellm.litellm_core_utils.litellm_logging import Logging, set_callbacks
 
     supplied: Final = arguments.get("litellm_logging_obj")
     if supplied is not None:
         return supplied
-    callbacks: Final = tuple(dict.fromkeys((*litellm.callbacks, *cast(list, arguments.get("callbacks") or []))))
+    callbacks: Final = tuple(dict.fromkeys(utils.get_dynamic_callbacks(cast(list, arguments.get("callbacks")))))
     success: Final = tuple(dict.fromkeys((*callbacks, *cast(list, arguments.get("success_callback") or []))))
     failure: Final = tuple(dict.fromkeys((*callbacks, *cast(list, arguments.get("failure_callback") or []))))
+    configured: Final = tuple(
+        dict.fromkeys(
+            (
+                *litellm.input_callback,
+                *litellm.success_callback,
+                *litellm.failure_callback,
+                *litellm._async_success_callback,
+                *litellm._async_failure_callback,
+                *success,
+                *failure,
+            )
+        )
+    )
+    uninitialized: Final = [
+        cb
+        for cb in configured
+        if isinstance(cb, str)
+        and (
+            cb not in litellm._known_custom_logger_compatible_callbacks
+            or cb in litellm.input_callback + litellm.success_callback + litellm.failure_callback
+        )
+        and cb not in (utils.callback_list or [])
+    ]
+    if uninitialized:
+        set_callbacks(uninitialized, function_id=arguments.get("id"))
+        utils.callback_list = list(dict.fromkeys((*(utils.callback_list or []), *uninitialized)))
+    if litellm_logging.customLogger is None:
+        set_callbacks([cb for cb in configured if callable(cb)], function_id=arguments.get("id"))
+    for event, registered, add_async in (
+        ("input", litellm.input_callback, litellm.logging_callback_manager.add_litellm_input_callback),
+        ("success", litellm.success_callback, litellm.logging_callback_manager.add_litellm_async_success_callback),
+        ("failure", litellm.failure_callback, litellm.logging_callback_manager.add_litellm_async_failure_callback),
+    ):
+        for cb in tuple(registered):
+            if coroutine_checker.is_async_callable(cb) or (event == "success" and cb in ("dynamodb", "openmeter")):
+                if cb not in getattr(litellm, f"_async_{event}_callback"):
+                    add_async(cb)
+                registered.remove(cb)
+            elif event != "input" and isinstance(cb, str) and cb in litellm._known_custom_logger_compatible_callbacks:
+                utils._add_custom_logger_callback_to_specific_event(cb, event)
+    for event, registered, add_sync in (
+        ("success", litellm._async_success_callback, litellm.logging_callback_manager.add_litellm_success_callback),
+        ("failure", litellm._async_failure_callback, litellm.logging_callback_manager.add_litellm_failure_callback),
+    ):
+        for cb in tuple(registered):
+            if callable(cb) and not isinstance(cb, CustomLogger) and not coroutine_checker.is_async_callable(cb):
+                if cb not in getattr(litellm, f"{event}_callback"):
+                    add_sync(cb)
+                registered.remove(cb)
     call_id: Final = str(arguments.get("litellm_call_id") or uuid4())
     logger: Final = Logging(
         model=str(arguments["model"]),
@@ -76,14 +128,27 @@ def initialize_logging(arguments: dict[str, object], asynchronous: bool) -> obje
         litellm_call_id=call_id,
         function_id=str(arguments.get("id") or ""),
         litellm_trace_id=cast(str | None, arguments.get("litellm_trace_id")),
-        dynamic_input_callbacks=[cb for cb in callbacks if cb not in litellm.input_callback],
-        dynamic_success_callbacks=[cb for cb in success if not inspect.iscoroutinefunction(cb)],
-        dynamic_async_success_callbacks=list(success),
-        dynamic_failure_callbacks=[cb for cb in failure if not inspect.iscoroutinefunction(cb)],
-        dynamic_async_failure_callbacks=list(failure),
+        dynamic_input_callbacks=[
+            cb for cb in callbacks if cb not in litellm.input_callback and not coroutine_checker.is_async_callable(cb)
+        ],
+        dynamic_success_callbacks=[
+            cb for cb in success if not coroutine_checker.is_async_callable(cb) and cb not in ("dynamodb", "s3")
+        ],
+        dynamic_async_success_callbacks=[
+            cb
+            for cb in success
+            if coroutine_checker.is_async_callable(cb) or isinstance(cb, CustomLogger) or cb in ("dynamodb", "s3")
+        ],
+        dynamic_failure_callbacks=[cb for cb in failure if not coroutine_checker.is_async_callable(cb)],
+        dynamic_async_failure_callbacks=[
+            cb for cb in failure if coroutine_checker.is_async_callable(cb) or isinstance(cb, CustomLogger)
+        ],
         kwargs=arguments,
         supports_correlation_logging=asynchronous,
     )
+    logger.dynamic_input_callbacks = [
+        cb for cb in dict.fromkeys(logger.dynamic_input_callbacks or []) if cb not in litellm.input_callback
+    ]
     arguments["litellm_call_id"] = call_id
     arguments["litellm_logging_obj"] = logger
     return logger
@@ -110,7 +175,13 @@ def invoke_terminal(
             _retained: Final = roots
             await logging.async_success_handler(value, start_time, end_time)
 
-        GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=run_async())
+        def enqueue() -> None:
+            GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=run_async())
+
+        if getattr(logging, "_defer_async_logging", False) is True:
+            logging._enqueue_deferred_logging = enqueue
+        else:
+            enqueue()
         return None
     if action == "sync_success_if_needed":
         if logging._should_run_sync_callbacks_for_async_calls():

tests/test_litellm/rust_bridge/native_route_wheel_test.py

@@ -10,6 +10,8 @@ import sys
 import tempfile
 import threading
 import zipfile
+from collections.abc import Awaitable
+from contextvars import ContextVar
 from dataclasses import dataclass
 from http.client import HTTPMessage
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
@@ -200,10 +202,35 @@ def initialize_ocr_logging(arguments: dict[str, object], asynchronous: bool) ->
 
 def invoke_ocr_terminal(
     action: str, roots: object, logger: OCRLogging, value: object, start: object, end: object
-) -> None:
+) -> Awaitable[None] | None:
     assert action in {"sync_success", "async_success", "sync_success_if_needed", "sync_failure", "async_failure"}
     assert isinstance(roots, tuple) and roots[0] is logger.arguments
     assert logger.calls == ("update", "pre")
+    if action == "async_failure":
+        return observe_ocr_failure({}, value, "aocr")
+    return None
+
+
+async def pre_ocr_deployment(arguments: dict[str, object], call_type: str) -> dict[str, object]:
+    assert call_type == "aocr"
+    return arguments
+
+
+async def post_ocr_deployment(arguments: dict[str, object], response: object, call_type: str) -> object:
+    assert call_type == "aocr"
+    return response
+
+
+async def observe_ocr_failure(arguments: dict[str, object], error: object, call_type: str) -> None:
+    assert call_type == "aocr"
+
+
+def restore_ocr_context(logger: object) -> None:
+    pass
+
+
+class WheelCallTypes:
+    aocr = "aocr"
 
 
 def route_kwargs(route: str, api_base: str, outcome: str) -> dict[str, object]:
@@ -369,11 +396,20 @@ def exercise_routes(native_path: Path, api_base: str) -> object:
     ocr_bridge: Final = ModuleType("litellm.rust_bridge.ocr")
     ocr_bridge.initialize_logging = initialize_ocr_logging
     ocr_bridge.invoke_terminal = invoke_ocr_terminal
+    utils: Final = ModuleType("litellm.utils")
+    utils.is_internal_call = ContextVar("wheel_internal_call", default=False)
+    utils.async_pre_call_deployment_hook = pre_ocr_deployment
+    utils.async_post_call_success_deployment_hook = post_ocr_deployment
+    utils.async_post_call_failure_deployment_hook = observe_ocr_failure
+    utils._restore_correlation_context_if_supported = restore_ocr_context
+    types_utils: Final = ModuleType("litellm.types.utils")
+    types_utils.CallTypes = WheelCallTypes
     packages: Final = {
         name: ModuleType(name)
         for name in (
             "litellm",
             "litellm.rust_bridge",
+            "litellm.types",
             "litellm.llms",
             "litellm.llms.base_llm",
             "litellm.llms.base_llm.ocr",
@@ -381,7 +417,8 @@ def exercise_routes(native_path: Path, api_base: str) -> object:
     }
     with patch.dict(
         sys.modules,
-        packages | {module.__name__: module for module in (transformation, exceptions, httpx, ocr_bridge)},
+        packages
+        | {module.__name__: module for module in (transformation, exceptions, httpx, ocr_bridge, utils, types_utils)},
     ):
         exercise_sync(native, api_base)
         asyncio.run(exercise_async(native, api_base))