wip

litellm-rust/Cargo.lock

@@ -1482,12 +1482,15 @@ dependencies = [
 name = "litellm-python-interop"
 version = "0.1.0"
 dependencies = [
+ "futures-util",
  "pyo3",
+ "pyo3-async-runtimes",
  "pythonize",
  "rstest",
  "serde",
  "serde_json",
  "serial_test",
+ "tokio",
 ]
 
 [[package]]

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

@@ -13,7 +13,7 @@ dials OpenAI upstream, and splices the two sockets frame-by-frame.
 | litellm-core | The LiteLLM SDK in Rust — per-route entrypoints (`messages::messages()`) that resolve the provider, transform, and make the call; plus types, provider transforms, and the router. |
 | litellm-config | Config-loading boundary. Returns resolved deployments and optionally delegates loading to Python. |
 | litellm-ai-gateway | The Axum server (behind the `server` feature) and WebSocket hosts. Translates HTTP/WS to core entrypoints; no provider handlers. |
-| litellm-python-interop | Domain-neutral PyO3 foundation for GIL handling and typed Python/Serde conversion. |
+| litellm-python-interop | Domain-neutral PyO3 foundation: typed Python/Serde conversion, retained callbacks, and sync/async Python↔Tokio execution. |
 | litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. |
 
 Dependency direction is acyclic: config depends on core, the gateway depends on config and core, and the Python bridge depends on the domain layers and Python interop.

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

@@ -2,6 +2,7 @@ use std::future::Future;
 use std::time::{Instant, SystemTime, UNIX_EPOCH};
 
 use crate::Error;
+use crate::lifecycle::ActionResult;
 
 pub mod types;
 
@@ -10,6 +11,39 @@ pub use types::{
     CallLifecycleTiming,
 };
 
+pub trait RequestPolicy<InitialReq, ProviderReq>: Send + Sync {
+    type PreCallFuture<'a>: Future<Output = ActionResult<InitialReq, Error>> + Send + 'a
+    where
+        Self: 'a,
+        InitialReq: 'a,
+        ProviderReq: 'a;
+
+    type DuringCallFuture<'a>: Future<Output = ActionResult<ProviderReq, Error>> + Send + 'a
+    where
+        Self: 'a,
+        InitialReq: 'a,
+        ProviderReq: 'a;
+
+    fn async_pre_call_hook<'a>(
+        &'a self,
+        context: &'a CallLifecycleContext,
+        request: InitialReq,
+    ) -> Self::PreCallFuture<'a>;
+
+    fn async_during_call_hook<'a>(
+        &'a self,
+        context: &'a CallLifecycleContext,
+        request: InitialReq,
+    ) -> Self::DuringCallFuture<'a>;
+}
+
+pub trait TerminalDispatcher: Send + Sync {
+    fn dispatch<'a>(
+        &'a self,
+        terminal: &'a crate::lifecycle::TerminalRecord,
+    ) -> crate::integrations::custom_logger::LogFuture<'a>;
+}
+
 pub trait CallLifecycleHooks<InitialReq, ProviderReq, Resp>: Send + Sync {
     type PreCallFuture<'a>: Future<Output = Result<InitialReq, Error>> + Send + 'a
     where
@@ -97,6 +131,21 @@ impl<'a> CallLifecycle<'a> {
         self.run(context, request, hooks, provider_call).await
     }
 
+    pub async fn run_result<InitialReq, ProviderReq, Resp, Hooks, ProviderCall, ProviderFuture>(
+        &self,
+        context: CallLifecycleContext,
+        request: InitialReq,
+        hooks: &Hooks,
+        provider_call: ProviderCall,
+    ) -> Result<Resp, Error>
+    where
+        Hooks: CallLifecycleHooks<InitialReq, ProviderReq, Resp>,
+        ProviderCall: FnOnce(ProviderReq) -> ProviderFuture,
+        ProviderFuture: Future<Output = Result<Resp, Error>>,
+    {
+        self.run(context, request, hooks, provider_call).await
+    }
+
     pub async fn run<InitialReq, ProviderReq, Resp, Hooks, ProviderCall, ProviderFuture>(
         &self,
         context: CallLifecycleContext,

litellm-rust/crates/core/src/integrations/custom_logger/mod.rs

@@ -96,6 +96,36 @@ impl CustomLoggerRunner {
     }
 }
 
+impl crate::call_lifecycle::TerminalDispatcher for CustomLoggerRunner {
+    fn dispatch<'a>(&'a self, terminal: &'a crate::lifecycle::TerminalRecord) -> LogFuture<'a> {
+        Box::pin(async move {
+            let details = ModelCallDetails::from(terminal);
+            let response = CallbackValue::new(
+                match &terminal.projection {
+                    crate::lifecycle::RouteProjection::Ocr { .. } => "ocr",
+                    crate::lifecycle::RouteProjection::Messages { .. } => "messages",
+                    crate::lifecycle::RouteProjection::ChatCompletions { .. } => "chat_completion",
+                    crate::lifecycle::RouteProjection::Audio { .. } => "audio",
+                    crate::lifecycle::RouteProjection::Realtime { .. } => "realtime",
+                    crate::lifecycle::RouteProjection::ResponsesWs { .. } => "responses_websocket",
+                },
+                terminal.projection.value().clone(),
+            );
+            match terminal.classification {
+                crate::lifecycle::TerminalClassification::Success => {
+                    self.async_log_success_event(&details, &response, terminal.timing)
+                        .await;
+                }
+                crate::lifecycle::TerminalClassification::Failure { .. } => {
+                    self.async_log_failure_event(&details, Some(&response), terminal.timing)
+                        .await;
+                }
+            }
+            Ok(())
+        })
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -303,6 +333,38 @@ mod tests {
         assert_eq!(report, CallbackDispatchReport::default());
     }
 
+    #[tokio::test]
+    async fn terminal_dispatcher_fans_out_shared_record() {
+        use crate::call_lifecycle::TerminalDispatcher;
+        use crate::integrations::types::Usage;
+        use crate::lifecycle::terminal::CostInputs;
+        use crate::lifecycle::{RouteProjection, TerminalClassification, TerminalRecord};
+
+        let logger = Arc::new(RecordingCustomLogger::default());
+        let runner = CustomLoggerRunner::new(vec![logger.clone()]);
+        let terminal = TerminalRecord {
+            call_id: "call-ocr".to_string(),
+            trace_id: Some("req-ocr".to_string()),
+            attempt: 1,
+            call_type: "ocr".to_string(),
+            model: "mistral-ocr-latest".to_string(),
+            provider: "mistral".to_string(),
+            timing: CallbackTiming::new(10.0, 11.5),
+            usage: Usage::default(),
+            cost_inputs: CostInputs::default(),
+            classification: TerminalClassification::Success,
+            projection: RouteProjection::Ocr {
+                value: json!({"pages": []}),
+            },
+        };
+
+        runner.dispatch(&terminal).await.expect("dispatch succeeds");
+
+        assert_eq!(logger.events().len(), 1);
+        assert_eq!(logger.events()[0].hook, "async_log_success_event");
+        assert_eq!(logger.events()[0].response_object.as_deref(), Some("ocr"));
+    }
+
     #[test]
     fn with_standard_logging_payload_keeps_top_level_fields_in_sync() {
         let details = ModelCallDetails::new("old-model", "old-provider", CallType::Completion)

litellm-rust/crates/core/src/integrations/custom_logger/types.rs

@@ -2,6 +2,7 @@ use std::collections::HashMap;
 use std::future::Future;
 use std::pin::Pin;
 
+use serde::Serialize;
 use serde_json::Value;
 
 use crate::integrations::types::{StandardLoggingMetadata, StandardLoggingPayload};
@@ -56,7 +57,7 @@ impl std::fmt::Display for CallType {
     }
 }
 
-#[derive(Clone, Copy, Debug, PartialEq)]
+#[derive(Clone, Copy, Debug, PartialEq, Serialize)]
 pub struct CallbackTiming {
     pub start_time: f64,
     pub end_time: f64,

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

@@ -13,7 +13,7 @@ use serde_json::Value;
 use std::collections::HashMap;
 
 /// Cumulative token usage for a realtime session.
-#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize)]
 pub struct Usage {
     pub prompt_tokens: u64,
     pub completion_tokens: u64,
@@ -65,7 +65,7 @@ pub struct StandardLoggingPayload {
 
 /// Cost-attribution keys. The replayer maps these into litellm_params.metadata,
 /// which the spend-logs builder reads to set user / team_id / organization_id.
-#[derive(Clone, Debug, Serialize, Default)]
+#[derive(Clone, Debug, Serialize, Default, PartialEq)]
 pub struct StandardLoggingMetadata {
     pub user_api_key_hash: Option<String>, // -> SpendLogs.api_key
     pub user_api_key_user_id: Option<String>, // -> SpendLogs.user

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

@@ -6,6 +6,7 @@ pub mod constants;
 pub mod error;
 pub mod http_utils;
 pub mod integrations;
+pub mod lifecycle;
 pub mod messages;
 #[cfg(any(feature = "observability", test))]
 pub mod observability;

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

@@ -0,0 +1,26 @@
+use serde::Serialize;
+
+use super::{ActionKind, Delivery, FailurePolicy};
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)]
+pub enum ResultPolicy {
+    Continue,
+    Replace,
+    Reject,
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)]
+pub enum Owner {
+    Core,
+    Route,
+    Host,
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)]
+pub struct ActionBinding {
+    pub kind: ActionKind,
+    pub delivery: Delivery,
+    pub on_result: ResultPolicy,
+    pub on_error: FailurePolicy,
+    pub owner: Owner,
+}

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

@@ -0,0 +1,64 @@
+use super::TerminalRecord;
+
+#[derive(Clone, Debug)]
+pub enum ExecutedCall<R, E> {
+    Success {
+        response: R,
+        terminal: TerminalRecord,
+    },
+    Failure {
+        error: E,
+        terminal: TerminalRecord,
+    },
+}
+
+impl<R, E> ExecutedCall<R, E> {
+    pub fn into_result(self) -> Result<R, E> {
+        match self {
+            Self::Success { response, .. } => Ok(response),
+            Self::Failure { error, .. } => Err(error),
+        }
+    }
+
+    pub fn terminal(&self) -> &TerminalRecord {
+        match self {
+            Self::Success { terminal, .. } | Self::Failure { terminal, .. } => terminal,
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::integrations::custom_logger::CallbackTiming;
+    use crate::integrations::types::Usage;
+    use crate::lifecycle::terminal::CostInputs;
+    use crate::lifecycle::{RouteProjection, TerminalClassification};
+    use serde_json::json;
+
+    #[test]
+    fn failure_retains_its_terminal_record() {
+        let call = ExecutedCall::<(), _>::Failure {
+            error: "provider failed",
+            terminal: TerminalRecord {
+                call_id: "call-1".to_string(),
+                trace_id: None,
+                attempt: 1,
+                call_type: "ocr".to_string(),
+                model: "model".to_string(),
+                provider: "provider".to_string(),
+                timing: CallbackTiming::new(1.0, 2.0),
+                usage: Usage::default(),
+                cost_inputs: CostInputs::default(),
+                classification: TerminalClassification::Failure {
+                    kind: "ProviderError".to_string(),
+                    message: "provider failed".to_string(),
+                },
+                projection: RouteProjection::Ocr { value: json!({}) },
+            },
+        };
+
+        assert_eq!(call.terminal().call_id, "call-1");
+        assert_eq!(call.into_result(), Err("provider failed"));
+    }
+}

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

@@ -0,0 +1,68 @@
+use std::marker::PhantomData;
+
+use super::ActionBinding;
+
+pub trait LifecycleRoute: Sized {
+    type Admission;
+    type Options;
+    type Context;
+    type Operation: Copy;
+    type Observation;
+    type Outcome;
+    type Transition;
+    type Error;
+    type Decline;
+    type State;
+
+    fn admit(
+        admission: &Self::Admission,
+        options: Self::Options,
+    ) -> Result<Result<Self::State, Self::Decline>, Self::Error>;
+
+    fn operation(state: &Self::State) -> Self::Operation;
+
+    fn advance(
+        state: &mut Self::State,
+        outcome: Self::Outcome,
+        observations: Self::Observation,
+    ) -> Result<Self::Transition, Self::Error>;
+
+    fn actions_for(operation: Self::Operation, context: &Self::Context)
+    -> &'static [ActionBinding];
+}
+
+#[derive(Debug)]
+pub struct Lifecycle<Route: LifecycleRoute> {
+    pub(crate) state: Route::State,
+    route: PhantomData<Route>,
+}
+
+impl<Route: LifecycleRoute> Lifecycle<Route> {
+    pub fn admit(
+        admission: &Route::Admission,
+        options: Route::Options,
+    ) -> Result<Result<Self, Route::Decline>, Route::Error> {
+        Route::admit(admission, options).map(|admission| {
+            admission.map(|state| Self {
+                state,
+                route: PhantomData,
+            })
+        })
+    }
+
+    pub fn operation(&self) -> Route::Operation {
+        Route::operation(&self.state)
+    }
+
+    pub fn advance(
+        &mut self,
+        outcome: Route::Outcome,
+        observations: Route::Observation,
+    ) -> Result<Route::Transition, Route::Error> {
+        Route::advance(&mut self.state, outcome, observations)
+    }
+
+    pub fn actions_for(&self, context: &Route::Context) -> &'static [ActionBinding] {
+        Route::actions_for(self.operation(), context)
+    }
+}

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

@@ -0,0 +1,12 @@
+pub mod action;
+pub mod executed;
+pub mod machine;
+pub mod ocr;
+pub mod terminal;
+pub mod types;
+
+pub use action::{ActionBinding, Owner, ResultPolicy};
+pub use executed::ExecutedCall;
+pub use machine::{Lifecycle, LifecycleRoute};
+pub use terminal::{RouteProjection, TerminalClassification, TerminalRecord};
+pub use types::{ActionKind, ActionResult, Delivery, ErrorDisposition, FailurePolicy, Outcome};

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

@@ -0,0 +1,264 @@
+use crate::Error;
+use crate::ocr::{OcrRequest, prepare};
+
+use super::{
+    ActionBinding, ActionKind, Delivery, ErrorDisposition, FailurePolicy, LifecycleRoute, Outcome,
+    Owner, ResultPolicy,
+};
+
+#[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 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 struct Transition {
+    pub operation: Operation,
+    pub error: ErrorDisposition,
+}
+
+#[derive(Debug)]
+pub struct OcrState {
+    operation: Operation,
+    outcome: Outcome,
+    asynchronous: bool,
+    internal_call: bool,
+    identity: Identity,
+}
+
+#[derive(Debug)]
+pub struct OcrRoute;
+
+pub type Lifecycle = super::Lifecycle<OcrRoute>;
+
+impl Lifecycle {
+    pub fn new(admission: &OcrRequest, options: Options) -> Result<NativeOutcome<Self>, Error> {
+        <super::Lifecycle<OcrRoute>>::admit(admission, options).map(|admission| match admission {
+            Ok(lifecycle) => NativeOutcome::Completed(lifecycle),
+            Err(decline) => NativeOutcome::Declined(decline),
+        })
+    }
+
+    pub fn identity(&self) -> &Identity {
+        &self.state.identity
+    }
+}
+
+impl LifecycleRoute for OcrRoute {
+    type Admission = OcrRequest;
+    type Options = Options;
+    type Context = Observations;
+    type Operation = Operation;
+    type Observation = Observations;
+    type Outcome = Outcome;
+    type Transition = Transition;
+    type Error = Error;
+    type Decline = Decline;
+    type State = OcrState;
+
+    fn admit(
+        admission: &Self::Admission,
+        options: Self::Options,
+    ) -> Result<Result<Self::State, Self::Decline>, Self::Error> {
+        match prepare::admission_capabilities(admission) {
+            Err(Error::Unsupported(reason)) => return Ok(Err(Decline(reason))),
+            Err(error) => return Err(error),
+            Ok(()) => {}
+        }
+        if options.credential_method == CredentialMethod::Acquisition {
+            return Ok(Err(Decline("OCR credential acquisition")));
+        }
+        let generated_call_id = options.call_id.is_none();
+        let call_id = options.call_id.unwrap_or_else(generate_call_id);
+        Ok(Ok(OcrState {
+            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,
+            },
+        }))
+    }
+
+    fn operation(state: &Self::State) -> Self::Operation {
+        state.operation
+    }
+
+    fn advance(
+        state: &mut Self::State,
+        outcome: Self::Outcome,
+        observations: Self::Observation,
+    ) -> Result<Self::Transition, Self::Error> {
+        use Operation::*;
+
+        if matches!(state.operation, Complete(_)) {
+            return Err(Error::InvalidRequest(
+                "OCR lifecycle is already complete".into(),
+            ));
+        }
+        let failure =
+            if observations.logger_available && !(state.asynchronous && state.internal_call) {
+                SyncFailure
+            } else {
+                Restore
+            };
+        let error = if outcome != Outcome::Success && state.operation != DeploymentFailure {
+            state.outcome = outcome;
+            ErrorDisposition::Replace
+        } else {
+            ErrorDisposition::Preserve
+        };
+        state.operation = match (state.operation, outcome) {
+            (Restore, _) => Complete(state.outcome),
+            (DeploymentFailure, _) => failure,
+            (_, Outcome::Abort) => Restore,
+            (SyncFailure | AsyncFailure, Outcome::Failure) => Restore,
+            (Prepare | Send, Outcome::Failure) if state.asynchronous => DeploymentFailure,
+            (_, Outcome::Failure) => failure,
+            (Setup, Outcome::Success) if state.asynchronous => DeploymentPre,
+            (Setup | DeploymentPre, Outcome::Success) => Prepare,
+            (Prepare, Outcome::Success) => Send,
+            (Send, Outcome::Success) if state.asynchronous => DeploymentSuccess,
+            (Send, Outcome::Success) => SyncSuccess,
+            (DeploymentSuccess, Outcome::Success) => {
+                if state.internal_call || observations.has_fallbacks {
+                    SyncSuccessIfNeeded
+                } else {
+                    AsyncSuccess
+                }
+            }
+            (AsyncSuccess, Outcome::Success) => SyncSuccessIfNeeded,
+            (SyncFailure, Outcome::Success) if state.asynchronous => AsyncFailure,
+            (SyncSuccess | SyncSuccessIfNeeded | SyncFailure | AsyncFailure, Outcome::Success) => {
+                Restore
+            }
+            (Complete(_), _) => unreachable!(),
+        };
+        Ok(Transition {
+            operation: state.operation,
+            error,
+        })
+    }
+
+    fn actions_for(
+        operation: Self::Operation,
+        _context: &Self::Context,
+    ) -> &'static [ActionBinding] {
+        match operation {
+            Operation::Prepare | Operation::Send => &PROVIDER_ACTION,
+            Operation::SyncFailure | Operation::AsyncFailure | Operation::DeploymentFailure => {
+                &FAILURE_ACTION
+            }
+            Operation::Restore => &RESTORE_ACTION,
+            Operation::Complete(_) => &[],
+            _ => &CALLBACK_ACTION,
+        }
+    }
+}
+
+const PROVIDER_ACTION: [ActionBinding; 1] = [ActionBinding {
+    kind: ActionKind::ProviderCall,
+    delivery: Delivery::InlineAwaited,
+    on_result: ResultPolicy::Replace,
+    on_error: FailurePolicy::Propagate,
+    owner: Owner::Core,
+}];
+
+const CALLBACK_ACTION: [ActionBinding; 1] = [ActionBinding {
+    kind: ActionKind::TerminalSuccess,
+    delivery: Delivery::InlineAwaited,
+    on_result: ResultPolicy::Continue,
+    on_error: FailurePolicy::RecordAndContinue,
+    owner: Owner::Route,
+}];
+
+const FAILURE_ACTION: [ActionBinding; 1] = [ActionBinding {
+    kind: ActionKind::TerminalFailure,
+    delivery: Delivery::InlineAwaited,
+    on_result: ResultPolicy::Continue,
+    on_error: FailurePolicy::PreserveOriginalFailure,
+    owner: Owner::Route,
+}];
+
+const RESTORE_ACTION: [ActionBinding; 1] = [ActionBinding {
+    kind: ActionKind::Restore,
+    delivery: Delivery::InlineDirect,
+    on_result: ResultPolicy::Continue,
+    on_error: FailurePolicy::Propagate,
+    owner: Owner::Core,
+}];
+
+fn generate_call_id() -> String {
+    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..]
+    )
+}

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

@@ -0,0 +1,156 @@
+use serde::Serialize;
+use serde_json::Value;
+
+use crate::integrations::custom_logger::{
+    CallType, CallbackTiming, LoggingError, ModelCallDetails,
+};
+use crate::integrations::types::{StandardLoggingMetadata, StandardLoggingPayload, Usage};
+
+#[derive(Clone, Debug, Default, PartialEq, Serialize)]
+pub struct CostInputs {
+    pub response_cost: f64,
+    pub metadata: StandardLoggingMetadata,
+}
+
+#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
+pub enum TerminalClassification {
+    Success,
+    Failure { kind: String, message: String },
+}
+
+#[derive(Clone, Debug, PartialEq, Serialize)]
+pub enum RouteProjection {
+    Ocr { value: Value },
+    Messages { value: Value },
+    ChatCompletions { value: Value },
+    Audio { value: Value },
+    Realtime { value: Value },
+    ResponsesWs { value: Value },
+}
+
+impl RouteProjection {
+    pub(crate) fn value(&self) -> &Value {
+        match self {
+            Self::Ocr { value }
+            | Self::Messages { value }
+            | Self::ChatCompletions { value }
+            | Self::Audio { value }
+            | Self::Realtime { value }
+            | Self::ResponsesWs { value } => value,
+        }
+    }
+}
+
+#[derive(Clone, Debug, PartialEq, Serialize)]
+pub struct TerminalRecord {
+    pub call_id: String,
+    pub trace_id: Option<String>,
+    pub attempt: u32,
+    pub call_type: String,
+    pub model: String,
+    pub provider: String,
+    pub timing: CallbackTiming,
+    pub usage: Usage,
+    pub cost_inputs: CostInputs,
+    pub classification: TerminalClassification,
+    pub projection: RouteProjection,
+}
+
+impl From<&TerminalRecord> for StandardLoggingPayload {
+    fn from(record: &TerminalRecord) -> Self {
+        Self {
+            id: record
+                .trace_id
+                .clone()
+                .unwrap_or_else(|| record.call_id.clone()),
+            litellm_call_id: record.call_id.clone(),
+            call_type: record.call_type.clone(),
+            model: record.model.clone(),
+            custom_llm_provider: record.provider.clone(),
+            response_cost: record.cost_inputs.response_cost,
+            prompt_tokens: record.usage.prompt_tokens,
+            completion_tokens: record.usage.completion_tokens,
+            total_tokens: record.usage.total_tokens,
+            start_time: record.timing.start_time,
+            end_time: record.timing.end_time,
+            stream: matches!(
+                record.projection,
+                RouteProjection::Realtime { .. } | RouteProjection::ResponsesWs { .. }
+            ),
+            metadata: record.cost_inputs.metadata.clone(),
+            messages: Some(record.projection.value().clone()),
+        }
+    }
+}
+
+impl From<&TerminalRecord> for ModelCallDetails {
+    fn from(record: &TerminalRecord) -> Self {
+        let details = Self::from_standard_logging_payload(record.into());
+        match &record.classification {
+            TerminalClassification::Success => details,
+            TerminalClassification::Failure { kind, message } => {
+                details.with_failure_error(LoggingError {
+                    kind: kind.clone(),
+                    message: message.clone(),
+                })
+            }
+        }
+    }
+}
+
+impl TerminalRecord {
+    pub fn call_type(&self) -> CallType {
+        CallType::from(self.call_type.as_str())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use serde_json::json;
+
+    fn record() -> TerminalRecord {
+        TerminalRecord {
+            call_id: "call-1".to_string(),
+            trace_id: Some("trace-1".to_string()),
+            attempt: 2,
+            call_type: "ocr".to_string(),
+            model: "mistral-ocr-latest".to_string(),
+            provider: "mistral".to_string(),
+            timing: CallbackTiming::new(10.0, 11.5),
+            usage: Usage {
+                prompt_tokens: 3,
+                completion_tokens: 4,
+                total_tokens: 7,
+            },
+            cost_inputs: CostInputs {
+                response_cost: 0.25,
+                metadata: StandardLoggingMetadata {
+                    user_api_key_user_id: Some("user-1".to_string()),
+                    ..Default::default()
+                },
+            },
+            classification: TerminalClassification::Success,
+            projection: RouteProjection::Ocr {
+                value: json!({"pages": [{"markdown": "ok"}]}),
+            },
+        }
+    }
+
+    #[test]
+    fn terminal_record_projects_existing_logging_payload() {
+        let record = record();
+        let payload = StandardLoggingPayload::from(&record);
+        let details = ModelCallDetails::from(&record);
+
+        assert_eq!(payload.id, "trace-1");
+        assert_eq!(payload.litellm_call_id, "call-1");
+        assert_eq!(payload.prompt_tokens, 3);
+        assert_eq!(payload.response_cost, 0.25);
+        assert_eq!(payload.messages, Some(record.projection.value().clone()));
+        assert_eq!(
+            details.standard_logging_payload.unwrap().litellm_call_id,
+            "call-1"
+        );
+    }
+}

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

@@ -0,0 +1,56 @@
+use serde::Serialize;
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)]
+pub enum ActionKind {
+    RequestPolicy,
+    ProviderCall,
+    Deployment,
+    TerminalSuccess,
+    TerminalFailure,
+    Restore,
+}
+
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub enum ActionResult<T, E> {
+    Continue(T),
+    Replace(T),
+    Reject(E),
+}
+
+impl<T, E> ActionResult<T, E> {
+    pub fn into_result(self) -> Result<T, E> {
+        match self {
+            Self::Continue(value) | Self::Replace(value) => Ok(value),
+            Self::Reject(error) => Err(error),
+        }
+    }
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)]
+pub enum FailurePolicy {
+    Propagate,
+    RecordAndContinue,
+    PreserveOriginalFailure,
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)]
+pub enum Delivery {
+    InlineDirect,
+    InlineAwaited,
+    BlockingWorker,
+    BackgroundTask,
+    Deferred,
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)]
+pub enum Outcome {
+    Success,
+    Failure,
+    Abort,
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)]
+pub enum ErrorDisposition {
+    Preserve,
+    Replace,
+}

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

@@ -1,201 +1,10 @@
+pub use crate::lifecycle::ocr::*;
+pub use crate::lifecycle::{ErrorDisposition, Outcome};
+
+#[cfg(test)]
 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)]
+use crate::ocr::{OcrRequest, prepare};
 
 #[cfg(test)]
 mod tests {

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

@@ -38,7 +38,7 @@ fn request_config(
     Ok((provider, config))
 }
 
-pub(super) fn admission_capabilities(request: &OcrRequest) -> Result<(), Error> {
+pub(crate) fn admission_capabilities(request: &OcrRequest) -> Result<(), Error> {
     check_admission_capabilities(request, &|key| std::env::var(key).ok())
 }
 

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

@@ -4,7 +4,7 @@ litellm-python-bridge is the PyO3 cdylib exposing LiteLLM Rust APIs to the Pytho
 - Keep it thin: no business logic, no transforms, no I/O orchestration; just marshal in/out and call the core entrypoint
 - One stable method per top-level route (`ocr`/`aocr`, `messages`/`amessages`, ...); do not add per-provider helpers
 - Provider dispatch lives in `litellm-core`, never here
-- Put domain-neutral Python/Serde conversion and GIL primitives in `litellm-python-interop`
+- Put domain-neutral Python/Serde conversion, retained callbacks, and Python↔Tokio execution in `litellm-python-interop`
 - Two logical parts, kept as modules: the domain adapter (`src/routes/*`, `marshal`, `errors`) and the binding artifact (`#[pymodule]`, `#[pyfunction]`, registration in `lib.rs`)
 - Data handling: do not log OCR payloads or provider responses; avoid copying large payloads; sanitize errors before they cross the boundary
 - Tests: `cargo test --workspace` compiles here; Python tests cover disabled, enabled, and module-missing fallback for every exposed route

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

@@ -21,7 +21,6 @@ trace-parity = [
 ]
 
 [dependencies]
-futures-util.workspace = true
 tracing = { workspace = true, optional = true }
 litellm-core = { workspace = true, features = ["bedrock-auth"] }
 litellm-ai-gateway = { workspace = true, default-features = false }
@@ -30,10 +29,11 @@ pyo3.workspace = true
 pyo3-async-runtimes.workspace = true
 serde.workspace = true
 serde_json.workspace = true
-tokio.workspace = true
 
 [dev-dependencies]
 criterion = "0.8.2"
+futures-util.workspace = true
+tokio.workspace = true
 tokio-tungstenite.workspace = true
 tracing.workspace = true
 

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

@@ -1,13 +1,4 @@
-use litellm_python_interop::release_count;
 use pyo3::prelude::*;
-use pyo3::types::PyDict;
-
-#[pyfunction]
-fn gil_stats(py: Python<'_>) -> PyResult<Py<PyAny>> {
-    let stats = PyDict::new(py);
-    stats.set_item("releases", release_count())?;
-    Ok(stats.into_any().unbind())
-}
 
 #[cfg(feature = "panic-test")]
 #[pyfunction]
@@ -15,9 +6,8 @@ fn _panic_for_test() {
     panic!("intentional PyO3 panic smoke test");
 }
 
-pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
-    module.add_function(wrap_pyfunction!(gil_stats, module)?)?;
+pub(crate) fn register(_module: &Bound<'_, PyModule>) -> PyResult<()> {
     #[cfg(feature = "panic-test")]
-    module.add_function(wrap_pyfunction!(_panic_for_test, module)?)?;
+    _module.add_function(wrap_pyfunction!(_panic_for_test, _module)?)?;
     Ok(())
 }

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

@@ -1,6 +1,5 @@
 mod diagnostics;
 mod errors;
-mod execution;
 #[cfg(feature = "trace-parity")]
 mod function_trace;
 mod marshal;
@@ -106,7 +105,6 @@ mod tests {
                 "chat_completions",
                 "achat_completions",
                 "ResponsesWebSocketConnection",
-                "gil_stats",
             ];
 
             let public_names: Vec<String> = module

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

@@ -31,7 +31,7 @@ macro_rules! bridge_route {
                 $($required_name,)*
                 $($optional_name),*
             })?;
-            $crate::execution::run_sync(py, future, $map_error)
+            litellm_python_interop::run_sync(py, future, $map_error)
         }
 
         #[pyfunction]
@@ -46,7 +46,7 @@ macro_rules! bridge_route {
                 $($required_name,)*
                 $($optional_name),*
             })?;
-            $crate::execution::run_async(py, future, $map_error)
+            litellm_python_interop::run_async(py, future, $map_error)
         }
 
         pub(super) fn register(
@@ -75,7 +75,7 @@ macro_rules! bridge_route {
                     $($required_name,)*
                     $($optional_name),*
                 })?;
-                $crate::execution::run_sync(
+                litellm_python_interop::run_sync(
                     py,
                     $crate::function_trace::capture(future),
                     $map_error,
@@ -94,7 +94,7 @@ macro_rules! bridge_route {
                     $($required_name,)*
                     $($optional_name),*
                 })?;
-                $crate::execution::run_async(
+                litellm_python_interop::run_async(
                     py,
                     $crate::function_trace::capture(future),
                     $map_error,

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

@@ -2,6 +2,7 @@ use pyo3::prelude::*;
 use serde_json::Value;
 
 use crate::errors::core_error_to_pyerr;
+use litellm_python_interop::run_async;
 
 #[pyfunction]
 fn gateway_messages<'py>(
@@ -17,7 +18,7 @@ fn gateway_messages<'py>(
         api_base,
         body,
     );
-    crate::execution::run_async(
+    run_async(
         py,
         crate::function_trace::capture(future),
         core_error_to_pyerr,

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

@@ -8,18 +8,16 @@ use litellm_core::ocr::lifecycle::{
 };
 use litellm_core::ocr::types::{OcrDocumentProjection, OcrRequest, PreparedOcr};
 use litellm_core::routing_utils::provider::get_custom_llm_provider;
-use litellm_python_interop::{
-    InvocationMode, InvocationOutcome, PreparedCall, Pythonized, from_py, to_py,
-};
+use litellm_python_interop::{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, PyTuple};
+use pyo3::types::PyDict;
 use serde_json::Value;
 
 use crate::errors::core_error_to_pyerr;
-use crate::execution::{run_async_value, run_sync_value};
+use litellm_python_interop::{run_async_value, run_sync_value};
 
 #[pyclass]
 struct OcrState {
@@ -305,32 +303,24 @@ fn invoke(
         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),
+    let (method, awaiting) = match operation {
+        Operation::Setup => ("setup", false),
+        Operation::DeploymentPre => ("deployment_pre", true),
+        Operation::Prepare => ("prepare", false),
+        Operation::Send if asynchronous => ("send", true),
+        Operation::Send => ("send_sync", false),
+        Operation::DeploymentSuccess => ("deployment_success", true),
+        Operation::DeploymentFailure => ("deployment_failure", true),
+        Operation::SyncSuccess => ("sync_success", false),
+        Operation::AsyncSuccess => ("async_success", false),
+        Operation::SyncSuccessIfNeeded => ("sync_success_if_needed", false),
+        Operation::SyncFailure => ("sync_failure", false),
+        Operation::AsyncFailure => ("async_failure", true),
+        Operation::Restore => ("restore", false),
         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)),
-    }
+    let value = host.getattr(py, method)?.call0(py)?;
+    Ok((awaiting, value))
 }
 
 #[pyfunction]

litellm-rust/crates/python-bridge/tests/marshal_boundary.rs

@@ -1,49 +0,0 @@
-use std::fs;
-use std::path::{Path, PathBuf};
-
-const DISALLOWED_OUTSIDE_INTEROP: &[&str] = &[
-    "py.import(\"json\")",
-    "pythonize::",
-    "serde_json::to_string",
-    "serde_json::from_str",
-];
-
-fn source_root() -> PathBuf {
-    Path::new(env!("CARGO_MANIFEST_DIR")).join("src")
-}
-
-fn rust_sources(directory: &Path) -> Vec<PathBuf> {
-    fs::read_dir(directory)
-        .expect("bridge source directory should be readable")
-        .map(|entry| {
-            entry
-                .expect("bridge source entry should be readable")
-                .path()
-        })
-        .flat_map(|path| {
-            if path.is_dir() {
-                rust_sources(&path)
-            } else if path.extension().is_some_and(|extension| extension == "rs") {
-                vec![path]
-            } else {
-                Vec::new()
-            }
-        })
-        .collect()
-}
-
-#[test]
-fn serialization_uses_the_interop_boundary() {
-    let root = source_root();
-
-    for path in rust_sources(&root) {
-        let source = fs::read_to_string(&path).expect("bridge source should be readable");
-        for disallowed in DISALLOWED_OUTSIDE_INTEROP {
-            assert!(
-                !source.contains(disallowed),
-                "{} bypasses litellm-python-interop with `{disallowed}`",
-                path.display()
-            );
-        }
-    }
-}

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

@@ -1,5 +1,6 @@
 litellm-python-interop is the domain-neutral PyO3 foundation.
 
-- Owns generic Python/Serde conversion and interpreter primitives (`gil`, `marshal`)
+- Owns generic Python/Serde conversion (`marshal`) and Python↔Tokio execution (`execution`)
 - Depends on PyO3 but no LiteLLM domain crate; no route types, no API registration, no cdylib
+- `execution` is generic over the caller's error type (`map_error: fn(E) -> PyErr`); the bridge passes its own mapper
 - Keep it free of `litellm-core`, `OcrRequest`, `CallServices`, LiteLLM exceptions or `_native` surface

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

@@ -10,14 +10,13 @@ autotests = false
 name = "synthetic"
 path = "tests/synthetic/mod.rs"
 
-[[test]]
-name = "integration"
-path = "tests/integration/mod.rs"
-
 [dependencies]
+futures-util.workspace = true
 pyo3.workspace = true
+pyo3-async-runtimes.workspace = true
 pythonize.workspace = true
 serde.workspace = true
+tokio.workspace = true
 
 [dev-dependencies]
 rstest.workspace = true

litellm-rust/crates/python-interop/src/callback.rs

@@ -1,113 +0,0 @@
-use pyo3::prelude::*;
-use pyo3::pyclass::{PyTraverseError, PyVisit};
-use pyo3::sync::PyOnceLock;
-use pyo3::types::{PyDict, PyTuple};
-
-use crate::constants::{
-    AWAIT_ADAPTER_FILENAME, AWAIT_ADAPTER_FUNCTION, AWAIT_ADAPTER_MODULE, AWAIT_ADAPTER_SOURCE,
-};
-
-static AWAIT_ADAPTER: PyOnceLock<Py<PyAny>> = PyOnceLock::new();
-
-/// How a retained callback is bound to its caller.
-///
-/// `Direct` mirrors `callable(*args, **kwargs)`: a coroutine returned by the
-/// callable is handed back untouched and never awaited. `Await` mirrors
-/// `await callable(*args, **kwargs)` inline in the caller's task.
-#[derive(Clone, Copy, Debug, PartialEq, Eq)]
-pub enum InvocationMode {
-    Direct,
-    Await,
-}
-
-/// Result of [`PreparedCall::invoke`].
-///
-/// `Awaitable` carries an adapter coroutine that has not yet called the
-/// callback. The callback runs, and any exception it raises surfaces, only
-/// when Python drives that coroutine.
-#[derive(Debug)]
-pub enum InvocationOutcome {
-    Returned(Py<PyAny>),
-    Awaitable(Py<PyAny>),
-}
-
-/// A callback plus its arguments, retained as owning Python references.
-///
-/// Arguments are passed to the callback by identity, never copied, so the
-/// callback observes and may mutate the caller's objects. Dropping the value
-/// releases the references; a Python-visible owner must also expose them to
-/// the cycle collector via [`PreparedCall::traverse`].
-pub struct PreparedCall {
-    mode: InvocationMode,
-    callable: Py<PyAny>,
-    positional: Py<PyTuple>,
-    keywords: Option<Py<PyDict>>,
-}
-
-impl PreparedCall {
-    pub fn new(
-        mode: InvocationMode,
-        callable: Py<PyAny>,
-        positional: Py<PyTuple>,
-        keywords: Option<Py<PyDict>>,
-    ) -> Self {
-        Self {
-            mode,
-            callable,
-            positional,
-            keywords,
-        }
-    }
-
-    pub fn invoke(&self, py: Python<'_>) -> PyResult<InvocationOutcome> {
-        match self.mode {
-            InvocationMode::Direct => self
-                .callable
-                .call(
-                    py,
-                    self.positional.bind(py),
-                    self.keywords.as_ref().map(|kwargs| kwargs.bind(py)),
-                )
-                .map(InvocationOutcome::Returned),
-            InvocationMode::Await => await_adapter(py)?
-                .call1(py, (&self.callable, &self.positional, &self.keywords))
-                .map(InvocationOutcome::Awaitable),
-        }
-    }
-
-    pub fn clone_ref(&self, py: Python<'_>) -> Self {
-        Self {
-            mode: self.mode,
-            callable: self.callable.clone_ref(py),
-            positional: self.positional.clone_ref(py),
-            keywords: self.keywords.as_ref().map(|value| value.clone_ref(py)),
-        }
-    }
-
-    pub fn traverse(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
-        visit.call(&self.callable)?;
-        visit.call(&self.positional)?;
-        if let Some(keywords) = &self.keywords {
-            visit.call(keywords)?;
-        }
-        Ok(())
-    }
-}
-
-/// Compiling the adapter runs Python, which may re-enter this function through
-/// audit hooks. `PyOnceLock` forbids re-entrant initialization, so compile
-/// first and only publish a finished adapter into the cell.
-fn await_adapter(py: Python<'_>) -> PyResult<&Py<PyAny>> {
-    if let Some(adapter) = AWAIT_ADAPTER.get(py) {
-        return Ok(adapter);
-    }
-    let compiled = PyModule::from_code(
-        py,
-        AWAIT_ADAPTER_SOURCE,
-        AWAIT_ADAPTER_FILENAME,
-        AWAIT_ADAPTER_MODULE,
-    )?
-    .getattr(AWAIT_ADAPTER_FUNCTION)?
-    .unbind();
-    Ok(AWAIT_ADAPTER.get_or_init(py, || compiled))
-}

litellm-rust/crates/python-interop/src/constants.rs

@@ -1,18 +0,0 @@
-use std::ffi::CStr;
-
-/// Python source of the coroutine adapter that awaits a retained callback
-/// inline in the caller's task. It is compiled once per interpreter.
-pub(crate) const AWAIT_ADAPTER_SOURCE: &CStr =
-    c"async def invoke_awaited(callable, positional, keywords):
-    if keywords is None:
-        return await callable(*positional)
-    return await callable(*positional, **keywords)
-";
-
-/// Filename recorded on the adapter's code object. Visible to Python
-/// `compile` audit hooks and tracebacks.
-pub const AWAIT_ADAPTER_FILENAME: &CStr = c"retained_callback.py";
-
-pub(crate) const AWAIT_ADAPTER_MODULE: &CStr = c"_retained_callback";
-
-pub(crate) const AWAIT_ADAPTER_FUNCTION: &str = "invoke_awaited";

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

@@ -3,35 +3,34 @@ use std::panic::AssertUnwindSafe;
 use std::time::Duration;
 
 use futures_util::FutureExt;
-use litellm_core::error::Error;
-use litellm_python_interop::{Pythonized, panic_to_pyerr, release_gil};
 use pyo3::exceptions::PyRuntimeError;
 use pyo3::prelude::*;
 use serde::Serialize;
 use tokio::runtime::{Handle, Runtime};
 use tokio::time::{self, MissedTickBehavior};
 
-pub(crate) fn run_sync<T, F>(
+use crate::marshal::Pythonized;
+use crate::marshal::panic_to_pyerr;
+
+pub fn run_sync<T, E, F>(
     py: Python<'_>,
     future: F,
-    map_error: fn(Error) -> PyErr,
+    map_error: fn(E) -> PyErr,
 ) -> PyResult<Py<PyAny>>
 where
     T: Serialize + Send + 'static,
-    F: Future<Output = Result<T, Error>> + Send + 'static,
+    E: Send + 'static,
+    F: Future<Output = Result<T, E>> + Send + 'static,
 {
     let result = run_sync_value(py, future, map_error)?;
     Pythonized(result).into_pyobject(py).map(Bound::unbind)
 }
 
-pub(crate) fn run_sync_value<T, F>(
-    py: Python<'_>,
-    future: F,
-    map_error: fn(Error) -> PyErr,
-) -> PyResult<T>
+pub fn run_sync_value<T, E, F>(py: Python<'_>, future: F, map_error: fn(E) -> PyErr) -> PyResult<T>
 where
     T: Send + 'static,
-    F: Future<Output = Result<T, Error>> + Send + 'static,
+    E: Send + 'static,
+    F: Future<Output = Result<T, E>> + Send + 'static,
 {
     run_sync_value_on(
         py,
@@ -41,15 +40,16 @@ where
     )
 }
 
-fn run_sync_value_on<T, F>(
+fn run_sync_value_on<T, E, F>(
     py: Python<'_>,
     runtime: &Runtime,
     future: F,
-    map_error: fn(Error) -> PyErr,
+    map_error: fn(E) -> PyErr,
 ) -> PyResult<T>
 where
     T: Send + 'static,
-    F: Future<Output = Result<T, Error>> + Send + 'static,
+    E: Send + 'static,
+    F: Future<Output = Result<T, E>> + Send + 'static,
 {
     if Handle::try_current().is_ok() {
         return Err(PyRuntimeError::new_err(
@@ -57,18 +57,19 @@ where
         ));
     }
 
-    let result = release_gil(py, move || runtime.block_on(wait_for_sync_result(future)))?;
+    let result = py.detach(move || runtime.block_on(wait_for_sync_result(future)))?;
     map_core_result(result, map_error)
 }
 
-pub(crate) fn run_async<T, F>(
+pub fn run_async<T, E, F>(
     py: Python<'_>,
     future: F,
-    map_error: fn(Error) -> PyErr,
+    map_error: fn(E) -> PyErr,
 ) -> PyResult<Bound<'_, PyAny>>
 where
     T: Serialize + Send + 'static,
-    F: Future<Output = Result<T, Error>> + Send + 'static,
+    E: Send + 'static,
+    F: Future<Output = Result<T, E>> + Send + 'static,
 {
     pyo3_async_runtimes::tokio::future_into_py(py, async move {
         let result = catch_future_panic(future).await?;
@@ -77,16 +78,20 @@ where
     })
 }
 
-pub(crate) async fn run_async_value<T, F>(future: F, map_error: fn(Error) -> PyErr) -> PyResult<T>
+pub async fn run_async_value<T, E, F>(future: F, map_error: fn(E) -> PyErr) -> PyResult<T>
 where
     T: Send + 'static,
-    F: Future<Output = Result<T, Error>> + Send + 'static,
+    E: Send + 'static,
+    F: Future<Output = Result<T, E>> + Send + 'static,
 {
     let result = catch_future_panic(future).await?;
     map_core_result(result, map_error)
 }
 
-fn map_core_result<T>(result: Result<T, Error>, map_error: fn(Error) -> PyErr) -> PyResult<T> {
+fn map_core_result<T, E>(result: Result<T, E>, map_error: fn(E) -> PyErr) -> PyResult<T>
+where
+    E: Send + 'static,
+{
     match result {
         Ok(value) => Ok(value),
         Err(error) => Err(
@@ -96,9 +101,9 @@ fn map_core_result<T>(result: Result<T, Error>, map_error: fn(Error) -> PyErr) -
     }
 }
 
-async fn catch_future_panic<T, F>(future: F) -> PyResult<Result<T, Error>>
+async fn catch_future_panic<T, E, F>(future: F) -> PyResult<Result<T, E>>
 where
-    F: Future<Output = Result<T, Error>>,
+    F: Future<Output = Result<T, E>>,
 {
     AssertUnwindSafe(future)
         .catch_unwind()
@@ -106,9 +111,9 @@ where
         .map_err(panic_to_pyerr)
 }
 
-async fn wait_for_sync_result<T, F>(future: F) -> PyResult<Result<T, Error>>
+async fn wait_for_sync_result<T, E, F>(future: F) -> PyResult<Result<T, E>>
 where
-    F: Future<Output = Result<T, Error>>,
+    F: Future<Output = Result<T, E>>,
 {
     let future = catch_future_panic(future);
     tokio::pin!(future);
@@ -142,11 +147,18 @@ mod tests {
 
     use super::*;
 
-    fn runtime_error(error: Error) -> PyErr {
-        PyRuntimeError::new_err(error.to_string())
+    #[derive(Debug)]
+    enum TestError {
+        InvalidRequest(&'static str),
     }
 
-    fn panicking_error_mapper(_error: Error) -> PyErr {
+    fn runtime_error(error: TestError) -> PyErr {
+        PyRuntimeError::new_err(match error {
+            TestError::InvalidRequest(message) => message,
+        })
+    }
+
+    fn panicking_error_mapper(_error: TestError) -> PyErr {
         panic!("error mapper panicked")
     }
 
@@ -258,7 +270,7 @@ mod tests {
 
         let error = runtime.block_on(async {
             Python::attach(|py| {
-                run_sync::<bool, _>(py, async { Ok(true) }, runtime_error)
+                run_sync::<bool, _, _>(py, async { Ok(true) }, runtime_error)
                     .expect_err("sync route should reject a nested Tokio runtime")
             })
         });
@@ -294,9 +306,9 @@ mod tests {
     fn sync_runner_maps_a_panicked_future() {
         Python::initialize();
         Python::attach(|py| {
-            let error = run_sync::<bool, _>(
+            let error = run_sync::<bool, _, _>(
                 py,
-                poll_fn(|_| -> Poll<Result<bool, Error>> { panic!("route future panicked") }),
+                poll_fn(|_| -> Poll<Result<bool, TestError>> { panic!("route future panicked") }),
                 runtime_error,
             )
             .expect_err("panicked route should become a Python exception");
@@ -310,9 +322,9 @@ mod tests {
     fn sync_runner_maps_a_panicked_error_mapper() {
         Python::initialize();
         Python::attach(|py| {
-            let error = run_sync::<bool, _>(
+            let error = run_sync::<bool, _, _>(
                 py,
-                async { Err(Error::InvalidRequest("invalid".to_string())) },
+                async { Err(TestError::InvalidRequest("invalid")) },
                 panicking_error_mapper,
             )
             .expect_err("panicked mapper should become a Python exception");
@@ -337,7 +349,7 @@ mod tests {
     #[test]
     fn sync_runner_supports_concurrent_callers_on_the_shared_runtime() {
         Python::initialize();
-        let barrier = Arc::new(tokio::sync::Barrier::new(2));
+        let barrier = Arc::new(std::sync::Barrier::new(2));
         let callers: Vec<_> = (0..2)
             .map(|_| {
                 let barrier = Arc::clone(&barrier);
@@ -348,9 +360,11 @@ mod tests {
                             run_sync(
                                 py,
                                 async move {
-                                    Ok(tokio::time::timeout(Duration::from_secs(2), barrier.wait())
+                                    let barrier = barrier;
+                                    tokio::task::spawn_blocking(move || barrier.wait())
                                         .await
-                                        .is_ok())
+                                        .map(|_| true)
+                                        .map_err(|_| TestError::InvalidRequest("join"))
                                 },
                                 runtime_error,
                             ),

litellm-rust/crates/python-interop/src/gil.rs

@@ -1,21 +0,0 @@
-use std::sync::atomic::{AtomicU64, Ordering};
-
-use pyo3::prelude::*;
-
-static GIL_RELEASES: AtomicU64 = AtomicU64::new(0);
-
-/// Runs work detached from the interpreter and records the release.
-///
-/// `f` must not access Python state while the interpreter is detached.
-pub fn release_gil<T, F>(py: Python<'_>, f: F) -> T
-where
-    F: FnOnce() -> T + Send,
-    T: Send,
-{
-    GIL_RELEASES.fetch_add(1, Ordering::Relaxed);
-    py.detach(f)
-}
-
-pub fn release_count() -> u64 {
-    GIL_RELEASES.load(Ordering::Relaxed)
-}

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

@@ -1,9 +1,5 @@
-mod callback;
-mod constants;
-mod gil;
+mod execution;
 mod marshal;
 
-pub use callback::{InvocationMode, InvocationOutcome, PreparedCall};
-pub use constants::AWAIT_ADAPTER_FILENAME;
-pub use gil::{release_count, release_gil};
+pub use execution::{run_async, run_async_value, run_sync, run_sync_value};
 pub use marshal::{Pythonized, from_py, panic_to_pyerr, to_py};

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

@@ -18,7 +18,8 @@ pub fn to_py<T>(py: Python<'_>, value: &T) -> PyResult<Py<PyAny>>
 where
     T: Serialize + ?Sized,
 {
-    pythonize::pythonize(py, value)
+    catch_unwind(AssertUnwindSafe(|| pythonize::pythonize(py, value)))
+        .map_err(panic_to_pyerr)?
         .map(Bound::unbind)
         .map_err(|error| PyValueError::new_err(error.to_string()))
 }

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

@@ -1,848 +0,0 @@
-import asyncio
-import atexit
-import contextvars
-import gc
-import json
-import threading
-import weakref
-from datetime import datetime
-from unittest import TestCase
-
-import litellm
-from litellm.integrations.custom_logger import CustomLogger
-from litellm.litellm_core_utils.core_helpers import safe_deep_copy
-from litellm.litellm_core_utils.litellm_logging import Logging
-from litellm.litellm_core_utils.logging_worker import LoggingWorker
-from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
-from litellm.types.utils import Delta, ModelResponse, ModelResponseStream, StreamingChoices, Usage
-
-
-def logger_for(
-    callbacks=(),
-    stream=False,
-    input_callbacks=(),
-    sync_callbacks=(),
-    failure_callbacks=(),
-    async_failure_callbacks=(),
-    call_type="acompletion",
-):
-    return Logging(
-        model="test",
-        messages=[{"role": "user", "content": "test"}],
-        stream=stream,
-        call_type=call_type,
-        start_time=datetime.now(),
-        litellm_call_id="retained-test",
-        function_id="retained-test",
-        dynamic_async_success_callbacks=list(callbacks),
-        dynamic_input_callbacks=list(input_callbacks),
-        dynamic_success_callbacks=list(sync_callbacks),
-        dynamic_failure_callbacks=list(failure_callbacks),
-        dynamic_async_failure_callbacks=list(async_failure_callbacks),
-    )
-
-
-def invoke_pre_call(owners, logger, additional):
-    owner = owners.prepare(logger.pre_call, (logger.messages, "test-key"), {"additional_args": additional})
-    try:
-        return owner.invoke()
-    finally:
-        owner.close()
-
-
-async def pre_call_identity_and_ignored_returns(owners):
-    saved = []
-    ignored = {"replacement": True}
-
-    class Retain(CustomLogger):
-        def log_pre_api_call(self, model, messages, kwargs):
-            saved.append((kwargs, messages))
-            return ignored
-
-    logger = logger_for(input_callbacks=[Retain(), Retain()])
-    details = logger.model_call_details
-    additional = {"headers": {"test": "header"}}
-    assert invoke_pre_call(owners, logger, additional) is None
-    assert len(saved) == 2
-    assert saved[0][0] is saved[1][0] is details
-    assert saved[0][1] is saved[1][1] is logger.messages is details["input"]
-    assert details["additional_args"] is additional
-    assert "replacement" not in details
-
-
-async def pre_call_mutations_visible_to_later_callbacks(owners):
-    saved, observed, order = [], [], []
-    metadata = {"secret": "private", "keep": []}
-    removed = object()
-
-    class Retain(CustomLogger):
-        def log_pre_api_call(self, model, messages, kwargs):
-            order.append("retain")
-            saved.append(kwargs)
-
-    class Mutate(CustomLogger):
-        def log_pre_api_call(self, model, messages, kwargs):
-            order.append("mutate")
-            kwargs["normalized"] = "normalized"
-            assert kwargs.pop("remove") is removed
-            kwargs["retained_metadata"]["secret"] = "masked"
-            return {"replacement": True}
-
-    class Observe(CustomLogger):
-        def log_pre_api_call(self, model, messages, kwargs):
-            order.append("observe")
-            observed.append((kwargs["normalized"], "remove" in kwargs, kwargs["retained_metadata"]["secret"]))
-
-    logger = logger_for(input_callbacks=[Retain(), Mutate(), Observe()])
-    details = logger.model_call_details
-    details.update(retained_metadata=metadata, normalized=None, remove=removed)
-    assert invoke_pre_call(owners, logger, {}) is None
-    assert order == ["retain", "mutate", "observe"]
-    assert observed == [("normalized", False, "masked")]
-    assert len(saved) == 1 and saved[0] is details
-    assert details["retained_metadata"] is metadata
-    assert metadata == {"secret": "masked", "keep": []}
-    assert details["normalized"] == "normalized" and "remove" not in details
-    assert "replacement" not in details
-
-
-async def pre_call_mutation_survives_failure(owners):
-    observed, order = [], []
-    metadata = {"keep": []}
-    lock = threading.Lock()
-
-    class Fail(CustomLogger):
-        def log_pre_api_call(self, model, messages, kwargs):
-            order.append("fail")
-            kwargs["lock"] = lock
-            kwargs["retained_metadata"]["keep"].append("before failure")
-            raise RuntimeError("expected pre-call callback failure")
-
-    class Observe(CustomLogger):
-        def log_pre_api_call(self, model, messages, kwargs):
-            order.append("observe")
-            observed.append((kwargs, tuple(kwargs["retained_metadata"]["keep"]), kwargs["lock"]))
-
-    logger = logger_for(input_callbacks=[Fail(), Observe()])
-    details = logger.model_call_details
-    details["retained_metadata"] = metadata
-    assert invoke_pre_call(owners, logger, {}) is None
-    assert order == ["fail", "observe"]
-    assert len(observed) == 1 and observed[0][0] is details
-    assert observed[0][1] == ("before failure",) and observed[0][2] is lock
-    assert details["retained_metadata"] is metadata
-    assert metadata == {"keep": ["before failure"]} and details["lock"] is lock
-    with TestCase().assertRaises(TypeError):
-        json.dumps({"lock": details["lock"]})
-
-
-async def real_post_call_logging(owners):
-    saved, observed, order = [], [], []
-    shared = {"values": []}
-    response = ModelResponse(model="test")
-    ignored = {"replacement": True}
-    error = RuntimeError("expected post-call callback failure")
-
-    class Retain(CustomLogger):
-        def log_pre_api_call(self, model, messages, kwargs):
-            kwargs["stash"] = shared
-            saved.append(kwargs)
-
-        def log_post_api_call(self, kwargs, response_obj, start_time, end_time):
-            order.append("retain")
-            saved.append(kwargs)
-            return ignored
-
-    class MutateThenFail(CustomLogger):
-        def log_post_api_call(self, kwargs, response_obj, start_time, end_time):
-            order.append("fail")
-            kwargs["stash"]["values"].append("post")
-            kwargs["callback_error"] = error
-            kwargs["original_response"].choices[0].message.content = "mutated"
-            raise error
-
-    class Observe(CustomLogger):
-        def log_post_api_call(self, kwargs, response_obj, start_time, end_time):
-            order.append("observe")
-            observed.append((kwargs, response_obj, start_time, end_time, tuple(kwargs["stash"]["values"])))
-            return ignored
-
-    logger = logger_for(input_callbacks=[Retain(), MutateThenFail(), Observe()])
-    details = logger.model_call_details
-    assert invoke_pre_call(owners, logger, {}) is None
-    additional = {"headers": {"test": "post"}}
-    owner = owners.prepare(logger.post_call, (response, logger.messages, "test-key"), {"additional_args": additional})
-    try:
-        assert owner.invoke() is None
-    finally:
-        owner.close()
-    assert order == ["retain", "fail", "observe"]
-    assert len(saved) == 2 and saved[0] is saved[1] is details
-    assert len(observed) == 1 and observed[0][0] is details
-    assert observed[0][1] is None and observed[0][2] is logger.start_time and observed[0][3] is None
-    assert observed[0][4] == ("post",)
-    assert details["original_response"] is response and response.choices[0].message.content == "mutated"
-    assert details["input"] is logger.messages and details["additional_args"] is additional
-    assert details["log_event_type"] == "post_api_call" and details["api_key"] == "test-key"
-    assert details["stash"] is shared and details["callback_error"] is error
-    assert "replacement" not in details
-    shared["values"].append("later")
-    assert saved[0]["stash"]["values"] == ["post", "later"]
-
-
-async def real_post_call_dict_response(owners):
-    observed = []
-
-    class Observe(CustomLogger):
-        def log_post_api_call(self, kwargs, response_obj, start_time, end_time):
-            observed.append(kwargs["original_response"])
-
-    response = {"content": ["original"], "timestamp": datetime(2026, 1, 1)}
-    logger = logger_for(input_callbacks=[Observe()])
-    owner = owners.prepare(logger.post_call, (response,))
-    try:
-        assert owner.invoke() is None
-    finally:
-        owner.close()
-    assert len(observed) == 1 and observed[0] is logger.model_call_details["original_response"]
-    assert isinstance(observed[0], str)
-    assert json.loads(observed[0]) == {"content": ["original"], "timestamp": "2026-01-01 00:00:00"}
-    response["content"].append("later")
-    assert json.loads(observed[0])["content"] == ["original"]
-
-
-async def real_sync_logging(owners):
-    saved, observations, replacements = [], [], []
-    shared = {"values": []}
-    result = ModelResponse(model="test", choices=[{"message": {"role": "assistant", "content": "original"}}])
-    replacement = ModelResponse(model="test", choices=[{"message": {"role": "assistant", "content": "replacement"}}])
-    ignored = {"ignored": True}, result
-
-    class Retain(CustomLogger):
-        def log_pre_api_call(self, model, messages, kwargs):
-            kwargs["stash"] = shared
-            saved.append(kwargs)
-
-        def logging_hook(self, kwargs, result, call_type):
-            observations.append(("retain", kwargs, result, call_type))
-            kwargs["stash"]["values"].append("hook")
-            result.choices[0].message.content = "mutated"
-            return kwargs, result
-
-        def log_success_event(self, kwargs, response_obj, start_time, end_time):
-            kwargs["stash"]["values"].append("event")
-            observations.append(("event", kwargs, response_obj, tuple(kwargs["stash"]["values"])))
-            return ignored
-
-    class Replace(CustomLogger):
-        def logging_hook(self, kwargs, result, call_type):
-            observations.append(("replace", kwargs, result, result.choices[0].message.content))
-            updated = {**kwargs, "adopted": True}
-            replacements.append(updated)
-            return updated, replacement
-
-    class Observe(CustomLogger):
-        def logging_hook(self, kwargs, result, call_type):
-            observations.append(("observe", kwargs, result, tuple(kwargs["stash"]["values"])))
-            return kwargs, result
-
-        def log_success_event(self, kwargs, response_obj, start_time, end_time):
-            observations.append(("success", kwargs, response_obj, tuple(kwargs["stash"]["values"])))
-
-    retain = Retain()
-    logger = logger_for(input_callbacks=[retain], sync_callbacks=[retain, Replace(), Observe()], call_type="completion")
-    assert invoke_pre_call(owners, logger, {}) is None
-    owner = owners.prepare(logger.success_handler, (result,))
-    try:
-        assert owner.invoke() is None
-    finally:
-        owner.close()
-    assert [entry[0] for entry in observations] == ["retain", "replace", "observe", "event", "success"]
-    assert len(saved) == len(replacements) == 1
-    assert observations[0][1] is observations[1][1] is saved[0]
-    assert observations[0][2] is observations[1][2] is result
-    assert observations[0][3] == "completion" and observations[1][3] == "mutated"
-    assert observations[2][3] == ("hook",) and observations[3][3] == observations[4][3] == ("hook", "event")
-    assert all(entry[1] is replacements[0] is logger.model_call_details for entry in observations[2:])
-    assert all(entry[2] is replacement for entry in observations[2:])
-    assert logger.model_call_details is not saved[0]
-    assert logger.model_call_details["adopted"] and "adopted" not in saved[0]
-    assert "ignored" not in logger.model_call_details
-    assert result.choices[0].message.content == "mutated"
-    assert replacement.choices[0].message.content == "replacement"
-    assert saved[0]["stash"] is logger.model_call_details["stash"] is shared
-    shared["values"].append("later")
-    assert observations[-1][1]["stash"]["values"] == ["hook", "event", "later"]
-
-
-async def real_sync_logging_hook_failure(owners):
-    order, saved = [], []
-    result = ModelResponse(model="test")
-    error = RuntimeError("expected sync logging hook failure")
-
-    class MutateThenFail(CustomLogger):
-        def logging_hook(self, kwargs, result, call_type):
-            order.append("fail")
-            saved.append(kwargs)
-            kwargs["callback_error"] = error
-            result.choices[0].message.content = "before failure"
-            raise error
-
-        def log_success_event(self, kwargs, response_obj, start_time, end_time):
-            order.append("unexpected success")
-
-    class Observe(CustomLogger):
-        def logging_hook(self, kwargs, result, call_type):
-            order.append("unexpected hook")
-            return kwargs, result
-
-        def log_success_event(self, kwargs, response_obj, start_time, end_time):
-            order.append("unexpected later success")
-
-    logger = logger_for(sync_callbacks=[MutateThenFail(), Observe()], call_type="completion")
-    details = logger.model_call_details
-    owner = owners.prepare(logger.success_handler, (result,))
-    try:
-        assert owner.invoke() is None
-    finally:
-        owner.close()
-    assert order == ["fail"]
-    assert len(saved) == 1 and saved[0] is logger.model_call_details is details
-    assert details["callback_error"] is error and error.__traceback__ is not None
-    assert result.choices[0].message.content == "before failure"
-
-
-async def real_sync_failure_chain(owners):
-    await real_failure_chain(owners, awaited=False)
-
-
-async def real_async_failure_chain(owners):
-    await real_failure_chain(owners, awaited=True)
-
-
-async def real_failure_chain(owners, awaited):
-    saved, observations, hooks = [], [], []
-    shared = {"values": []}
-    error = ValueError("provider failure")
-    callback_error = RuntimeError("expected failure callback error")
-    ignored = {"replacement": True}, object()
-    task = asyncio.current_task()
-    end = datetime.now()
-
-    class Stage(CustomLogger):
-        def __init__(self, name):
-            super().__init__()
-            self.name = name
-
-        def log_pre_api_call(self, model, messages, kwargs):
-            kwargs["stash"] = shared
-            saved.append(kwargs)
-
-        def logging_hook(self, kwargs, result, call_type):
-            hooks.append("sync")
-            return ignored
-
-        async def async_logging_hook(self, kwargs, result, call_type):
-            hooks.append("async")
-            return ignored
-
-        def record(self, kwargs, response_obj, start_time, end_time):
-            observations.append(
-                (
-                    self.name,
-                    kwargs,
-                    response_obj,
-                    kwargs["exception"],
-                    tuple(kwargs["stash"]["values"]),
-                    start_time,
-                    end_time,
-                    asyncio.current_task(),
-                )
-            )
-            if self.name == "fail":
-                kwargs["stash"]["values"].append("before failure")
-                kwargs["callback_error"] = callback_error
-                raise callback_error
-            return ignored
-
-        def log_failure_event(self, kwargs, response_obj, start_time, end_time):
-            return self.record(kwargs, response_obj, start_time, end_time)
-
-        async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
-            await asyncio.sleep(0)
-            return self.record(kwargs, response_obj, start_time, end_time)
-
-    retain, fail, observe = Stage("retain"), Stage("fail"), Stage("observe")
-    callbacks = [retain, fail, observe]
-    logger = logger_for(
-        input_callbacks=[retain],
-        failure_callbacks=() if awaited else callbacks,
-        async_failure_callbacks=callbacks if awaited else (),
-        call_type="acompletion" if awaited else "completion",
-    )
-    logger.model_call_details["litellm_params"]["acompletion"] = awaited
-    details = logger.model_call_details
-    assert invoke_pre_call(owners, logger, {}) is None
-    owner = owners.prepare(
-        logger.async_failure_handler if awaited else logger.failure_handler,
-        (error, "provider traceback"),
-        {"start_time": logger.start_time, "end_time": end},
-        awaited=awaited,
-    )
-    try:
-        if awaited:
-            assert await owner.invoke() is None
-        else:
-            assert owner.invoke() is None
-    finally:
-        owner.close()
-    assert [entry[0] for entry in observations] == ["retain", "fail", "observe"]
-    assert len(saved) == 1 and saved[0] is logger.model_call_details is details
-    assert all(entry[1] is details and entry[2] is None and entry[3] is error for entry in observations)
-    assert [entry[4] for entry in observations] == [(), (), ("before failure",)]
-    assert all(entry[5] is logger.start_time and entry[6] is end and entry[7] is task for entry in observations)
-    assert details["exception"] is error and details["callback_error"] is callback_error
-    assert callback_error.__traceback__ is not None
-    assert details["traceback_exception"] == "provider traceback" and details["log_event_type"] == "failed_api_call"
-    assert details["stash"] is shared and "replacement" not in details and hooks == []
-    shared["values"].append("later")
-    assert saved[0]["stash"]["values"] == observations[-1][1]["stash"]["values"] == ["before failure", "later"]
-
-
-async def real_async_logging(owners):
-    observations = []
-    task = asyncio.current_task()
-    gate = asyncio.Event()
-    result = ModelResponse(model="test", choices=[{"message": {"role": "assistant", "content": "original"}}])
-    replacement = ModelResponse(model="test", choices=[{"message": {"role": "assistant", "content": "replacement"}}])
-    shared = {}
-    side_channel = {}
-    replaced_kwargs = []
-
-    class Retain(CustomLogger):
-        async def async_logging_hook(self, kwargs, result, call_type):
-            observations.append(("retained", kwargs, result, asyncio.current_task()))
-            return kwargs, result
-
-    class MutateThenFail(CustomLogger):
-        async def async_logging_hook(self, kwargs, result, call_type):
-            asyncio.get_running_loop().call_soon(gate.set)
-            await gate.wait()
-            kwargs["retained_shared"]["changed"] = True
-            side_channel["failed_hook"] = kwargs
-            result.choices[0].message.content = "mutated"
-            raise RuntimeError("expected async callback failure")
-
-    class Replace(CustomLogger):
-        async def async_logging_hook(self, kwargs, result, call_type):
-            observations.append(("replace", kwargs, result))
-            updated = {**kwargs, "adopted": True}
-            replaced_kwargs.append(updated)
-            side_channel["replacement"] = replacement
-            return updated, replacement
-
-    class Observe(CustomLogger):
-        async def async_logging_hook(self, kwargs, result, call_type):
-            observations.append(("observe", kwargs, result, asyncio.current_task()))
-            return kwargs, result
-
-        async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
-            observations.append(("success", kwargs, response_obj))
-
-    logger = logger_for([Retain(), MutateThenFail(), Replace(), Observe()])
-    logger.model_call_details["retained_shared"] = shared
-    owner = owners.prepare(logger.async_success_handler, (), {"result": result}, awaited=True)
-    await owner.invoke()
-    owner.close()
-    assert [entry[0] for entry in observations] == ["retained", "replace", "observe", "success"]
-    assert observations[0][3] is task and observations[2][3] is task
-    assert observations[0][2] is result and observations[1][2] is result
-    assert observations[2][2] is replacement and observations[3][2] is replacement
-    assert observations[0][1]["retained_shared"] is shared and shared["changed"]
-    assert observations[0][1] is observations[1][1] is side_channel["failed_hook"]
-    assert observations[2][1] is observations[3][1] is logger.model_call_details is replaced_kwargs[0]
-    assert logger.model_call_details is not observations[0][1]
-    assert logger.model_call_details["adopted"] and "adopted" not in observations[0][1]
-    assert logger.model_call_details["retained_shared"] is shared
-    assert side_channel["replacement"] is replacement
-    assert result.choices[0].message.content == "mutated"
-    observations[0][1]["retained_shared"]["after_replacement"] = True
-    assert observations[3][1]["retained_shared"]["after_replacement"]
-
-
-async def real_copy_boundaries(owners):
-    lock = threading.Lock()
-    shared = {"values": []}
-    standard = {
-        "messages": [{"role": "user", "content": "private"}],
-        "response": {"choices": [{"message": {"content": "private"}}]},
-        "metadata": shared,
-    }
-    details = {"standard_logging_object": standard, "shared": shared, "lock": lock}
-    passthrough = owners.prepare(CustomLogger().redact_standard_logging_payload_from_model_call_details, (details,))
-    try:
-        assert passthrough.invoke() is details
-    finally:
-        passthrough.close()
-    logger = CustomLogger(turn_off_message_logging=True)
-    redact = owners.prepare(logger.redact_standard_logging_payload_from_model_call_details, (details,))
-    try:
-        redacted = redact.invoke()
-    finally:
-        redact.close()
-    assert redacted is not details
-    assert redacted["standard_logging_object"] is not standard
-    assert redacted["shared"] is shared and redacted["lock"] is lock
-    assert redacted["standard_logging_object"]["metadata"] is shared
-    redacted["standard_logging_object"]["metadata"]["values"].append("shared mutation")
-    assert shared["values"] == ["shared mutation"]
-    assert redacted["standard_logging_object"]["messages"][0]["content"] == "redacted-by-litellm"
-    assert redacted["standard_logging_object"]["response"]["choices"][0]["message"]["content"] == "redacted-by-litellm"
-    assert standard["messages"][0]["content"] == "private"
-    assert standard["response"]["choices"][0]["message"]["content"] == "private"
-
-    original_mode = litellm.safe_memory_mode
-    try:
-        for safe_mode in (False, True):
-            litellm.safe_memory_mode = safe_mode
-            uncopyable = {"lock": lock, "values": []}
-            values = []
-            data = {"copyable": {"values": values, "alias": values}, "uncopyable": uncopyable}
-            owner = owners.prepare(safe_deep_copy, (data,))
-            try:
-                copied = owner.invoke()
-            finally:
-                owner.close()
-            assert (copied is data) is safe_mode
-            assert (copied["copyable"] is data["copyable"]) is safe_mode
-            assert copied["copyable"]["values"] is copied["copyable"]["alias"]
-            assert (copied["copyable"]["values"] is values) is safe_mode
-            assert copied["uncopyable"] is uncopyable and copied["uncopyable"]["lock"] is lock
-            copied["copyable"]["values"].append("copy")
-            assert copied["copyable"]["alias"] == ["copy"]
-            copied["uncopyable"]["values"].append("fallback")
-            assert data["copyable"]["values"] == (["copy"] if safe_mode else [])
-            assert uncopyable["values"] == ["fallback"]
-    finally:
-        litellm.safe_memory_mode = original_mode
-
-
-async def real_logging_worker(owners):
-    context = contextvars.ContextVar("component_worker_context", default="outside")
-    entered, release = asyncio.Event(), asyncio.Event()
-    observations = []
-    worker = LoggingWorker(timeout=5, concurrency=1)
-
-    class Payload:
-        pass
-
-    async def upload(value, *, alias):
-        assert value is alias
-        assert context.get() == "submitted"
-        entered.set()
-        await release.wait()
-        observations.append((value.changed, context.get()))
-        context.set("worker only")
-
-    payload = Payload()
-    payload.changed = False
-    reference = weakref.ref(payload)
-    invocation = owners.prepare(upload, (payload,), {"alias": payload}, awaited=True)
-    pending = invocation.invoke()
-    invocation.close()
-    enqueue = owners.prepare(worker.ensure_initialized_and_enqueue, (pending,))
-    stop = owners.prepare(worker.stop, (), awaited=True)
-    flush = owners.prepare(worker.flush, (), awaited=True)
-    token = context.set("submitted")
-    try:
-        enqueue.invoke()
-        enqueue.close()
-        del pending, payload
-        context.set("consumer")
-        await entered.wait()
-        assert reference() is not None
-        reference().changed = True
-        release.set()
-        await flush.invoke()
-        assert observations == [(True, "submitted")]
-        assert context.get() == "consumer"
-    finally:
-        release.set()
-        enqueue.close()
-        flush.close()
-        await stop.invoke()
-        stop.close()
-        context.reset(token)
-        atexit.unregister(worker._flush_on_exit)
-    assert worker._worker_task is None and not worker._running_tasks and not worker._dequeued_tasks
-    assert worker._queue.empty()
-    gc.collect()
-    assert reference() is None
-
-
-class ControlledStream:
-    def __init__(self):
-        self.originals = [
-            ModelResponseStream(model="test", choices=[StreamingChoices(delta=Delta(content="hello"), index=0)]),
-            ModelResponseStream(
-                model="test", choices=[StreamingChoices(delta=Delta(content=""), index=0, finish_reason="stop")]
-            ),
-            ModelResponseStream(
-                model="test", choices=[], usage=Usage(prompt_tokens=3, completion_tokens=5, total_tokens=8)
-            ),
-        ]
-        self.chunks = iter(self.originals)
-        self.closed = 0
-
-    def __aiter__(self):
-        return self
-
-    async def __anext__(self):
-        try:
-            return next(self.chunks)
-        except StopIteration:
-            raise StopAsyncIteration
-
-    async def aclose(self):
-        self.closed += 1
-
-
-async def real_stream_completion(owners):
-    logger = logger_for(stream=True)
-    completions = []
-    cached = []
-    cache_done = asyncio.Event()
-
-    class CacheRecorder:
-        async def _add_streaming_response_to_cache(self, response):
-            cached.append(response)
-            cache_done.set()
-
-    logger._llm_caching_handler = CacheRecorder()
-
-    async def complete(response, cache_hit):
-        completions.append(response)
-
-    logger._on_deferred_stream_complete = complete
-    stream = ControlledStream()
-    stream.originals[1].usage = Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2)
-    wrapper = CustomStreamWrapper(
-        completion_stream=stream, model="test", logging_obj=logger, custom_llm_provider="bedrock"
-    )
-    pull = owners.prepare(wrapper.__anext__, (), awaited=True)
-    chunks = []
-    retained_hidden = None
-    hidden_owner = None
-    while True:
-        try:
-            chunk = await pull.invoke()
-            chunks.append(chunk)
-            if len(chunks) == 1:
-                assert wrapper.chunks[-1] is chunk
-                chunk.choices[0].delta.content = "retained hello"
-            if chunk.choices and chunk.choices[0].finish_reason:
-                stored = wrapper.chunks[-1]
-                assert stored is not chunk and stored is not stream.originals[1]
-                assert stored.usage is stream.originals[1].usage
-                assert getattr(chunk, "usage", None) is None and stored.usage.total_tokens == 2
-                retained_hidden = chunk._hidden_params
-                hidden_owner = owners.prepare(lambda value: value, (retained_hidden,))
-                assert not completions
-        except StopAsyncIteration:
-            break
-    pull.close()
-    assert retained_hidden is not None
-    assert retained_hidden["usage"].total_tokens == 8
-    assert completions == []
-    response, cache_hit = logger._deferred_stream_complete_args
-    assert response.usage.total_tokens == 8
-    assert response.choices[0].message.content == "retained hello"
-    assert retained_hidden["usage"] is response.usage
-    usage_chunk = wrapper.chunks[-1]
-    assert usage_chunk is not stream.originals[-1]
-    assert usage_chunk.usage is stream.originals[-1].usage
-    stream.originals[-1].usage.total_tokens = 13
-    assert usage_chunk.usage.total_tokens == 13
-    assert response.usage.total_tokens == 8
-    deferred = owners.prepare(logger._on_deferred_stream_complete, (response, cache_hit), awaited=True)
-    logger._on_deferred_stream_complete = None
-    logger._deferred_stream_complete_args = None
-    close = owners.prepare(wrapper.aclose, (), awaited=True)
-    await close.invoke()
-    await close.invoke()
-    close.close()
-    assert stream.closed == 1
-    del wrapper, logger
-    await deferred.invoke()
-    deferred.close()
-    assert completions == [response]
-    assert retained_hidden["usage"].total_tokens == 8
-    assert hidden_owner.invoke() is retained_hidden
-    hidden_owner.close()
-    await cache_done.wait()
-    assert len(cached) == 1 and cached[0] is not response
-    assert cached[0].choices[0] is not response.choices[0]
-    cached[0].choices[0].message.content = "cache only"
-    assert response.choices[0].message.content == "retained hello"
-
-
-async def real_sync_stream_copies(owners):
-    original_disable = litellm.disable_streaming_logging
-    copy_attempts = []
-
-    class Uncopyable:
-        def __deepcopy__(self, memo):
-            copy_attempts.append(True)
-            raise RuntimeError("expected streaming deepcopy failure")
-
-    class CacheRecorder:
-        def __init__(self, responses):
-            self.responses = responses
-
-        def _sync_add_streaming_response_to_cache(self, response):
-            self.responses.append(response)
-
-    class Observe(CustomLogger):
-        def __init__(self, responses, finished):
-            super().__init__()
-            self.responses = responses
-            self.finished = finished
-
-        def log_success_event(self, kwargs, response_obj, start_time, end_time):
-            self.responses.append(response_obj)
-            self.finished.set()
-
-    try:
-        litellm.disable_streaming_logging = True
-        for fallback in (False, True):
-            cached, logged = [], []
-            finished = threading.Event()
-
-            logger = logger_for(stream=True, sync_callbacks=[Observe(logged, finished)])
-            logger._llm_caching_handler = CacheRecorder(cached)
-            source = ControlledStream()
-            source.originals[1].usage = source.originals[2].usage
-            wrapper = CustomStreamWrapper(
-                completion_stream=iter(source.originals[:2]),
-                model="test",
-                logging_obj=logger,
-                custom_llm_provider="bedrock",
-            )
-            pull = owners.prepare(wrapper.__next__, ())
-            close = owners.prepare(wrapper.aclose, (), awaited=True)
-            try:
-                first = pull.invoke()
-                assert wrapper.chunks[0] is first
-                first.choices[0].delta.content = "consumer mutation"
-                shared = {"values": []}
-                last = pull.invoke()
-                assert last.choices[0].finish_reason == "stop"
-                for chunk in wrapper.chunks:
-                    chunk._hidden_params["retained_shared"] = shared
-                    if fallback:
-                        chunk._hidden_params["uncopyable"] = Uncopyable()
-                retained_hidden = last._hidden_params
-                with TestCase().assertRaises(StopIteration):
-                    pull.invoke()
-                assert await asyncio.to_thread(finished.wait, 5)
-            finally:
-                pull.close()
-                await close.invoke()
-                close.close()
-            assert len(cached) == len(logged) == 1
-            cache_response, log_response = cached[0], logged[0]
-            assert cache_response is not log_response
-            assert cache_response.choices[0].message.content == "consumer mutation"
-            assert log_response.choices[0].message.content == "consumer mutation"
-            assert retained_hidden["usage"].total_tokens == 8
-            assert (cache_response.choices is log_response.choices) is fallback
-            assert (cache_response.usage is log_response.usage) is fallback
-            assert (cache_response.usage is retained_hidden["usage"]) is fallback
-            assert (cache_response._hidden_params is log_response._hidden_params) is fallback
-            assert (cache_response._hidden_params["retained_shared"] is shared) is fallback
-            cache_response.choices[0].message.content = "cache mutation"
-            cache_response._hidden_params["retained_shared"]["values"].append("cache mutation")
-            assert log_response.choices[0].message.content == ("cache mutation" if fallback else "consumer mutation")
-            assert shared["values"] == (["cache mutation"] if fallback else [])
-            assert log_response._hidden_params["retained_shared"]["values"] == (["cache mutation"] if fallback else [])
-        assert copy_attempts == [True]
-    finally:
-        litellm.disable_streaming_logging = original_disable
-
-
-async def real_stream_close(owners):
-    source = ControlledStream()
-    logger = logger_for(stream=True)
-    wrapper = CustomStreamWrapper(
-        completion_stream=source, model="test", logging_obj=logger, custom_llm_provider="bedrock"
-    )
-    pull = owners.prepare(wrapper.__anext__, (), awaited=True)
-    chunk = await pull.invoke()
-    assert wrapper.chunks[0] is chunk
-    pull.close()
-    retained = owners.prepare(lambda value: value, (chunk,))
-    close = owners.prepare(wrapper.aclose, (), awaited=True)
-    await close.invoke()
-    await close.invoke()
-    close.close()
-    assert source.closed == 1
-    assert wrapper.completion_stream is None
-    assert not getattr(logger, "_deferred_stream_complete_args", None)
-    assert retained.invoke() is chunk
-    assert chunk.choices[0].delta.content == "hello"
-    retained.close()
-
-
-async def real_stream_cancellation(owners):
-    entered = asyncio.Event()
-
-    class SuspendedStream(ControlledStream):
-        async def __anext__(self):
-            if self.chunks is not None:
-                chunk = next(self.chunks)
-                self.chunks = None
-                return chunk
-            entered.set()
-            await asyncio.Event().wait()
-
-    source = SuspendedStream()
-    source.originals[0].usage = Usage(prompt_tokens=3, completion_tokens=2, total_tokens=5)
-    logger = logger_for(stream=True)
-    wrapper = CustomStreamWrapper(
-        completion_stream=source, model="test", logging_obj=logger, custom_llm_provider="bedrock"
-    )
-    pull = owners.prepare(wrapper.__anext__, (), awaited=True)
-    chunk = await pull.invoke()
-    retained = owners.prepare(lambda value: value, (chunk,))
-    assert wrapper.chunks[0] is not chunk
-    assert wrapper.chunks[0].usage is source.originals[0].usage
-    task = asyncio.create_task(pull.invoke())
-    pull.close()
-    await entered.wait()
-    task.cancel()
-    with TestCase().assertRaises(asyncio.CancelledError):
-        await task
-    assert logger.model_call_details.get("combined_usage_object") is None
-    recover = owners.prepare(wrapper._record_partial_usage_for_failure, ())
-    try:
-        assert recover.invoke() is None
-    finally:
-        recover.close()
-    usage = logger.model_call_details["combined_usage_object"]
-    assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (3, 2, 5)
-    assert usage is not source.originals[0].usage
-    source.originals[0].usage.total_tokens = 99
-    assert usage.total_tokens == 5
-    retained_usage = owners.prepare(lambda value: value, (usage,))
-    close = owners.prepare(wrapper.aclose, (), awaited=True)
-    await close.invoke()
-    await close.invoke()
-    close.close()
-    assert source.closed == 1
-    assert wrapper.completion_stream is None and len(wrapper.chunks) == 1
-    assert not getattr(logger, "_deferred_stream_complete_args", None)
-    assert retained.invoke() is chunk and chunk.choices[0].delta.content == "hello"
-    retained.close()
-    del wrapper, logger, usage
-    assert retained_usage.invoke().total_tokens == 5
-    retained_usage.close()

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

@@ -1,346 +0,0 @@
-import asyncio
-import gc
-import inspect
-import weakref
-from collections.abc import Awaitable, Callable
-from dataclasses import dataclass
-from typing import Protocol, cast
-
-from callback_lifecycle import ReferenceFactory, run_checked, settle
-
-
-class PreparedInvocation(Protocol):
-    def invoke(self) -> object: ...
-
-    def close(self) -> None: ...
-
-
-class CallFactory(Protocol):
-    def prepare(
-        self,
-        callable: Callable[..., object],
-        positional: tuple[object, ...],
-        keywords: dict[str, object] | None = None,
-        awaited: bool = False,
-    ) -> PreparedInvocation: ...
-
-
-class LiveCallFactory(CallFactory, Protocol):
-    @property
-    def live(self) -> int: ...
-
-
-def unchanged_result(value: object) -> object:
-    return value
-
-
-@dataclass(frozen=True, slots=True)
-class ResultTransformInvocation:
-    inner: PreparedInvocation
-    transform: Callable[[object], object]
-    awaited: bool
-
-    def invoke(self) -> object:
-        result = self.inner.invoke()
-        transform = self.transform
-        if not self.awaited:
-            return transform(result)
-
-        async def run() -> object:
-            return transform(await cast(Awaitable[object], result))
-
-        return run()
-
-    def close(self) -> None:
-        self.inner.close()
-
-
-@dataclass(frozen=True, slots=True)
-class ResultTransformFactory:
-    inner: CallFactory
-    transform: Callable[[object], object]
-
-    def prepare(
-        self,
-        callable: Callable[..., object],
-        positional: tuple[object, ...],
-        keywords: dict[str, object] | None = None,
-        awaited: bool = False,
-    ) -> PreparedInvocation:
-        return ResultTransformInvocation(
-            self.inner.prepare(callable, positional, keywords, awaited=awaited), self.transform, awaited
-        )
-
-
-@dataclass(frozen=True, slots=True)
-class ExpiredBorrow:
-    edges: tuple[str, ...]
-
-
-@dataclass(frozen=True, slots=True)
-class CheckedWeakInvocation:
-    callback: weakref.ReferenceType[Callable[..., object]]
-    positional: tuple[weakref.ReferenceType[object], ...]
-    keywords: tuple[tuple[str, weakref.ReferenceType[object]], ...]
-    awaited: bool
-
-    def resolve(self) -> tuple[Callable[..., object], tuple[object, ...], dict[str, object]] | ExpiredBorrow:
-        callback = self.callback()
-        positional = tuple(reference() for reference in self.positional)
-        keywords = {name: reference() for name, reference in self.keywords}
-        expired = (
-            *(("callable",) if callback is None else ()),
-            *(f"positional:{index}" for index, value in enumerate(positional) if value is None),
-            *(f"keyword:{name}" for name, value in keywords.items() if value is None),
-        )
-        if expired:
-            return ExpiredBorrow(expired)
-        assert callback is not None
-        return callback, positional, keywords
-
-    def invoke(self) -> object:
-        if not self.awaited:
-            resolved = self.resolve()
-            if isinstance(resolved, ExpiredBorrow):
-                return resolved
-            callback, positional, keywords = resolved
-            return callback(*positional, **keywords)
-
-        async def run() -> object:
-            resolved = self.resolve()
-            if isinstance(resolved, ExpiredBorrow):
-                return resolved
-            callback, positional, keywords = resolved
-            return await cast(Awaitable[object], callback(*positional, **keywords))
-
-        return run()
-
-    def close(self) -> None:
-        pass
-
-
-class CheckedWeakFactory:
-    def prepare(
-        self,
-        callable: Callable[..., object],
-        positional: tuple[object, ...],
-        keywords: dict[str, object] | None = None,
-        awaited: bool = False,
-    ) -> PreparedInvocation:
-        return CheckedWeakInvocation(
-            weakref.ref(callable),
-            tuple(weakref.ref(value) for value in positional),
-            tuple((name, weakref.ref(value)) for name, value in (keywords or {}).items()),
-            awaited,
-        )
-
-
-@dataclass(frozen=True, slots=True)
-class MissingHandoffInvocation:
-    inner: PreparedInvocation
-    borrowed: PreparedInvocation
-
-    def invoke(self) -> object:
-        return self.borrowed.invoke()
-
-    def close(self) -> None:
-        try:
-            self.inner.close()
-        finally:
-            self.borrowed.close()
-
-
-@dataclass(frozen=True, slots=True)
-class MissingHandoffFactory:
-    inner: CallFactory
-
-    def prepare(
-        self,
-        callable: Callable[..., object],
-        positional: tuple[object, ...],
-        keywords: dict[str, object] | None = None,
-        awaited: bool = False,
-    ) -> PreparedInvocation:
-        borrowed = CheckedWeakFactory().prepare(callable, positional, keywords, awaited=awaited)
-        return MissingHandoffInvocation(self.inner.prepare(callable, positional, keywords, awaited=awaited), borrowed)
-
-
-def control_factory(control: str, inner: CallFactory) -> CallFactory:
-    if control == "identity":
-        return inner
-    if control == "weak":
-        return CheckedWeakFactory()
-    if control == "missing_handoff":
-        return MissingHandoffFactory(inner)
-    return {"result_passthrough": ResultTransformFactory(inner, unchanged_result)}[control]
-
-
-@dataclass
-class ControlNode:
-    stage: int = 0
-
-
-@dataclass
-class LifetimeCallback:
-    awaited: bool
-
-    def __call__(self, value: ControlNode, *, alias: ControlNode) -> object:
-        if self.awaited:
-            return self.run(value, alias=alias)
-        return value.stage + alias.stage
-
-    async def run(self, value: ControlNode, *, alias: ControlNode) -> int:
-        await asyncio.sleep(0)
-        return value.stage + alias.stage
-
-
-def prepare_released(
-    owners: CallFactory, awaited: bool
-) -> tuple[
-    PreparedInvocation,
-    tuple[
-        weakref.ReferenceType[LifetimeCallback], weakref.ReferenceType[ControlNode], weakref.ReferenceType[ControlNode]
-    ],
-]:
-    callback = LifetimeCallback(awaited)
-    value = ControlNode(13)
-    alias = ControlNode(29)
-    return (
-        owners.prepare(callback, (value,), {"alias": alias}, awaited=awaited),
-        (weakref.ref(callback), weakref.ref(value), weakref.ref(alias)),
-    )
-
-
-@dataclass(frozen=True, slots=True)
-class LifetimeObservation:
-    alive: tuple[bool, bool, bool]
-    result: int | ExpiredBorrow
-
-
-async def deferred_lifetime(owners: CallFactory, awaited: bool) -> LifetimeObservation:
-    owner, references = prepare_released(owners, awaited)
-    try:
-        gc.collect()
-        alive = tuple(reference() is not None for reference in references)
-        pending = owner.invoke()
-        result = await settle(pending, awaited)
-        observation = LifetimeObservation(alive, result)
-    finally:
-        owner.close()
-    gc.collect()
-    assert all(reference() is None for reference in references)
-    return observation
-
-
-@dataclass(frozen=True, slots=True)
-class BorrowedObservation:
-    positional_identity: bool
-    keyword_identity: bool
-    positional_stage: int
-    keyword_stage: int
-
-
-async def borrowed_lifetime(owners: CallFactory, awaited: bool) -> BorrowedObservation:
-    value, alias = ControlNode(), ControlNode()
-    value_ref, alias_ref = weakref.ref(value), weakref.ref(alias)
-
-    def observe(value: ControlNode, *, alias: ControlNode) -> BorrowedObservation:
-        return BorrowedObservation(value is value_ref(), alias is alias_ref(), value.stage, alias.stage)
-
-    async def observe_async(value: ControlNode, *, alias: ControlNode) -> BorrowedObservation:
-        return observe(value, alias=alias)
-
-    owner = owners.prepare(observe_async if awaited else observe, (value,), {"alias": alias}, awaited=awaited)
-    try:
-        value.stage, alias.stage = 13, 29
-        pending = owner.invoke()
-        value.stage, alias.stage = 17, 31
-        return await settle(pending, awaited)
-    finally:
-        owner.close()
-
-
-async def pending_handoff(owners: CallFactory, awaited: bool) -> LifetimeObservation:
-    assert awaited
-    owner, references = prepare_released(owners, awaited)
-    try:
-        pending = owner.invoke()
-        try:
-            owner.close()
-            gc.collect()
-            alive = tuple(reference() is not None for reference in references)
-            observation = LifetimeObservation(alive, await pending)
-        finally:
-            pending.close()
-    finally:
-        owner.close()
-    gc.collect()
-    assert all(reference() is None for reference in references)
-    return observation
-
-
-async def direct_coroutine(owners: CallFactory, awaited: bool) -> bool:
-    async def body() -> int:
-        return 73
-
-    original = body()
-    try:
-        owner = owners.prepare(lambda: original, (), awaited=awaited)
-        try:
-            pending = owner.invoke()
-            if awaited:
-                assert await pending == 73
-                nested = body()
-                try:
-
-                    async def returns_coroutine() -> Awaitable[int]:
-                        return nested
-
-                    nested_owner = owners.prepare(returns_coroutine, (), awaited=True)
-                    try:
-                        result = await nested_owner.invoke()
-                        return result is nested and inspect.getcoroutinestate(nested) == inspect.CORO_CREATED
-                    finally:
-                        nested_owner.close()
-                finally:
-                    nested.close()
-            try:
-                return pending is original and inspect.getcoroutinestate(original) == inspect.CORO_CREATED
-            finally:
-                if inspect.iscoroutine(pending):
-                    pending.close()
-        finally:
-            owner.close()
-    finally:
-        original.close()
-
-
-def expected_control(witness: str, control: str, awaited: bool) -> object:
-    if witness in ("deferred_lifetime", "pending_handoff"):
-        if control == "weak" or (witness == "pending_handoff" and control == "missing_handoff"):
-            return LifetimeObservation(
-                (False, False, False), ExpiredBorrow(("callable", "positional:0", "keyword:alias"))
-            )
-        return LifetimeObservation((True, True, True), 42)
-    if witness == "borrowed_lifetime":
-        return BorrowedObservation(True, True, 17 if awaited else 13, 31 if awaited else 29)
-    return {"direct_coroutine": True}[witness]
-
-
-def run_control(witness: str, control: str, retained: bool, awaited: bool, factory: LiveCallFactory) -> None:
-    inner = factory if retained else ReferenceFactory()
-    owners = control_factory(control, inner)
-
-    async def run() -> None:
-        observed = await WITNESSES[witness](owners, awaited)
-        assert observed == expected_control(witness, control, awaited), (witness, control, awaited, observed)
-
-    run_checked(inner, run())
-
-
-WITNESSES: dict[str, Callable[[CallFactory, bool], object]] = {
-    "deferred_lifetime": deferred_lifetime,
-    "borrowed_lifetime": borrowed_lifetime,
-    "pending_handoff": pending_handoff,
-    "direct_coroutine": direct_coroutine,
-}

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

@@ -1,667 +0,0 @@
-import asyncio
-import copy
-import gzip
-import json
-import threading
-from collections import OrderedDict
-from dataclasses import dataclass
-from datetime import datetime
-from typing import Literal
-from unittest import TestCase
-
-import httpx
-from fastapi import HTTPException
-
-import litellm
-from litellm.integrations.custom_batch_logger import CustomBatchLogger
-from litellm.integrations.custom_guardrail import CustomGuardrail, ModifyResponseException
-from litellm.integrations.custom_logger import CustomLogger
-from litellm.integrations.datadog.datadog import DataDogLogger
-from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger
-from litellm.integrations.literal_ai import LiteralAILogger
-from litellm.integrations.rubrik import RubrikLogger
-from litellm.litellm_core_utils.litellm_logging import Logging, create_dummy_standard_logging_payload
-from litellm.llms.openai.chat.guardrail_translation.handler import OpenAIChatCompletionsHandler
-from litellm.proxy._types import UserAPIKeyAuth
-from litellm.proxy.guardrails.guardrail_hooks.crowdstrike_aidr.crowdstrike_aidr import CrowdStrikeAIDRHandler
-from litellm.proxy.guardrails.guardrail_hooks.microsoft_purview.purview_dlp import MicrosoftPurviewDLPGuardrail
-from litellm.types.guardrails import GuardrailEventHooks
-from litellm.types.utils import ModelResponse
-
-
-async def integration_invoke(owners, callback, *args, **kwargs):
-    owner = owners.prepare(callback, args, kwargs, awaited=True)
-    pending = owner.invoke()
-    owner.close()
-    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))
-
-
-@dataclass(frozen=True, slots=True)
-class QueueObservation:
-    gcs_model_parameters: str
-    datadog_snapshot: str
-    literal_prepared_settings: str
-
-
-async def real_logging_queue_copy_control(owners):
-    baseline = await integration_logging_queue_case(owners)
-    copied = await integration_logging_queue_case(owners, literal_copy="payload")
-    envelope = await integration_logging_queue_case(owners, literal_copy="envelope")
-    assert envelope == baseline
-    assert json.loads(baseline.gcs_model_parameters) == {"stream": True, "temperature": 0.25}
-    assert copied.datadog_snapshot == baseline.datadog_snapshot
-    assert copied.literal_prepared_settings == baseline.literal_prepared_settings
-    assert json.loads(copied.literal_prepared_settings) == {"stream": True}
-    assert json.loads(copied.gcs_model_parameters) == {
-        **json.loads(baseline.gcs_model_parameters),
-        "tools": [{"type": "function", "function": {"name": "lookup"}}],
-    }
-    return copied
-
-
-def queue_loggers(entered, release, uploads):
-    class VertexTransport:
-        async def _ensure_access_token_async(self, **kwargs):
-            entered.set()
-            await release.wait()
-            return "fixture-token", "fixture-project"
-
-        def _get_token_and_url(self, **kwargs):
-            return kwargs["auth_header"], None
-
-    class Transport:
-        async def post(self, url, **kwargs):
-            wire = {**kwargs, "json": json.loads(json.dumps(kwargs["json"]))} if "json" in kwargs else kwargs
-            uploads.append((url, wire))
-            return integration_response(url, {}, 202 if "datadog" in url else 200)
-
-    datadog = DataDogLogger.__new__(DataDogLogger)
-    CustomBatchLogger.__init__(datadog, batch_size=100, flush_lock=asyncio.Lock())
-    datadog.intake_url, datadog.DD_API_KEY, datadog.is_mock_mode = "https://datadog.invalid/logs", "test", False
-    datadog.async_client = Transport()
-    gcs = GCSBucketLogger.__new__(GCSBucketLogger)
-    CustomBatchLogger.__init__(gcs, batch_size=100)
-    gcs.log_queue = asyncio.Queue()
-    gcs.BUCKET_NAME, gcs.path_service_account_json = "fixture-bucket", None
-    gcs.vertex_instances = {"IAM_AUTH": VertexTransport()}
-    gcs.use_batched_logging = True
-    gcs.async_httpx_client = Transport()
-    literal = LiteralAILogger.__new__(LiteralAILogger)
-    CustomBatchLogger.__init__(literal, batch_size=100, flush_lock=asyncio.Lock())
-    literal.literalai_api_url, literal.headers = "https://literal.invalid", {}
-    literal.async_httpx_client = Transport()
-    return datadog, gcs, literal
-
-
-async def integration_logging_queue_case(owners, *, literal_copy: Literal["direct", "envelope", "payload"] = "direct"):
-    entered, release = asyncio.Event(), asyncio.Event()
-    uploads = []
-    datadog, gcs, literal = queue_loggers(entered, release, uploads)
-    copy_payload = literal_copy == "payload"
-
-    class LiteralCallback(CustomLogger):
-        def __init__(self, delegate):
-            super().__init__()
-            self.delegate = delegate
-            self.calls = self.completed = 0
-
-        async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
-            self.calls += 1
-            self.received = kwargs
-            self.forwarded = {
-                **kwargs,
-                "standard_logging_object": (
-                    copy.deepcopy(kwargs["standard_logging_object"])
-                    if copy_payload
-                    else kwargs["standard_logging_object"]
-                ),
-            }
-            await self.delegate.async_log_failure_event(self.forwarded, response_obj, start_time, end_time)
-            self.completed += 1
-
-    literal_callback = literal if literal_copy == "direct" else LiteralCallback(literal)
-
-    payload = create_dummy_standard_logging_payload()
-    payload.update(status="failure", error_str="x" * 10001)
-    messages, settings, metadata = payload["messages"], payload["model_parameters"], payload["metadata"]
-    completion = payload["response"]["choices"][0]["message"]
-    tools = [{"type": "function", "function": {"name": "lookup"}}]
-    settings["tools"] = tools
-    now = datetime.now()
-    logging = Logging(
-        model="fixture-model",
-        messages=messages,
-        stream=False,
-        call_type="acompletion",
-        start_time=now,
-        litellm_call_id="fixture-queue",
-        function_id="fixture",
-        dynamic_async_failure_callbacks=[datadog, gcs, literal_callback],
-    )
-    error = RuntimeError("fixture failure")
-    kwargs = logging.model_call_details
-    kwargs.update(standard_logging_object=payload, model="fixture-model", exception=error, end_time=now)
-    await integration_invoke(owners, logging.async_failure_handler, error, "fixture traceback", now, now)
-    if literal_copy != "direct":
-        assert literal_callback.calls == literal_callback.completed == 1
-        assert literal_callback.received is kwargs and literal_callback.forwarded is not kwargs
-        assert (literal_callback.forwarded["standard_logging_object"] is payload) is (not copy_payload)
-    assert len(datadog.log_queue) == gcs.log_queue.qsize() == len(literal.log_queue) == 1
-    assert kwargs["standard_logging_object"] is payload
-    assert payload["messages"] is messages and payload["model_parameters"] is settings
-    assert payload["error_str"].endswith("truncated by litellm, this logger does not support large content")
-    assert ("tools" in settings) is copy_payload
-    dd_snapshot = json.loads(datadog.log_queue[0]["message"])
-    assert dd_snapshot["model_parameters"]["tools"] == tools
-    queued = gcs.log_queue.get_nowait()
-    assert queued["payload"] is payload and queued["kwargs"] is kwargs and queued["response_obj"] is None
-    gcs.log_queue.put_nowait(queued)
-    generation = literal.log_queue[0]["generation"]
-    prepared_settings = json.dumps(generation["settings"], sort_keys=True)
-    assert "tools" not in generation["settings"] and generation["tools"] == tools
-    if copy_payload:
-        assert generation["settings"] is not settings and generation["tools"] is not tools
-        assert generation["messages"] is not messages and generation["messageCompletion"] is not completion
-        assert literal.log_queue[0]["metadata"] is not metadata
-    else:
-        assert generation["settings"] is settings and generation["tools"] is tools
-        assert generation["messages"] is messages and generation["messageCompletion"] is completion
-        assert literal.log_queue[0]["metadata"] is metadata
-
-    flush = asyncio.create_task(integration_invoke(owners, gcs.flush_queue))
-    try:
-        await entered.wait()
-        assert not uploads and not flush.done() and gcs.log_queue.empty()
-        messages[0]["content"] = "mutated before serialization"
-        settings["temperature"] = 0.25
-        completion["content"] = "late completion"
-        payload["messages"] = [{"role": "user", "content": "replacement field"}]
-        kwargs["standard_logging_object"] = {"replacement": True}
-        release.set()
-        await flush
-    finally:
-        release.set()
-        if not flush.done():
-            flush.cancel()
-        await asyncio.gather(flush, return_exceptions=True)
-    assert len(uploads) == 1
-    gcs_snapshot = json.loads(uploads[0][1]["data"])
-    assert gcs_snapshot["messages"] == payload["messages"]
-    assert gcs_snapshot["model_parameters"] == settings
-    assert "replacement" not in gcs_snapshot
-    await integration_invoke(owners, datadog.flush_queue)
-    await integration_invoke(owners, literal.flush_queue)
-    assert len(uploads) == 3 and not datadog.log_queue and not literal.log_queue
-    sent_dd = json.loads(gzip.decompress(uploads[1][1]["data"]))
-    assert json.loads(sent_dd[0]["message"]) == dd_snapshot
-    literal_wire = uploads[2][1]["json"]
-    sent_generation = literal_wire["variables"]["generation_0"]
-    assert sent_generation["messages"] != gcs_snapshot["messages"]
-    if copy_payload:
-        assert sent_generation["messages"] == dd_snapshot["messages"]
-        assert json.dumps(sent_generation["settings"], sort_keys=True) == prepared_settings
-        assert sent_generation["messageCompletion"] == dd_snapshot["response"]["choices"][0]["message"]
-    else:
-        assert sent_generation["messages"] == messages
-        assert sent_generation["settings"]["temperature"] == 0.25
-        assert sent_generation["messageCompletion"]["content"] == "late completion"
-    messages[0]["content"] = "after serialization"
-    assert sent_generation["messages"][0]["content"] == (
-        "Hello, world!" if copy_payload else "mutated before serialization"
-    )
-    assert dd_snapshot["messages"][0]["content"] == "Hello, world!"
-    return QueueObservation(
-        gcs_model_parameters=json.dumps(gcs_snapshot["model_parameters"], sort_keys=True),
-        datadog_snapshot=json.dumps(dd_snapshot, sort_keys=True),
-        literal_prepared_settings=prepared_settings,
-    )
-
-
-async def real_crowdstrike_translator_identity(owners):
-    entered, release = asyncio.Event(), asyncio.Event()
-    calls = []
-
-    class Transport:
-        async def post(self, url, json, **kwargs):
-            calls.append(json)
-            entered.set()
-            await release.wait()
-            return integration_response(
-                url,
-                {
-                    "result": {
-                        "blocked": False,
-                        "transformed": True,
-                        "guard_output": {"messages": [{"role": "user", "content": "redacted"}]},
-                    }
-                },
-            )
-
-    guardrail = CrowdStrikeAIDRHandler.__new__(CrowdStrikeAIDRHandler)
-    CustomGuardrail.__init__(guardrail, guardrail_name="fixture-crowdstrike", event_hook=GuardrailEventHooks.pre_call)
-    guardrail.api_base, guardrail.api_key, guardrail.fail_on_error = "https://crowdstrike.invalid", "test", True
-    guardrail.skip_system_message_in_guardrail = True
-    guardrail.async_handler = Transport()
-    system = {"role": "system", "content": "internal policy"}
-    user = {"role": "user", "content": "private text", "extra": {"retained": True}}
-    messages = [system, user]
-    data = {"model": "fixture-model", "messages": messages}
-    task = asyncio.create_task(
-        integration_invoke(
-            owners,
-            OpenAIChatCompletionsHandler().process_input_messages,
-            data,
-            guardrail,
-        )
-    )
-    try:
-        await entered.wait()
-        assert data["messages"] is messages and not task.done()
-        assert calls[0]["guard_input"]["messages"] == [{"role": "user", "content": "private text"}]
-        user["extra"]["during_http"] = True
-        release.set()
-        assert await task is data
-    finally:
-        release.set()
-        if not task.done():
-            task.cancel()
-        await asyncio.gather(task, return_exceptions=True)
-    assert data["messages"] is not messages
-    assert data["messages"][0] is system
-    assert data["messages"][1] is not user
-    assert data["messages"][1]["extra"] is user["extra"]
-    assert data["messages"][1]["content"] == "redacted"
-    assert user["content"] == "private text" and messages[1] is user
-    detached = copy.deepcopy(messages[1:])
-    inputs = {"texts": ["private text"], "structured_messages": detached}
-    control = await integration_invoke(owners, guardrail.apply_guardrail, inputs, {"messages": messages}, "request")
-    assert len(calls) == 2
-    assert control["structured_messages"] is detached and detached[0] is not user
-    assert control["texts"] == ["redacted"]
-    assert detached[0]["content"] == user["content"] == "private text"
-
-
-async def real_rubrik_block_lifecycle(owners):
-    for input_type, populated in (("request", False), ("response", True)):
-        entered, release = asyncio.Event(), asyncio.Event()
-        moderation, uploads = [], []
-
-        class Transport:
-            def __init__(self, moderation, uploads, entered, release):
-                self.moderation, self.uploads = moderation, uploads
-                self.entered, self.release = entered, release
-
-            async def post(self, url, json, **kwargs):
-                if url.endswith("/batch"):
-                    self.uploads.append(json)
-                    return integration_response(url, {})
-                self.moderation.append(json)
-                self.entered.set()
-                await self.release.wait()
-                return integration_response(url, {"choices": [{"message": {"content": "blocked by policy"}}]})
-
-        rubrik = RubrikLogger.__new__(RubrikLogger)
-        CustomGuardrail.__init__(
-            rubrik,
-            guardrail_name="fixture-rubrik",
-            event_hook=GuardrailEventHooks.post_call,
-            flush_lock=asyncio.Lock(),
-            batch_size=100,
-        )
-        rubrik._periodic_flush_task = None
-        rubrik.sampling_rate, rubrik._headers = 1.0, {}
-        rubrik._dropped_since_warning, rubrik._last_drop_warning_time = 0, 0.0
-        rubrik.prompt_moderation_endpoint = "https://rubrik.invalid/before"
-        rubrik.response_moderation_endpoint = "https://rubrik.invalid/after"
-        rubrik.logging_endpoint = "https://rubrik.invalid/batch"
-        rubrik.moderation_client = rubrik.async_httpx_client = Transport(moderation, uploads, entered, release)
-        other = RubrikLogger.__new__(RubrikLogger)
-        CustomGuardrail.__init__(other, guardrail_name="fixture-other-rubrik", event_hook=GuardrailEventHooks.post_call)
-        logging = Logging(
-            model="fixture-model",
-            messages=[{"role": "user", "content": "original prompt"}],
-            stream=False,
-            call_type="acompletion",
-            start_time=datetime.now(),
-            litellm_call_id="fixture-correlation",
-            function_id="fixture",
-        )
-        details = logging.model_call_details
-        details.update(messages=logging.messages, model="fixture-model", litellm_call_id="fixture-correlation")
-        details["system"] = "system scaffold"
-        if populated:
-            details["standard_logging_object"] = create_dummy_standard_logging_payload()
-        messages = details["messages"]
-        request = {"model": "fixture-model", "litellm_call_id": "fixture-correlation", "messages": messages}
-        inputs = {"texts": ["original response"], "structured_messages": messages}
-        success = owners.prepare(rubrik.async_log_success_event, (details, None, None, None), awaited=True)
-        task = asyncio.create_task(
-            integration_invoke(owners, rubrik.apply_guardrail, inputs, request, input_type, logging)
-        )
-        try:
-            await entered.wait()
-            assert not task.done() and "_rubrik_logging_obj" not in request
-            assert "_rubrik_blocked" not in details
-            if input_type == "request":
-                assert moderation[0]["correlation_key"] == "fixture-correlation"
-                assert moderation[0]["messages"][0]["content"] == "original prompt"
-            else:
-                assert moderation[0]["request"]["messages"] is messages
-                assert moderation[0]["response"]["id"] == "fixture-correlation"
-            release.set()
-            with TestCase().assertRaises(ModifyResponseException) as caught:
-                await task
-            error = caught.exception
-            assert error.request_data is request and error.message == "blocked by policy"
-            assert request["_rubrik_logging_obj"] is logging and details["_rubrik_blocked"] is True
-            await integration_invoke(
-                owners, other.async_post_call_failure_hook, request, error, UserAPIKeyAuth(user_id="fixture-user")
-            )
-            assert request["_rubrik_logging_obj"] is logging and details["_rubrik_blocked"] is True
-            assert not other.log_queue and not rubrik.log_queue and not uploads
-            await integration_invoke(
-                owners, rubrik.async_post_call_failure_hook, request, error, UserAPIKeyAuth(user_id="fixture-user")
-            )
-            assert "_rubrik_logging_obj" not in request and details["_rubrik_blocked"] is True
-            assert len(rubrik.log_queue) == 1
-            queued = rubrik.log_queue[0]
-            assert queued["id"] == "fixture-correlation"
-            assert queued["response"] == "ModifyResponseException: blocked by policy"
-            assert queued["messages"][0] == {"role": "system", "content": "system scaffold"}
-            assert messages[0] == {"role": "user", "content": "original prompt"}
-            if populated:
-                base = details["standard_logging_object"]
-                assert queued["metadata"] is not base["metadata"]
-                assert queued["messages"][1] is not base["messages"][0]
-                assert isinstance(base["response"], dict)
-            else:
-                assert queued["messages"][1] is messages[0]
-                assert queued["status"] == "failure"
-                assert queued["metadata"]["user_api_key_user_id"] == "fixture-user"
-            pending = success.invoke()
-            success.close()
-            assert await pending is None
-            assert len(rubrik.log_queue) == 1 and rubrik.log_queue[0] is queued
-            await integration_invoke(owners, rubrik.flush_queue)
-            assert not rubrik.log_queue and len(uploads) == 1 and uploads[0][0] is queued
-        finally:
-            success.close()
-            release.set()
-            if not task.done():
-                task.cancel()
-            await asyncio.gather(task, return_exceptions=True)
-            await rubrik.aclose()
-            if rubrik._periodic_flush_task is not None:
-                await asyncio.gather(rubrik._periodic_flush_task, return_exceptions=True)
-
-
-async def real_parallel_guardrail_snapshots(owners):
-    original_mode = litellm.safe_memory_mode
-    try:
-        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, *, 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
-
-    class Uncopyable:
-        def __init__(self):
-            self.attempts = 0
-            self.observed = []
-
-        def __deepcopy__(self, memo):
-            self.attempts += 1
-            raise TypeError("fixture cannot be copied")
-
-    arrived, release, mutated = asyncio.Event(), asyncio.Event(), asyncio.Event()
-    observations = {}
-    sentinel = Uncopyable()
-    live = {"messages": [{"role": "user", "content": "original"}], "uncopyable": sentinel}
-    raw = independent_snapshot(live)
-    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):
-            observations[self.guardrail_name] = data
-            if len(observations) == 4:
-                arrived.set()
-            await release.wait()
-            if self.guardrail_name == "fixture-live-writer":
-                data["messages"][0]["content"] = "shared mutation"
-                mutated.set()
-            await mutated.wait()
-            if self.scan_raw_request:
-                assert data["messages"][0]["content"] == "original"
-                data["messages"][0]["content"] = self.guardrail_name
-                assert data["uncopyable"] is sentinel
-                data["uncopyable"].observed.append(self.guardrail_name)
-            else:
-                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(
-        Inspect(
-            guardrail_name=name,
-            event_hook=GuardrailEventHooks.pre_call,
-            default_on=True,
-            run_in_parallel=True,
-            scan_raw_request="raw" in name,
-        )
-        for name in ("fixture-live-writer", "fixture-live-reader", "fixture-raw-a", "fixture-raw-b")
-    )
-    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,
-            dispatched,
-            raw,
-            UserAPIKeyAuth(),
-            "acompletion",
-        )
-    )
-    try:
-        await arrived.wait()
-        assert not task.done()
-        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 + copy_live and not sentinel.observed
-        release.set()
-        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 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 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(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()
-        if not task.done():
-            task.cancel()
-        await asyncio.gather(task, return_exceptions=True)
-
-
-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()
-
-    class Transport:
-        async def post(self, url, **kwargs):
-            workers.append(threading.current_thread())
-            calls.append((url, kwargs))
-            if url.endswith("/token"):
-                entered.set()
-                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"):
-                return integration_response(url, {}, headers={"etag": "fixture-etag"})
-            assert url.endswith("/processContent")
-            assert kwargs["headers"]["If-None-Match"] == "fixture-etag"
-            return integration_response(
-                url, {"policyActions": [{"action": "restrictAccess", "restrictionAction": "block"}]}
-            )
-
-    purview = MicrosoftPurviewDLPGuardrail.__new__(MicrosoftPurviewDLPGuardrail)
-    CustomGuardrail.__init__(purview, guardrail_name="fixture-purview", event_hook=GuardrailEventHooks.logging_only)
-    purview.async_handler = Transport()
-    purview.tenant_id, purview.client_id, purview.client_secret = "fixture-tenant", "fixture-client", "test"
-    purview.purview_app_name, purview.user_id_field, purview.guardrail_provider = (
-        "fixture",
-        "user_id",
-        "microsoft_purview",
-    )
-    purview._token_cache, purview._scope_cache = None, OrderedDict()
-    purview._scope_cache_maxsize, purview._cache_lock = 1000, threading.Lock()
-    metadata = {"user_api_key_user_id": "fixture-user"}
-    kwargs = {
-        "messages": [{"role": "user", "content": "prompt at dispatch"}],
-        "litellm_params": {"metadata": metadata},
-        "litellm_call_id": "before-http",
-    }
-    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:
-        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
-        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()
-        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"]
-        assert entries[0]["contentEntries"][0]["content"]["data"] == "prompt at dispatch"
-        assert entries[1]["contentEntries"][0]["content"]["data"] == "response mutated while audit waits"
-        assert all(entry["contentEntries"][0]["correlationId"] == "after-http" for entry in entries)
-        assert kwargs["litellm_params"]["metadata"] is metadata
-        info = kwargs["metadata"]["standard_logging_guardrail_information"]
-        assert len(info) == 2 and all(item["guardrail_status"] == "guardrail_intervened" for item in info)
-        assert all(item["end_time"] >= item["start_time"] and item["duration"] >= 0 for item in info)
-    finally:
-        owner.close()
-        release.set()
-        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

@@ -1,633 +0,0 @@
-import asyncio
-import contextvars
-import copy
-import gc
-import inspect
-import json
-import sys
-import threading
-import weakref
-from unittest import TestCase
-
-
-async def checkpoint():
-    ready = asyncio.Event()
-    asyncio.get_running_loop().call_soon(ready.set)
-    await ready.wait()
-
-
-async def settle(pending, awaited):
-    return await pending if awaited else pending
-
-
-class Value:
-    pass
-
-
-def cold_awaited_adapter_reentry(factory):
-    events = []
-    compilations = []
-
-    def invoke_failure():
-        error = LookupError("cold-cache callback error")
-
-        async def failing():
-            raise error
-
-        owner = factory.prepare(failing, (), awaited=True)
-        pending = owner.invoke()
-        try:
-            assert inspect.getcoroutinestate(pending) == inspect.CORO_CREATED
-            with TestCase().assertRaises(LookupError) as caught:
-                pending.send(None)
-            assert caught.exception is error
-        finally:
-            pending.close()
-            owner.close()
-            error.__traceback__ = None
-
-    def audit(event, args):
-        if event != "compile" or args[1] != "retained_callback.py":
-            return
-        compilations.append(args[1])
-        if len(compilations) == 1:
-            events.append("entered")
-            invoke_failure()
-            events.append("nested completed")
-
-    sys.addaudithook(audit)
-    invoke_failure()
-    events.append("outer completed")
-    assert events == ["entered", "nested completed", "outer completed"]
-    assert len(compilations) == 2
-    assert factory.live == 0
-
-    result = object()
-    direct = factory.prepare(lambda value: value, (result,))
-    try:
-        assert direct.invoke() is result
-    finally:
-        direct.close()
-
-    async def successful(value, *, alias):
-        assert value is alias
-        await checkpoint()
-        return value
-
-    owner = factory.prepare(successful, (result,), {"alias": result}, awaited=True)
-    try:
-        assert asyncio.run(owner.invoke()) is result
-    finally:
-        owner.close()
-    assert len(compilations) == 2
-    assert factory.live == 0
-
-
-class ReferenceFactory:
-    def __init__(self):
-        self.live = 0
-
-    def prepare(self, callable, positional, keywords=None, awaited=False):
-        return ReferenceOwner(self, callable, positional, keywords, awaited)
-
-
-class ReferenceOwner:
-    def __init__(self, factory, callable, positional, keywords, awaited):
-        self.factory = factory
-        self.call = (callable, positional, keywords, awaited)
-        factory.live += 1
-
-    def invoke(self):
-        if self.call is None:
-            raise RuntimeError("invocation owner released")
-        callable, positional, keywords, awaited = self.call
-        if not awaited:
-            return callable(*positional, **(keywords if keywords is not None else {}))
-
-        async def run():
-            return await callable(*positional, **(keywords if keywords is not None else {}))
-
-        return run()
-
-    def clone_owner(self):
-        return ReferenceOwner(self.factory, *self.call)
-
-    def close(self):
-        if self.call is not None:
-            released, self.call = self.call, None
-            self.factory.live -= 1
-            del released
-
-    def __del__(self):
-        self.close()
-
-
-async def awaitable_kinds(owners):
-    payload = Value()
-    calls = []
-
-    async def coroutine(value, *, alias):
-        assert value is alias
-        calls.append("called")
-        return value
-
-    class CustomAwaitable:
-        def __await__(self):
-            return coroutine(payload, alias=payload).__await__()
-
-    for kind in ("async", "sync_coroutine", "custom", "future"):
-        future = asyncio.get_running_loop().create_future()
-        future.set_result(payload)
-        callback = {
-            "async": coroutine,
-            "sync_coroutine": lambda value, *, alias: coroutine(value, alias=alias),
-            "custom": lambda value, *, alias: CustomAwaitable(),
-            "future": lambda value, *, alias, future=future: future,
-        }[kind]
-        owner = owners.prepare(callback, (payload,), {"alias": payload}, awaited=True)
-        pending = owner.invoke()
-        before = len(calls)
-        owner.close()
-        assert await pending is payload
-        assert len(calls) == before + (kind != "future")
-
-    owner = owners.prepare(lambda: payload, (), awaited=True)
-    with TestCase().assertRaises(TypeError):
-        await owner.invoke()
-    owner.close()
-
-    inner = coroutine(payload, alias=payload)
-
-    async def returns_coroutine():
-        return inner
-
-    owner = owners.prepare(returns_coroutine, (), awaited=True)
-    assert await owner.invoke() is inner
-    assert inner.cr_frame is not None
-    inner.close()
-    owner.close()
-
-    direct = owners.prepare(coroutine, (payload,), {"alias": payload})
-    untouched = direct.invoke()
-    assert untouched.cr_frame is not None
-    assert untouched.cr_await is None
-    untouched.close()
-    direct.close()
-
-
-async def identity_and_context(owners):
-    context = contextvars.ContextVar("retained_context", default="outside")
-    task = asyncio.current_task()
-    loop = asyncio.get_running_loop()
-    thread = threading.get_ident()
-    payload = {"nested": {}}
-    alias = payload["nested"]
-    saved = []
-    gate = asyncio.Event()
-
-    async def nested(value):
-        assert asyncio.current_task() is task
-        assert context.get() == "inside"
-        value["nested"]["nested_call"] = True
-        context.set("nested")
-        return value
-
-    async def callback(value, *, shared):
-        assert value is payload and shared is alias
-        assert asyncio.current_task() is task
-        assert asyncio.get_running_loop() is loop
-        assert threading.get_ident() == thread
-        assert context.get() == "at_await"
-        owner.close()
-        payload["closure_mutation"] = True
-        context.set("inside")
-        saved.append(value)
-        shared["before"] = True
-        loop.call_soon(gate.set)
-        await gate.wait()
-        inner = owners.prepare(nested, (value,), awaited=True)
-        try:
-            assert await inner.invoke() is value
-        finally:
-            inner.close()
-        return value
-
-    owner = owners.prepare(callback, (payload,), {"shared": alias}, awaited=True)
-    pending = owner.invoke()
-    context.set("at_await")
-    assert await pending is payload
-    assert context.get() == "nested"
-    assert payload["closure_mutation"] is True
-    assert owners.live == 0
-    alias["after"] = True
-    assert saved[0]["nested"] == {"before": True, "nested_call": True, "after": True}
-
-
-async def exceptions(owners):
-    for error in (RuntimeError("original"), KeyboardInterrupt("original"), asyncio.CancelledError("original")):
-        cause = ValueError("cause")
-        payload = {}
-
-        async def callback(payload=payload, error=error, cause=cause):
-            payload["changed"] = True
-            raise error from cause
-
-        owner = owners.prepare(callback, (), awaited=True)
-        caught_error = None
-        try:
-            await owner.invoke()
-        except BaseException as caught:
-            caught_error = caught
-        finally:
-            owner.close()
-        assert caught_error is error and caught_error.__cause__ is cause
-        frames = []
-        tb = caught_error.__traceback__
-        while tb:
-            frames.append(tb.tb_frame.f_code.co_name)
-            tb = tb.tb_next
-        assert "callback" in frames
-        assert payload["changed"] is True
-
-
-async def exception_ownership(owners):
-    class Callback:
-        def __init__(self, error):
-            self.error = error
-
-        async def __call__(self, value):
-            value.changed = True
-            raise self.error
-
-    value = Value()
-    error = RuntimeError("retained exception")
-    callback = Callback(error)
-    value_ref, callback_ref = weakref.ref(value), weakref.ref(callback)
-    owner = owners.prepare(callback, (value,), awaited=True)
-    del value, callback
-    caught_error = None
-    try:
-        await owner.invoke()
-    except RuntimeError as caught:
-        caught_error = caught
-    assert caught_error is error
-    owner.close()
-    assert value_ref().changed and callback_ref() is not None
-    del error, caught_error
-    gc.collect()
-    assert value_ref() is None and callback_ref() is None
-
-
-async def cancellation_before_start(owners):
-    started = []
-
-    async def callback(value):
-        started.append(value)
-
-    for operation in ("close", "cancel"):
-        value = Value()
-        ref = weakref.ref(value)
-        owner = owners.prepare(callback, (value,), awaited=True)
-        pending = owner.invoke()
-        owner.close()
-        del value
-        assert ref() is not None
-        if operation == "close":
-            pending.close()
-        else:
-            task = asyncio.create_task(pending)
-            task.cancel()
-            try:
-                await task
-            except asyncio.CancelledError:
-                pass
-            del task
-        del pending
-        await checkpoint()
-        gc.collect()
-        assert ref() is None
-    assert started == []
-
-
-async def cancellation_case(owners, repeated=False, suppress=False):
-    started, cleaning, finish = asyncio.Event(), asyncio.Event(), asyncio.Event()
-    value = Value()
-    ref = weakref.ref(value)
-    observed = []
-
-    async def callback(argument):
-        try:
-            started.set()
-            await asyncio.Event().wait()
-        except asyncio.CancelledError:
-            cleaning.set()
-            try:
-                await finish.wait()
-            except asyncio.CancelledError:
-                observed.append("second cancellation")
-                await finish.wait()
-            argument.cleaned = True
-            if suppress:
-                return argument
-            raise
-
-    owner = owners.prepare(callback, (value,), awaited=True)
-    task = asyncio.create_task(owner.invoke())
-    owner.close()
-    del value
-    try:
-        await started.wait()
-        task.cancel()
-        await cleaning.wait()
-        assert ref() is not None and not task.done()
-        if repeated:
-            task.cancel()
-            barrier = asyncio.Event()
-            asyncio.get_running_loop().call_soon(barrier.set)
-            await barrier.wait()
-            assert observed == ["second cancellation"] and not task.done()
-        finish.set()
-        try:
-            result = await task
-            assert suppress and result is ref() and result.cleaned
-            del result
-        except asyncio.CancelledError:
-            assert not suppress
-            assert ref().cleaned
-    finally:
-        finish.set()
-        if not task.done():
-            task.cancel()
-            await asyncio.gather(task, return_exceptions=True)
-    del task
-    await checkpoint()
-    gc.collect()
-    assert ref() is None
-
-
-async def cancellation_unwinds(owners):
-    await cancellation_case(owners)
-
-
-async def cancellation_during_cleanup(owners):
-    await cancellation_case(owners, repeated=True)
-
-
-async def cancellation_suppressed(owners):
-    await cancellation_case(owners, suppress=True)
-
-
-async def registration_and_gc(owners):
-    original = Value()
-    original_ref = weakref.ref(original)
-    registered = owners.prepare(lambda value: value, (original,))
-    active = registered.clone_owner()
-    registered.close()
-    registered = owners.prepare(lambda: "replacement", ())
-    del original
-    assert active.invoke() is original_ref()
-    active.close()
-    gc.collect()
-    assert original_ref() is None
-    assert registered.invoke() == "replacement"
-    registered.close()
-
-    for edge in ("callable", "positional", "keywords"):
-
-        class Callback:
-            def __call__(self, *args, **kwargs):
-                pass
-
-        value = Callback()
-        owner = owners.prepare(
-            value if edge == "callable" else lambda *a, **k: None,
-            (value,) if edge == "positional" else (),
-            {"value": value} if edge == "keywords" else None,
-        )
-        value.owner = owner
-        value_ref, owner_ref = weakref.ref(value), weakref.ref(owner)
-        del value, owner
-        gc.collect()
-        assert value_ref() is None and owner_ref() is None
-        assert owners.live == 0
-
-    finalized = []
-
-    class Reenter:
-        def __del__(self):
-            try:
-                reentrant.close()
-                another = owners.prepare(lambda: 42, ())
-                finalized.append(another.invoke())
-                another.close()
-            except BaseException as error:
-                finalized.append(type(error).__name__)
-
-    value = Reenter()
-    reentrant = owners.prepare(lambda value: None, (value,))
-    del value
-    reentrant.close()
-    reentrant.close()
-    assert finalized == [42]
-    assert owners.live == 0
-
-
-async def background_and_session(owners):
-    context = contextvars.ContextVar("background_context", default="initial")
-    start, finish = asyncio.Event(), asyncio.Event()
-    payload = {"nested": {"value": "queued"}}
-    saved = []
-
-    async def upload(value):
-        assert context.get() == "submission"
-        start.set()
-        await finish.wait()
-        saved.append(json.dumps(value))
-
-    session = owners.prepare(lambda value: value, (payload,))
-    first_response, second_response = session.clone_owner(), session.clone_owner()
-    upload_owner = owners.prepare(upload, (payload,), awaited=True)
-    context.set("submission")
-    task = asyncio.create_task(upload_owner.invoke())
-    context.set("consumer")
-    upload_owner.close()
-    first_response.close()
-    await start.wait()
-    assert second_response.invoke() is payload
-    payload["nested"]["value"] = "later"
-    consumed = json.dumps(payload)
-    payload["nested"]["after_consumption"] = True
-    assert "after_consumption" not in consumed
-    second_response.close()
-    session.close()
-    assert owners.live == 0 and not task.done()
-    finish.set()
-    await task
-    assert json.loads(saved[0]) == payload
-    assert context.get() == "consumer"
-
-
-async def stream_lifecycle(owners):
-    for terminal in ("exhaustion", "failure", "close"):
-        nested = {"usage": 0}
-        item = {"nested": nested}
-        closed = []
-
-        async def source(item=item, nested=nested, terminal=terminal, closed=closed):
-            try:
-                yield item
-                nested["usage"] = 12
-                if terminal == "failure":
-                    raise ValueError("stream failure")
-            finally:
-                closed.append(True)
-
-        stream = source()
-        pull = owners.prepare(stream.__anext__, (), awaited=True)
-        yielded = await pull.invoke()
-        assert yielded is item
-        shallow, deep = copy.copy(yielded), copy.deepcopy(yielded)
-        retained = owners.prepare(lambda value: value, (yielded["nested"],))
-        yielded["nested"] = {"replacement": True}
-        if terminal == "close":
-            close = owners.prepare(stream.aclose, (), awaited=True)
-            await close.invoke()
-            await close.invoke()
-            close.close()
-        else:
-            with TestCase().assertRaises(ValueError if terminal == "failure" else StopAsyncIteration):
-                await pull.invoke()
-        pull.close()
-        assert closed == [True]
-        assert retained.invoke() is nested
-        assert shallow["nested"] is nested and deep["nested"]["usage"] == 0
-        assert nested["usage"] == (0 if terminal == "close" else 12)
-        retained.close()
-
-
-async def sync_stream_lifecycle(owners):
-    for terminal in ("exhaustion", "failure", "close"):
-        value = {"nested": {"usage": 0}}
-        closed = []
-
-        def source(value=value, terminal=terminal, closed=closed):
-            try:
-                yield value
-                value["nested"]["usage"] = 12
-                if terminal == "failure":
-                    raise ValueError("stream failure")
-            finally:
-                closed.append(True)
-
-        stream = source()
-        pull = owners.prepare(stream.__next__, ())
-        assert pull.invoke() is value
-        saved = owners.prepare(lambda value: value, (value,))
-        if terminal == "close":
-            close = owners.prepare(stream.close, ())
-            close.invoke()
-            close.invoke()
-            close.close()
-        else:
-            with TestCase().assertRaises(ValueError if terminal == "failure" else StopIteration):
-                pull.invoke()
-        pull.close()
-        assert saved.invoke() is value
-        assert value["nested"]["usage"] == (0 if terminal == "close" else 12)
-        assert closed == [True]
-        saved.close()
-
-
-async def repeated_ownership(owners):
-    refs = []
-    for batch in range(8):
-        gate = asyncio.Event()
-
-        async def work(value, gate=gate):
-            await gate.wait()
-            return None
-
-        tasks = []
-        for index in range(4):
-            value = Value()
-            refs.append(weakref.ref(value))
-            owner = owners.prepare(work, (value,), awaited=True)
-            tasks.append(asyncio.create_task(owner.invoke()))
-            owner.close()
-            del value
-        gate.set()
-        await asyncio.gather(*tasks)
-        del tasks
-        gc.collect()
-        assert owners.live == 0
-        assert all(ref() is None for ref in refs)
-
-
-async def detached_work_after_error(owners):
-    for raises in (False, True):
-        entered, release = asyncio.Event(), asyncio.Event()
-        tasks, observed = [], []
-        value = Value()
-        value.status = "before return"
-        reference = weakref.ref(value)
-
-        async def consume(argument, entered=entered, release=release, observed=observed):
-            entered.set()
-            await release.wait()
-            observed.append(argument.status)
-
-        def callback(argument, tasks=tasks, consume=consume, raises=raises):
-            tasks.append(asyncio.create_task(consume(argument)))
-            if raises:
-                raise ValueError("after task creation")
-
-        owner = owners.prepare(callback, (value,))
-        try:
-            if raises:
-                with TestCase().assertRaisesRegex(ValueError, "after task creation"):
-                    owner.invoke()
-            else:
-                assert owner.invoke() is None
-        finally:
-            owner.close()
-        del value
-        try:
-            await entered.wait()
-            assert owners.live == 0 and reference() is not None
-            reference().status = "after invocation"
-            release.set()
-            await tasks[0]
-            assert observed == ["after invocation"]
-        finally:
-            release.set()
-            await asyncio.gather(*tasks, return_exceptions=True)
-        tasks.clear()
-        await checkpoint()
-        gc.collect()
-        assert reference() is None
-
-
-def run_checked(owners, scenario):
-    baseline = owners.live
-    background_failures = []
-
-    async def run():
-        asyncio.get_running_loop().set_exception_handler(lambda loop, context: background_failures.append(context))
-        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
-
-    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):
-    owners = factory if retained else ReferenceFactory()
-    assert owners.live == 0
-    run_checked(owners, globals()[name](owners))

litellm-rust/crates/python-interop/tests/integration/lifecycle.rs

@@ -1,96 +0,0 @@
-use pyo3::prelude::*;
-use pyo3::types::PyDict;
-use rstest::{fixture, rstest};
-use serial_test::serial;
-
-use crate::support::Backend;
-use crate::support::python::run_fixture;
-use crate::support::scenarios::{run_scenario_fixture, scenario_scope};
-
-#[fixture]
-fn component_scope(scenario_scope: Py<PyDict>) -> Py<PyDict> {
-    Python::attach(|py| {
-        run_fixture(
-            py,
-            scenario_scope.bind(py),
-            include_str!("../fixtures/callback_components.py"),
-            concat!(
-                env!("CARGO_MANIFEST_DIR"),
-                "/tests/fixtures/callback_components.py"
-            ),
-        )
-        .unwrap();
-    });
-    scenario_scope
-}
-
-#[fixture]
-fn integration_scope(scenario_scope: Py<PyDict>) -> Py<PyDict> {
-    Python::attach(|py| {
-        run_fixture(
-            py,
-            scenario_scope.bind(py),
-            include_str!("../fixtures/callback_integrations.py"),
-            concat!(
-                env!("CARGO_MANIFEST_DIR"),
-                "/tests/fixtures/callback_integrations.py"
-            ),
-        )
-        .unwrap();
-    });
-    scenario_scope
-}
-
-#[rstest]
-#[case::identity_and_ignored_returns("pre_call_identity_and_ignored_returns")]
-#[case::mutations_visible_to_later_callbacks("pre_call_mutations_visible_to_later_callbacks")]
-#[case::mutation_survives_failure("pre_call_mutation_survives_failure")]
-#[ignore = "requires the repository Python environment and LiteLLM on PYTHONPATH"]
-#[serial(python_interpreter)]
-fn pre_call_contract(
-    component_scope: Py<PyDict>,
-    #[case] scenario: &str,
-    #[values(Backend::Python, Backend::PreparedCall)] backend: Backend,
-) -> PyResult<()> {
-    run_scenario_fixture(component_scope, scenario, backend)
-}
-
-#[rstest]
-#[case::real_post_call_logging("real_post_call_logging")]
-#[case::real_post_call_dict_response("real_post_call_dict_response")]
-#[case::real_sync_logging("real_sync_logging")]
-#[case::real_sync_logging_hook_failure("real_sync_logging_hook_failure")]
-#[case::real_sync_failure_chain("real_sync_failure_chain")]
-#[case::real_async_failure_chain("real_async_failure_chain")]
-#[case::real_async_logging("real_async_logging")]
-#[case::real_copy_boundaries("real_copy_boundaries")]
-#[case::real_logging_worker("real_logging_worker")]
-#[case::real_sync_stream_copies("real_sync_stream_copies")]
-#[case::real_stream_completion("real_stream_completion")]
-#[case::real_stream_close("real_stream_close")]
-#[case::real_stream_cancellation("real_stream_cancellation")]
-#[ignore = "requires the repository Python environment and LiteLLM on PYTHONPATH"]
-#[serial(python_interpreter)]
-fn component_contract(
-    component_scope: Py<PyDict>,
-    #[case] scenario: &str,
-    #[values(Backend::Python, Backend::PreparedCall)] backend: Backend,
-) -> PyResult<()> {
-    run_scenario_fixture(component_scope, scenario, backend)
-}
-
-#[rstest]
-#[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_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(
-    integration_scope: Py<PyDict>,
-    #[case] scenario: &str,
-    #[values(Backend::Python, Backend::PreparedCall)] backend: Backend,
-) -> PyResult<()> {
-    run_scenario_fixture(integration_scope, scenario, backend)
-}

litellm-rust/crates/python-interop/tests/integration/mod.rs

@@ -1,5 +0,0 @@
-#[path = "../support/mod.rs"]
-mod support;
-
-mod lifecycle;
-mod ocr;

litellm-rust/crates/python-interop/tests/integration/ocr.rs

@@ -1,175 +0,0 @@
-use litellm_python_interop::{InvocationMode, InvocationOutcome, PreparedCall};
-use pyo3::prelude::*;
-use pyo3::types::{PyDict, PyList, PyTuple};
-use rstest::rstest;
-use serial_test::serial;
-
-use crate::support::python::{InitializedPython, initialized_python, item, scope};
-
-fn prepare_pre_call(
-    py: Python<'_>,
-    logger: &Bound<'_, PyAny>,
-    view: &Bound<'_, PyDict>,
-) -> PyResult<PreparedCall> {
-    let keywords = PyDict::new(py);
-    keywords.set_item("input", "OCR document processing")?;
-    keywords.set_item("api_key", py.None())?;
-    keywords.set_item("additional_args", view)?;
-    Ok(PreparedCall::new(
-        InvocationMode::Direct,
-        logger.getattr("pre_call")?.unbind(),
-        PyTuple::empty(py).unbind(),
-        Some(keywords.unbind()),
-    ))
-}
-
-#[rstest]
-#[ignore = "requires the repository Python environment and LiteLLM on PYTHONPATH"]
-#[serial(python_interpreter)]
-fn real_ocr_logging_preserves_execution_roots_and_continues_after_error(
-    initialized_python: &InitializedPython,
-) -> PyResult<()> {
-    let _ = initialized_python;
-    Python::attach(|py| {
-        let globals = scope(
-            py,
-            c"
-from litellm.integrations.custom_logger import CustomLogger
-
-class Retain(CustomLogger):
-    def log_pre_api_call(self, model, messages, kwargs):
-        order.append('retain')
-        self.view = kwargs['additional_args']
-        self.headers = self.view['headers']
-        self.body = self.view['complete_input_dict']
-        self.snapshot = (self.headers['X-Trace'], self.body['document']['value'])
-        return {'ignored_replacement': True}
-
-class MutateThenFail(CustomLogger):
-    def log_pre_api_call(self, model, messages, kwargs):
-        order.append('mutate_then_fail')
-        view = kwargs['additional_args']
-        view['headers']['X-Trace'] = 'mutated'
-        view['complete_input_dict']['document']['value'] = 'mutated'
-        view['headers'] = {'X-Trace': 'replacement'}
-        view['complete_input_dict'] = {'replacement': True}
-        raise RuntimeError('expected callback failure')
-
-class Observe(CustomLogger):
-    def log_pre_api_call(self, model, messages, kwargs):
-        order.append('observe')
-        self.view = kwargs['additional_args']
-        self.snapshot = (
-            tuple(sorted(self.view['headers'].items())),
-            self.view['complete_input_dict'].get('replacement'),
-            'document' in self.view['complete_input_dict'],
-        )
-
-",
-        )?;
-        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)?;
-        match invocation.invoke(py)? {
-            InvocationOutcome::Returned(value) => assert!(value.is_none(py)),
-            InvocationOutcome::Awaitable(_) => {
-                panic!("direct binding produced an awaitable outcome")
-            }
-        }
-        drop(invocation);
-        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!(
-            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"
-        );
-        Ok(())
-    })
-}

litellm-rust/crates/python-interop/tests/support/callback_owner.rs

@@ -1,114 +0,0 @@
-use std::sync::{
-    Arc,
-    atomic::{AtomicUsize, Ordering},
-};
-
-use litellm_python_interop::{InvocationMode, InvocationOutcome, PreparedCall};
-use pyo3::exceptions::PyRuntimeError;
-use pyo3::prelude::*;
-use pyo3::pyclass::{PyTraverseError, PyVisit};
-use pyo3::types::{PyDict, PyTuple};
-
-#[pyclass]
-#[derive(Default)]
-pub struct OwnerFactory {
-    live: Arc<AtomicUsize>,
-}
-
-#[pymethods]
-impl OwnerFactory {
-    #[getter]
-    fn live(&self) -> usize {
-        self.live.load(Ordering::SeqCst)
-    }
-
-    #[pyo3(signature = (callable, positional, keywords=None, awaited=false))]
-    fn prepare(
-        &self,
-        callable: Py<PyAny>,
-        positional: Py<PyTuple>,
-        keywords: Option<Py<PyDict>>,
-        awaited: bool,
-    ) -> Owner {
-        self.live.fetch_add(1, Ordering::SeqCst);
-        Owner {
-            call: Some(PreparedCall::new(
-                if awaited {
-                    InvocationMode::Await
-                } else {
-                    InvocationMode::Direct
-                },
-                callable,
-                positional,
-                keywords,
-            )),
-            live: self.live.clone(),
-        }
-    }
-}
-
-#[pyclass(weakref)]
-struct Owner {
-    call: Option<PreparedCall>,
-    live: Arc<AtomicUsize>,
-}
-
-impl Owner {
-    fn release(&mut self) -> Option<PreparedCall> {
-        let call = self.call.take();
-        if call.is_some() {
-            self.live.fetch_sub(1, Ordering::SeqCst);
-        }
-        call
-    }
-}
-
-#[pymethods]
-impl Owner {
-    fn invoke(slf: &Bound<'_, Self>) -> PyResult<Py<PyAny>> {
-        let py = slf.py();
-        let call = slf
-            .borrow()
-            .call
-            .as_ref()
-            .map(|call| call.clone_ref(py))
-            .ok_or_else(|| PyRuntimeError::new_err("invocation owner released"))?;
-        match call.invoke(py)? {
-            InvocationOutcome::Returned(value) | InvocationOutcome::Awaitable(value) => Ok(value),
-        }
-    }
-
-    fn clone_owner(&self, py: Python<'_>) -> PyResult<Self> {
-        let call = self
-            .call
-            .as_ref()
-            .ok_or_else(|| PyRuntimeError::new_err("invocation owner released"))?;
-        self.live.fetch_add(1, Ordering::SeqCst);
-        Ok(Self {
-            call: Some(call.clone_ref(py)),
-            live: self.live.clone(),
-        })
-    }
-
-    fn close(slf: &Bound<'_, Self>) {
-        let released = slf.borrow_mut().release();
-        drop(released);
-    }
-
-    fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
-        if let Some(call) = &self.call {
-            call.traverse(visit)?;
-        }
-        Ok(())
-    }
-
-    fn __clear__(slf: &Bound<'_, Self>) {
-        Self::close(slf);
-    }
-}
-
-impl Drop for Owner {
-    fn drop(&mut self) {
-        drop(self.release());
-    }
-}

litellm-rust/crates/python-interop/tests/support/mod.rs

@@ -1,9 +1 @@
-mod callback_owner;
 pub mod python;
-pub mod scenarios;
-
-#[derive(Clone, Copy, Debug)]
-pub enum Backend {
-    Python,
-    PreparedCall,
-}

litellm-rust/crates/python-interop/tests/support/scenarios.rs

@@ -1,54 +0,0 @@
-use pyo3::prelude::*;
-use pyo3::types::PyDict;
-use rstest::fixture;
-
-use super::python::{InitializedPython, initialized_python, run_fixture};
-use super::{Backend, callback_owner};
-
-#[fixture]
-pub fn scenario_scope(initialized_python: &InitializedPython) -> Py<PyDict> {
-    initialized_python.attach(|py| {
-        let globals = PyDict::new(py);
-        globals
-            .set_item(
-                "factory",
-                Py::new(py, callback_owner::OwnerFactory::default()).unwrap(),
-            )
-            .unwrap();
-        globals
-            .set_item(
-                "AWAIT_ADAPTER_FILENAME",
-                litellm_python_interop::AWAIT_ADAPTER_FILENAME
-                    .to_str()
-                    .unwrap(),
-            )
-            .unwrap();
-        run_fixture(
-            py,
-            &globals,
-            include_str!("../fixtures/callback_lifecycle.py"),
-            concat!(
-                env!("CARGO_MANIFEST_DIR"),
-                "/tests/fixtures/callback_lifecycle.py"
-            ),
-        )
-        .unwrap();
-        globals.unbind()
-    })
-}
-
-pub fn run_scenario_fixture(
-    scenario_scope: Py<PyDict>,
-    scenario: &str,
-    backend: Backend,
-) -> PyResult<()> {
-    Python::attach(|py| {
-        let globals = scenario_scope.bind(py);
-        globals.get_item("run_scenario")?.unwrap().call1((
-            scenario,
-            matches!(backend, Backend::PreparedCall),
-            globals.get_item("factory")?.unwrap(),
-        ))?;
-        Ok(())
-    })
-}

litellm-rust/crates/python-interop/tests/synthetic/controls.rs

@@ -1,225 +0,0 @@
-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;
-
-use crate::support::Backend;
-use crate::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>>,
-        backend: Backend,
-        mode: InvocationMode,
-    ) -> PyResult<Self> {
-        if matches!(backend, Backend::PreparedCall) {
-            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(Backend::Python, Backend::PreparedCall)] backend: Backend,
-    #[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>()?),
-            backend,
-            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(Backend::Python, Backend::PreparedCall)] backend: Backend,
-    #[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,
-            backend,
-            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/synthetic/lifecycle.rs

@@ -1,411 +0,0 @@
-use std::io::Read;
-use std::process::{Command, Stdio};
-use std::time::{Duration, Instant};
-
-use litellm_python_interop::{InvocationMode, InvocationOutcome, PreparedCall};
-use pyo3::prelude::*;
-use pyo3::types::{PyDict, PyTuple};
-use rstest::rstest;
-use serial_test::{parallel, serial};
-
-use crate::support::Backend;
-use crate::support::python::{InitializedPython, initialized_python, item, run_fixture};
-use crate::support::scenarios::{run_scenario_fixture, scenario_scope};
-
-#[test]
-fn cold_awaited_adapter_initialization_allows_reentry() -> PyResult<()> {
-    let test = concat!(
-        module_path!(),
-        "::cold_awaited_adapter_initialization_allows_reentry"
-    )
-    .split_once("::")
-    .unwrap()
-    .1;
-    let child_env = "LITELLM_INTEROP_COLD_REENTRY_CHILD";
-    if std::env::var(child_env).as_deref() != Ok(test) {
-        let mut child = Command::new(std::env::current_exe().unwrap())
-            .args(["--exact", test, "--nocapture"])
-            .env(child_env, test)
-            .stdout(Stdio::piped())
-            .spawn()
-            .unwrap();
-        let deadline = Instant::now() + Duration::from_secs(15);
-        loop {
-            if let Some(status) = child.try_wait().unwrap() {
-                let mut output = String::new();
-                child
-                    .stdout
-                    .take()
-                    .unwrap()
-                    .read_to_string(&mut output)
-                    .unwrap();
-                assert!(
-                    status.success(),
-                    "awaited adapter reentry child failed: {status}\n{output}"
-                );
-                assert!(
-                    output.contains("test result: ok. 1 passed; 0 failed; 0 ignored;"),
-                    "awaited adapter reentry child did not run exactly one test:\n{output}"
-                );
-                return Ok(());
-            }
-            if Instant::now() >= deadline {
-                child.kill().unwrap();
-                child.wait().unwrap();
-                panic!("awaited adapter reentry did not complete within 15 seconds");
-            }
-            std::thread::sleep(Duration::from_millis(10));
-        }
-    }
-
-    let globals = scenario_scope(&initialized_python());
-    Python::attach(|py| {
-        globals
-            .bind(py)
-            .get_item("cold_awaited_adapter_reentry")?
-            .unwrap()
-            .call1((globals.bind(py).get_item("factory")?.unwrap(),))?;
-        Ok(())
-    })
-}
-
-#[rstest]
-#[case::awaitable_kinds("awaitable_kinds")]
-#[case::identity_and_context("identity_and_context")]
-#[case::exceptions("exceptions")]
-#[case::exception_ownership("exception_ownership")]
-#[case::cancellation_before_start("cancellation_before_start")]
-#[case::cancellation_unwinds("cancellation_unwinds")]
-#[case::cancellation_during_cleanup("cancellation_during_cleanup")]
-#[case::cancellation_suppressed("cancellation_suppressed")]
-#[case::registration_and_gc("registration_and_gc")]
-#[case::background_and_session("background_and_session")]
-#[case::stream_lifecycle("stream_lifecycle")]
-#[case::sync_stream_lifecycle("sync_stream_lifecycle")]
-#[case::repeated_ownership("repeated_ownership")]
-#[case::detached_work_after_error("detached_work_after_error")]
-#[serial(python_interpreter)]
-fn lifecycle_contract(
-    scenario_scope: Py<PyDict>,
-    #[case] scenario: &str,
-    #[values(Backend::Python, Backend::PreparedCall)] backend: Backend,
-) -> PyResult<()> {
-    run_scenario_fixture(scenario_scope, scenario, backend)
-}
-
-#[rstest]
-#[case::retained_lifetime("deferred_lifetime", "identity")]
-#[case::prepared_ownership("deferred_lifetime", "missing_handoff")]
-#[case::externally_owned_retained("borrowed_lifetime", "identity")]
-#[case::original_coroutine("direct_coroutine", "identity")]
-#[case::passthrough_coroutine("direct_coroutine", "result_passthrough")]
-#[serial(python_interpreter)]
-fn control_contract(
-    scenario_scope: Py<PyDict>,
-    #[case] witness: &str,
-    #[case] control: &str,
-    #[values(Backend::Python, Backend::PreparedCall)] backend: Backend,
-    #[values(false, true)] awaited: bool,
-) -> PyResult<()> {
-    run_control_fixture(scenario_scope, witness, control, backend, awaited)
-}
-
-#[rstest]
-#[case::expired_borrow("deferred_lifetime")]
-#[case::externally_owned_borrow("borrowed_lifetime")]
-#[serial(python_interpreter)]
-fn weak_control(
-    scenario_scope: Py<PyDict>,
-    #[case] witness: &str,
-    #[values(false, true)] awaited: bool,
-) -> PyResult<()> {
-    run_control_fixture(scenario_scope, witness, "weak", Backend::Python, awaited)
-}
-
-#[rstest]
-#[case::retained("identity")]
-#[case::missing_handoff("missing_handoff")]
-#[serial(python_interpreter)]
-fn pending_handoff_control(
-    scenario_scope: Py<PyDict>,
-    #[case] control: &str,
-    #[values(Backend::Python, Backend::PreparedCall)] backend: Backend,
-) -> PyResult<()> {
-    run_control_fixture(scenario_scope, "pending_handoff", control, backend, true)
-}
-
-fn run_control_fixture(
-    scenario_scope: Py<PyDict>,
-    witness: &str,
-    control: &str,
-    backend: Backend,
-    awaited: bool,
-) -> PyResult<()> {
-    Python::attach(|py| {
-        let globals = scenario_scope.bind(py);
-        run_fixture(
-            py,
-            globals,
-            include_str!("../fixtures/callback_controls.py"),
-            concat!(
-                env!("CARGO_MANIFEST_DIR"),
-                "/tests/fixtures/callback_controls.py"
-            ),
-        )?;
-        globals.get_item("run_control")?.unwrap().call1((
-            witness,
-            control,
-            matches!(backend, Backend::PreparedCall),
-            awaited,
-            globals.get_item("factory")?.unwrap(),
-        ))?;
-        Ok(())
-    })
-}
-
-fn invoke_direct_callback(
-    py: Python<'_>,
-    globals: &Bound<'_, PyDict>,
-    callback: &str,
-    argument: &str,
-    backend: Backend,
-) -> PyResult<Py<PyAny>> {
-    let args = PyTuple::new(py, [item(globals, argument)])?;
-    if matches!(backend, Backend::Python) {
-        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(Backend::Python, Backend::PreparedCall)] backend: Backend,
-) -> 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", backend)?.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(Backend::Python, Backend::PreparedCall)] backend: Backend,
-) -> 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", backend)?.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<()> {
-    use litellm_python_interop::{InvocationMode, PreparedCall};
-    use pyo3::types::PyTuple;
-    let _ = initialized_python;
-    let (call, reference) = Python::attach(|py| -> PyResult<_> {
-        let globals = PyDict::new(py);
-        py.run(
-            c"import weakref\nclass Value: pass\nvalue = Value()\nreference = weakref.ref(value)",
-            Some(&globals),
-            None,
-        )?;
-        let value = globals.get_item("value")?.unwrap();
-        let call = PreparedCall::new(
-            InvocationMode::Direct,
-            py.None(),
-            PyTuple::new(py, [value])?.unbind(),
-            None,
-        );
-        let reference = globals.get_item("reference")?.unwrap().unbind();
-        globals.del_item("value")?;
-        Ok((call, reference))
-    })?;
-    drop(call);
-    Python::attach(|py| {
-        assert!(reference.call0(py)?.is_none(py));
-        Ok(())
-    })
-}
-
-#[rstest]
-#[case::direct(false)]
-#[case::awaited(true)]
-#[parallel(python_interpreter)]
-fn outcome_identifies_binding(scenario_scope: Py<PyDict>, #[case] awaited: bool) -> PyResult<()> {
-    use litellm_python_interop::{InvocationMode, InvocationOutcome, PreparedCall};
-    use pyo3::types::PyTuple;
-    Python::attach(|py| {
-        let callback = py.eval(c"lambda: None", Some(scenario_scope.bind(py)), None)?;
-        let call = PreparedCall::new(
-            if awaited {
-                InvocationMode::Await
-            } else {
-                InvocationMode::Direct
-            },
-            callback.unbind(),
-            PyTuple::empty(py).unbind(),
-            None,
-        );
-        match call.invoke(py)? {
-            InvocationOutcome::Returned(value) => {
-                assert!(!awaited);
-                assert!(value.is_none(py));
-            }
-            InvocationOutcome::Awaitable(value) => {
-                assert!(awaited);
-                value.call_method0(py, "close")?;
-            }
-        }
-        Ok(())
-    })
-}
-
-#[rstest]
-#[parallel(python_interpreter)]
-fn awaited_raise_surfaces_when_driven(scenario_scope: Py<PyDict>) -> PyResult<()> {
-    use litellm_python_interop::{InvocationMode, InvocationOutcome, PreparedCall};
-    use pyo3::types::PyTuple;
-    Python::attach(|py| {
-        let globals = scenario_scope.bind(py);
-        py.run(
-            c"events = []\nerror = ValueError('await failure')\nasync def callback():\n    events.append('started')\n    raise error\ndef drive(coroutine):\n    try:\n        coroutine.send(None)\n        return False\n    except BaseException as caught:\n        return caught is error\n",
-            Some(globals),
-            None,
-        )?;
-        let started = || -> PyResult<usize> { globals.get_item("events")?.unwrap().len() };
-        let call = PreparedCall::new(
-            InvocationMode::Await,
-            globals.get_item("callback")?.unwrap().unbind(),
-            PyTuple::empty(py).unbind(),
-            None,
-        );
-        let pending = match call.invoke(py)? {
-            InvocationOutcome::Awaitable(value) => value,
-            InvocationOutcome::Returned(_) => panic!("await binding produced a settled outcome"),
-        };
-        assert_eq!(started()?, 0);
-        assert!(
-            globals
-                .get_item("drive")?
-                .unwrap()
-                .call1((pending,))?
-                .extract::<bool>()?
-        );
-        assert_eq!(started()?, 1);
-        Ok(())
-    })
-}

litellm-rust/crates/python-interop/tests/synthetic/mod.rs

@@ -1,8 +1,4 @@
 #[path = "../support/mod.rs"]
 mod support;
 
-mod controls;
-mod lifecycle;
-mod patterns;
-mod prepared_call;
 mod primitives;

litellm-rust/crates/python-interop/tests/synthetic/patterns.rs

@@ -1,712 +0,0 @@
-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;
-
-use crate::support::Backend;
-use crate::support::python::{InitializedPython, initialized_python, item, run_fixture, scope};
-
-#[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/synthetic/prepared_call.rs

@@ -1,350 +0,0 @@
-use litellm_python_interop::{InvocationMode, InvocationOutcome, PreparedCall};
-use pyo3::exceptions::{PyAssertionError, PyKeyboardInterrupt, PyValueError};
-use pyo3::prelude::*;
-use pyo3::types::{PyDict, PyList, PyTuple};
-use rstest::rstest;
-use serial_test::parallel;
-
-use crate::support::python::{InitializedPython, initialized_python, item, run_fixture, scope};
-
-#[rstest]
-#[parallel(python_interpreter)]
-fn retains_aliases_mutations_and_original_result(
-    initialized_python: &InitializedPython,
-) -> PyResult<()> {
-    let _ = initialized_python;
-    Python::attach(|py| {
-        let globals = scope(
-            py,
-            c"
-def callback(data, *, alias):
-    observed.append(data['nested'] is alias)
-    saved.append(data)
-    alias['value'] = 'during'
-    return data
-",
-        )?;
-        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", &shared)?;
-        let invocation = PreparedCall::new(
-            InvocationMode::Direct,
-            item(&globals, "callback").unbind(),
-            PyTuple::new(py, [&payload])?.unbind(),
-            Some(keywords.unbind()),
-        );
-        let result = invoke_direct(&invocation, py)?;
-        assert!(result.bind(py).is(&payload));
-        drop(invocation);
-        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(())
-    })
-}
-
-#[rstest]
-#[parallel(python_interpreter)]
-fn preserves_exception_identity_cause_traceback_and_prior_mutation(
-    initialized_python: &InitializedPython,
-) -> PyResult<()> {
-    let _ = initialized_python;
-    Python::attach(|py| {
-        let globals = scope(
-            py,
-            c"
-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, [&payload])?.unbind(),
-            None,
-        );
-        let error = invoke_direct(&invocation, py).unwrap_err();
-        assert!(error.value(py).is(original.value(py)));
-        drop(invocation);
-        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(())
-    })
-}
-
-#[rstest]
-#[parallel(python_interpreter)]
-fn returns_coroutine_without_executing_it(initialized_python: &InitializedPython) -> PyResult<()> {
-    let _ = initialized_python;
-    Python::attach(|py| {
-        let globals = scope(
-            py,
-            c"
-async def work():
-    started.append(True)
-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(),
-            PyTuple::empty(py).unbind(),
-            None,
-        );
-        let result = invoke_direct(&invocation, py)?;
-        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(())
-    })
-}
-
-#[pyfunction]
-fn reenter(py: Python<'_>, callback: Py<PyAny>, payload: Py<PyAny>) -> PyResult<Py<PyAny>> {
-    invoke_direct(
-        &PreparedCall::new(
-            InvocationMode::Direct,
-            callback,
-            PyTuple::new(py, [payload])?.unbind(),
-            None,
-        ),
-        py,
-    )
-}
-
-#[rstest]
-#[parallel(python_interpreter)]
-fn preserves_current_context_thread_and_reentry(
-    initialized_python: &InitializedPython,
-) -> PyResult<()> {
-    let _ = initialized_python;
-    Python::attach(|py| {
-        let globals = scope(
-            py,
-            c"
-import threading
-def inner(data):
-    observed.append((context.get(), threading.get_ident()))
-    data['inner'] = True
-    context.set('inner')
-    return data
-def outer():
-    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,
-            item(&globals, "outer").unbind(),
-            PyTuple::empty(py).unbind(),
-            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(())
-    })
-}
-
-#[rstest]
-#[parallel(python_interpreter)]
-fn owns_arguments_until_release_and_preserves_callback_retention(
-    initialized_python: &InitializedPython,
-) -> PyResult<()> {
-    let _ = initialized_python;
-    let (invocation, globals) = Python::attach(|py| {
-        let globals = scope(
-            py,
-            c"
-class Value:
-    pass
-class Callback:
-    def __call__(self, value, *, other):
-        saved.append(value)
-        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", other)?;
-        let invocation = PreparedCall::new(
-            InvocationMode::Direct,
-            callback.unbind(),
-            PyTuple::new(py, [value])?.unbind(),
-            Some(keywords.unbind()),
-        );
-        Ok::<_, PyErr>((invocation, globals.unbind()))
-    })?;
-    Python::attach(|py| {
-        let globals = globals.bind(py);
-        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);
-        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(())
-    })
-}
-
-#[rstest]
-#[parallel(python_interpreter)]
-fn checked_runner_rejects_unhandled_background_failures(
-    initialized_python: &InitializedPython,
-) -> PyResult<()> {
-    let _ = initialized_python;
-    Python::attach(|py| {
-        let globals = PyDict::new(py);
-        run_fixture(
-            py,
-            &globals,
-            include_str!("../fixtures/callback_lifecycle.py"),
-            concat!(
-                env!("CARGO_MANIFEST_DIR"),
-                "/tests/fixtures/callback_lifecycle.py"
-            ),
-        )?;
-        py.run(
-            c"
-async def fail():
-    raise RuntimeError('background task regression')
-
-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(())
-    })
-}
-
-fn invoke_direct(call: &PreparedCall, py: Python<'_>) -> PyResult<Py<PyAny>> {
-    match call.invoke(py)? {
-        InvocationOutcome::Returned(value) => Ok(value),
-        InvocationOutcome::Awaitable(_) => panic!("direct binding produced an awaitable outcome"),
-    }
-}

litellm-rust/crates/python-interop/tests/synthetic/primitives.rs

@@ -2,7 +2,7 @@ use rstest::rstest;
 use serde_json::{Value, json};
 use serial_test::parallel;
 
-use litellm_python_interop::{from_py, release_count, release_gil, to_py};
+use litellm_python_interop::{from_py, to_py};
 
 use crate::support::python::{InitializedPython, initialized_python};
 
@@ -18,13 +18,3 @@ fn serde_values_round_trip_through_python(#[from(initialized_python)] python: &I
         assert_eq!(actual, expected);
     });
 }
-
-#[rstest]
-#[parallel(python_interpreter)]
-fn release_gil_runs_work_and_records_it(#[from(initialized_python)] python: &InitializedPython) {
-    let before = release_count();
-    let result = python.attach(|py| release_gil(py, || 42));
-
-    assert_eq!(result, 42);
-    assert_eq!(release_count(), before + 1);
-}