refactor(rust): unify call lifecycle execution

litellm-rust/Cargo.lock

@@ -1925,6 +1925,26 @@ dependencies = [
  "tracing",
 ]
 
+[[package]]
+name = "litellm-cache"
+version = "0.1.0"
+dependencies = [
+ "rstest",
+ "serde",
+ "serde_json",
+ "sha2 0.10.9",
+ "thiserror 2.0.19",
+]
+
+[[package]]
+name = "litellm-cache-memory"
+version = "0.1.0"
+dependencies = [
+ "litellm-cache",
+ "rstest",
+ "serde_json",
+]
+
 [[package]]
 name = "litellm-config"
 version = "0.1.0"
@@ -1952,6 +1972,8 @@ dependencies = [
  "data-url",
  "futures-util",
  "gcp_auth",
+ "litellm-cache",
+ "litellm-cache-memory",
  "mime_guess",
  "moka",
  "rand 0.8.7",
@@ -1980,6 +2002,8 @@ version = "0.1.0"
 dependencies = [
  "criterion",
  "futures-util",
+ "litellm-cache",
+ "litellm-cache-memory",
  "litellm-core",
  "litellm-python-interop",
  "litellm-token-counter",

litellm-rust/Cargo.toml

@@ -1,5 +1,7 @@
 [workspace]
 members = [
+    "crates/cache",
+    "crates/cache-memory",
     "crates/core",
     "crates/token-counter",
     "crates/config",
@@ -20,6 +22,8 @@ bytes = "1"
 tracing = "0.1"
 tracing-subscriber = { version = "0.3", default-features = false, features = ["registry", "std"] }
 litellm-core = { path = "crates/core" }
+litellm-cache = { path = "crates/cache" }
+litellm-cache-memory = { path = "crates/cache-memory" }
 litellm-token-counter = { path = "crates/token-counter" }
 litellm-config = { path = "crates/config" }
 litellm-ai-gateway = { path = "crates/ai-gateway", default-features = false }

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

@@ -1,12 +1,10 @@
-use litellm_core::audio_transcription::{
-    AudioTranscriptionRequest as CoreAudioTranscriptionRequest, ProviderAudioTranscriptionRequest,
-    prepare_audio_transcription_provider_call,
-};
-use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming};
-use litellm_core::error::Error;
 use serde_json::{Map, Value, json};
-use std::future::Future;
-use std::pin::Pin;
+
+use litellm_core::call_lifecycle::provider::{
+    ProviderHookFuture, ProviderHooks, ProviderRequest, ProviderResponse,
+};
+use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleTiming};
+use litellm_core::error::Error;
 
 use super::types::PreparedAudioTranscriptionRequest;
 use crate::integrations::custom_guardrail::{
@@ -23,25 +21,25 @@ pub(crate) struct AudioTranscriptionLifecycleHooks {
     logger_runner: CustomLoggerRunner,
     guardrail_runner: CustomGuardrailRunner,
     request_metadata: RequestMetadata,
+    provider: String,
 }
 
-type AudioFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
-type AudioLogFuture<'a> = Pin<Box<dyn Future<Output = ()> + Send + 'a>>;
-
 impl AudioTranscriptionLifecycleHooks {
     pub(crate) fn new(
         logger_runner: CustomLoggerRunner,
         guardrail_runner: CustomGuardrailRunner,
         request_metadata: RequestMetadata,
+        provider: String,
     ) -> Self {
         Self {
             logger_runner,
             guardrail_runner,
             request_metadata,
+            provider,
         }
     }
 
-    async fn run_pre_call_guardrails(
+    pub(crate) async fn run_pre_call_guardrails(
         &self,
         request: PreparedAudioTranscriptionRequest,
     ) -> Result<PreparedAudioTranscriptionRequest, Error> {
@@ -85,39 +83,10 @@ impl AudioTranscriptionLifecycleHooks {
         })
     }
 
-    async fn prepare_provider_request(
-        &self,
-        request: PreparedAudioTranscriptionRequest,
-    ) -> Result<ProviderAudioTranscriptionRequest, Error> {
-        let PreparedAudioTranscriptionRequest {
-            model,
-            custom_llm_provider,
-            audio,
-            api_key,
-            api_base,
-            extra_headers,
-            optional_params,
-            timeout,
-            ..
-        } = request;
-        let provider_request =
-            prepare_audio_transcription_provider_call(CoreAudioTranscriptionRequest {
-                model: &model,
-                audio,
-                api_key: api_key.as_deref(),
-                api_base: api_base.as_deref(),
-                custom_llm_provider: Some(&custom_llm_provider),
-                extra_headers,
-                optional_params,
-                timeout,
-            })?;
-        self.run_during_call_guardrails(provider_request).await
-    }
-
     async fn run_during_call_guardrails(
         &self,
-        request: ProviderAudioTranscriptionRequest,
-    ) -> Result<ProviderAudioTranscriptionRequest, Error> {
+        request: ProviderRequest,
+    ) -> Result<ProviderRequest, Error> {
         if self.guardrail_runner.is_empty() {
             return Ok(request);
         }
@@ -126,10 +95,10 @@ impl AudioTranscriptionLifecycleHooks {
             .run_during_call(
                 &guardrail_context(&self.request_metadata),
                 GuardrailRequest::new(json!({
-                    "model": request.model(),
-                    "custom_llm_provider": request.custom_llm_provider(),
-                    "url": request.url(),
-                    "body": request.body(),
+                    "model": &request.model,
+                    "custom_llm_provider": &self.provider,
+                    "url": &request.url,
+                    "body": &request.body,
                 })),
             )
             .await
@@ -142,7 +111,7 @@ impl AudioTranscriptionLifecycleHooks {
         let body = data.remove("body").ok_or_else(|| {
             Error::InvalidRequest("audio transcription guardrail removed body".to_string())
         })?;
-        Ok(request.with_body(body))
+        Ok(ProviderRequest { body, ..request })
     }
 
     fn logging_payload(
@@ -172,82 +141,65 @@ impl AudioTranscriptionLifecycleHooks {
             messages: None,
         }
     }
+    pub(crate) async fn log_success(
+        &self,
+        context: &CallLifecycleContext,
+        response: &Value,
+        timing: &CallLifecycleTiming,
+    ) {
+        if self.logger_runner.is_empty() {
+            return;
+        }
+        self.logger_runner
+            .async_log_success_event(
+                &ModelCallDetails::from_standard_logging_payload(
+                    self.logging_payload(context, timing),
+                ),
+                &CallbackValue::new("audio_transcription", response.clone()),
+                CallbackTiming::new(timing.start_time, timing.end_time),
+            )
+            .await;
+    }
+
+    pub(crate) async fn log_failure(
+        &self,
+        context: &CallLifecycleContext,
+        error: &Error,
+        timing: &CallLifecycleTiming,
+    ) {
+        if self.logger_runner.is_empty() {
+            return;
+        }
+        let logging_error = LoggingError {
+            message: error.to_string(),
+            kind: core_error_kind(error).to_string(),
+        };
+        self.logger_runner
+            .async_log_failure_event(
+                &ModelCallDetails::from_standard_logging_payload(
+                    self.logging_payload(context, timing),
+                )
+                .with_failure_error(logging_error.clone()),
+                Some(&CallbackValue::new(
+                    "error",
+                    json!({"message": logging_error.message, "kind": logging_error.kind}),
+                )),
+                CallbackTiming::new(timing.start_time, timing.end_time),
+            )
+            .await;
+    }
 }
 
-impl CallLifecycleHooks<PreparedAudioTranscriptionRequest, ProviderAudioTranscriptionRequest, Value>
-    for AudioTranscriptionLifecycleHooks
-{
-    type PreCallFuture<'a> = AudioFuture<'a, PreparedAudioTranscriptionRequest>;
-    type DuringCallFuture<'a> = AudioFuture<'a, ProviderAudioTranscriptionRequest>;
-    type SuccessFuture<'a> = AudioLogFuture<'a>;
-    type FailureFuture<'a> = AudioLogFuture<'a>;
-
-    fn async_pre_call_hook<'a>(
-        &'a self,
-        _context: &'a CallLifecycleContext,
-        request: PreparedAudioTranscriptionRequest,
-    ) -> Self::PreCallFuture<'a> {
-        Box::pin(async move { self.run_pre_call_guardrails(request).await })
+impl ProviderHooks for AudioTranscriptionLifecycleHooks {
+    fn before_request(&self, request: ProviderRequest) -> ProviderHookFuture<'_, ProviderRequest> {
+        Box::pin(async move { self.run_during_call_guardrails(request).await })
     }
 
-    fn async_during_call_hook<'a>(
-        &'a self,
-        _context: &'a CallLifecycleContext,
-        request: PreparedAudioTranscriptionRequest,
-    ) -> Self::DuringCallFuture<'a> {
-        Box::pin(async move { self.prepare_provider_request(request).await })
-    }
-
-    fn async_log_success_event<'a>(
-        &'a self,
-        context: &'a CallLifecycleContext,
-        response: &'a Value,
-        timing: &'a CallLifecycleTiming,
-    ) -> Self::SuccessFuture<'a> {
-        Box::pin(async move {
-            if self.logger_runner.is_empty() {
-                return;
-            }
-            self.logger_runner
-                .async_log_success_event(
-                    &ModelCallDetails::from_standard_logging_payload(
-                        self.logging_payload(context, timing),
-                    ),
-                    &CallbackValue::new("audio_transcription", response.clone()),
-                    CallbackTiming::new(timing.start_time, timing.end_time),
-                )
-                .await;
-        })
-    }
-
-    fn async_log_failure_event<'a>(
-        &'a self,
-        context: &'a CallLifecycleContext,
-        error: &'a Error,
-        timing: &'a CallLifecycleTiming,
-    ) -> Self::FailureFuture<'a> {
-        Box::pin(async move {
-            if self.logger_runner.is_empty() {
-                return;
-            }
-            let logging_error = LoggingError {
-                message: error.to_string(),
-                kind: core_error_kind(error).to_string(),
-            };
-            self.logger_runner
-                .async_log_failure_event(
-                    &ModelCallDetails::from_standard_logging_payload(
-                        self.logging_payload(context, timing),
-                    )
-                    .with_failure_error(logging_error.clone()),
-                    Some(&CallbackValue::new(
-                        "error",
-                        json!({"message": logging_error.message, "kind": logging_error.kind}),
-                    )),
-                    CallbackTiming::new(timing.start_time, timing.end_time),
-                )
-                .await;
-        })
+    fn after_response(
+        &self,
+        response: ProviderResponse,
+    ) -> ProviderHookFuture<'_, ProviderResponse> {
+        Box::pin(async move { Ok(response) })
     }
 }
 

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

@@ -1,22 +1,137 @@
-use litellm_core::Error;
-use litellm_core::audio_transcription::execute_audio_transcription_provider_call;
-use litellm_core::call_lifecycle::CallLifecycle;
+use std::sync::Mutex;
+use std::time::{SystemTime, UNIX_EPOCH};
+
 use serde_json::Value;
 
+use litellm_core::Error;
+use litellm_core::audio_transcription::lifecycle::{
+    AudioTranscriptionCall, OwnedAudioTranscriptionRequest,
+};
+use litellm_core::call_lifecycle::CallLifecycleRequest;
+use litellm_core::call_lifecycle::host::{LifecycleBackend, LifecycleBackendFuture, drive};
+use litellm_core::call_lifecycle::provider::{
+    CompletedOperation, CompletedReply, CompletedWorkflow, ProviderHooks, ProviderOptions,
+};
+use litellm_core::call_lifecycle::workflow::LifecycleOperation;
+
 mod hooks;
 mod prepare;
 mod types;
 
 pub use types::AudioTranscriptionRequest;
 
+use hooks::AudioTranscriptionLifecycleHooks;
 use prepare::{PreparedAudioTranscriptionCall, prepare_audio_transcription_call};
+use types::PreparedAudioTranscriptionRequest;
 
 pub async fn audio_transcription(request: AudioTranscriptionRequest<'_>) -> Result<Value, Error> {
     let PreparedAudioTranscriptionCall { request, hooks } =
         prepare_audio_transcription_call(request);
-    CallLifecycle::default()
-        .run_request(request, &hooks, execute_audio_transcription_provider_call)
-        .await
+    let mut call = AudioTranscriptionCall::new(CompletedWorkflow::default(), false);
+    drive(
+        &mut call,
+        &AudioTranscriptionBackend {
+            request: Mutex::new(Some(request)),
+            hooks,
+        },
+    )
+    .await
+}
+
+struct AudioTranscriptionBackend {
+    request: Mutex<Option<PreparedAudioTranscriptionRequest>>,
+    hooks: AudioTranscriptionLifecycleHooks,
+}
+
+impl LifecycleBackend<CompletedOperation<Value>, CompletedReply<OwnedAudioTranscriptionRequest>>
+    for AudioTranscriptionBackend
+{
+    fn invoke(
+        &self,
+        operation: CompletedOperation<Value>,
+    ) -> LifecycleBackendFuture<'_, CompletedReply<OwnedAudioTranscriptionRequest>> {
+        Box::pin(async move {
+            match operation {
+                CompletedOperation::Lifecycle(LifecycleOperation::ProjectRequest) => {
+                    let request = self
+                        .request
+                        .lock()
+                        .unwrap_or_else(|error| error.into_inner())
+                        .take()
+                        .ok_or_else(|| {
+                            Error::InvalidRequest("request was already projected".into())
+                        });
+                    CompletedReply::Request(match request {
+                        Ok(request) => {
+                            let context = request.lifecycle_context();
+                            let started = epoch_seconds();
+                            match self.hooks.run_pre_call_guardrails(request).await {
+                                Ok(request) => Ok(owned_request(request)),
+                                Err(error) => {
+                                    let timing =
+                                        litellm_core::call_lifecycle::CallLifecycleTiming::new(
+                                            started,
+                                            epoch_seconds(),
+                                        );
+                                    self.hooks.log_failure(&context, &error, &timing).await;
+                                    Err(error)
+                                }
+                            }
+                        }
+                        Err(error) => Err(error),
+                    })
+                }
+                CompletedOperation::Lifecycle(LifecycleOperation::Success {
+                    context,
+                    response,
+                    timing,
+                }) => {
+                    self.hooks
+                        .log_success(&context, response.as_ref(), &timing)
+                        .await;
+                    CompletedReply::Lifecycle(Ok(()))
+                }
+                CompletedOperation::Lifecycle(LifecycleOperation::Failure {
+                    context,
+                    error,
+                    timing,
+                }) => {
+                    self.hooks.log_failure(&context, &error, &timing).await;
+                    CompletedReply::Lifecycle(Ok(()))
+                }
+                CompletedOperation::Lifecycle(_) => CompletedReply::Lifecycle(Ok(())),
+                CompletedOperation::BeforeRequest(request) => {
+                    CompletedReply::BeforeRequest(self.hooks.before_request(request).await)
+                }
+                CompletedOperation::AfterResponse(response) => {
+                    CompletedReply::AfterResponse(self.hooks.after_response(response).await)
+                }
+            }
+        })
+    }
+}
+
+fn owned_request(request: PreparedAudioTranscriptionRequest) -> OwnedAudioTranscriptionRequest {
+    OwnedAudioTranscriptionRequest {
+        options: ProviderOptions {
+            model: request.model,
+            litellm_call_id: Some(request.litellm_call_id),
+            api_key: request.api_key,
+            api_base: request.api_base,
+            custom_llm_provider: Some(request.custom_llm_provider),
+            extra_headers: request.extra_headers,
+            timeout: request.timeout,
+        },
+        audio: request.audio,
+        optional_params: request.optional_params,
+    }
+}
+
+fn epoch_seconds() -> f64 {
+    SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .map(|duration| duration.as_secs_f64())
+        .unwrap_or(0.0)
 }
 
 #[cfg(test)]

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

@@ -25,10 +25,12 @@ pub(crate) fn prepare_audio_transcription_call(
             model: request.model,
             custom_llm_provider: "bedrock",
         });
+    let model = provider_info.model.to_string();
+    let provider = provider_info.custom_llm_provider.to_string();
     PreparedAudioTranscriptionCall {
         request: PreparedAudioTranscriptionRequest {
-            model: provider_info.model.to_string(),
-            custom_llm_provider: provider_info.custom_llm_provider.to_string(),
+            model,
+            custom_llm_provider: provider.clone(),
             litellm_call_id: call_id,
             audio: request.audio,
             api_key: request.api_key.map(str::to_string),
@@ -41,6 +43,7 @@ pub(crate) fn prepare_audio_transcription_call(
             CustomLoggerRunner::new(request.callbacks),
             CustomGuardrailRunner::new(request.guardrails),
             request.request_metadata,
+            provider,
         ),
     }
 }

litellm-rust/crates/ai-gateway/src/routes/responses/service.rs

@@ -3,7 +3,6 @@ use std::time::Duration;
 
 use futures_util::{Sink, Stream};
 use litellm_core::Error;
-use litellm_core::call_lifecycle::{CallLifecycle, CallLifecycleContext};
 use litellm_core::responses::instrumentation::{
     ResponsesWsCallbackPayload, ResponsesWsInstrumentation, ResponsesWsLogOutcome,
     ResponsesWsMetadata,
@@ -55,24 +54,20 @@ where
         },
     ));
     let observer_instrumentation = Arc::clone(&instrumentation);
-    let context = CallLifecycleContext::new("responses_websocket", model, "openai", call_id);
-    let result = CallLifecycle::default()
-        .run(context, (), instrumentation.as_ref(), |_| async move {
-            crate::io::responses_ws::async_responses_websocket(
-                provider_model,
-                params.api_key.as_deref(),
-                params.api_base.as_deref(),
-                first_frame,
-                idle_timeout,
-                move |event| {
-                    observer_instrumentation.observe(event);
-                },
-                client_in,
-                client_out,
-            )
-            .await
-        })
-        .await;
+    let result = crate::io::responses_ws::async_responses_websocket(
+        provider_model,
+        params.api_key.as_deref(),
+        params.api_base.as_deref(),
+        first_frame,
+        idle_timeout,
+        move |event| {
+            observer_instrumentation.observe(event);
+        },
+        client_in,
+        client_out,
+    )
+    .await;
+    instrumentation.record_outcome(result.is_ok());
     let outcome = instrumentation.take_or_build_outcome(result.is_ok());
     dispatch_outcome(loggers, outcome).await;
     result

litellm-rust/crates/core/Cargo.toml

@@ -8,6 +8,8 @@ autotests = false
 
 [dependencies]
 bytes.workspace = true
+litellm-cache.workspace = true
+litellm-cache-memory.workspace = true
 futures-util.workspace = true
 base64.workspace = true
 azure_core.workspace = true

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

@@ -1,5 +1,8 @@
 use serde_json::Value;
 
+use crate::call_lifecycle::provider::{
+    NoopProviderHooks, ProviderHooks, ProviderRequest, ProviderResponse,
+};
 use crate::error::Error;
 use crate::http_utils::{http_request, truncate_error_body};
 
@@ -10,6 +13,28 @@ use super::types::ProviderAudioTranscriptionRequest;
 pub async fn execute_audio_transcription_provider_call(
     request: ProviderAudioTranscriptionRequest,
 ) -> Result<Value, Error> {
+    execute_with_hooks(request, &NoopProviderHooks).await
+}
+
+pub(super) async fn execute_with_hooks(
+    request: ProviderAudioTranscriptionRequest,
+    hooks: &dyn ProviderHooks,
+) -> Result<Value, Error> {
+    let changed = hooks
+        .before_request(ProviderRequest {
+            model: request.model.clone(),
+            url: request.url.clone(),
+            headers: request.upstream_headers.clone(),
+            body: request.body.clone(),
+        })
+        .await?;
+    let request = ProviderAudioTranscriptionRequest {
+        model: changed.model,
+        url: changed.url,
+        upstream_headers: changed.headers,
+        body: changed.body,
+        ..request
+    };
     let body = serde_json::to_vec(&request.body)
         .map_err(|error| Error::InvalidRequest(format!("invalid audio request body: {error}")))?;
     let headers = signed_headers(&request, &body).await?;
@@ -28,9 +53,17 @@ pub async fn execute_audio_transcription_provider_call(
         .text()
         .await
         .map_err(|error| Error::Network(error.to_string()))?;
-    if !status.is_success() {
-        return Err(Error::Http {
+    let observed = hooks
+        .after_response(ProviderResponse {
             status: status.as_u16(),
+            body: text,
+        })
+        .await?;
+    let observed_status = observed.status;
+    let text = observed.body;
+    if !(200..300).contains(&observed_status) {
+        return Err(Error::Http {
+            status: observed_status,
             body: truncate_error_body(&text),
         });
     }

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

@@ -0,0 +1,69 @@
+use std::sync::Arc;
+
+use serde_json::Value;
+
+use super::types::AudioTranscriptionRequest;
+use crate::call_lifecycle::provider::{
+    CompletedCall, CompletedRoute, ProviderHooks, ProviderOptions,
+};
+use crate::call_lifecycle::workflow::WorkflowFuture;
+
+pub struct OwnedAudioTranscriptionRequest {
+    pub options: ProviderOptions,
+    pub audio: Value,
+    pub optional_params: serde_json::Map<String, Value>,
+}
+
+impl From<AudioTranscriptionRequest<'_>> for OwnedAudioTranscriptionRequest {
+    fn from(request: AudioTranscriptionRequest<'_>) -> Self {
+        Self {
+            options: ProviderOptions {
+                model: request.model.to_owned(),
+                litellm_call_id: None,
+                api_key: request.api_key.map(str::to_owned),
+                api_base: request.api_base.map(str::to_owned),
+                custom_llm_provider: request.custom_llm_provider.map(str::to_owned),
+                extra_headers: request.extra_headers,
+                timeout: request.timeout,
+            },
+            audio: request.audio,
+            optional_params: request.optional_params,
+        }
+    }
+}
+
+pub struct AudioTranscriptionRoute;
+pub type AudioTranscriptionCall = CompletedCall<AudioTranscriptionRoute>;
+
+impl CompletedRoute for AudioTranscriptionRoute {
+    type Request = OwnedAudioTranscriptionRequest;
+    type Response = Value;
+
+    fn run(
+        request: Self::Request,
+        hooks: Arc<dyn ProviderHooks>,
+    ) -> WorkflowFuture<Self::Response> {
+        Box::pin(async move {
+            let options = request.options;
+            let request = AudioTranscriptionRequest {
+                model: &options.model,
+                api_key: options.api_key.as_deref(),
+                api_base: options.api_base.as_deref(),
+                custom_llm_provider: options.custom_llm_provider.as_deref(),
+                extra_headers: options.extra_headers,
+                timeout: options.timeout,
+                audio: request.audio,
+                optional_params: request.optional_params,
+            };
+            super::handler::execute_with_hooks(
+                super::prepare::prepare_audio_transcription_provider_call(request)?,
+                hooks.as_ref(),
+            )
+            .await
+        })
+    }
+
+    fn context(request: &Self::Request) -> crate::call_lifecycle::CallLifecycleContext {
+        request.options.lifecycle_context("audio_transcription")
+    }
+}

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

@@ -1,6 +1,7 @@
 use crate::Error;
 mod client;
 mod handler;
+pub mod lifecycle;
 mod prepare;
 pub mod transformation;
 pub mod types;
@@ -13,8 +14,10 @@ pub use types::{AudioTranscriptionRequest, ProviderAudioTranscriptionRequest};
 
 #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
 pub async fn audio_transcription(request: AudioTranscriptionRequest<'_>) -> Result<Value, Error> {
-    execute_audio_transcription_provider_call(prepare_audio_transcription_provider_call(request)?)
-        .await
+    crate::call_lifecycle::provider::run_completed::<lifecycle::AudioTranscriptionRoute>(
+        request.into(),
+    )
+    .await
 }
 
 pub fn admit(

litellm-rust/crates/core/src/caching/in_memory_cache.rs

@@ -1,258 +0,0 @@
-use std::cmp::Reverse;
-use std::collections::{BinaryHeap, HashMap};
-use std::time::{Duration, SystemTime, UNIX_EPOCH};
-
-const DEFAULT_MAX_SIZE_IN_MEMORY: usize = 200;
-const DEFAULT_TTL: Duration = Duration::from_secs(600);
-
-pub struct InMemoryCache<V: Clone> {
-    pub cache_dict: HashMap<String, V>,
-    pub ttl_dict: HashMap<String, Duration>,
-    pub expiration_heap: BinaryHeap<Reverse<(Duration, String)>>,
-    pub max_size_in_memory: usize,
-    pub default_ttl: Duration,
-    now: Box<dyn Fn() -> Duration + Send + Sync>,
-}
-
-impl<V: Clone> Default for InMemoryCache<V> {
-    fn default() -> Self {
-        Self::new(None, None)
-    }
-}
-
-impl<V: Clone> InMemoryCache<V> {
-    pub fn new(max_size_in_memory: Option<usize>, default_ttl: Option<Duration>) -> Self {
-        Self::with_clock(max_size_in_memory, default_ttl, || {
-            SystemTime::now()
-                .duration_since(UNIX_EPOCH)
-                .unwrap_or_default()
-        })
-    }
-
-    pub fn with_clock(
-        max_size_in_memory: Option<usize>,
-        default_ttl: Option<Duration>,
-        now: impl Fn() -> Duration + Send + Sync + 'static,
-    ) -> Self {
-        Self {
-            cache_dict: HashMap::new(),
-            ttl_dict: HashMap::new(),
-            expiration_heap: BinaryHeap::new(),
-            max_size_in_memory: max_size_in_memory.unwrap_or(DEFAULT_MAX_SIZE_IN_MEMORY),
-            default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
-            now: Box::new(now),
-        }
-    }
-
-    pub fn evict_cache(&mut self) {
-        if self.max_size_in_memory == 0 {
-            return;
-        }
-
-        let current_time = (self.now)();
-        while let Some(Reverse((expiration_time, key))) = self.expiration_heap.peek().cloned() {
-            if self.ttl_dict.get(&key).copied() != Some(expiration_time) {
-                self.expiration_heap.pop();
-            } else if expiration_time <= current_time {
-                self.expiration_heap.pop();
-                self.remove_key(&key);
-            } else {
-                break;
-            }
-        }
-
-        while self.cache_dict.len() >= self.max_size_in_memory {
-            let Some(Reverse((expiration_time, key))) = self.expiration_heap.pop() else {
-                break;
-            };
-            if self.ttl_dict.get(&key).copied() == Some(expiration_time) {
-                self.remove_key(&key);
-            }
-        }
-    }
-
-    pub fn allow_ttl_override(&self, key: &str) -> bool {
-        match self.ttl_dict.get(key).copied() {
-            None => true,
-            Some(expiration_time) => expiration_time < (self.now)(),
-        }
-    }
-
-    pub fn set_cache(&mut self, key: impl Into<String>, value: V, ttl: Option<Duration>) {
-        if self.max_size_in_memory == 0 {
-            return;
-        }
-
-        self.evict_cache();
-        let key = key.into();
-        self.cache_dict.insert(key.clone(), value);
-        if self.allow_ttl_override(&key) {
-            let expiration_time = (self.now)() + ttl.unwrap_or(self.default_ttl);
-            self.ttl_dict.insert(key.clone(), expiration_time);
-            self.expiration_heap.push(Reverse((expiration_time, key)));
-        }
-    }
-
-    // Generic values intentionally omit Python's per-item size check.
-    pub fn get_cache(&mut self, key: &str) -> Option<V> {
-        if self.cache_dict.contains_key(key) {
-            if self.is_key_expired(key) {
-                self.remove_key(key);
-                return None;
-            }
-            return self.cache_dict.get(key).cloned();
-        }
-        None
-    }
-
-    pub fn get_ttl(&self, key: &str) -> Option<Duration> {
-        self.ttl_dict.get(key).copied()
-    }
-
-    pub fn delete_cache(&mut self, key: &str) {
-        self.remove_key(key);
-    }
-
-    pub fn flush_cache(&mut self) {
-        self.cache_dict.clear();
-        self.ttl_dict.clear();
-        self.expiration_heap.clear();
-    }
-
-    fn is_key_expired(&self, key: &str) -> bool {
-        self.ttl_dict
-            .get(key)
-            .is_some_and(|expiration_time| *expiration_time < (self.now)())
-    }
-
-    fn remove_key(&mut self, key: &str) {
-        self.cache_dict.remove(key);
-        self.ttl_dict.remove(key);
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use std::sync::{
-        Arc,
-        atomic::{AtomicU64, Ordering},
-    };
-
-    use super::InMemoryCache;
-    use std::time::Duration;
-
-    fn cache(now: Arc<AtomicU64>, max_size: usize, default_ttl: Duration) -> InMemoryCache<String> {
-        InMemoryCache::with_clock(Some(max_size), Some(default_ttl), move || {
-            Duration::from_secs(now.load(Ordering::Relaxed))
-        })
-    }
-
-    #[test]
-    fn ttl_expiry_is_deterministic() {
-        let now = Arc::new(AtomicU64::new(100));
-        let mut cache = cache(now.clone(), 10, Duration::from_secs(60));
-        cache.set_cache("key", "value".to_string(), None);
-        assert_eq!(cache.get_cache("key"), Some("value".to_string()));
-        now.store(161, Ordering::Relaxed);
-        assert_eq!(cache.get_cache("key"), None);
-        assert_eq!(cache.get_ttl("key"), None);
-    }
-
-    #[test]
-    fn default_and_per_set_ttl_are_applied() {
-        let now = Arc::new(AtomicU64::new(100));
-        let mut cache = cache(now.clone(), 10, Duration::from_secs(60));
-        cache.set_cache("default", "value".to_string(), None);
-        cache.set_cache("custom", "value".to_string(), Some(Duration::from_secs(20)));
-        assert_eq!(cache.get_ttl("default"), Some(Duration::from_secs(160)));
-        assert_eq!(cache.get_ttl("custom"), Some(Duration::from_secs(120)));
-    }
-
-    #[test]
-    fn unexpired_entries_do_not_allow_ttl_override() {
-        let now = Arc::new(AtomicU64::new(100));
-        let mut cache = cache(now.clone(), 10, Duration::from_secs(60));
-        cache.set_cache("key", "first".to_string(), Some(Duration::from_secs(20)));
-        cache.set_cache("key", "second".to_string(), Some(Duration::from_secs(80)));
-        assert_eq!(cache.get_cache("key"), Some("second".to_string()));
-        assert_eq!(cache.get_ttl("key"), Some(Duration::from_secs(120)));
-        now.store(121, Ordering::Relaxed);
-        cache.set_cache("key", "third".to_string(), Some(Duration::from_secs(80)));
-        assert_eq!(cache.get_ttl("key"), Some(Duration::from_secs(201)));
-    }
-
-    #[test]
-    fn max_size_evicts_earliest_expiration() {
-        let now = Arc::new(AtomicU64::new(100));
-        let mut cache = cache(now, 2, Duration::from_secs(60));
-        cache.set_cache("early", "value".to_string(), Some(Duration::from_secs(10)));
-        cache.set_cache("late", "value".to_string(), Some(Duration::from_secs(20)));
-        cache.set_cache("new", "value".to_string(), Some(Duration::from_secs(30)));
-        assert_eq!(cache.get_cache("early"), None);
-        assert!(cache.get_cache("late").is_some());
-        assert!(cache.get_cache("new").is_some());
-    }
-
-    #[test]
-    fn expired_entries_are_evicted_before_live_entries() {
-        let now = Arc::new(AtomicU64::new(100));
-        let mut cache = cache(now.clone(), 3, Duration::from_secs(60));
-        cache.set_cache(
-            "expired-one",
-            "value".to_string(),
-            Some(Duration::from_secs(10)),
-        );
-        cache.set_cache(
-            "expired-two",
-            "value".to_string(),
-            Some(Duration::from_secs(20)),
-        );
-        cache.set_cache("live", "value".to_string(), Some(Duration::from_secs(100)));
-        now.store(121, Ordering::Relaxed);
-        cache.set_cache("new", "value".to_string(), Some(Duration::from_secs(100)));
-        assert_eq!(cache.get_cache("expired-one"), None);
-        assert_eq!(cache.get_cache("expired-two"), None);
-        assert!(cache.get_cache("live").is_some());
-        assert!(cache.get_cache("new").is_some());
-    }
-
-    #[test]
-    fn stale_heap_entries_are_skipped() {
-        let now = Arc::new(AtomicU64::new(100));
-        let mut cache = cache(now, 1, Duration::from_secs(60));
-        cache.set_cache(
-            "removed",
-            "value".to_string(),
-            Some(Duration::from_secs(10)),
-        );
-        cache.delete_cache("removed");
-        cache.set_cache("kept", "value".to_string(), Some(Duration::from_secs(20)));
-        cache.set_cache("new", "value".to_string(), Some(Duration::from_secs(30)));
-        assert_eq!(cache.get_cache("removed"), None);
-        assert_eq!(cache.get_cache("kept"), None);
-        assert!(cache.get_cache("new").is_some());
-    }
-
-    #[test]
-    fn delete_and_flush_remove_values_and_ttls() {
-        let now = Arc::new(AtomicU64::new(100));
-        let mut cache = cache(now, 10, Duration::from_secs(60));
-        cache.set_cache("one", "value".to_string(), None);
-        cache.set_cache("two", "value".to_string(), None);
-        cache.delete_cache("one");
-        assert_eq!(cache.get_cache("one"), None);
-        cache.flush_cache();
-        assert!(cache.cache_dict.is_empty());
-        assert!(cache.ttl_dict.is_empty());
-        assert!(cache.expiration_heap.is_empty());
-    }
-
-    #[test]
-    fn zero_max_size_does_not_cache() {
-        let now = Arc::new(AtomicU64::new(100));
-        let mut cache = cache(now, 0, Duration::from_secs(60));
-        cache.set_cache("key", "value".to_string(), None);
-        assert_eq!(cache.get_cache("key"), None);
-        assert!(cache.cache_dict.is_empty());
-    }
-}

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

@@ -1 +0,0 @@
-pub mod in_memory_cache;

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

@@ -0,0 +1,71 @@
+use std::sync::Arc;
+use std::time::{Duration, SystemTime, UNIX_EPOCH};
+
+use litellm_cache::{BaseCache, CacheControls, CacheEntry, CacheKwargs};
+use serde_json::Value;
+
+use super::admission::AdmissionDecline;
+
+#[derive(Clone, Default)]
+pub struct ResponseCachePlan {
+    pub controls: CacheControls,
+    pub key: String,
+    pub backend: Option<Arc<dyn BaseCache<Value = CacheEntry>>>,
+    pub ttl: Option<Duration>,
+    pub max_age: Option<Duration>,
+}
+
+impl ResponseCachePlan {
+    pub fn admit(&self) -> Result<(), AdmissionDecline> {
+        if !self.controls.native_backend && (self.controls.reads() || self.controls.writes()) {
+            return Err(AdmissionDecline::Feature(
+                "response cache backend is not implemented in Rust",
+            ));
+        }
+        Ok(())
+    }
+}
+
+pub async fn lookup(plan: &ResponseCachePlan) -> Option<Value> {
+    let backend = plan.backend.as_ref().filter(|_| plan.controls.reads())?;
+    let entry = backend
+        .async_get_cache(&plan.key, &CacheKwargs::default())
+        .await
+        .ok()??;
+    let now = SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .unwrap_or_default();
+    if !entry.fresh(now, plan.max_age) {
+        return None;
+    }
+    match entry.response {
+        Value::String(text) => serde_json::from_str(&text).ok(),
+        value => Some(value),
+    }
+}
+
+pub async fn store(plan: &ResponseCachePlan, value: Option<Value>) {
+    let (Some(value), Some(backend)) = (
+        value,
+        plan.backend.as_ref().filter(|_| plan.controls.writes()),
+    ) else {
+        return;
+    };
+    let timestamp = SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .unwrap_or_default()
+        .as_secs_f64();
+    let _ = backend
+        .async_set_cache(
+            &plan.key,
+            CacheEntry {
+                timestamp,
+                response: value,
+            },
+            CacheKwargs {
+                ttl: plan.ttl,
+                ..Default::default()
+            },
+        )
+        .await;
+}

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

@@ -0,0 +1,220 @@
+use std::future::Future;
+
+use tokio::sync::{mpsc, oneshot};
+
+use super::host::HostCallStep;
+use crate::Error;
+
+struct PendingOperation<O, R> {
+    operation: O,
+    reply: oneshot::Sender<R>,
+}
+
+pub struct HostExchange<O, R> {
+    sender: mpsc::UnboundedSender<PendingOperation<O, R>>,
+}
+
+impl<O, R> Clone for HostExchange<O, R> {
+    fn clone(&self) -> Self {
+        Self {
+            sender: self.sender.clone(),
+        }
+    }
+}
+
+impl<O, R> std::fmt::Debug for HostExchange<O, R> {
+    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        formatter
+            .debug_struct("HostExchange")
+            .finish_non_exhaustive()
+    }
+}
+
+impl<O, R> HostExchange<O, R> {
+    pub async fn invoke(&self, operation: O) -> Result<R, Error> {
+        let (reply, receiver) = oneshot::channel();
+        self.sender
+            .send(PendingOperation { operation, reply })
+            .map_err(|_| Error::InvalidRequest("host driver was abandoned".into()))?;
+        receiver
+            .await
+            .map_err(|_| Error::InvalidRequest("host operation was abandoned".into()))
+    }
+}
+
+pub struct HostExecution<O, R, T> {
+    exchange: HostExchange<O, R>,
+    operations: mpsc::UnboundedReceiver<PendingOperation<O, R>>,
+    pending: Option<PendingOperation<O, R>>,
+    task: Option<tokio::task::JoinHandle<Result<T, Error>>>,
+    completed: bool,
+    accepts: fn(&O, &R) -> bool,
+}
+
+impl<O: Clone, R, T: Send + 'static> HostExecution<O, R, T> {
+    pub fn new(accepts: fn(&O, &R) -> bool) -> Self {
+        let (sender, operations) = mpsc::unbounded_channel();
+        Self {
+            exchange: HostExchange { sender },
+            operations,
+            pending: None,
+            task: None,
+            completed: false,
+            accepts,
+        }
+    }
+
+    pub fn exchange(&self) -> HostExchange<O, R> {
+        self.exchange.clone()
+    }
+
+    pub fn started(&self) -> bool {
+        self.task.is_some() || self.completed
+    }
+
+    pub fn start(
+        &mut self,
+        future: impl Future<Output = Result<T, Error>> + Send + 'static,
+    ) -> Result<(), Error> {
+        if self.started() {
+            return Err(Error::InvalidRequest(
+                "host execution already started".into(),
+            ));
+        }
+        self.task = Some(tokio::spawn(future));
+        Ok(())
+    }
+
+    pub async fn resume(&mut self, result: Option<R>) -> Result<HostCallStep<O, T>, Error> {
+        if self.completed {
+            return Err(Error::InvalidRequest(
+                "call cannot be resumed after completion".into(),
+            ));
+        }
+        match (&self.pending, &result) {
+            (Some(pending), Some(reply)) if (self.accepts)(&pending.operation, reply) => {}
+            (None, None) => {}
+            _ => {
+                return Err(Error::InvalidRequest(
+                    "host reply does not match pending operation".into(),
+                ));
+            }
+        }
+        if let (Some(pending), Some(reply)) = (self.pending.take(), result) {
+            pending
+                .reply
+                .send(reply)
+                .map_err(|_| Error::InvalidRequest("host operation was abandoned".into()))?;
+        }
+        let task = self
+            .task
+            .as_mut()
+            .ok_or_else(|| Error::InvalidRequest("host execution has not started".into()))?;
+        tokio::select! {
+            operation = self.operations.recv() => {
+                let pending = operation.ok_or_else(|| Error::InvalidRequest("host operation channel closed".into()))?;
+                let operation = pending.operation.clone();
+                self.pending = Some(pending);
+                Ok(HostCallStep::Host(operation))
+            }
+            result = task => {
+                self.task = None;
+                self.completed = true;
+                result.map_err(|error| Error::Network(format!("execution task failed: {error}")))?.map(HostCallStep::Complete)
+            }
+        }
+    }
+
+    pub fn cancel(&mut self) {
+        self.pending = None;
+        if let Some(task) = &self.task {
+            task.abort();
+        }
+    }
+
+    pub async fn stop(&mut self) {
+        self.cancel();
+        if let Some(task) = self.task.as_mut() {
+            let _ = task.await;
+        }
+        self.task = None;
+        self.completed = true;
+    }
+}
+
+impl<O, R, T> Drop for HostExecution<O, R, T> {
+    fn drop(&mut self) {
+        if let Some(task) = &self.task {
+            task.abort();
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::sync::{
+        Arc,
+        atomic::{AtomicBool, Ordering},
+    };
+
+    #[tokio::test]
+    async fn wrong_and_missing_replies_preserve_the_pending_exchange() {
+        let mut execution =
+            HostExecution::new(|operation: &u32, reply: &u32| *reply == *operation + 1);
+        let exchange = execution.exchange();
+        execution
+            .start(async move { exchange.invoke(40).await })
+            .unwrap();
+        assert!(matches!(
+            execution.resume(None).await.unwrap(),
+            HostCallStep::Host(40)
+        ));
+        assert!(execution.resume(Some(99)).await.is_err());
+        assert!(execution.resume(None).await.is_err());
+        assert!(matches!(
+            execution.resume(Some(41)).await.unwrap(),
+            HostCallStep::Complete(41)
+        ));
+        assert!(execution.resume(None).await.is_err());
+    }
+
+    struct Capture(Arc<AtomicBool>);
+    impl Drop for Capture {
+        fn drop(&mut self) {
+            self.0.store(true, Ordering::SeqCst);
+        }
+    }
+
+    #[tokio::test]
+    async fn stop_waits_for_provider_captures_to_drop() {
+        let dropped = Arc::new(AtomicBool::new(false));
+        let capture = Capture(dropped.clone());
+        let mut execution = HostExecution::<u32, u32, u32>::new(|_, _| true);
+        let exchange = execution.exchange();
+        execution
+            .start(async move {
+                let _capture = capture;
+                exchange.invoke(1).await
+            })
+            .unwrap();
+        assert!(matches!(
+            execution.resume(None).await.unwrap(),
+            HostCallStep::Host(1)
+        ));
+        execution.stop().await;
+        assert!(dropped.load(Ordering::SeqCst));
+        assert!(execution.resume(Some(2)).await.is_err());
+    }
+
+    #[tokio::test]
+    async fn provider_panics_are_terminal_errors() {
+        let mut execution = HostExecution::<u32, u32, u32>::new(|_, _| true);
+        execution.start(async { panic!("provider panic") }).unwrap();
+        assert!(matches!(
+            execution.resume(None).await,
+            Err(Error::Network(_))
+        ));
+        assert!(execution.resume(None).await.is_err());
+    }
+}

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

@@ -25,6 +25,26 @@ pub trait HostCall: Send + Sync {
     ) -> HostCallFuture<'_, Self::Operation, Self::Complete>;
 }
 
+pub type LifecycleBackendFuture<'a, R> = Pin<Box<dyn Future<Output = R> + Send + 'a>>;
+
+pub trait LifecycleBackend<O, R>: Send + Sync {
+    fn invoke(&self, operation: O) -> LifecycleBackendFuture<'_, R>;
+}
+
+pub async fn drive<C, B>(call: &mut C, backend: &B) -> Result<C::Complete, crate::Error>
+where
+    C: HostCall,
+    B: LifecycleBackend<C::Operation, C::Result> + ?Sized,
+{
+    let mut result = None;
+    loop {
+        match call.resume(result.take()).await? {
+            HostCallStep::Host(operation) => result = Some(backend.invoke(operation).await),
+            HostCallStep::Complete(response) => return Ok(response),
+        }
+    }
+}
+
 pub enum HostStep<V, S> {
     Ready(V),
     Suspend(S),
@@ -35,9 +55,12 @@ pub enum HostPhase {
     Setup,
     DeploymentPreCall,
     Prepare,
+    CacheLookup,
     Execute,
     ConstructResponse,
+    PostProcess,
     DeploymentPostCall,
+    CacheStore,
     Finalize,
     Success,
     MapFailure,
@@ -56,6 +79,7 @@ pub enum HostFailure {
 pub struct HostLifecycle {
     phase: HostPhase,
     asynchronous: bool,
+    cached: bool,
 }
 
 impl HostLifecycle {
@@ -63,9 +87,15 @@ impl HostLifecycle {
         Self {
             phase: HostPhase::Setup,
             asynchronous,
+            cached: false,
         }
     }
 
+    pub fn cache_hit(&mut self) {
+        self.cached = true;
+        self.phase = HostPhase::ConstructResponse;
+    }
+
     pub fn phase(&self) -> HostPhase {
         self.phase
     }
@@ -104,10 +134,16 @@ impl HostLifecycle {
         self.phase = match self.phase {
             HostPhase::Setup if self.asynchronous => HostPhase::DeploymentPreCall,
             HostPhase::Setup | HostPhase::DeploymentPreCall => HostPhase::Prepare,
-            HostPhase::Prepare => HostPhase::Execute,
+            HostPhase::Prepare => HostPhase::CacheLookup,
+            HostPhase::CacheLookup => HostPhase::Execute,
             HostPhase::Execute => HostPhase::ConstructResponse,
-            HostPhase::ConstructResponse if self.asynchronous => HostPhase::DeploymentPostCall,
-            HostPhase::ConstructResponse | HostPhase::DeploymentPostCall => HostPhase::Finalize,
+            HostPhase::ConstructResponse => HostPhase::PostProcess,
+            HostPhase::PostProcess if self.cached => HostPhase::Finalize,
+            HostPhase::PostProcess if self.asynchronous => HostPhase::DeploymentPostCall,
+            HostPhase::PostProcess => HostPhase::CacheStore,
+            HostPhase::DeploymentPostCall if self.cached => HostPhase::Finalize,
+            HostPhase::DeploymentPostCall => HostPhase::CacheStore,
+            HostPhase::CacheStore => HostPhase::Finalize,
             HostPhase::Finalize => HostPhase::Success,
             HostPhase::MapFailure if self.asynchronous => HostPhase::DeploymentFailure,
             HostPhase::MapFailure | HostPhase::DeploymentFailure => HostPhase::Failure,

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

@@ -1,420 +1,13 @@
-use std::future::Future;
-use std::time::{Instant, SystemTime, UNIX_EPOCH};
-
-use crate::Error;
-
 pub mod admission;
+pub mod cache;
 pub mod dispatch;
+pub mod execution;
 pub mod host;
 #[cfg(test)]
 #[path = "../../tests/host_lifecycle.rs"]
 mod host_tests;
+pub mod provider;
 pub mod types;
+pub mod workflow;
 
-pub use types::{
-    CallLifecycleContext, CallLifecyclePhase, CallLifecyclePhaseTiming, CallLifecycleRequest,
-    CallLifecycleTiming,
-};
-
-pub trait CallLifecycleHooks<InitialReq, ProviderReq, Resp>: Send + Sync {
-    type PreCallFuture<'a>: Future<Output = Result<InitialReq, Error>> + Send + 'a
-    where
-        Self: 'a,
-        InitialReq: 'a,
-        ProviderReq: 'a,
-        Resp: 'a;
-
-    type DuringCallFuture<'a>: Future<Output = Result<ProviderReq, Error>> + Send + 'a
-    where
-        Self: 'a,
-        InitialReq: 'a,
-        ProviderReq: 'a,
-        Resp: 'a;
-
-    type SuccessFuture<'a>: Future<Output = ()> + Send + 'a
-    where
-        Self: 'a,
-        Resp: 'a;
-
-    type FailureFuture<'a>: Future<Output = ()> + Send + 'a
-    where
-        Self: '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>;
-
-    fn async_log_success_event<'a>(
-        &'a self,
-        context: &'a CallLifecycleContext,
-        response: &'a Resp,
-        timing: &'a CallLifecycleTiming,
-    ) -> Self::SuccessFuture<'a>;
-
-    fn async_log_failure_event<'a>(
-        &'a self,
-        context: &'a CallLifecycleContext,
-        error: &'a Error,
-        timing: &'a CallLifecycleTiming,
-    ) -> Self::FailureFuture<'a>;
-}
-
-pub trait CallLifecycleObserver: Send + Sync {
-    fn on_phase_start(&self, _context: &CallLifecycleContext, _phase: CallLifecyclePhase) {}
-
-    fn on_phase_end(&self, _context: &CallLifecycleContext, _timing: &CallLifecyclePhaseTiming) {}
-}
-
-#[derive(Default)]
-pub struct NoopCallLifecycleObserver;
-
-impl CallLifecycleObserver for NoopCallLifecycleObserver {}
-
-pub struct CallLifecycle<'a> {
-    observer: &'a dyn CallLifecycleObserver,
-}
-
-impl<'a> CallLifecycle<'a> {
-    pub fn new(observer: &'a dyn CallLifecycleObserver) -> Self {
-        Self { observer }
-    }
-
-    pub async fn run_request<InitialReq, ProviderReq, Resp, Hooks, ProviderCall, ProviderFuture>(
-        &self,
-        request: InitialReq,
-        hooks: &Hooks,
-        provider_call: ProviderCall,
-    ) -> Result<Resp, Error>
-    where
-        InitialReq: CallLifecycleRequest,
-        Hooks: CallLifecycleHooks<InitialReq, ProviderReq, Resp>,
-        ProviderCall: FnOnce(ProviderReq) -> ProviderFuture,
-        ProviderFuture: Future<Output = Result<Resp, Error>>,
-    {
-        let context = request.lifecycle_context();
-        self.run(context, request, hooks, provider_call).await
-    }
-
-    pub async fn run<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>>,
-    {
-        let call_start = epoch_seconds();
-        let mut phases = Vec::new();
-
-        let pre_call = self.start_phase(&context, CallLifecyclePhase::PreCall);
-        let request = match hooks.async_pre_call_hook(&context, request).await {
-            Ok(request) => {
-                phases.push(self.finish_phase(&context, pre_call));
-                request
-            }
-            Err(error) => {
-                phases.push(self.finish_phase(&context, pre_call));
-                self.log_failure(&context, hooks, &error, call_start, &mut phases)
-                    .await;
-                return Err(error);
-            }
-        };
-
-        let during_call = self.start_phase(&context, CallLifecyclePhase::DuringCall);
-        let provider_request = match hooks.async_during_call_hook(&context, request).await {
-            Ok(request) => {
-                phases.push(self.finish_phase(&context, during_call));
-                request
-            }
-            Err(error) => {
-                phases.push(self.finish_phase(&context, during_call));
-                self.log_failure(&context, hooks, &error, call_start, &mut phases)
-                    .await;
-                return Err(error);
-            }
-        };
-
-        let provider_phase = self.start_phase(&context, CallLifecyclePhase::ProviderCall);
-        let result = provider_call(provider_request).await;
-        phases.push(self.finish_phase(&context, provider_phase));
-
-        match &result {
-            Ok(response) => {
-                let success_phase = self.start_phase(&context, CallLifecyclePhase::SuccessCallback);
-                let timing = CallLifecycleTiming::new(call_start, epoch_seconds(), phases.clone());
-                hooks
-                    .async_log_success_event(&context, response, &timing)
-                    .await;
-                phases.push(self.finish_phase(&context, success_phase));
-            }
-            Err(error) => {
-                self.log_failure(&context, hooks, error, call_start, &mut phases)
-                    .await;
-            }
-        }
-
-        result
-    }
-
-    async fn log_failure<InitialReq, ProviderReq, Resp, Hooks>(
-        &self,
-        context: &CallLifecycleContext,
-        hooks: &Hooks,
-        error: &Error,
-        call_start: f64,
-        phases: &mut Vec<CallLifecyclePhaseTiming>,
-    ) where
-        Hooks: CallLifecycleHooks<InitialReq, ProviderReq, Resp>,
-    {
-        let failure_phase = self.start_phase(context, CallLifecyclePhase::FailureCallback);
-        let timing = CallLifecycleTiming::new(call_start, epoch_seconds(), phases.clone());
-        hooks.async_log_failure_event(context, error, &timing).await;
-        phases.push(self.finish_phase(context, failure_phase));
-    }
-
-    fn start_phase(&self, context: &CallLifecycleContext, phase: CallLifecyclePhase) -> PhaseStart {
-        self.observer.on_phase_start(context, phase);
-        PhaseStart {
-            phase,
-            start_time: epoch_seconds(),
-            started_at: Instant::now(),
-        }
-    }
-
-    fn finish_phase(
-        &self,
-        context: &CallLifecycleContext,
-        phase_start: PhaseStart,
-    ) -> CallLifecyclePhaseTiming {
-        let timing = CallLifecyclePhaseTiming {
-            phase: phase_start.phase,
-            start_time: phase_start.start_time,
-            end_time: epoch_seconds(),
-            duration: phase_start.started_at.elapsed(),
-        };
-        self.observer.on_phase_end(context, &timing);
-        timing
-    }
-}
-
-impl Default for CallLifecycle<'static> {
-    fn default() -> Self {
-        static OBSERVER: NoopCallLifecycleObserver = NoopCallLifecycleObserver;
-        Self::new(&OBSERVER)
-    }
-}
-
-struct PhaseStart {
-    phase: CallLifecyclePhase,
-    start_time: f64,
-    started_at: Instant,
-}
-
-fn epoch_seconds() -> f64 {
-    SystemTime::now()
-        .duration_since(UNIX_EPOCH)
-        .map(|duration| duration.as_secs_f64())
-        .unwrap_or(0.0)
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-    use std::pin::Pin;
-    use std::sync::Mutex;
-
-    type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
-
-    #[derive(Default)]
-    struct RecordingHooks {
-        events: Mutex<Vec<&'static str>>,
-    }
-
-    struct RecordingRequest(String);
-
-    impl CallLifecycleRequest for RecordingRequest {
-        fn lifecycle_context(&self) -> CallLifecycleContext {
-            CallLifecycleContext::new("ocr", "mistral-ocr-latest", "mistral", "call_1")
-        }
-    }
-
-    impl RecordingHooks {
-        fn events(&self) -> Vec<&'static str> {
-            self.events.lock().unwrap().clone()
-        }
-    }
-
-    impl CallLifecycleHooks<String, String, String> for RecordingHooks {
-        type PreCallFuture<'a> = BoxFuture<'a, Result<String, Error>>;
-        type DuringCallFuture<'a> = BoxFuture<'a, Result<String, Error>>;
-        type SuccessFuture<'a> = BoxFuture<'a, ()>;
-        type FailureFuture<'a> = BoxFuture<'a, ()>;
-
-        fn async_pre_call_hook<'a>(
-            &'a self,
-            _context: &'a CallLifecycleContext,
-            request: String,
-        ) -> Self::PreCallFuture<'a> {
-            Box::pin(async move {
-                self.events.lock().unwrap().push("pre_call");
-                Ok(format!("{request}:pre"))
-            })
-        }
-
-        fn async_during_call_hook<'a>(
-            &'a self,
-            _context: &'a CallLifecycleContext,
-            request: String,
-        ) -> Self::DuringCallFuture<'a> {
-            Box::pin(async move {
-                self.events.lock().unwrap().push("during_call");
-                Ok(format!("{request}:during"))
-            })
-        }
-
-        fn async_log_success_event<'a>(
-            &'a self,
-            _context: &'a CallLifecycleContext,
-            _response: &'a String,
-            timing: &'a CallLifecycleTiming,
-        ) -> Self::SuccessFuture<'a> {
-            Box::pin(async move {
-                assert!(timing.end_time >= timing.start_time);
-                assert_eq!(timing.phases.len(), 3);
-                self.events.lock().unwrap().push("success");
-            })
-        }
-
-        fn async_log_failure_event<'a>(
-            &'a self,
-            _context: &'a CallLifecycleContext,
-            _error: &'a Error,
-            _timing: &'a CallLifecycleTiming,
-        ) -> Self::FailureFuture<'a> {
-            Box::pin(async move {
-                self.events.lock().unwrap().push("failure");
-            })
-        }
-    }
-
-    impl CallLifecycleHooks<RecordingRequest, String, String> for RecordingHooks {
-        type PreCallFuture<'a> = BoxFuture<'a, Result<RecordingRequest, Error>>;
-        type DuringCallFuture<'a> = BoxFuture<'a, Result<String, Error>>;
-        type SuccessFuture<'a> = BoxFuture<'a, ()>;
-        type FailureFuture<'a> = BoxFuture<'a, ()>;
-
-        fn async_pre_call_hook<'a>(
-            &'a self,
-            _context: &'a CallLifecycleContext,
-            request: RecordingRequest,
-        ) -> Self::PreCallFuture<'a> {
-            Box::pin(async move {
-                self.events.lock().unwrap().push("pre_call");
-                Ok(RecordingRequest(format!("{}:pre", request.0)))
-            })
-        }
-
-        fn async_during_call_hook<'a>(
-            &'a self,
-            _context: &'a CallLifecycleContext,
-            request: RecordingRequest,
-        ) -> Self::DuringCallFuture<'a> {
-            Box::pin(async move {
-                self.events.lock().unwrap().push("during_call");
-                Ok(format!("{}:during", request.0))
-            })
-        }
-
-        fn async_log_success_event<'a>(
-            &'a self,
-            _context: &'a CallLifecycleContext,
-            _response: &'a String,
-            _timing: &'a CallLifecycleTiming,
-        ) -> Self::SuccessFuture<'a> {
-            Box::pin(async move {
-                self.events.lock().unwrap().push("success");
-            })
-        }
-
-        fn async_log_failure_event<'a>(
-            &'a self,
-            _context: &'a CallLifecycleContext,
-            _error: &'a Error,
-            _timing: &'a CallLifecycleTiming,
-        ) -> Self::FailureFuture<'a> {
-            Box::pin(async move {
-                self.events.lock().unwrap().push("failure");
-            })
-        }
-    }
-
-    #[tokio::test]
-    async fn lifecycle_runs_hooks_around_provider_call() {
-        let hooks = RecordingHooks::default();
-        let response = CallLifecycle::default()
-            .run(
-                CallLifecycleContext::new("ocr", "mistral-ocr-latest", "mistral", "call_1"),
-                "request".to_string(),
-                &hooks,
-                |request| async move {
-                    assert_eq!(request, "request:pre:during");
-                    Ok("response".to_string())
-                },
-            )
-            .await
-            .expect("call succeeds");
-
-        assert_eq!(response, "response");
-        assert_eq!(hooks.events(), vec!["pre_call", "during_call", "success"]);
-    }
-
-    #[tokio::test]
-    async fn lifecycle_logs_failure_when_provider_fails() {
-        let hooks = RecordingHooks::default();
-        let error = CallLifecycle::default()
-            .run(
-                CallLifecycleContext::new("ocr", "mistral-ocr-latest", "mistral", "call_1"),
-                "request".to_string(),
-                &hooks,
-                |_request| async move {
-                    Err::<String, Error>(Error::Network("provider down".to_string()))
-                },
-            )
-            .await
-            .expect_err("call fails");
-
-        assert_eq!(error, Error::Network("provider down".to_string()));
-        assert_eq!(hooks.events(), vec!["pre_call", "during_call", "failure"]);
-    }
-
-    #[tokio::test]
-    async fn lifecycle_can_run_any_request_with_embedded_context() {
-        let hooks = RecordingHooks::default();
-        let response = CallLifecycle::default()
-            .run_request(
-                RecordingRequest("request".to_string()),
-                &hooks,
-                |request| async move {
-                    assert_eq!(request, "request:pre:during");
-                    Ok("response".to_string())
-                },
-            )
-            .await
-            .expect("call succeeds");
-
-        assert_eq!(response, "response");
-        assert_eq!(hooks.events(), vec!["pre_call", "during_call", "success"]);
-    }
-}
+pub use types::{CallLifecycleContext, CallLifecycleRequest, CallLifecycleTiming};

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

@@ -0,0 +1,516 @@
+use std::future::Future;
+use std::marker::PhantomData;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::sync::Mutex;
+use std::time::{Duration, SystemTime, UNIX_EPOCH};
+
+use serde::{Deserialize, Serialize};
+use serde_json::{Map, Value};
+
+use super::cache::{ResponseCachePlan, lookup, store};
+use super::execution::HostExchange;
+use super::host::{HostFailure, LifecycleBackend, LifecycleBackendFuture, drive};
+use super::workflow::{LifecycleCall, LifecycleOperation, Workflow, WorkflowFuture, WorkflowReply};
+use crate::Error;
+use crate::call_lifecycle::{CallLifecycleContext, CallLifecycleTiming};
+
+#[derive(Clone, Default)]
+pub struct ProviderOptions {
+    pub model: String,
+    pub litellm_call_id: Option<String>,
+    pub api_key: Option<String>,
+    pub api_base: Option<String>,
+    pub custom_llm_provider: Option<String>,
+    pub extra_headers: Option<Map<String, Value>>,
+    pub timeout: Option<Duration>,
+}
+
+impl ProviderOptions {
+    pub fn lifecycle_context(&self, call_type: &str) -> CallLifecycleContext {
+        CallLifecycleContext::new(
+            call_type,
+            self.model.clone(),
+            self.custom_llm_provider.clone().unwrap_or_default(),
+            self.litellm_call_id
+                .clone()
+                .unwrap_or_else(|| format!("native-{:032x}", rand::random::<u128>())),
+        )
+    }
+}
+
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct ProviderRequest {
+    pub model: String,
+    pub url: String,
+    pub headers: Vec<(String, String)>,
+    pub body: Value,
+}
+
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct ProviderResponse {
+    pub status: u16,
+    pub body: String,
+}
+
+pub type ProviderHookFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
+
+pub trait ProviderHooks: Send + Sync {
+    fn before_request(&self, request: ProviderRequest) -> ProviderHookFuture<'_, ProviderRequest>;
+    fn after_response(
+        &self,
+        response: ProviderResponse,
+    ) -> ProviderHookFuture<'_, ProviderResponse>;
+}
+
+pub struct NoopProviderHooks;
+
+impl ProviderHooks for NoopProviderHooks {
+    fn before_request(&self, request: ProviderRequest) -> ProviderHookFuture<'_, ProviderRequest> {
+        Box::pin(async move { Ok(request) })
+    }
+
+    fn after_response(
+        &self,
+        response: ProviderResponse,
+    ) -> ProviderHookFuture<'_, ProviderResponse> {
+        Box::pin(async move { Ok(response) })
+    }
+}
+
+pub struct ProviderHookChain {
+    hooks: Box<[Arc<dyn ProviderHooks>]>,
+}
+
+impl ProviderHookChain {
+    pub fn new(hooks: impl IntoIterator<Item = Arc<dyn ProviderHooks>>) -> Self {
+        Self {
+            hooks: hooks.into_iter().collect(),
+        }
+    }
+
+    fn before<'a>(
+        hooks: &'a [Arc<dyn ProviderHooks>],
+        request: ProviderRequest,
+    ) -> ProviderHookFuture<'a, ProviderRequest> {
+        Box::pin(async move {
+            let Some((hook, remaining)) = hooks.split_first() else {
+                return Ok(request);
+            };
+            let request = hook.before_request(request).await?;
+            Self::before(remaining, request).await
+        })
+    }
+
+    fn after<'a>(
+        hooks: &'a [Arc<dyn ProviderHooks>],
+        response: ProviderResponse,
+    ) -> ProviderHookFuture<'a, ProviderResponse> {
+        Box::pin(async move {
+            let Some((hook, remaining)) = hooks.split_first() else {
+                return Ok(response);
+            };
+            let response = hook.after_response(response).await?;
+            Self::after(remaining, response).await
+        })
+    }
+}
+
+impl ProviderHooks for ProviderHookChain {
+    fn before_request(&self, request: ProviderRequest) -> ProviderHookFuture<'_, ProviderRequest> {
+        Self::before(&self.hooks, request)
+    }
+
+    fn after_response(
+        &self,
+        response: ProviderResponse,
+    ) -> ProviderHookFuture<'_, ProviderResponse> {
+        Self::after(&self.hooks, response)
+    }
+}
+
+#[derive(Clone, Debug)]
+pub enum CompletedOperation<T> {
+    Lifecycle(LifecycleOperation<T>),
+    BeforeRequest(ProviderRequest),
+    AfterResponse(ProviderResponse),
+}
+
+pub enum CompletedReply<Q> {
+    Prepared(Result<ResponseCachePlan, Error>),
+    CacheStore(Result<Option<Value>, Error>),
+    Request(Result<Q, Error>),
+    Lifecycle(Result<(), HostFailure>),
+    BeforeRequest(Result<ProviderRequest, Error>),
+    AfterResponse(Result<ProviderResponse, Error>),
+}
+
+pub trait CompletedRoute: Send + Sync + 'static {
+    type Request: Send + Sync + 'static;
+    type Response: Clone + Send + Sync + serde::de::DeserializeOwned + 'static;
+
+    fn run(request: Self::Request, hooks: Arc<dyn ProviderHooks>)
+    -> WorkflowFuture<Self::Response>;
+    fn context(request: &Self::Request) -> CallLifecycleContext;
+}
+
+pub struct CompletedWorkflow<R: CompletedRoute> {
+    route: PhantomData<R>,
+    hooks: Arc<dyn ProviderHooks>,
+    cache: ResponseCachePlan,
+    cached: Arc<Mutex<Option<Value>>>,
+    pending_write: Option<Value>,
+    terminal: Arc<Mutex<Option<(CallLifecycleContext, CallLifecycleTiming)>>>,
+}
+
+impl<R: CompletedRoute> Default for CompletedWorkflow<R> {
+    fn default() -> Self {
+        Self {
+            route: PhantomData,
+            hooks: Arc::new(NoopProviderHooks),
+            cache: ResponseCachePlan::default(),
+            cached: Arc::default(),
+            pending_write: None,
+            terminal: Arc::default(),
+        }
+    }
+}
+
+impl<R: CompletedRoute> CompletedWorkflow<R> {
+    pub fn with_hooks(hooks: Arc<dyn ProviderHooks>) -> Self {
+        Self {
+            hooks,
+            ..Self::default()
+        }
+    }
+}
+
+impl<R: CompletedRoute> Workflow for CompletedWorkflow<R> {
+    type Request = R::Request;
+    type Operation = CompletedOperation<R::Response>;
+    type Reply = CompletedReply<R::Request>;
+    type Response = R::Response;
+
+    fn operation(operation: LifecycleOperation<Self::Response>) -> Self::Operation {
+        CompletedOperation::Lifecycle(operation)
+    }
+
+    fn accepts(operation: &Self::Operation, reply: &Self::Reply) -> bool {
+        match (operation, reply) {
+            (_, CompletedReply::Lifecycle(Err(_))) => true,
+            (
+                CompletedOperation::Lifecycle(LifecycleOperation::Phase(
+                    super::host::HostPhase::Prepare,
+                )),
+                CompletedReply::Prepared(_),
+            ) => true,
+            (
+                CompletedOperation::Lifecycle(LifecycleOperation::Phase(
+                    super::host::HostPhase::CacheStore,
+                )),
+                CompletedReply::CacheStore(_),
+            ) => true,
+            (
+                CompletedOperation::Lifecycle(LifecycleOperation::ProjectRequest),
+                CompletedReply::Request(_),
+            ) => true,
+            (CompletedOperation::Lifecycle(LifecycleOperation::ProjectRequest), _) => false,
+            (CompletedOperation::Lifecycle(_), CompletedReply::Lifecycle(_)) => true,
+            (CompletedOperation::BeforeRequest(_), CompletedReply::BeforeRequest(_)) => true,
+            (CompletedOperation::AfterResponse(_), CompletedReply::AfterResponse(_)) => true,
+            _ => false,
+        }
+    }
+
+    fn reply(&mut self, reply: Self::Reply) -> WorkflowReply<Self::Request, Self::Reply> {
+        match reply {
+            CompletedReply::Prepared(result) => WorkflowReply::Lifecycle(
+                result
+                    .map(|cache| self.cache = cache)
+                    .map_err(HostFailure::Error),
+            ),
+            CompletedReply::CacheStore(result) => WorkflowReply::Lifecycle(
+                result
+                    .map(|value| self.pending_write = value)
+                    .map_err(HostFailure::Error),
+            ),
+            CompletedReply::Request(request) => WorkflowReply::Request(request),
+            CompletedReply::Lifecycle(result) => WorkflowReply::Lifecycle(result),
+            reply => WorkflowReply::Operation(reply),
+        }
+    }
+
+    fn start(
+        &mut self,
+        request: Self::Request,
+        host: HostExchange<Self::Operation, Self::Reply>,
+    ) -> WorkflowFuture<Self::Response> {
+        let context = R::context(&request);
+        let started = epoch_seconds();
+        let terminal = self.terminal.clone();
+        let host_hooks: Arc<dyn ProviderHooks> = Arc::new(ExchangeHooks { host });
+        let future = R::run(
+            request,
+            Arc::new(ProviderHookChain::new([self.hooks.clone(), host_hooks])),
+        );
+        Box::pin(async move {
+            let result = future.await;
+            let timing = CallLifecycleTiming::new(started, epoch_seconds());
+            *terminal.lock().unwrap_or_else(|error| error.into_inner()) = Some((context, timing));
+            result
+        })
+    }
+
+    fn terminal(&self) -> Option<(CallLifecycleContext, CallLifecycleTiming)> {
+        self.terminal
+            .lock()
+            .unwrap_or_else(|error| error.into_inner())
+            .clone()
+    }
+
+    fn cache_lookup(&mut self) -> WorkflowFuture<Option<Self::Response>> {
+        let cached = self.cached.clone();
+        let cache = self.cache.clone();
+        Box::pin(async move {
+            let Some(value) = lookup(&cache).await else {
+                return Ok(None);
+            };
+            let Ok(response) = serde_json::from_value(value.clone()) else {
+                return Ok(None);
+            };
+            *cached.lock().unwrap_or_else(|error| error.into_inner()) = Some(value);
+            Ok(Some(response))
+        })
+    }
+
+    fn cached_public_response(&self) -> Option<Value> {
+        self.cached
+            .lock()
+            .unwrap_or_else(|error| error.into_inner())
+            .clone()
+    }
+
+    fn flush_cache(&mut self) -> WorkflowFuture<()> {
+        let value = self.pending_write.take();
+        let cache = self.cache.clone();
+        Box::pin(async move {
+            store(&cache, value).await;
+            Ok(())
+        })
+    }
+}
+
+fn epoch_seconds() -> f64 {
+    SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .map(|duration| duration.as_secs_f64())
+        .unwrap_or(0.0)
+}
+
+struct ExchangeHooks<Q, T> {
+    host: HostExchange<CompletedOperation<T>, CompletedReply<Q>>,
+}
+
+impl<Q: Send + 'static, T: Send + Sync + 'static> ProviderHooks for ExchangeHooks<Q, T> {
+    fn before_request(&self, request: ProviderRequest) -> ProviderHookFuture<'_, ProviderRequest> {
+        Box::pin(async move {
+            match self
+                .host
+                .invoke(CompletedOperation::BeforeRequest(request))
+                .await?
+            {
+                CompletedReply::BeforeRequest(result) => result,
+                _ => Err(Error::InvalidRequest(
+                    "unexpected provider request reply".into(),
+                )),
+            }
+        })
+    }
+
+    fn after_response(
+        &self,
+        response: ProviderResponse,
+    ) -> ProviderHookFuture<'_, ProviderResponse> {
+        Box::pin(async move {
+            match self
+                .host
+                .invoke(CompletedOperation::AfterResponse(response))
+                .await?
+            {
+                CompletedReply::AfterResponse(result) => result,
+                _ => Err(Error::InvalidRequest(
+                    "unexpected provider response reply".into(),
+                )),
+            }
+        })
+    }
+}
+
+pub type CompletedCall<R> = LifecycleCall<CompletedWorkflow<R>>;
+
+pub async fn run_completed<R: CompletedRoute>(request: R::Request) -> Result<R::Response, Error> {
+    run_completed_with_hooks::<R>(request, Arc::new(NoopProviderHooks)).await
+}
+
+pub async fn run_completed_with_hooks<R: CompletedRoute>(
+    request: R::Request,
+    hooks: Arc<dyn ProviderHooks>,
+) -> Result<R::Response, Error> {
+    let mut call = CompletedCall::<R>::new(CompletedWorkflow::with_hooks(hooks), false);
+    drive(
+        &mut call,
+        &CompletedBackend::<R> {
+            request: Mutex::new(Some(request)),
+            route: PhantomData,
+        },
+    )
+    .await
+}
+
+struct CompletedBackend<R: CompletedRoute> {
+    request: Mutex<Option<R::Request>>,
+    route: PhantomData<R>,
+}
+
+impl<R: CompletedRoute>
+    LifecycleBackend<CompletedOperation<R::Response>, CompletedReply<R::Request>>
+    for CompletedBackend<R>
+{
+    fn invoke(
+        &self,
+        operation: CompletedOperation<R::Response>,
+    ) -> LifecycleBackendFuture<'_, CompletedReply<R::Request>> {
+        Box::pin(async move {
+            match operation {
+                CompletedOperation::Lifecycle(LifecycleOperation::ProjectRequest) => {
+                    CompletedReply::Request(
+                        self.request
+                            .lock()
+                            .unwrap_or_else(|error| error.into_inner())
+                            .take()
+                            .ok_or_else(|| {
+                                Error::InvalidRequest("request was already projected".into())
+                            }),
+                    )
+                }
+                CompletedOperation::Lifecycle(_) => CompletedReply::Lifecycle(Ok(())),
+                CompletedOperation::BeforeRequest(request) => {
+                    CompletedReply::BeforeRequest(Ok(request))
+                }
+                CompletedOperation::AfterResponse(response) => {
+                    CompletedReply::AfterResponse(Ok(response))
+                }
+            }
+        })
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    struct RecordingHook {
+        name: &'static str,
+        events: Arc<Mutex<Vec<String>>>,
+    }
+
+    impl ProviderHooks for RecordingHook {
+        fn before_request(
+            &self,
+            request: ProviderRequest,
+        ) -> ProviderHookFuture<'_, ProviderRequest> {
+            Box::pin(async move {
+                self.events
+                    .lock()
+                    .unwrap()
+                    .push(format!("{}:before", self.name));
+                Ok(ProviderRequest {
+                    body: Value::String(format!(
+                        "{}:{}",
+                        request.body.as_str().unwrap(),
+                        self.name
+                    )),
+                    ..request
+                })
+            })
+        }
+
+        fn after_response(
+            &self,
+            response: ProviderResponse,
+        ) -> ProviderHookFuture<'_, ProviderResponse> {
+            Box::pin(async move {
+                self.events
+                    .lock()
+                    .unwrap()
+                    .push(format!("{}:after", self.name));
+                Ok(ProviderResponse {
+                    body: format!("{}:{}", response.body, self.name),
+                    ..response
+                })
+            })
+        }
+    }
+
+    #[derive(Clone, Deserialize)]
+    struct TestResponse(String);
+
+    struct TestRoute;
+
+    impl CompletedRoute for TestRoute {
+        type Request = ();
+        type Response = TestResponse;
+
+        fn run((): Self::Request, hooks: Arc<dyn ProviderHooks>) -> WorkflowFuture<Self::Response> {
+            Box::pin(async move {
+                let request = hooks
+                    .before_request(ProviderRequest {
+                        model: "model".into(),
+                        url: "https://example.com".into(),
+                        headers: Vec::new(),
+                        body: Value::String("request".into()),
+                    })
+                    .await?;
+                let response = hooks
+                    .after_response(ProviderResponse {
+                        status: 200,
+                        body: request.body.as_str().unwrap().to_owned(),
+                    })
+                    .await?;
+                Ok(TestResponse(response.body))
+            })
+        }
+
+        fn context(_: &Self::Request) -> CallLifecycleContext {
+            CallLifecycleContext::new("test", "model", "provider", "call")
+        }
+    }
+
+    #[tokio::test]
+    async fn native_hook_chain_runs_in_order_through_the_lifecycle() {
+        let events = Arc::new(Mutex::new(Vec::new()));
+        let hooks = ProviderHookChain::new([
+            Arc::new(RecordingHook {
+                name: "first",
+                events: events.clone(),
+            }) as Arc<dyn ProviderHooks>,
+            Arc::new(RecordingHook {
+                name: "second",
+                events: events.clone(),
+            }) as Arc<dyn ProviderHooks>,
+        ]);
+        let response = run_completed_with_hooks::<TestRoute>((), Arc::new(hooks))
+            .await
+            .unwrap();
+        assert_eq!(response.0, "request:first:second:first:second");
+        assert_eq!(
+            *events.lock().unwrap(),
+            [
+                "first:before",
+                "second:before",
+                "first:after",
+                "second:after"
+            ]
+        );
+    }
+}

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

@@ -1,5 +1,3 @@
-use std::time::Duration;
-
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub struct CallLifecycleContext {
     pub call_type: String,
@@ -28,48 +26,17 @@ pub trait CallLifecycleRequest {
     fn lifecycle_context(&self) -> CallLifecycleContext;
 }
 
-#[derive(Clone, Copy, Debug, PartialEq, Eq)]
-pub enum CallLifecyclePhase {
-    PreCall,
-    DuringCall,
-    ProviderCall,
-    SuccessCallback,
-    FailureCallback,
-}
-
-impl CallLifecyclePhase {
-    pub fn as_str(self) -> &'static str {
-        match self {
-            Self::PreCall => "pre_call",
-            Self::DuringCall => "during_call",
-            Self::ProviderCall => "provider_call",
-            Self::SuccessCallback => "success_callback",
-            Self::FailureCallback => "failure_callback",
-        }
-    }
-}
-
-#[derive(Clone, Copy, Debug, PartialEq)]
-pub struct CallLifecyclePhaseTiming {
-    pub phase: CallLifecyclePhase,
-    pub start_time: f64,
-    pub end_time: f64,
-    pub duration: Duration,
-}
-
 #[derive(Clone, Debug, PartialEq)]
 pub struct CallLifecycleTiming {
     pub start_time: f64,
     pub end_time: f64,
-    pub phases: Vec<CallLifecyclePhaseTiming>,
 }
 
 impl CallLifecycleTiming {
-    pub fn new(start_time: f64, end_time: f64, phases: Vec<CallLifecyclePhaseTiming>) -> Self {
+    pub fn new(start_time: f64, end_time: f64) -> Self {
         Self {
             start_time,
             end_time,
-            phases,
         }
     }
 }

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

@@ -0,0 +1,250 @@
+use std::future::Future;
+use std::pin::Pin;
+use std::sync::Arc;
+
+use super::execution::{HostExchange, HostExecution};
+use super::host::{HostCall, HostCallFuture, HostCallStep, HostFailure, HostLifecycle, HostPhase};
+use super::{CallLifecycleContext, CallLifecycleTiming};
+use crate::Error;
+
+pub type WorkflowFuture<T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'static>>;
+
+#[derive(Clone, Debug)]
+pub enum LifecycleOperation<T> {
+    ProjectRequest,
+    Phase(HostPhase),
+    ConstructResponse(Arc<T>),
+    ConstructCachedResponse(serde_json::Value),
+    MapFailure(Error),
+    Success {
+        context: CallLifecycleContext,
+        response: Arc<T>,
+        timing: CallLifecycleTiming,
+    },
+    Failure {
+        context: CallLifecycleContext,
+        error: Error,
+        timing: CallLifecycleTiming,
+    },
+}
+
+pub enum WorkflowReply<Q, R> {
+    Request(Result<Q, Error>),
+    Lifecycle(Result<(), HostFailure>),
+    Operation(R),
+}
+
+pub trait Workflow: Send + Sync {
+    type Request: Send + 'static;
+    type Operation: Clone + Send + Sync + 'static;
+    type Reply: Send + Sync + 'static;
+    type Response: Clone + Send + Sync + 'static;
+
+    fn operation(operation: LifecycleOperation<Self::Response>) -> Self::Operation;
+    fn accepts(operation: &Self::Operation, reply: &Self::Reply) -> bool;
+    fn reply(&mut self, reply: Self::Reply) -> WorkflowReply<Self::Request, Self::Reply>;
+    fn start(
+        &mut self,
+        request: Self::Request,
+        host: HostExchange<Self::Operation, Self::Reply>,
+    ) -> WorkflowFuture<Self::Response>;
+    fn terminal(&self) -> Option<(CallLifecycleContext, CallLifecycleTiming)> {
+        None
+    }
+    fn cache_lookup(&mut self) -> WorkflowFuture<Option<Self::Response>> {
+        Box::pin(async { Ok(None) })
+    }
+    fn cached_public_response(&self) -> Option<serde_json::Value> {
+        None
+    }
+    fn flush_cache(&mut self) -> WorkflowFuture<()> {
+        Box::pin(async { Ok(()) })
+    }
+}
+
+pub struct LifecycleCall<W: Workflow> {
+    workflow: W,
+    lifecycle: HostLifecycle,
+    execution: HostExecution<W::Operation, W::Reply, W::Response>,
+    pending: Option<W::Operation>,
+    response: Option<Arc<W::Response>>,
+    error: Option<Error>,
+    completed: bool,
+}
+
+impl<W: Workflow> LifecycleCall<W> {
+    pub fn new(workflow: W, asynchronous: bool) -> Self {
+        Self {
+            workflow,
+            lifecycle: HostLifecycle::new(asynchronous),
+            execution: HostExecution::new(W::accepts),
+            pending: None,
+            response: None,
+            error: None,
+            completed: false,
+        }
+    }
+
+    pub async fn resume(
+        &mut self,
+        result: Option<W::Reply>,
+    ) -> Result<HostCallStep<W::Operation, W::Response>, Error> {
+        if self.completed {
+            return Err(Error::InvalidRequest(
+                "call cannot be resumed after completion".into(),
+            ));
+        }
+        match (&self.pending, &result) {
+            (Some(operation), Some(reply)) if W::accepts(operation, reply) => {}
+            (None, None) => {}
+            _ => {
+                return Err(Error::InvalidRequest(
+                    "host reply does not match pending operation".into(),
+                ));
+            }
+        }
+        self.pending = None;
+        let provider_reply = match result.map(|reply| self.workflow.reply(reply)) {
+            Some(WorkflowReply::Request(Ok(request))) => {
+                let future = self.workflow.start(request, self.execution.exchange());
+                self.execution.start(future)?;
+                None
+            }
+            Some(WorkflowReply::Request(Err(error))) => {
+                self.accept(Err(HostFailure::Error(error)));
+                None
+            }
+            Some(WorkflowReply::Lifecycle(result)) => {
+                self.accept(result);
+                None
+            }
+            Some(WorkflowReply::Operation(reply)) => Some(reply),
+            None => None,
+        };
+        self.workflow.flush_cache().await?;
+        if self.lifecycle.phase() == HostPhase::CacheLookup {
+            match self.workflow.cache_lookup().await {
+                Ok(Some(response)) => {
+                    self.response = Some(Arc::new(response));
+                    self.lifecycle.cache_hit();
+                }
+                Ok(None) => self.accept(Ok(())),
+                Err(error) => self.accept(Err(HostFailure::Error(error))),
+            }
+        }
+        if self.lifecycle.phase() == HostPhase::Execute {
+            if !self.execution.started() {
+                return Ok(self.host_step(W::operation(LifecycleOperation::ProjectRequest)));
+            }
+            match self.execution.resume(provider_reply).await {
+                Ok(HostCallStep::Host(operation)) => return Ok(self.host_step(operation)),
+                Ok(HostCallStep::Complete(response)) => {
+                    self.response = Some(Arc::new(response));
+                    self.accept(Ok(()));
+                }
+                Err(error) => self.accept(Err(HostFailure::Error(error))),
+            }
+        }
+        if self.error.is_some() {
+            self.execution.stop().await;
+        }
+        let operation = match self.lifecycle.phase() {
+            HostPhase::Complete => {
+                self.completed = true;
+                return match self.error.take() {
+                    Some(error) => Err(error),
+                    None => self
+                        .response
+                        .take()
+                        .map(Arc::unwrap_or_clone)
+                        .map(HostCallStep::Complete)
+                        .ok_or_else(|| {
+                            Error::InvalidRequest("call completed without a response".into())
+                        }),
+                };
+            }
+            HostPhase::ConstructResponse => match self.workflow.cached_public_response() {
+                Some(response) => LifecycleOperation::ConstructCachedResponse(response),
+                None => LifecycleOperation::ConstructResponse(self.response()?),
+            },
+            HostPhase::MapFailure => LifecycleOperation::MapFailure(self.error()?),
+            phase => match (phase, self.workflow.terminal()) {
+                (HostPhase::Success, Some((context, timing))) => LifecycleOperation::Success {
+                    context,
+                    response: self.response()?,
+                    timing,
+                },
+                (HostPhase::Failure, Some((context, timing))) => LifecycleOperation::Failure {
+                    context,
+                    error: self.error()?,
+                    timing,
+                },
+                _ => LifecycleOperation::Phase(phase),
+            },
+        };
+        Ok(self.host_step(W::operation(operation)))
+    }
+
+    fn response(&self) -> Result<Arc<W::Response>, Error> {
+        self.response
+            .clone()
+            .ok_or_else(|| Error::InvalidRequest("missing response".into()))
+    }
+
+    fn error(&self) -> Result<Error, Error> {
+        self.error
+            .clone()
+            .ok_or_else(|| Error::InvalidRequest("missing failure".into()))
+    }
+
+    fn host_step(&mut self, operation: W::Operation) -> HostCallStep<W::Operation, W::Response> {
+        self.pending = Some(operation.clone());
+        HostCallStep::Host(operation)
+    }
+
+    fn accept(&mut self, result: Result<(), HostFailure>) {
+        let cancelled = matches!(&result, Err(HostFailure::Cancelled(_)));
+        if let Some(error) = self.lifecycle.accept(result) {
+            if cancelled {
+                self.error = Some(error);
+            } else {
+                self.error.get_or_insert(error);
+            }
+            self.execution.cancel();
+        }
+    }
+
+    pub async fn interrupt(
+        &mut self,
+        failure: HostFailure,
+    ) -> Result<HostCallStep<W::Operation, W::Response>, Error> {
+        if self.completed {
+            return Err(Error::InvalidRequest(
+                "call cannot be interrupted after completion".into(),
+            ));
+        }
+        self.pending = None;
+        self.accept(Err(failure));
+        self.resume(None).await
+    }
+}
+
+impl<W: Workflow> HostCall for LifecycleCall<W> {
+    type Operation = W::Operation;
+    type Result = W::Reply;
+    type Complete = W::Response;
+
+    fn resume(
+        &mut self,
+        result: Option<Self::Result>,
+    ) -> HostCallFuture<'_, Self::Operation, Self::Complete> {
+        Box::pin(LifecycleCall::resume(self, result))
+    }
+
+    fn interrupt(
+        &mut self,
+        failure: HostFailure,
+    ) -> HostCallFuture<'_, Self::Operation, Self::Complete> {
+        Box::pin(LifecycleCall::interrupt(self, failure))
+    }
+}

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

@@ -1,5 +1,6 @@
 use serde_json::Value;
 
+use crate::call_lifecycle::provider::{ProviderHooks, ProviderRequest, ProviderResponse};
 use crate::error::Error;
 use crate::http_utils::{http_request, truncate_error_body};
 
@@ -14,8 +15,24 @@ use super::types::{
 #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
 pub(super) async fn execute_chat_completions_provider_call(
     request: ResolvedChatCompletionsRequest<'_>,
+    hooks: &dyn ProviderHooks,
 ) -> Result<ChatCompletionsResponse, Error> {
     let request = prepare_provider_request(request)?;
+    let changed = hooks
+        .before_request(ProviderRequest {
+            model: request.model.clone(),
+            url: request.url.clone(),
+            headers: request.upstream_headers.clone(),
+            body: request.body.clone(),
+        })
+        .await?;
+    let request = ProviderChatCompletionsRequest {
+        model: changed.model,
+        url: changed.url,
+        upstream_headers: changed.headers,
+        body: changed.body,
+        ..request
+    };
     let body = serde_json::to_vec(&request.body).map_err(|err| {
         Error::InvalidRequest(format!(
             "failed to serialize chat completions request: {err}"
@@ -48,9 +65,17 @@ pub(super) async fn execute_chat_completions_provider_call(
         .await
         .map_err(|err| Error::Network(err.to_string()))?;
 
-    if !status.is_success() {
-        return Err(Error::Http {
+    let observed = hooks
+        .after_response(ProviderResponse {
             status: status.as_u16(),
+            body: text,
+        })
+        .await?;
+    let observed_status = observed.status;
+    let text = observed.body;
+    if !(200..300).contains(&observed_status) {
+        return Err(Error::Http {
+            status: observed_status,
             body: truncate_error_body(&text),
         });
     }

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

@@ -0,0 +1,69 @@
+use std::sync::Arc;
+
+use serde_json::Value;
+
+use super::types::{ChatCompletionsRequest, ChatCompletionsResponse};
+use crate::call_lifecycle::provider::{
+    CompletedCall, CompletedRoute, ProviderHooks, ProviderOptions,
+};
+use crate::call_lifecycle::workflow::WorkflowFuture;
+
+pub struct OwnedChatCompletionsRequest {
+    pub options: ProviderOptions,
+    pub messages: Value,
+    pub optional_params: serde_json::Map<String, Value>,
+}
+
+impl From<ChatCompletionsRequest<'_>> for OwnedChatCompletionsRequest {
+    fn from(request: ChatCompletionsRequest<'_>) -> Self {
+        Self {
+            options: ProviderOptions {
+                model: request.model.to_owned(),
+                litellm_call_id: None,
+                api_key: request.api_key.map(str::to_owned),
+                api_base: request.api_base.map(str::to_owned),
+                custom_llm_provider: request.custom_llm_provider.map(str::to_owned),
+                extra_headers: request.extra_headers,
+                timeout: request.timeout,
+            },
+            messages: request.messages,
+            optional_params: request.optional_params,
+        }
+    }
+}
+
+pub struct ChatCompletionsRoute;
+pub type ChatCompletionsCall = CompletedCall<ChatCompletionsRoute>;
+
+impl CompletedRoute for ChatCompletionsRoute {
+    type Request = OwnedChatCompletionsRequest;
+    type Response = ChatCompletionsResponse;
+
+    fn run(
+        request: Self::Request,
+        hooks: Arc<dyn ProviderHooks>,
+    ) -> WorkflowFuture<Self::Response> {
+        Box::pin(async move {
+            let options = request.options;
+            let request = ChatCompletionsRequest {
+                model: &options.model,
+                api_key: options.api_key.as_deref(),
+                api_base: options.api_base.as_deref(),
+                custom_llm_provider: options.custom_llm_provider.as_deref(),
+                extra_headers: options.extra_headers,
+                timeout: options.timeout,
+                messages: request.messages,
+                optional_params: request.optional_params,
+            };
+            super::handler::execute_chat_completions_provider_call(
+                super::prepare::resolve_request(request)?,
+                hooks.as_ref(),
+            )
+            .await
+        })
+    }
+
+    fn context(request: &Self::Request) -> crate::call_lifecycle::CallLifecycleContext {
+        request.options.lifecycle_context("completion")
+    }
+}

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

@@ -11,6 +11,7 @@ mod client;
 mod common_utils;
 pub mod conversation;
 pub(crate) mod handler;
+pub mod lifecycle;
 mod prepare;
 pub mod response_utils;
 pub mod transformation;
@@ -18,15 +19,17 @@ pub mod types;
 
 use serde_json::{Map, Value};
 
-use handler::execute_chat_completions_provider_call;
-use prepare::{parse_messages, resolve_provider_config, resolve_request};
+use prepare::{parse_messages, resolve_provider_config};
 use types::{ChatCompletionsRequest, ChatCompletionsResponse};
 
 #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
 pub async fn chat_completions(
     request: ChatCompletionsRequest<'_>,
 ) -> Result<ChatCompletionsResponse, Error> {
-    execute_chat_completions_provider_call(resolve_request(request)?).await
+    crate::call_lifecycle::provider::run_completed::<lifecycle::ChatCompletionsRoute>(
+        request.into(),
+    )
+    .await
 }
 
 /// Whether the core would accept this request, without resolving credentials or

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

@@ -1,6 +1,5 @@
 pub mod audio_transcription;
 pub mod auth;
-pub mod caching;
 pub mod call_lifecycle;
 pub mod chat_completions;
 pub mod constants;

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

@@ -1,3 +1,4 @@
+use crate::call_lifecycle::provider::{ProviderHooks, ProviderRequest, ProviderResponse};
 use crate::constants::ANTHROPIC_MESSAGES_PROVIDER;
 use crate::error::Error;
 use crate::http_utils::http_request;
@@ -10,8 +11,24 @@ use super::types::{AnthropicMessagesResponse, MessagesRequest};
 #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
 pub(super) async fn execute_messages_provider_call(
     request: MessagesRequest<'_>,
+    hooks: &dyn ProviderHooks,
 ) -> Result<AnthropicMessagesResponse, Error> {
     let request = prepare_provider_request(request)?;
+    let changed = hooks
+        .before_request(ProviderRequest {
+            model: request.model.clone(),
+            url: request.url.clone(),
+            headers: request.upstream_headers.clone(),
+            body: request.body.clone(),
+        })
+        .await?;
+    let request = super::types::ProviderMessagesRequest {
+        model: changed.model,
+        url: changed.url,
+        upstream_headers: changed.headers,
+        body: changed.body,
+        ..request
+    };
     let mut request_builder = http_client().post(&request.url).json(&request.body);
     for (key, value) in &request.upstream_headers {
         request_builder = request_builder.header(key, value);
@@ -30,9 +47,17 @@ pub(super) async fn execute_messages_provider_call(
         .await
         .map_err(|err| Error::Network(err.to_string()))?;
 
-    if !status.is_success() {
-        return Err(Error::Http {
+    let observed = hooks
+        .after_response(ProviderResponse {
             status: status.as_u16(),
+            body: text,
+        })
+        .await?;
+    let observed_status = observed.status;
+    let text = observed.body;
+    if !(200..300).contains(&observed_status) {
+        return Err(Error::Http {
+            status: observed_status,
             body: truncate_error_body(&text),
         });
     }

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

@@ -0,0 +1,62 @@
+use std::sync::Arc;
+
+use serde_json::Value;
+
+use super::types::{AnthropicMessagesResponse, MessagesRequest};
+use crate::call_lifecycle::provider::{
+    CompletedCall, CompletedRoute, ProviderHooks, ProviderOptions,
+};
+use crate::call_lifecycle::workflow::WorkflowFuture;
+
+pub struct OwnedMessagesRequest {
+    pub options: ProviderOptions,
+    pub body: Value,
+}
+
+impl From<MessagesRequest<'_>> for OwnedMessagesRequest {
+    fn from(request: MessagesRequest<'_>) -> Self {
+        Self {
+            options: ProviderOptions {
+                model: request.model.to_owned(),
+                litellm_call_id: None,
+                api_key: request.api_key.map(str::to_owned),
+                api_base: request.api_base.map(str::to_owned),
+                custom_llm_provider: request.custom_llm_provider.map(str::to_owned),
+                extra_headers: request.extra_headers,
+                timeout: request.timeout,
+            },
+            body: request.body,
+        }
+    }
+}
+
+pub struct MessagesRoute;
+pub type MessagesCall = CompletedCall<MessagesRoute>;
+
+impl CompletedRoute for MessagesRoute {
+    type Request = OwnedMessagesRequest;
+    type Response = AnthropicMessagesResponse;
+
+    fn run(
+        request: Self::Request,
+        hooks: Arc<dyn ProviderHooks>,
+    ) -> WorkflowFuture<Self::Response> {
+        Box::pin(async move {
+            let options = request.options;
+            let request = MessagesRequest {
+                model: &options.model,
+                api_key: options.api_key.as_deref(),
+                api_base: options.api_base.as_deref(),
+                custom_llm_provider: options.custom_llm_provider.as_deref(),
+                extra_headers: options.extra_headers,
+                timeout: options.timeout,
+                body: request.body,
+            };
+            super::handler::execute_messages_provider_call(request, hooks.as_ref()).await
+        })
+    }
+
+    fn context(request: &Self::Request) -> crate::call_lifecycle::CallLifecycleContext {
+        request.options.lifecycle_context("anthropic_messages")
+    }
+}

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

@@ -11,16 +11,17 @@ use crate::Error;
 mod client;
 mod common_utils;
 mod handler;
+pub mod lifecycle;
 mod prepare;
 pub mod transformation;
 pub mod types;
 
-use handler::{execute_messages_provider_call, execute_messages_provider_stream};
+use handler::execute_messages_provider_stream;
 use types::{AnthropicMessagesResponse, MessagesRequest};
 
 #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
 pub async fn messages(request: MessagesRequest<'_>) -> Result<AnthropicMessagesResponse, Error> {
-    execute_messages_provider_call(request).await
+    crate::call_lifecycle::provider::run_completed::<lifecycle::MessagesRoute>(request.into()).await
 }
 
 pub async fn messages_stream(request: MessagesRequest<'_>) -> Result<reqwest::Response, Error> {

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

@@ -1,10 +1,9 @@
 use super::OcrClient;
 use super::adapters::OcrAdapter;
-use super::hooks::{OcrHooks, OcrLifecycleHooks, OcrPostCallRequest};
+use super::hooks::{OcrHooks, OcrPostCallRequest, OcrPreCallRequest};
 use super::registry::OcrAdapterKind;
 use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse};
 use crate::Error;
-use crate::call_lifecycle::{CallLifecycle, CallLifecycleContext};
 use std::sync::Arc;
 
 pub(crate) async fn perform_ocr_request(
@@ -12,28 +11,38 @@ pub(crate) async fn perform_ocr_request(
     request: LiteLLMOcrRequest,
 ) -> Result<LiteLLMOcrResponse, Error> {
     request.response_format()?;
-    let context = CallLifecycleContext::new(
-        "ocr",
-        request.model.clone(),
-        request.adapter.provider().as_str(),
-        request
-            .litellm_call_id
-            .clone()
-            .unwrap_or_else(|| format!("ocr-{:032x}", rand::random::<u128>())),
-    );
-    let hooks = OcrLifecycleHooks {
-        hooks: request.hooks.clone(),
-        provider_name: context.custom_llm_provider.clone(),
-    };
-    CallLifecycle::default()
-        .run(context, request, &hooks, |request| async move {
-            PreparedOcrCall::prepare(client.clone(), request)
-                .await?
-                .execute()
-                .await?
-                .normalize()
+    let request = pre_call(request).await?;
+    PreparedOcrCall::prepare(client.clone(), request)
+        .await?
+        .execute()
+        .await?
+        .normalize()
+}
+
+async fn pre_call(request: LiteLLMOcrRequest) -> Result<LiteLLMOcrRequest, Error> {
+    if !request.hooks.intercepts_requests() {
+        return Ok(request);
+    }
+    let changed = request
+        .hooks
+        .pre_call(OcrPreCallRequest {
+            model: request.model.clone(),
+            custom_llm_provider: request.adapter.provider().as_str().to_owned(),
+            document: request.document,
+            optional_params: serde_json::Value::Object(request.optional_params),
         })
-        .await
+        .await?;
+    let serde_json::Value::Object(optional_params) = changed.optional_params else {
+        return Err(super::error::OcrRequestError::RequestField {
+            path: "guardrail.optional_params".into(),
+        }
+        .into());
+    };
+    Ok(LiteLLMOcrRequest {
+        document: changed.document,
+        optional_params,
+        ..request
+    })
 }
 
 pub(crate) struct PreparedOcrCall {

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

@@ -1,10 +1,9 @@
 use std::future::Future;
 use std::pin::Pin;
-use std::sync::Arc;
 
-use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrDocument};
+use super::types::{LiteLLMOcrResponse, OcrDocument};
 use crate::Error;
-use crate::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming};
+use crate::call_lifecycle::{CallLifecycleContext, CallLifecycleTiming};
 use serde::Serialize;
 use serde_json::Value;
 
@@ -71,87 +70,3 @@ pub trait OcrHooks: Send + Sync {
 
 pub struct NoopOcrHooks;
 impl OcrHooks for NoopOcrHooks {}
-
-pub(crate) struct OcrLifecycleHooks {
-    pub hooks: Arc<dyn OcrHooks>,
-    pub provider_name: String,
-}
-
-impl CallLifecycleHooks<LiteLLMOcrRequest, LiteLLMOcrRequest, LiteLLMOcrResponse>
-    for OcrLifecycleHooks
-{
-    type PreCallFuture<'a> = OcrHookFuture<'a, LiteLLMOcrRequest>;
-    type DuringCallFuture<'a> = OcrHookFuture<'a, LiteLLMOcrRequest>;
-    type SuccessFuture<'a> = OcrLogFuture<'a>;
-    type FailureFuture<'a> = OcrLogFuture<'a>;
-
-    fn async_pre_call_hook<'a>(
-        &'a self,
-        _context: &'a CallLifecycleContext,
-        request: LiteLLMOcrRequest,
-    ) -> Self::PreCallFuture<'a> {
-        Box::pin(async move {
-            if !self.hooks.intercepts_requests() {
-                return Ok(request);
-            }
-            let changed = self
-                .hooks
-                .pre_call(OcrPreCallRequest {
-                    model: request.model.clone(),
-                    custom_llm_provider: self.provider_name.clone(),
-                    document: request.document,
-                    optional_params: Value::Object(request.optional_params),
-                })
-                .await?;
-            let Value::Object(optional_params) = changed.optional_params else {
-                return Err(super::error::OcrRequestError::RequestField {
-                    path: "guardrail.optional_params".into(),
-                }
-                .into());
-            };
-            Ok(LiteLLMOcrRequest {
-                document: changed.document,
-                optional_params,
-                ..request
-            })
-        })
-    }
-
-    fn async_during_call_hook<'a>(
-        &'a self,
-        _context: &'a CallLifecycleContext,
-        request: LiteLLMOcrRequest,
-    ) -> Self::DuringCallFuture<'a> {
-        Box::pin(async move { Ok(request) })
-    }
-
-    #[tracing::instrument(
-        name = "success_callback",
-        target = "litellm::function_trace",
-        level = "trace",
-        skip_all
-    )]
-    fn async_log_success_event<'a>(
-        &'a self,
-        context: &'a CallLifecycleContext,
-        response: &'a LiteLLMOcrResponse,
-        timing: &'a CallLifecycleTiming,
-    ) -> Self::SuccessFuture<'a> {
-        self.hooks.success(context, response, timing)
-    }
-
-    #[tracing::instrument(
-        name = "failure_callback",
-        target = "litellm::function_trace",
-        level = "trace",
-        skip_all
-    )]
-    fn async_log_failure_event<'a>(
-        &'a self,
-        context: &'a CallLifecycleContext,
-        error: &'a Error,
-        timing: &'a CallLifecycleTiming,
-    ) -> Self::FailureFuture<'a> {
-        self.hooks.failure(context, error, timing)
-    }
-}

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

@@ -1,20 +1,21 @@
-use std::future::Future;
-use std::pin::Pin;
 use std::sync::Arc;
-
-use tokio::sync::{mpsc, oneshot};
+use std::time::{SystemTime, UNIX_EPOCH};
 
 use super::handler::perform_ocr_request;
 use super::hooks::{
-    OcrDuringCallRequest, OcrHookFuture, OcrHooks, OcrLogFuture, OcrPostCallRequest,
-    OcrPreCallRequest,
+    OcrDuringCallRequest, OcrHookFuture, OcrHooks, OcrPostCallRequest, OcrPreCallRequest,
 };
 use super::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrClient};
 use crate::AuthError;
 use crate::Error;
 use crate::auth::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle};
+use crate::call_lifecycle::execution::HostExchange;
 use crate::call_lifecycle::host::{
-    HostCall, HostCallFuture, HostCallStep, HostFailure, HostLifecycle, HostPhase,
+    HostCall, HostCallFuture, HostCallStep, HostFailure, HostPhase, LifecycleBackend,
+    LifecycleBackendFuture,
+};
+use crate::call_lifecycle::workflow::{
+    LifecycleCall, LifecycleOperation, Workflow, WorkflowFuture, WorkflowReply,
 };
 use crate::call_lifecycle::{CallLifecycleContext, CallLifecycleTiming};
 
@@ -93,15 +94,7 @@ pub enum OcrHostResult {
 
 pub type OcrCallStep = HostCallStep<OcrHostOperation, LiteLLMOcrResponse>;
 
-pub struct OcrCall {
-    lifecycle: HostLifecycle,
-    execution: OcrExecution,
-    response: Option<Arc<LiteLLMOcrResponse>>,
-    error: Option<Error>,
-    pending: bool,
-    completed: bool,
-    projecting: bool,
-}
+pub struct OcrCall(LifecycleCall<OcrWorkflow>);
 
 impl OcrCall {
     pub fn admit(client: OcrClient, admission: OcrAdmission) -> NativeOutcome<Self> {
@@ -111,177 +104,21 @@ impl OcrCall {
         if !admission.host_operations {
             return NativeOutcome::Declined(OcrDecline::HostOperations);
         }
-        NativeOutcome::Completed(Self {
-            lifecycle: HostLifecycle::new(admission.asynchronous),
-            execution: OcrExecution::new(client),
-            response: None,
-            error: None,
-            pending: false,
-            completed: false,
-            projecting: false,
-        })
+        NativeOutcome::Completed(Self(LifecycleCall::new(
+            OcrWorkflow {
+                client,
+                terminal: Arc::default(),
+            },
+            admission.asynchronous,
+        )))
     }
 
     pub async fn resume(&mut self, result: Option<OcrHostResult>) -> Result<OcrCallStep, Error> {
-        if self.completed {
-            return Err(Error::InvalidRequest(
-                "OCR call cannot be resumed after completion".into(),
-            ));
-        }
-        if self.pending != result.is_some() {
-            return Err(Error::InvalidRequest(
-                "OCR host operation result does not match pending state".into(),
-            ));
-        }
-        match &result {
-            Some(OcrHostResult::Lifecycle(Ok(())))
-                if self.lifecycle.phase() == HostPhase::Execute =>
-            {
-                return Err(Error::InvalidRequest(
-                    "OCR provider operation requires a typed result".into(),
-                ));
-            }
-            Some(result)
-                if !matches!(result, OcrHostResult::Lifecycle(_))
-                    && self.lifecycle.phase() != HostPhase::Execute =>
-            {
-                return Err(Error::InvalidRequest(
-                    "unexpected OCR provider operation result".into(),
-                ));
-            }
-            _ => {}
-        }
-        self.pending = false;
-        let provider_result = match result {
-            Some(OcrHostResult::Request(result)) if self.projecting => {
-                self.projecting = false;
-                match result {
-                    Ok((request, azure_ad_token_provider)) => {
-                        self.execution.request = Some(*request);
-                        self.execution.azure_ad_token_provider = azure_ad_token_provider;
-                    }
-                    Err(error) => self.accept(Err(HostFailure::Error(error))),
-                }
-                None
-            }
-            Some(OcrHostResult::Request(_)) => {
-                return Err(Error::InvalidRequest(
-                    "unexpected OCR request projection".into(),
-                ));
-            }
-            Some(OcrHostResult::Lifecycle(result)) => {
-                self.accept(result);
-                None
-            }
-            result => result,
-        };
-        if self.lifecycle.phase() == HostPhase::Execute {
-            if self.execution.request.is_none()
-                && self.execution.execution.is_none()
-                && !self.execution.completed
-            {
-                self.projecting = true;
-                return Ok(self.host_step(OcrHostOperation::ProjectRequest));
-            }
-            match self.execution.resume(provider_result).await {
-                Ok(OcrCallStep::Host(operation)) => return Ok(self.host_step(operation)),
-                Ok(OcrCallStep::Complete(response)) => {
-                    self.response = Some(Arc::new(response));
-                    self.accept(Ok(()));
-                }
-                Err(error) => self.accept(Err(HostFailure::Error(error))),
-            }
-        }
-        if self.error.is_some() {
-            self.execution.stop().await;
-        }
-        let operation = match self.lifecycle.phase() {
-            HostPhase::Complete => {
-                self.completed = true;
-                return match self.error.take() {
-                    Some(error) => Err(error),
-                    None => self
-                        .response
-                        .take()
-                        .map(Arc::unwrap_or_clone)
-                        .map(OcrCallStep::Complete)
-                        .ok_or_else(|| {
-                            Error::InvalidRequest("OCR completed without a response".into())
-                        }),
-                };
-            }
-            HostPhase::ConstructResponse => OcrHostOperation::ConstructResponse(
-                self.response
-                    .as_ref()
-                    .ok_or_else(|| Error::InvalidRequest("missing OCR response".into()))?
-                    .clone(),
-            ),
-            HostPhase::MapFailure => OcrHostOperation::MapFailure(
-                self.error
-                    .as_ref()
-                    .ok_or_else(|| Error::InvalidRequest("missing OCR failure".into()))?
-                    .clone(),
-            ),
-            HostPhase::Success | HostPhase::Failure => {
-                let snapshot = self
-                    .execution
-                    .terminal
-                    .lock()
-                    .unwrap_or_else(|error| error.into_inner())
-                    .clone();
-                match (self.lifecycle.phase(), snapshot) {
-                    (HostPhase::Success, Some((context, timing))) => OcrHostOperation::Success {
-                        context,
-                        response: self
-                            .response
-                            .as_ref()
-                            .ok_or_else(|| Error::InvalidRequest("missing OCR response".into()))?
-                            .clone(),
-                        timing,
-                    },
-                    (HostPhase::Failure, Some((context, timing))) => OcrHostOperation::Failure {
-                        context,
-                        error: self
-                            .error
-                            .as_ref()
-                            .ok_or_else(|| Error::InvalidRequest("missing OCR failure".into()))?
-                            .clone(),
-                        timing,
-                    },
-                    (phase, _) => OcrHostOperation::Lifecycle(phase),
-                }
-            }
-            phase => OcrHostOperation::Lifecycle(phase),
-        };
-        Ok(self.host_step(operation))
-    }
-
-    fn accept(&mut self, result: Result<(), HostFailure>) {
-        let cancelled = matches!(&result, Err(HostFailure::Cancelled(_)));
-        if let Some(error) = self.lifecycle.accept(result) {
-            if cancelled {
-                self.error = Some(error);
-            } else {
-                self.error.get_or_insert(error);
-            }
-            self.execution.cancel();
-        }
+        self.0.resume(result).await
     }
 
     pub async fn interrupt(&mut self, failure: HostFailure) -> Result<OcrCallStep, Error> {
-        if self.completed {
-            return Err(Error::InvalidRequest(
-                "OCR call cannot be interrupted after completion".into(),
-            ));
-        }
-        self.pending = false;
-        self.accept(Err(failure));
-        self.resume(None).await
-    }
-
-    fn host_step(&mut self, operation: OcrHostOperation) -> OcrCallStep {
-        self.pending = true;
-        OcrCallStep::Host(operation)
+        self.0.interrupt(failure).await
     }
 }
 
@@ -305,159 +142,148 @@ impl HostCall for OcrCall {
     }
 }
 
-struct PendingOperation {
-    operation: OcrHostOperation,
-    result: oneshot::Sender<OcrHostResult>,
-}
-
-struct OcrExecution {
-    client: Option<OcrClient>,
-    request: Option<LiteLLMOcrRequest>,
-    operations_tx: mpsc::UnboundedSender<PendingOperation>,
-    operations_rx: mpsc::UnboundedReceiver<PendingOperation>,
-    pending_result: Option<oneshot::Sender<OcrHostResult>>,
-    execution: Option<tokio::task::JoinHandle<Result<LiteLLMOcrResponse, Error>>>,
-    completed: bool,
-    azure_ad_token_provider: bool,
+struct OcrWorkflow {
+    client: OcrClient,
     terminal: Arc<std::sync::Mutex<Option<(CallLifecycleContext, CallLifecycleTiming)>>>,
 }
 
-impl OcrExecution {
-    fn new(client: OcrClient) -> Self {
-        let (operations_tx, operations_rx) = mpsc::unbounded_channel();
-        Self {
-            client: Some(client),
-            request: None,
-            operations_tx,
-            operations_rx,
-            pending_result: None,
-            execution: None,
-            completed: false,
-            azure_ad_token_provider: false,
-            terminal: Arc::default(),
+impl Workflow for OcrWorkflow {
+    type Request = (Box<LiteLLMOcrRequest>, bool);
+    type Operation = OcrHostOperation;
+    type Reply = OcrHostResult;
+    type Response = LiteLLMOcrResponse;
+
+    fn operation(operation: LifecycleOperation<Self::Response>) -> Self::Operation {
+        match operation {
+            LifecycleOperation::ProjectRequest => OcrHostOperation::ProjectRequest,
+            LifecycleOperation::Phase(phase) => OcrHostOperation::Lifecycle(phase),
+            LifecycleOperation::ConstructResponse(response) => {
+                OcrHostOperation::ConstructResponse(response)
+            }
+            LifecycleOperation::ConstructCachedResponse(_) => {
+                unreachable!("OCR does not cache responses")
+            }
+            LifecycleOperation::MapFailure(error) => OcrHostOperation::MapFailure(error),
+            LifecycleOperation::Success {
+                context,
+                response,
+                timing,
+            } => OcrHostOperation::Success {
+                context,
+                response,
+                timing,
+            },
+            LifecycleOperation::Failure {
+                context,
+                error,
+                timing,
+            } => OcrHostOperation::Failure {
+                context,
+                error,
+                timing,
+            },
         }
     }
 
-    pub async fn resume(&mut self, result: Option<OcrHostResult>) -> Result<OcrCallStep, Error> {
-        if self.completed {
-            return Err(Error::InvalidRequest(
-                "OCR call cannot be resumed after completion".into(),
-            ));
-        }
-        match (self.pending_result.take(), result) {
-            (Some(sender), Some(result)) => sender
-                .send(result)
-                .map_err(|_| Error::InvalidRequest("OCR host operation was abandoned".into()))?,
-            (None, None) if self.execution.is_none() => self.start(),
-            (Some(sender), None) => {
-                self.pending_result = Some(sender);
-                return Err(Error::InvalidRequest(
-                    "OCR host operation result is required".into(),
-                ));
-            }
-            (None, Some(_)) => {
-                return Err(Error::InvalidRequest(
-                    "unexpected OCR host operation result".into(),
-                ));
-            }
-            (None, None) => {}
-        }
+    fn accepts(operation: &OcrHostOperation, reply: &OcrHostResult) -> bool {
+        matches!(reply, OcrHostResult::Lifecycle(Err(_)))
+            || matches!(
+                (operation, reply),
+                (OcrHostOperation::ProjectRequest, OcrHostResult::Request(_))
+                    | (
+                        OcrHostOperation::AcquireAzureAdToken,
+                        OcrHostResult::AzureAdToken(_)
+                    )
+                    | (OcrHostOperation::PreCall(_), OcrHostResult::PreCall(_))
+                    | (
+                        OcrHostOperation::DuringCall(_),
+                        OcrHostResult::DuringCall(_)
+                    )
+                    | (OcrHostOperation::PostCall(_), OcrHostResult::PostCall(_))
+                    | (
+                        OcrHostOperation::Lifecycle(_)
+                            | OcrHostOperation::ConstructResponse(_)
+                            | OcrHostOperation::MapFailure(_)
+                            | OcrHostOperation::Success { .. }
+                            | OcrHostOperation::Failure { .. },
+                        OcrHostResult::Lifecycle(_)
+                    )
+            )
+    }
 
-        let execution = self.execution.as_mut().ok_or_else(|| {
-            Error::InvalidRequest("OCR call cannot be resumed after completion".into())
-        })?;
-        tokio::select! {
-            operation = self.operations_rx.recv() => {
-                let operation = operation.ok_or_else(|| Error::InvalidRequest("OCR operation channel closed".into()))?;
-                self.pending_result = Some(operation.result);
-                Ok(OcrCallStep::Host(operation.operation))
-            }
-            result = execution => {
-                self.execution = None;
-                self.completed = true;
-                result
-                    .map_err(|error| Error::Network(format!("OCR execution task failed: {error}")))?
-                    .map(OcrCallStep::Complete)
-            }
+    fn reply(&mut self, reply: OcrHostResult) -> WorkflowReply<Self::Request, OcrHostResult> {
+        match reply {
+            OcrHostResult::Request(request) => WorkflowReply::Request(request),
+            OcrHostResult::Lifecycle(result) => WorkflowReply::Lifecycle(result),
+            reply => WorkflowReply::Operation(reply),
         }
     }
 
-    fn start(&mut self) {
-        let client = self.client.take().expect("admitted OCR call has a client");
-        let mut request = self
-            .request
-            .take()
-            .expect("admitted OCR call has a request");
+    fn start(
+        &mut self,
+        (request, azure_ad_token_provider): Self::Request,
+        operations: HostExchange<OcrHostOperation, OcrHostResult>,
+    ) -> WorkflowFuture<LiteLLMOcrResponse> {
+        let client = self.client.clone();
+        let mut request = *request;
         let intercepts_requests = request.hooks.intercepts_requests();
-        if self.azure_ad_token_provider {
+        let context = CallLifecycleContext::new(
+            "ocr",
+            request.model.clone(),
+            request.adapter.provider().as_str(),
+            request
+                .litellm_call_id
+                .clone()
+                .unwrap_or_else(|| format!("ocr-{:032x}", rand::random::<u128>())),
+        );
+        let started = epoch_seconds();
+        if azure_ad_token_provider {
             request.azure_ad_token_provider = Some(TokenProviderHandle::new(Arc::new(
                 OcrAzureAdTokenProvider {
-                    operations: self.operations_tx.clone(),
+                    operations: operations.clone(),
                 },
             )));
         }
         request.hooks = Arc::new(ProtocolHooks {
-            operations: self.operations_tx.clone(),
+            operations,
             intercepts_requests,
-            terminal: self.terminal.clone(),
         });
-        self.execution = Some(tokio::spawn(async move {
-            perform_ocr_request(&client, request).await
-        }));
+        let terminal = self.terminal.clone();
+        Box::pin(async move {
+            let result = perform_ocr_request(&client, request).await;
+            let timing = CallLifecycleTiming::new(started, epoch_seconds());
+            *terminal.lock().unwrap_or_else(|error| error.into_inner()) = Some((context, timing));
+            result
+        })
     }
 
-    fn cancel(&mut self) {
-        self.pending_result = None;
-        if let Some(execution) = &self.execution {
-            execution.abort();
-        }
-    }
-
-    async fn stop(&mut self) {
-        self.cancel();
-        if let Some(execution) = self.execution.as_mut() {
-            let _ = execution.await;
-        }
-        self.execution = None;
-    }
-}
-
-impl Drop for OcrExecution {
-    fn drop(&mut self) {
-        if let Some(execution) = &self.execution {
-            execution.abort();
-        }
+    fn terminal(&self) -> Option<(CallLifecycleContext, CallLifecycleTiming)> {
+        self.terminal
+            .lock()
+            .unwrap_or_else(|error| error.into_inner())
+            .clone()
     }
 }
 
 struct ProtocolHooks {
-    operations: mpsc::UnboundedSender<PendingOperation>,
+    operations: HostExchange<OcrHostOperation, OcrHostResult>,
     intercepts_requests: bool,
-    terminal: Arc<std::sync::Mutex<Option<(CallLifecycleContext, CallLifecycleTiming)>>>,
 }
 
 #[derive(Debug)]
 struct OcrAzureAdTokenProvider {
-    operations: mpsc::UnboundedSender<PendingOperation>,
+    operations: HostExchange<OcrHostOperation, OcrHostResult>,
 }
 
 impl TokenProvider for OcrAzureAdTokenProvider {
     fn acquire(&self) -> TokenFuture<'_> {
         Box::pin(async move {
-            let (result, receiver) = oneshot::channel();
-            self.operations
-                .send(PendingOperation {
-                    operation: OcrHostOperation::AcquireAzureAdToken,
-                    result,
-                })
-                .map_err(|_| {
-                    AuthError::AzureTokenAcquisition("OCR host driver was abandoned".into())
-                })?;
-            match receiver.await.map_err(|_| {
-                AuthError::AzureTokenAcquisition(
-                    "OCR token provider operation was abandoned".into(),
-                )
-            })? {
+            match self
+                .operations
+                .invoke(OcrHostOperation::AcquireAzureAdToken)
+                .await
+                .map_err(|error| AuthError::AzureTokenAcquisition(error.to_string()))?
+            {
                 OcrHostResult::AzureAdToken(result) => result,
                 _ => Err(AuthError::AzureTokenAcquisition(
                     "invalid OCR token provider host result".into(),
@@ -469,13 +295,7 @@ impl TokenProvider for OcrAzureAdTokenProvider {
 
 impl ProtocolHooks {
     async fn invoke(&self, operation: OcrHostOperation) -> Result<OcrHostResult, Error> {
-        let (result, receiver) = oneshot::channel();
-        self.operations
-            .send(PendingOperation { operation, result })
-            .map_err(|_| Error::InvalidRequest("OCR host driver was abandoned".into()))?;
-        receiver
-            .await
-            .map_err(|_| Error::InvalidRequest("OCR host operation was abandoned".into()))
+        self.operations.invoke(operation).await
     }
 }
 
@@ -519,48 +339,19 @@ impl OcrHooks for ProtocolHooks {
             }
         })
     }
-
-    fn success<'a>(
-        &'a self,
-        context: &'a CallLifecycleContext,
-        _response: &'a LiteLLMOcrResponse,
-        timing: &'a CallLifecycleTiming,
-    ) -> OcrLogFuture<'a> {
-        Box::pin(async move {
-            *self
-                .terminal
-                .lock()
-                .unwrap_or_else(|error| error.into_inner()) =
-                Some((context.clone(), timing.clone()));
-        })
-    }
-
-    fn failure<'a>(
-        &'a self,
-        context: &'a CallLifecycleContext,
-        _error: &'a Error,
-        timing: &'a CallLifecycleTiming,
-    ) -> OcrLogFuture<'a> {
-        Box::pin(async move {
-            *self
-                .terminal
-                .lock()
-                .unwrap_or_else(|error| error.into_inner()) =
-                Some((context.clone(), timing.clone()));
-        })
-    }
 }
 
-pub type OcrHostFuture<'a> = Pin<Box<dyn Future<Output = OcrHostResult> + Send + 'a>>;
-
-pub trait OcrHost: Send + Sync {
-    fn invoke(&self, operation: OcrHostOperation) -> OcrHostFuture<'_>;
+fn epoch_seconds() -> f64 {
+    SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .map(|duration| duration.as_secs_f64())
+        .unwrap_or(0.0)
 }
 
 pub struct NoopOcrHost;
 
-impl OcrHost for NoopOcrHost {
-    fn invoke(&self, operation: OcrHostOperation) -> OcrHostFuture<'_> {
+impl LifecycleBackend<OcrHostOperation, OcrHostResult> for NoopOcrHost {
+    fn invoke(&self, operation: OcrHostOperation) -> LifecycleBackendFuture<'_, OcrHostResult> {
         Box::pin(async move {
             match operation {
                 OcrHostOperation::ProjectRequest => OcrHostResult::Request(Err(
@@ -594,8 +385,8 @@ impl OcrHookHost {
     }
 }
 
-impl OcrHost for OcrHookHost {
-    fn invoke(&self, operation: OcrHostOperation) -> OcrHostFuture<'_> {
+impl LifecycleBackend<OcrHostOperation, OcrHostResult> for OcrHookHost {
+    fn invoke(&self, operation: OcrHostOperation) -> LifecycleBackendFuture<'_, OcrHostResult> {
         Box::pin(async move {
             match operation {
                 OcrHostOperation::ProjectRequest => OcrHostResult::Request(Err(

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

@@ -11,11 +11,12 @@ mod registry;
 pub mod types;
 pub mod wire;
 
+pub use crate::call_lifecycle::host::LifecycleBackend as OcrHost;
 pub use client::{OcrClient, ocr};
 pub use document::{encode_file_document, mime_type_for_name, upload_mime_type};
 pub use lifecycle::{
     NativeOutcome, NativeResult, NoopOcrHost, OcrAdmission, OcrCall, OcrCallStep, OcrDecline,
-    OcrHookHost, OcrHost, OcrHostOperation, OcrHostResult,
+    OcrHookHost, OcrHostOperation, OcrHostResult,
 };
 pub use types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection, OcrDocument};
 

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

@@ -210,13 +210,18 @@ pub fn decode_request(wire: OcrWireRequest) -> Result<LiteLLMOcrRequest, Error>
 }
 
 fn decode_document(value: Value) -> Result<OcrDocument, OcrRequestError> {
+    validate_document_url(&value)?;
+    decode_request_value(value, "document")
+}
+
+pub fn validate_document_url(value: &Value) -> Result<(), OcrRequestError> {
     let kind = value.get("type").and_then(Value::as_str);
     let missing_url = matches!(kind, Some("document_url")) && value.get("document_url").is_none()
         || matches!(kind, Some("image_url")) && value.get("image_url").is_none();
     if missing_url {
         return Err(OcrRequestError::MissingDocumentUrl);
     }
-    decode_request_value(value, "document")
+    Ok(())
 }
 
 fn source_for(sources: &BTreeMap<String, InputSource>, name: &str) -> InputSource {

litellm-rust/crates/core/src/providers/bedrock/aws_base.rs

@@ -1,9 +1,8 @@
 use std::collections::BTreeMap;
-use std::sync::{Mutex, OnceLock};
+use std::sync::OnceLock;
 use std::time::Duration;
 use std::time::{SystemTime, UNIX_EPOCH};
 
-use crate::caching::in_memory_cache::InMemoryCache;
 use crate::error::Error;
 use aws_credential_types::Credentials;
 use aws_credential_types::provider::ProvideCredentials;
@@ -12,6 +11,7 @@ use aws_sigv4::http_request::{
 };
 use aws_sigv4::sign::v4;
 use aws_smithy_runtime_api::client::identity::Identity;
+use litellm_cache_memory::InMemoryCache;
 use serde_json::{Map, Value};
 use sha2::{Digest, Sha256};
 
@@ -26,7 +26,7 @@ use super::constants::{
 const STATIC_CREDENTIALS_TTL: Duration = Duration::from_secs(3600 - 60);
 const AMBIENT_CREDENTIALS_TTL: Duration = Duration::from_secs(600);
 
-static IAM_CREDENTIALS_CACHE: OnceLock<Mutex<InMemoryCache<Credentials>>> = OnceLock::new();
+static IAM_CREDENTIALS_CACHE: OnceLock<InMemoryCache<Credentials>> = OnceLock::new();
 
 fn credential_cache_ttl(flow: &AwsAuthFlow) -> Option<Duration> {
     match flow {
@@ -108,16 +108,17 @@ fn cache_key(config: &AwsAuthConfig, flow: &AwsAuthFlow) -> String {
 }
 
 fn get_cached_credentials(key: &str) -> Option<Credentials> {
-    let cache = IAM_CREDENTIALS_CACHE.get_or_init(|| Mutex::new(InMemoryCache::default()));
-    let mut entries = cache.lock().ok()?;
-    entries.get_cache(key)
+    IAM_CREDENTIALS_CACHE
+        .get_or_init(InMemoryCache::default)
+        .get_cache(key)
+        .ok()
+        .flatten()
 }
 
 fn set_cached_credentials(key: String, credentials: Credentials, ttl: Duration) {
-    let cache = IAM_CREDENTIALS_CACHE.get_or_init(|| Mutex::new(InMemoryCache::default()));
-    if let Ok(mut entries) = cache.lock() {
-        entries.set_cache(key, credentials, Some(ttl));
-    }
+    let _ = IAM_CREDENTIALS_CACHE
+        .get_or_init(InMemoryCache::default)
+        .set_cache(key, credentials, Some(ttl));
 }
 
 fn role_identity(arn: &str) -> Option<(&str, &str, &str)> {

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

@@ -1,12 +1,8 @@
-use std::future::Future;
-use std::pin::Pin;
 use std::sync::Mutex;
 use std::time::{SystemTime, UNIX_EPOCH};
 
 use serde_json::Value;
 
-use crate::Error;
-use crate::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming};
 use crate::responses::types::{ResponsesWsEvent, ResponsesWsEventType};
 
 #[derive(Clone, Debug, Default, PartialEq, Eq)]
@@ -203,58 +199,16 @@ impl ResponsesWsInstrumentation {
             }
         })
     }
-}
 
-type LifecycleFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
-
-impl CallLifecycleHooks<(), (), ()> for ResponsesWsInstrumentation {
-    type PreCallFuture<'a> = LifecycleFuture<'a, ()>;
-    type DuringCallFuture<'a> = LifecycleFuture<'a, ()>;
-    type SuccessFuture<'a> = Pin<Box<dyn Future<Output = ()> + Send + 'a>>;
-    type FailureFuture<'a> = Pin<Box<dyn Future<Output = ()> + Send + 'a>>;
-
-    fn async_pre_call_hook<'a>(
-        &'a self,
-        _context: &'a CallLifecycleContext,
-        request: (),
-    ) -> Self::PreCallFuture<'a> {
-        Box::pin(async move { Ok(request) })
-    }
-
-    fn async_during_call_hook<'a>(
-        &'a self,
-        _context: &'a CallLifecycleContext,
-        request: (),
-    ) -> Self::DuringCallFuture<'a> {
-        Box::pin(async move { Ok(request) })
-    }
-
-    fn async_log_success_event<'a>(
-        &'a self,
-        _context: &'a CallLifecycleContext,
-        _response: &'a (),
-        _timing: &'a CallLifecycleTiming,
-    ) -> Self::SuccessFuture<'a> {
-        Box::pin(async move {
-            let outcome = self.success_outcome();
-            if let Ok(mut state) = self.state.lock() {
-                state.outcome = Some(outcome);
-            }
-        })
-    }
-
-    fn async_log_failure_event<'a>(
-        &'a self,
-        _context: &'a CallLifecycleContext,
-        _error: &'a Error,
-        _timing: &'a CallLifecycleTiming,
-    ) -> Self::FailureFuture<'a> {
-        Box::pin(async move {
-            let outcome = self.failure_outcome();
-            if let Ok(mut state) = self.state.lock() {
-                state.outcome = Some(outcome);
-            }
-        })
+    pub fn record_outcome(&self, success: bool) {
+        let outcome = if success {
+            self.success_outcome()
+        } else {
+            self.failure_outcome()
+        };
+        if let Ok(mut state) = self.state.lock() {
+            state.outcome = Some(outcome);
+        }
     }
 }
 
@@ -328,25 +282,12 @@ mod tests {
         ));
     }
 
-    #[tokio::test]
-    async fn lifecycle_records_success_outcome_for_provider_completion() {
+    #[test]
+    fn records_success_outcome_for_provider_completion() {
         let instrumentation =
             ResponsesWsInstrumentation::new("call-1", "gpt-5", ResponsesWsMetadata::default());
-        let result = crate::call_lifecycle::CallLifecycle::default()
-            .run(
-                crate::call_lifecycle::CallLifecycleContext::new(
-                    "responses_websocket",
-                    "gpt-5",
-                    "openai",
-                    "call-1",
-                ),
-                (),
-                &instrumentation,
-                |_| async { Ok::<(), Error>(()) },
-            )
-            .await;
+        instrumentation.record_outcome(true);
 
-        assert!(result.is_ok());
         assert!(matches!(
             instrumentation.take_outcome(),
             Some(ResponsesWsLogOutcome::Success { .. })

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

@@ -27,6 +27,34 @@ fn public_outcome_is_finalized_before_a_single_terminal_dispatch() {
     for asynchronous in [false, true] {
         let (events, failures) = run(None, asynchronous);
         assert!(failures.is_empty());
+        let expected = if asynchronous {
+            vec![
+                HostPhase::Setup,
+                HostPhase::DeploymentPreCall,
+                HostPhase::Prepare,
+                HostPhase::CacheLookup,
+                HostPhase::Execute,
+                HostPhase::ConstructResponse,
+                HostPhase::PostProcess,
+                HostPhase::DeploymentPostCall,
+                HostPhase::CacheStore,
+                HostPhase::Finalize,
+                HostPhase::Success,
+            ]
+        } else {
+            vec![
+                HostPhase::Setup,
+                HostPhase::Prepare,
+                HostPhase::CacheLookup,
+                HostPhase::Execute,
+                HostPhase::ConstructResponse,
+                HostPhase::PostProcess,
+                HostPhase::CacheStore,
+                HostPhase::Finalize,
+                HostPhase::Success,
+            ]
+        };
+        assert_eq!(events, expected);
         assert_eq!(
             &events[events.len() - 2..],
             &[HostPhase::Finalize, HostPhase::Success]
@@ -51,9 +79,12 @@ fn only_provider_and_response_construction_failures_use_provider_mapping() {
         HostPhase::Setup,
         HostPhase::DeploymentPreCall,
         HostPhase::Prepare,
+        HostPhase::CacheLookup,
         HostPhase::Execute,
         HostPhase::ConstructResponse,
+        HostPhase::PostProcess,
         HostPhase::DeploymentPostCall,
+        HostPhase::CacheStore,
         HostPhase::Finalize,
     ] {
         let (events, failures) = run(Some(phase), true);
@@ -76,6 +107,35 @@ fn only_provider_and_response_construction_failures_use_provider_mapping() {
     }
 }
 
+#[test]
+fn cache_hit_uses_the_same_graph_without_provider_or_cache_store() {
+    let mut lifecycle = HostLifecycle::new(true);
+    let mut events = Vec::new();
+    while lifecycle.phase() != HostPhase::CacheLookup {
+        events.push(lifecycle.phase());
+        lifecycle.accept(Ok(()));
+    }
+    events.push(HostPhase::CacheLookup);
+    lifecycle.cache_hit();
+    while lifecycle.phase() != HostPhase::Complete {
+        events.push(lifecycle.phase());
+        lifecycle.accept(Ok(()));
+    }
+    assert_eq!(
+        events,
+        [
+            HostPhase::Setup,
+            HostPhase::DeploymentPreCall,
+            HostPhase::Prepare,
+            HostPhase::CacheLookup,
+            HostPhase::ConstructResponse,
+            HostPhase::PostProcess,
+            HostPhase::Finalize,
+            HostPhase::Success,
+        ]
+    );
+}
+
 #[test]
 fn failure_handler_errors_do_not_replace_selected_failure_or_suppress_async_dispatch() {
     let mut lifecycle = HostLifecycle::new(true);

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

@@ -490,7 +490,9 @@ async fn direct_native_host_drives_the_same_state_machine() {
             "DuringCall",
             "PostCall",
             "ConstructResponse",
+            "PostProcess",
             "DeploymentPostCall",
+            "CacheStore",
             "Finalize",
             "Success",
         ]

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

@@ -23,6 +23,8 @@ trace-parity = [
 futures-util.workspace = true
 tracing = { workspace = true, optional = true }
 litellm-core = { workspace = true, features = ["bedrock-auth"] }
+litellm-cache.workspace = true
+litellm-cache-memory.workspace = true
 litellm-token-counter.workspace = true
 litellm-python-interop.workspace = true
 pyo3.workspace = true

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

@@ -1,4 +1,5 @@
 mod auth;
+mod cache;
 mod constants;
 mod diagnostics;
 mod errors;

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

@@ -3,6 +3,8 @@ use pyo3::gc::{PyTraverseError, PyVisit};
 use pyo3::prelude::*;
 use pyo3::types::{PyDict, PyTuple};
 
+use super::contract::CallbackPhase;
+
 #[derive(FromPyObject)]
 pub(crate) struct PythonLogger(Py<PyAny>);
 
@@ -19,7 +21,7 @@ impl PythonLogger {
         visit.call(&self.0)
     }
 
-    pub(crate) fn callbacks_needed(&self, py: Python<'_>, phase: &str) -> PyResult<bool> {
+    pub(crate) fn callbacks_needed(&self, py: Python<'_>, phase: CallbackPhase) -> PyResult<bool> {
         if !self
             .object(py)
             .getattr("_native_callback_fast_path")
@@ -29,7 +31,7 @@ impl PythonLogger {
         }
         py.import("litellm.rust_bridge.lifecycle")?
             .getattr("callbacks_needed")?
-            .call1((self.object(py), phase))?
+            .call1((self.object(py), phase.as_str()))?
             .extract()
     }
 
@@ -68,7 +70,7 @@ impl PythonLogger {
         start: &Py<PyAny>,
         end: &Option<Py<PyAny>>,
     ) -> PyResult<()> {
-        if !self.callbacks_needed(py, "sync_success_async")? {
+        if !self.callbacks_needed(py, CallbackPhase::SyncSuccessForAsyncCall)? {
             return Ok(());
         }
         self.object(py).call_method1(
@@ -89,9 +91,9 @@ impl PythonLogger {
         if !self.callbacks_needed(
             py,
             if asynchronous {
-                "async_failure"
+                CallbackPhase::AsyncFailure
             } else {
-                "sync_failure"
+                CallbackPhase::SyncFailure
             },
         )? {
             py.import("litellm.rust_bridge.lifecycle")?
@@ -129,7 +131,7 @@ impl PythonLogger {
         start: &Py<PyAny>,
         end: &Option<Py<PyAny>>,
     ) -> PyResult<()> {
-        if !self.callbacks_needed(py, "sync_success")? {
+        if !self.callbacks_needed(py, CallbackPhase::SyncSuccess)? {
             return self.success_bookkeeping(py, response, start, end, false);
         }
         let context = py.import("contextvars")?.call_method0("copy_context")?;
@@ -155,7 +157,7 @@ impl PythonLogger {
         start: &Py<PyAny>,
         end: &Option<Py<PyAny>>,
     ) -> PyResult<()> {
-        if !self.callbacks_needed(py, "async_success")? {
+        if !self.callbacks_needed(py, CallbackPhase::AsyncSuccess)? {
             return self.success_bookkeeping(py, response, start, end, true);
         }
         let context = py.import("contextvars")?.call_method0("copy_context")?;

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

@@ -0,0 +1,281 @@
+use std::marker::PhantomData;
+
+use pyo3::gc::{PyTraverseError, PyVisit};
+use pyo3::prelude::*;
+use pyo3::types::{PyDict, PyTuple};
+use serde::Serialize;
+
+use litellm_core::call_lifecycle::host::{HostPhase, HostStep};
+use litellm_core::call_lifecycle::provider::{
+    CompletedCall, CompletedOperation, CompletedReply, CompletedRoute, CompletedWorkflow,
+    ProviderRequest, ProviderResponse,
+};
+use litellm_core::call_lifecycle::workflow::LifecycleOperation;
+use litellm_python_interop::{
+    from_py_preserving_errors as from_py, to_py_preserving_errors as to_py,
+};
+
+use super::contract::{AdapterOperation, CallMode, PythonCallType};
+use super::{OperationClass, PythonCallState, PythonRoute, missing_state, now, run_call};
+use crate::errors::execution_error_to_pyerr;
+
+pub(crate) trait PythonCompletedRoute: CompletedRoute {
+    const SYNC_CALL_TYPE: PythonCallType;
+    const ASYNC_CALL_TYPE: PythonCallType;
+
+    fn admit(request: &Bound<'_, PyDict>) -> PyResult<()>;
+    fn project(request: &Bound<'_, PyDict>) -> PyResult<Self::Request>;
+}
+
+struct PythonCompletedHost<R: PythonCompletedRoute> {
+    state: PythonCallState,
+    request: Py<PyDict>,
+    adapter: Py<PyAny>,
+    pending: Option<AdapterOperation>,
+    route: PhantomData<R>,
+}
+
+impl<R: PythonCompletedRoute> PythonRoute for PythonCompletedHost<R>
+where
+    R::Response: Serialize,
+{
+    type Call = CompletedCall<R>;
+
+    fn state(&self) -> &PythonCallState {
+        &self.state
+    }
+    fn state_mut(&mut self) -> &mut PythonCallState {
+        &mut self.state
+    }
+
+    fn classify(operation: &CompletedOperation<R::Response>) -> OperationClass {
+        match operation {
+            CompletedOperation::Lifecycle(LifecycleOperation::Phase(
+                HostPhase::Prepare | HostPhase::PostProcess | HostPhase::CacheStore,
+            )) => OperationClass::Route,
+            CompletedOperation::Lifecycle(LifecycleOperation::Phase(phase)) => {
+                OperationClass::Phase(*phase)
+            }
+            CompletedOperation::Lifecycle(LifecycleOperation::Success { .. }) => {
+                OperationClass::Phase(HostPhase::Success)
+            }
+            CompletedOperation::Lifecycle(LifecycleOperation::Failure { .. }) => {
+                OperationClass::Phase(HostPhase::Failure)
+            }
+            _ => OperationClass::Route,
+        }
+    }
+
+    fn lifecycle_result() -> CompletedReply<R::Request> {
+        CompletedReply::Lifecycle(Ok(()))
+    }
+    fn map_error(error: litellm_core::Error) -> PyErr {
+        execution_error_to_pyerr(error)
+    }
+
+    fn invoke(
+        &mut self,
+        py: Python<'_>,
+        operation: CompletedOperation<R::Response>,
+    ) -> PyResult<CompletedReply<R::Request>> {
+        match self.invoke_step(py, operation)? {
+            HostStep::Ready(reply) => Ok(reply),
+            HostStep::Suspend(_) => Err(pyo3::exceptions::PyRuntimeError::new_err(
+                "route operation requires async execution",
+            )),
+        }
+    }
+
+    fn invoke_step(
+        &mut self,
+        py: Python<'_>,
+        operation: CompletedOperation<R::Response>,
+    ) -> PyResult<HostStep<CompletedReply<R::Request>, Py<PyAny>>> {
+        let (operation, payload) = match operation {
+            CompletedOperation::Lifecycle(LifecycleOperation::Phase(HostPhase::Prepare)) => {
+                self.state.prepare(py)?;
+                return Ok(HostStep::Ready(CompletedReply::Prepared(Ok(
+                    crate::cache::plan(
+                        py,
+                        self.state.call_type.as_str(),
+                        self.state.kwargs.bind(py),
+                    )?,
+                ))));
+            }
+            CompletedOperation::Lifecycle(LifecycleOperation::Phase(HostPhase::PostProcess)) => (
+                AdapterOperation::PostProcess,
+                self.state
+                    .response
+                    .as_ref()
+                    .ok_or_else(missing_state)?
+                    .clone_ref(py),
+            ),
+            CompletedOperation::Lifecycle(LifecycleOperation::Phase(HostPhase::CacheStore)) => (
+                AdapterOperation::CacheStore,
+                self.state
+                    .response
+                    .as_ref()
+                    .ok_or_else(missing_state)?
+                    .clone_ref(py),
+            ),
+            CompletedOperation::Lifecycle(LifecycleOperation::ConstructCachedResponse(
+                response,
+            )) => {
+                self.state.end = Some(now(py)?);
+                (AdapterOperation::CachedResponse, to_py(py, &response)?)
+            }
+            CompletedOperation::Lifecycle(LifecycleOperation::ProjectRequest) => (
+                AdapterOperation::ProjectRequest,
+                self.request.clone_ref(py).into_any(),
+            ),
+            CompletedOperation::BeforeRequest(request) => {
+                (AdapterOperation::BeforeRequest, to_py(py, &request)?)
+            }
+            CompletedOperation::AfterResponse(response) => {
+                (AdapterOperation::AfterResponse, to_py(py, &response)?)
+            }
+            CompletedOperation::Lifecycle(LifecycleOperation::ConstructResponse(response)) => {
+                self.state.end = Some(now(py)?);
+                (
+                    AdapterOperation::ConstructResponse,
+                    to_py(py, response.as_ref())?,
+                )
+            }
+            CompletedOperation::Lifecycle(LifecycleOperation::MapFailure(error)) => {
+                if self.state.error.is_none() {
+                    self.state.retain_error(py, execution_error_to_pyerr(error));
+                }
+                if self.state.end.is_none() {
+                    self.state.end = Some(now(py)?);
+                }
+                (
+                    AdapterOperation::MapFailure,
+                    self.state
+                        .error
+                        .as_ref()
+                        .ok_or_else(missing_state)?
+                        .clone_ref(py)
+                        .into_any(),
+                )
+            }
+            _ => return Err(missing_state()),
+        };
+        self.pending = Some(operation);
+        let step = self.adapter.bind(py).call_method1(
+            "invoke",
+            (
+                operation.as_str(),
+                payload,
+                &self.request,
+                &self.state.kwargs,
+                self.state.logger()?.object(py),
+            ),
+        )?;
+        let protocol = py.import("litellm.rust_bridge.lifecycle")?;
+        if step.is_instance(&protocol.getattr("Await")?)? {
+            if !self.state.mode.is_async() {
+                return Err(pyo3::exceptions::PyRuntimeError::new_err(
+                    "sync route operation suspended",
+                ));
+            }
+            return Ok(HostStep::Suspend(step.getattr("awaitable")?.unbind()));
+        }
+        if !step.is_instance(&protocol.getattr("Complete")?)? {
+            return Err(missing_state());
+        }
+        self.accept_route(py, step.getattr("value")?.unbind())
+            .map(HostStep::Ready)
+    }
+
+    fn accept_route(
+        &mut self,
+        py: Python<'_>,
+        value: Py<PyAny>,
+    ) -> PyResult<CompletedReply<R::Request>> {
+        Ok(match self.pending.take().ok_or_else(missing_state)? {
+            AdapterOperation::PostProcess => Self::lifecycle_result(),
+            AdapterOperation::CacheStore => CompletedReply::CacheStore(Ok(if value.is_none(py) {
+                None
+            } else {
+                Some(from_py(value.bind(py))?)
+            })),
+            AdapterOperation::ProjectRequest => {
+                let request = value.into_bound(py).cast_into::<PyDict>()?;
+                CompletedReply::Request(Ok(R::project(&request)?))
+            }
+            AdapterOperation::BeforeRequest => {
+                CompletedReply::BeforeRequest(Ok(from_py::<ProviderRequest>(value.bind(py))?))
+            }
+            AdapterOperation::AfterResponse => {
+                CompletedReply::AfterResponse(Ok(from_py::<ProviderResponse>(value.bind(py))?))
+            }
+            AdapterOperation::CachedResponse | AdapterOperation::ConstructResponse => {
+                self.state.response = Some(value);
+                Self::lifecycle_result()
+            }
+            AdapterOperation::MapFailure => {
+                self.state
+                    .retain_error(py, PyErr::from_value(value.into_bound(py)));
+                Self::lifecycle_result()
+            }
+        })
+    }
+
+    fn cleanup(&mut self) {
+        self.pending = None;
+    }
+
+    fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
+        visit.call(&self.request)?;
+        visit.call(&self.adapter)
+    }
+}
+
+pub(crate) fn run<R: PythonCompletedRoute>(
+    py: Python<'_>,
+    request: Bound<'_, PyDict>,
+    args: Bound<'_, PyTuple>,
+    kwargs: Bound<'_, PyDict>,
+    asynchronous: bool,
+    host: Bound<'_, PyAny>,
+) -> PyResult<Py<PyAny>>
+where
+    R::Response: Serialize,
+{
+    R::admit(&request)?;
+    let controls = crate::cache::snapshot(
+        py,
+        if asynchronous {
+            R::ASYNC_CALL_TYPE.as_str()
+        } else {
+            R::SYNC_CALL_TYPE.as_str()
+        },
+        &request,
+    )?;
+    crate::errors::admit(
+        litellm_core::call_lifecycle::cache::ResponseCachePlan {
+            controls,
+            ..Default::default()
+        }
+        .admit(),
+    )?;
+    let call = CompletedCall::<R>::new(CompletedWorkflow::default(), asynchronous);
+    let host = PythonCompletedHost::<R> {
+        state: PythonCallState::new(
+            py,
+            args.unbind(),
+            kwargs.copy()?.unbind(),
+            CallMode::from_async(asynchronous),
+            if asynchronous {
+                R::ASYNC_CALL_TYPE
+            } else {
+                R::SYNC_CALL_TYPE
+            },
+        )?,
+        request: request.unbind(),
+        adapter: host.unbind(),
+        pending: None,
+        route: PhantomData,
+    };
+    run_call(py, call, host)
+}

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

@@ -0,0 +1,163 @@
+use pyo3::intern;
+use pyo3::prelude::*;
+use pyo3::types::PyString;
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub(crate) enum CallMode {
+    Sync,
+    Async,
+}
+
+impl CallMode {
+    pub(crate) fn from_async(asynchronous: bool) -> Self {
+        if asynchronous {
+            Self::Async
+        } else {
+            Self::Sync
+        }
+    }
+
+    pub(crate) fn is_async(self) -> bool {
+        matches!(self, Self::Async)
+    }
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub(crate) enum PythonCallType {
+    Ocr,
+    AsyncOcr,
+    Completion,
+    AsyncCompletion,
+    AnthropicMessages,
+    Transcription,
+    AsyncTranscription,
+    #[cfg(test)]
+    Synthetic,
+    #[cfg(test)]
+    Test,
+}
+
+impl PythonCallType {
+    pub(crate) const fn as_str(self) -> &'static str {
+        match self {
+            Self::Ocr => "ocr",
+            Self::AsyncOcr => "aocr",
+            Self::Completion => "completion",
+            Self::AsyncCompletion => "acompletion",
+            Self::AnthropicMessages => "anthropic_messages",
+            Self::Transcription => "transcription",
+            Self::AsyncTranscription => "atranscription",
+            #[cfg(test)]
+            Self::Synthetic => "synthetic",
+            #[cfg(test)]
+            Self::Test => "test",
+        }
+    }
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub(crate) enum CallbackPhase {
+    Input,
+    Payload,
+    SyncSuccess,
+    SyncSuccessForAsyncCall,
+    AsyncSuccess,
+    SyncFailure,
+    AsyncFailure,
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub(crate) enum AdapterOperation {
+    PostProcess,
+    CacheStore,
+    CachedResponse,
+    ProjectRequest,
+    BeforeRequest,
+    AfterResponse,
+    ConstructResponse,
+    MapFailure,
+}
+
+impl AdapterOperation {
+    pub(crate) const fn as_str(self) -> &'static str {
+        match self {
+            Self::PostProcess => "post_process",
+            Self::CacheStore => "cache_response",
+            Self::CachedResponse => "cached_response",
+            Self::ProjectRequest => "project",
+            Self::BeforeRequest => "before_request",
+            Self::AfterResponse => "after_response",
+            Self::ConstructResponse => "response",
+            Self::MapFailure => "map_failure",
+        }
+    }
+}
+
+impl CallbackPhase {
+    pub(crate) const fn as_str(self) -> &'static str {
+        match self {
+            Self::Input => "input",
+            Self::Payload => "payload",
+            Self::SyncSuccess => "sync_success",
+            Self::SyncSuccessForAsyncCall => "sync_success_async",
+            Self::AsyncSuccess => "async_success",
+            Self::SyncFailure => "sync_failure",
+            Self::AsyncFailure => "async_failure",
+        }
+    }
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub(crate) enum RequestField {
+    ApiBase,
+    ApiKey,
+    Audio,
+    Body,
+    CustomLlmProvider,
+    ExtraHeaders,
+    HasAgenticHook,
+    HostFacts,
+    LitellmCallId,
+    Messages,
+    Model,
+    OptionalParams,
+    Timeout,
+}
+
+impl RequestField {
+    pub(crate) fn key(self, py: Python<'_>) -> &Bound<'_, PyString> {
+        match self {
+            Self::ApiBase => intern!(py, "api_base"),
+            Self::ApiKey => intern!(py, "api_key"),
+            Self::Audio => intern!(py, "audio"),
+            Self::Body => intern!(py, "body"),
+            Self::CustomLlmProvider => intern!(py, "custom_llm_provider"),
+            Self::ExtraHeaders => intern!(py, "extra_headers"),
+            Self::HasAgenticHook => intern!(py, "has_agentic_hook"),
+            Self::HostFacts => intern!(py, "host_facts"),
+            Self::LitellmCallId => intern!(py, "litellm_call_id"),
+            Self::Messages => intern!(py, "messages"),
+            Self::Model => intern!(py, "model"),
+            Self::OptionalParams => intern!(py, "optional_params"),
+            Self::Timeout => intern!(py, "timeout"),
+        }
+    }
+
+    pub(crate) const fn as_str(self) -> &'static str {
+        match self {
+            Self::ApiBase => "api_base",
+            Self::ApiKey => "api_key",
+            Self::Audio => "audio",
+            Self::Body => "body",
+            Self::CustomLlmProvider => "custom_llm_provider",
+            Self::ExtraHeaders => "extra_headers",
+            Self::HasAgenticHook => "has_agentic_hook",
+            Self::HostFacts => "host_facts",
+            Self::LitellmCallId => "litellm_call_id",
+            Self::Messages => "messages",
+            Self::Model => "model",
+            Self::OptionalParams => "optional_params",
+            Self::Timeout => "timeout",
+        }
+    }
+}

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

@@ -6,6 +6,7 @@ use pyo3::exceptions::PyException;
 use pyo3::prelude::*;
 
 use super::bindings::PythonLogger;
+use super::contract::CallbackPhase;
 use super::state::PythonCallState;
 
 impl PythonCallState {
@@ -32,7 +33,7 @@ impl PythonCallState {
             state: self,
             logger,
         };
-        match success_dispatch(self.asynchronous, self.internal, &facts)? {
+        match success_dispatch(self.mode.is_async(), self.internal, &facts)? {
             SuccessDispatch::SyncBookkeeping => {
                 logger.success_bookkeeping(py, &self.response, &self.start, &self.end, false)
             }
@@ -62,7 +63,7 @@ impl PythonCallState {
         py: Python<'_>,
         asynchronous: bool,
     ) -> PyResult<Option<Py<PyAny>>> {
-        if !failure_dispatch(self.asynchronous, self.internal, self.logger.is_some()) {
+        if !failure_dispatch(self.mode.is_async(), self.internal, self.logger.is_some()) {
             return Ok(None);
         }
         let Some(error) = &self.error else {
@@ -95,9 +96,9 @@ impl SuccessFacts for PythonSuccessFacts<'_, '_> {
         self.logger.callbacks_needed(
             self.py,
             if asynchronous {
-                "async_success"
+                CallbackPhase::AsyncSuccess
             } else {
-                "sync_success"
+                CallbackPhase::SyncSuccess
             },
         )
     }

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

@@ -1,10 +1,11 @@
 use std::panic::{AssertUnwindSafe, catch_unwind};
 
-use litellm_python_interop::panic_to_pyerr;
 use pyo3::exceptions::{PyBaseException, PyRuntimeError};
 use pyo3::gc::{PyTraverseError, PyVisit};
 use pyo3::prelude::*;
 
+use litellm_python_interop::panic_to_pyerr;
+
 pub(super) enum ExecutionStep {
     Return(Py<PyAny>),
     Await(Py<PyAny>),

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

@@ -1,7 +1,10 @@
 mod bindings;
+pub(crate) mod completed;
+pub(crate) mod contract;
 mod dispatch;
 mod handle;
 mod preparation;
+pub(crate) mod request;
 mod runner;
 mod state;
 

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

@@ -0,0 +1,63 @@
+use litellm_core::call_lifecycle::provider::ProviderOptions;
+use litellm_python_interop::from_py_preserving_errors as from_py;
+use pyo3::prelude::*;
+use pyo3::types::PyDict;
+use serde_json::{Map, Value};
+
+use super::contract::RequestField;
+
+pub(crate) fn required<'py>(
+    request: &Bound<'py, PyDict>,
+    field: RequestField,
+) -> PyResult<Bound<'py, PyAny>> {
+    let name = field.as_str();
+    request
+        .get_item(field.key(request.py()))?
+        .ok_or_else(|| pyo3::exceptions::PyTypeError::new_err(format!("missing {name}")))
+}
+
+pub(crate) fn optional_string(
+    request: &Bound<'_, PyDict>,
+    field: RequestField,
+) -> PyResult<Option<String>> {
+    request
+        .get_item(field.key(request.py()))?
+        .map(|value| value.extract())
+        .transpose()
+        .map(Option::flatten)
+}
+
+pub(crate) fn object(
+    request: &Bound<'_, PyDict>,
+    field: RequestField,
+) -> PyResult<Map<String, Value>> {
+    request
+        .get_item(field.key(request.py()))?
+        .filter(|value| !value.is_none())
+        .map(|value| from_py(&value))
+        .transpose()
+        .map(Option::unwrap_or_default)
+}
+
+pub(crate) fn options(request: &Bound<'_, PyDict>) -> PyResult<ProviderOptions> {
+    Ok(ProviderOptions {
+        model: required(request, RequestField::Model)?.extract()?,
+        litellm_call_id: optional_string(request, RequestField::LitellmCallId)?,
+        api_key: optional_string(request, RequestField::ApiKey)?,
+        api_base: optional_string(request, RequestField::ApiBase)?,
+        custom_llm_provider: optional_string(request, RequestField::CustomLlmProvider)?,
+        extra_headers: request
+            .get_item(RequestField::ExtraHeaders.key(request.py()))?
+            .filter(|value| !value.is_none())
+            .map(|value| from_py(&value))
+            .transpose()?,
+        timeout: crate::marshal::optional_timeout(
+            request
+                .get_item(RequestField::Timeout.key(request.py()))?
+                .filter(|value| !value.is_none())
+                .map(|value| crate::marshal::python_timeout_seconds(request.py(), value.unbind()))
+                .transpose()?
+                .flatten(),
+        ),
+    })
+}

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

@@ -2,13 +2,15 @@ use std::sync::Arc;
 use std::task::Poll;
 
 use futures_util::future::{AbortHandle, Abortable};
-use litellm_core::call_lifecycle::host::{
-    HostCall as NativeCall, HostCallStep as NativeCallStep, HostFailure, HostPhase, HostStep,
-};
+use tokio::sync::Mutex;
+
 use pyo3::exceptions::{PyException, PyRuntimeError};
 use pyo3::gc::{PyTraverseError, PyVisit};
 use pyo3::prelude::*;
-use tokio::sync::Mutex;
+
+use litellm_core::call_lifecycle::host::{
+    HostCall as NativeCall, HostCallStep as NativeCallStep, HostFailure, HostPhase, HostStep,
+};
 
 use super::handle::{Execution, ExecutionBody, ExecutionStep};
 use super::state::{PythonCallState, missing_state, now};
@@ -32,6 +34,20 @@ pub(crate) trait PythonRoute: Send + Sync {
         py: Python<'_>,
         operation: <Self::Call as NativeCall>::Operation,
     ) -> PyResult<<Self::Call as NativeCall>::Result>;
+    fn invoke_step(
+        &mut self,
+        py: Python<'_>,
+        operation: <Self::Call as NativeCall>::Operation,
+    ) -> PyResult<HostStep<<Self::Call as NativeCall>::Result, Py<PyAny>>> {
+        self.invoke(py, operation).map(HostStep::Ready)
+    }
+    fn accept_route(
+        &mut self,
+        _py: Python<'_>,
+        _value: Py<PyAny>,
+    ) -> PyResult<<Self::Call as NativeCall>::Result> {
+        Err(missing_state())
+    }
     fn cleanup(&mut self);
     fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError>;
 }
@@ -48,6 +64,7 @@ struct NativeCallState<C: NativeCall> {
 enum PendingOperation {
     Native,
     Host(HostPhase),
+    Route,
 }
 
 struct PythonLifecycle<R: PythonRoute> {
@@ -62,7 +79,7 @@ pub(crate) fn run_call<R: PythonRoute + 'static>(
     call: R::Call,
     route: R,
 ) -> PyResult<Py<PyAny>> {
-    let asynchronous = route.state().asynchronous;
+    let asynchronous = route.state().mode.is_async();
     let mut lifecycle = PythonLifecycle {
         route,
         call: Some(Arc::new(Mutex::new(NativeCallState { call, result: None }))),
@@ -102,7 +119,7 @@ impl<R: PythonRoute> PythonLifecycle<R> {
             call.result = Some(result);
             Ok(())
         };
-        if self.route.state().asynchronous {
+        if self.route.state().mode.is_async() {
             let mut future = Box::pin(future);
             if let Poll::Ready(()) = poll_async_value(py, future.as_mut())? {
                 return Ok(HostStep::Ready(self.take_native_result()?));
@@ -182,6 +199,11 @@ impl<R: PythonRoute> PythonLifecycle<R> {
                 };
                 self.resume_core(py, Some(result))?
             }
+            (Some(PendingOperation::Route), Some(result)) => {
+                let result = result.and_then(|value| self.route.accept_route(py, value));
+                let result = result.map_err(|error| self.host_failure(py, error, None));
+                self.resume_core(py, Some(result))?
+            }
             _ => return Err(missing_state()),
         };
         loop {
@@ -215,7 +237,14 @@ impl<R: PythonRoute> PythonLifecycle<R> {
                         .map(|()| R::lifecycle_result()),
                     Err(error) => Err(error),
                 },
-                None => self.route.invoke(py, operation),
+                None => match self.route.invoke_step(py, operation) {
+                    Ok(HostStep::Ready(result)) => Ok(result),
+                    Ok(HostStep::Suspend(awaitable)) => {
+                        self.pending = Some(PendingOperation::Route);
+                        return Ok(ExecutionStep::Await(awaitable));
+                    }
+                    Err(error) => Err(error),
+                },
             };
             let result = match result {
                 Ok(result) => Ok(result),

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

@@ -5,6 +5,7 @@ use pyo3::prelude::*;
 use pyo3::types::{PyDict, PyTuple};
 
 use super::bindings::{self, DeploymentHooks, PythonLogger};
+use super::contract::{CallMode, PythonCallType};
 use super::preparation;
 
 pub(crate) fn missing_state() -> PyErr {
@@ -19,9 +20,9 @@ pub(crate) struct PythonCallState {
     pub end: Option<Py<PyAny>>,
     pub response: Option<Py<PyAny>>,
     pub error: Option<Py<PyBaseException>>,
-    pub asynchronous: bool,
+    pub mode: CallMode,
     pub internal: bool,
-    pub call_type: &'static str,
+    pub call_type: PythonCallType,
 }
 
 pub(crate) fn now(py: Python<'_>) -> PyResult<Py<PyAny>> {
@@ -46,7 +47,7 @@ impl PythonCallState {
                 return Ok(HostStep::Suspend(DeploymentHooks::before_call(
                     py,
                     &self.kwargs,
-                    self.call_type,
+                    self.call_type.as_str(),
                 )?));
             }
             HostPhase::Prepare => self.prepare(py)?,
@@ -62,7 +63,7 @@ impl PythonCallState {
                     py,
                     &self.kwargs,
                     &self.response,
-                    self.call_type,
+                    self.call_type.as_str(),
                 )?));
             }
             HostPhase::Finalize => self.finalize(py)?,
@@ -75,7 +76,7 @@ impl PythonCallState {
                         py,
                         &self.kwargs,
                         error,
-                        self.call_type,
+                        self.call_type.as_str(),
                     )?));
                 }
             }
@@ -86,7 +87,9 @@ impl PythonCallState {
                     return Ok(HostStep::Suspend(awaitable));
                 }
             }
+            HostPhase::PostProcess | HostPhase::CacheStore => {}
             HostPhase::Execute
+            | HostPhase::CacheLookup
             | HostPhase::ConstructResponse
             | HostPhase::MapFailure
             | HostPhase::Complete => return Err(missing_state()),
@@ -114,8 +117,8 @@ impl PythonCallState {
         py: Python<'_>,
         args: Py<PyTuple>,
         kwargs: Py<PyDict>,
-        asynchronous: bool,
-        call_type: &'static str,
+        mode: CallMode,
+        call_type: PythonCallType,
     ) -> PyResult<Self> {
         Ok(Self {
             args,
@@ -125,7 +128,7 @@ impl PythonCallState {
             end: None,
             response: None,
             error: None,
-            asynchronous,
+            mode,
             internal: false,
             call_type,
         })
@@ -142,11 +145,11 @@ impl PythonCallState {
         self.internal = bindings::is_internal_call(py)?;
         let result = bindings::setup(
             py,
-            self.call_type,
+            self.call_type.as_str(),
             &self.args,
             &self.kwargs,
             &self.start,
-            self.asynchronous,
+            self.mode.is_async(),
         )?;
         self.logger = Some(result.logger()?);
         self.kwargs = result.kwargs()?;

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

@@ -6,6 +6,7 @@ use pyo3::gc::{PyTraverseError, PyVisit};
 use pyo3::prelude::*;
 use pyo3::types::PyTuple;
 
+use super::contract::{CallMode, PythonCallType};
 use super::dispatch::{PendingLogging, PendingSuccess};
 use super::handle::{Execution, ExecutionBody, ExecutionStep};
 use super::*;
@@ -160,8 +161,8 @@ fn shared_runner_executes_a_non_ocr_adapter() {
                 py,
                 PyTuple::empty(py).unbind(),
                 PyDict::new(py).unbind(),
-                false,
-                "synthetic",
+                CallMode::Sync,
+                PythonCallType::Synthetic,
             )
             .unwrap(),
         );
@@ -196,8 +197,8 @@ fn ready_native_lifecycle_completes_without_scheduling() {
                 py,
                 PyTuple::empty(py).unbind(),
                 PyDict::new(py).unbind(),
-                true,
-                "synthetic",
+                CallMode::Async,
+                PythonCallType::Synthetic,
             )
             .unwrap(),
         );
@@ -280,8 +281,8 @@ fn error_execution(py: Python<'_>, error: Bound<'_, PyBaseException>) -> Executi
         py,
         PyTuple::empty(py).unbind(),
         PyDict::new(py).unbind(),
-        true,
-        "test",
+        CallMode::Async,
+        PythonCallType::Test,
     )
     .unwrap();
     state.retain_error(py, PyErr::from_value(error.into_any()));
@@ -342,9 +343,9 @@ fn state(
         end: Some(py.None()),
         response: Some(response),
         error: None,
-        asynchronous,
+        mode: CallMode::from_async(asynchronous),
         internal: false,
-        call_type: "test",
+        call_type: PythonCallType::Test,
     }
 }
 

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

@@ -1,2 +1,57 @@
-// TODO: implement chat_completions lifecycle checkpoints before replacing the Python lifecycle
-unimplemented_lifecycle_route!(ChatCompletions, _chat_completions_lifecycle);
+use litellm_core::chat_completions::lifecycle::{
+    ChatCompletionsRoute, OwnedChatCompletionsRequest,
+};
+use litellm_python_interop::from_py_preserving_errors as from_py;
+use pyo3::prelude::*;
+use pyo3::types::{PyDict, PyTuple};
+
+use crate::lifecycle::completed::{self, PythonCompletedRoute};
+use crate::lifecycle::contract::{PythonCallType, RequestField};
+use crate::lifecycle::request::{object, optional_string, options, required};
+
+impl PythonCompletedRoute for ChatCompletionsRoute {
+    const SYNC_CALL_TYPE: PythonCallType = PythonCallType::Completion;
+    const ASYNC_CALL_TYPE: PythonCallType = PythonCallType::AsyncCompletion;
+
+    fn admit(request: &Bound<'_, PyDict>) -> PyResult<()> {
+        crate::errors::admit(litellm_core::chat_completions::admit(
+            &required(request, RequestField::Model)?.extract::<String>()?,
+            optional_string(request, RequestField::CustomLlmProvider)?.as_deref(),
+            from_py(&required(request, RequestField::Messages)?)?,
+            &object(request, RequestField::OptionalParams)?,
+            Some(&object(request, RequestField::ExtraHeaders)?),
+            request
+                .get_item(RequestField::HostFacts.key(request.py()))?
+                .map(|value| from_py(&value))
+                .transpose()?
+                .unwrap_or_default(),
+        ))
+    }
+
+    fn project(request: &Bound<'_, PyDict>) -> PyResult<OwnedChatCompletionsRequest> {
+        Ok(OwnedChatCompletionsRequest {
+            options: options(request)?,
+            messages: from_py(&required(request, RequestField::Messages)?)?,
+            optional_params: object(request, RequestField::OptionalParams)?,
+        })
+    }
+}
+
+#[pyfunction]
+fn _chat_completions_lifecycle(
+    py: Python<'_>,
+    request: Bound<'_, PyDict>,
+    args: Bound<'_, PyTuple>,
+    kwargs: Bound<'_, PyDict>,
+    asynchronous: bool,
+    host: Bound<'_, PyAny>,
+) -> PyResult<Py<PyAny>> {
+    completed::run::<ChatCompletionsRoute>(py, request, args, kwargs, asynchronous, host)
+}
+
+pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
+    crate::routes::definition::add_function(
+        module,
+        wrap_pyfunction!(_chat_completions_lifecycle, module)?,
+    )
+}

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

@@ -88,7 +88,7 @@ bridge_route! {
         extra_headers: Option<serde_json::Value>,
         timeout_seconds: Option<f64>,
         #[pyo3(from_py_with = litellm_python_interop::from_py)]
-        host_facts: Option<Value>,
+        host_facts: Option<serde_json::Value>,
         on_request: Option<Py<PyAny>>,
     },
     prepare = prepare_chat_completions,

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

@@ -1,2 +1,49 @@
-// TODO: implement messages lifecycle checkpoints before replacing the Python lifecycle
-unimplemented_lifecycle_route!(Messages, _messages_lifecycle);
+use pyo3::prelude::*;
+use pyo3::types::{PyDict, PyTuple};
+
+use litellm_core::messages::lifecycle::{MessagesRoute, OwnedMessagesRequest};
+use litellm_python_interop::from_py_preserving_errors as from_py;
+
+use crate::lifecycle::completed::{self, PythonCompletedRoute};
+use crate::lifecycle::contract::{PythonCallType, RequestField};
+use crate::lifecycle::request::{optional_string, options, required};
+
+impl PythonCompletedRoute for MessagesRoute {
+    const SYNC_CALL_TYPE: PythonCallType = PythonCallType::AnthropicMessages;
+    const ASYNC_CALL_TYPE: PythonCallType = PythonCallType::AnthropicMessages;
+
+    fn admit(request: &Bound<'_, PyDict>) -> PyResult<()> {
+        crate::errors::admit(litellm_core::messages::admit(
+            &required(request, RequestField::Model)?.extract::<String>()?,
+            optional_string(request, RequestField::CustomLlmProvider)?.as_deref(),
+            request
+                .get_item(RequestField::HasAgenticHook.key(request.py()))?
+                .map(|value| value.extract())
+                .transpose()?
+                .unwrap_or(false),
+        ))
+    }
+
+    fn project(request: &Bound<'_, PyDict>) -> PyResult<OwnedMessagesRequest> {
+        Ok(OwnedMessagesRequest {
+            options: options(request)?,
+            body: from_py(&required(request, RequestField::Body)?)?,
+        })
+    }
+}
+
+#[pyfunction]
+fn _messages_lifecycle(
+    py: Python<'_>,
+    request: Bound<'_, PyDict>,
+    args: Bound<'_, PyTuple>,
+    kwargs: Bound<'_, PyDict>,
+    asynchronous: bool,
+    host: Bound<'_, PyAny>,
+) -> PyResult<Py<PyAny>> {
+    completed::run::<MessagesRoute>(py, request, args, kwargs, asynchronous, host)
+}
+
+pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
+    crate::routes::definition::add_function(module, wrap_pyfunction!(_messages_lifecycle, module)?)
+}

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

@@ -8,6 +8,7 @@ use litellm_core::ocr::hooks::OcrPreCallRequest;
 use litellm_python_interop::to_py_preserving_errors as to_py;
 
 use crate::lifecycle::PythonLogger;
+use crate::lifecycle::contract::CallbackPhase;
 
 pub(super) struct OcrLoggingFields {
     model: String,
@@ -89,7 +90,7 @@ impl PythonLogger {
         kwargs.set_item("input", "OCR document processing")?;
         kwargs.set_item("api_key", api_key)?;
         kwargs.set_item("additional_args", &additional)?;
-        if self.callbacks_needed(py, "input")? {
+        if self.callbacks_needed(py, CallbackPhase::Input)? {
             self.object(py).call_method("pre_call", (), Some(&kwargs))?;
         } else {
             self.object(py)
@@ -109,7 +110,7 @@ impl PythonLogger {
         let additional = PyDict::new(py);
         additional.set_item("complete_input_dict", body)?;
         additional.set_item("headers", headers)?;
-        if self.callbacks_needed(py, "input")? {
+        if self.callbacks_needed(py, CallbackPhase::Input)? {
             let kwargs = PyDict::new(py);
             kwargs.set_item("original_response", to_py(py, original_response)?)?;
             kwargs.set_item("additional_args", &additional)?;

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

@@ -11,6 +11,7 @@ use litellm_python_interop::{
 use super::callbacks;
 use super::errors::to_pyerr as ocr_error_to_pyerr;
 use super::project::{ProjectedOcrFields, admitted_call, project_request};
+use crate::lifecycle::contract::{CallMode, CallbackPhase, PythonCallType};
 use crate::lifecycle::{
     OperationClass, PythonCallState, PythonRoute, missing_state, now, run_call,
 };
@@ -99,7 +100,11 @@ impl PythonOcrHost {
             &projected.fields.secret_fields,
             &request.url,
         )?;
-        if !self.state.logger()?.callbacks_needed(py, "payload")? {
+        if !self
+            .state
+            .logger()?
+            .callbacks_needed(py, CallbackPhase::Payload)?
+        {
             self.state
                 .logger()?
                 .object(py)
@@ -147,7 +152,7 @@ impl PythonOcrHost {
         request: OcrPostCallRequest,
     ) -> PyResult<OcrPostCallRequest> {
         let logger = self.state.logger()?;
-        if logger.callbacks_needed(py, "payload")? {
+        if logger.callbacks_needed(py, CallbackPhase::Payload)? {
             let projected = self.projected()?;
             logger.post_ocr(
                 py,
@@ -301,8 +306,12 @@ fn _ocr_lifecycle(
             py,
             args.unbind(),
             kwargs.copy()?.unbind(),
-            asynchronous,
-            if asynchronous { "aocr" } else { "ocr" },
+            CallMode::from_async(asynchronous),
+            if asynchronous {
+                PythonCallType::AsyncOcr
+            } else {
+                PythonCallType::Ocr
+            },
         )?,
         adapter: host.unbind(),
         data: OcrHostData::Unprojected {

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

@@ -1,10 +1,11 @@
-use litellm_core::Error;
 use std::future::Future;
 
-use litellm_core::ocr::wire::{OcrWireRequest, decode_request};
 use pyo3::prelude::*;
 use serde_json::Value;
 
+use litellm_core::Error;
+use litellm_core::ocr::wire::{OcrWireRequest, decode_request, validate_document_url};
+
 use super::errors::to_pyerr as ocr_error_to_pyerr;
 use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty};
 
@@ -28,6 +29,9 @@ fn prepare_ocr(
         .map_err(|error| pyo3::exceptions::PyValueError::new_err(error.to_string()))?
         .unwrap_or_default();
 
+    validate_document_url(&document)
+        .map_err(Error::from)
+        .map_err(ocr_error_to_pyerr)?;
     crate::errors::admit(litellm_core::ocr::admit_value(
         &options.model,
         options.custom_llm_provider.as_deref(),

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

@@ -1,2 +1,52 @@
-// TODO: implement transcription lifecycle checkpoints before replacing the Python lifecycle
-unimplemented_lifecycle_route!(Transcription, _transcription_lifecycle);
+use pyo3::prelude::*;
+use pyo3::types::{PyDict, PyTuple};
+
+use litellm_core::audio_transcription::lifecycle::{
+    AudioTranscriptionRoute, OwnedAudioTranscriptionRequest,
+};
+use litellm_python_interop::from_py_preserving_errors as from_py;
+
+use crate::lifecycle::completed::{self, PythonCompletedRoute};
+use crate::lifecycle::contract::{PythonCallType, RequestField};
+use crate::lifecycle::request::{object, optional_string, options, required};
+
+impl PythonCompletedRoute for AudioTranscriptionRoute {
+    const SYNC_CALL_TYPE: PythonCallType = PythonCallType::Transcription;
+    const ASYNC_CALL_TYPE: PythonCallType = PythonCallType::AsyncTranscription;
+
+    fn admit(request: &Bound<'_, PyDict>) -> PyResult<()> {
+        let audio = from_py(&required(request, RequestField::Audio)?)?;
+        crate::errors::admit(litellm_core::audio_transcription::admit(
+            &required(request, RequestField::Model)?.extract::<String>()?,
+            optional_string(request, RequestField::CustomLlmProvider)?.as_deref(),
+            &audio,
+        ))
+    }
+
+    fn project(request: &Bound<'_, PyDict>) -> PyResult<OwnedAudioTranscriptionRequest> {
+        Ok(OwnedAudioTranscriptionRequest {
+            options: options(request)?,
+            audio: from_py(&required(request, RequestField::Audio)?)?,
+            optional_params: object(request, RequestField::OptionalParams)?,
+        })
+    }
+}
+
+#[pyfunction]
+fn _transcription_lifecycle(
+    py: Python<'_>,
+    request: Bound<'_, PyDict>,
+    args: Bound<'_, PyTuple>,
+    kwargs: Bound<'_, PyDict>,
+    asynchronous: bool,
+    host: Bound<'_, PyAny>,
+) -> PyResult<Py<PyAny>> {
+    completed::run::<AudioTranscriptionRoute>(py, request, args, kwargs, asynchronous, host)
+}
+
+pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
+    crate::routes::definition::add_function(
+        module,
+        wrap_pyfunction!(_transcription_lifecycle, module)?,
+    )
+}

litellm/rust_bridge/_native.pyi

@@ -1,12 +1,53 @@
 from asyncio import Future
 from collections.abc import Callable, Coroutine
-from typing import Literal, final, overload
+from typing import Literal, Protocol, TypedDict, final, overload
 
-from typing_extensions import Never
+from typing_extensions import Never, NotRequired, Required
 
 from litellm.llms.base_llm.ocr.transformation import OCRResponse
+from litellm.rust_bridge.lifecycle import Await, Complete
 from litellm.rust_bridge.ocr import LiteLLMOcrRequest
 
+class _CommonLifecycleRequest(TypedDict):
+    model: Required[str]
+    api_key: NotRequired[str | None]
+    api_base: NotRequired[str | None]
+    custom_llm_provider: NotRequired[str | None]
+    extra_headers: NotRequired[object]
+    timeout: NotRequired[object]
+
+class _MessagesLifecycleRequest(_CommonLifecycleRequest):
+    body: Required[object]
+    has_agentic_hook: NotRequired[bool | None]
+
+class _ChatCompletionsLifecycleRequest(_CommonLifecycleRequest):
+    messages: Required[object]
+    optional_params: NotRequired[object]
+    host_facts: NotRequired[object]
+
+class _TranscriptionLifecycleRequest(_CommonLifecycleRequest):
+    audio: Required[object]
+    optional_params: NotRequired[object]
+
+class _CompletedLifecycleHost(Protocol):
+    def invoke(
+        self,
+        operation: Literal[
+            "post_process",
+            "cache_response",
+            "cached_response",
+            "project",
+            "before_request",
+            "after_response",
+            "response",
+            "map_failure",
+        ],
+        payload: object,
+        request: object,
+        kwargs: dict[str, object],
+        logger: object,
+    ) -> Await | Complete: ...
+
 class RustBridgeDeclined(Exception): ...
 class RustBridgeUnavailable(Exception): ...
 class RustHostCallbackError(Exception): ...
@@ -28,15 +69,54 @@ def _ocr_lifecycle(
     asynchronous: Literal[True],
     host: object,
 ) -> Coroutine[object, object, OCRResponse]: ...
+@overload
 def _messages_lifecycle(
-    request: object, args: tuple[object, ...], kwargs: dict[str, object], asynchronous: bool, host: object
-) -> Never: ...
+    request: _MessagesLifecycleRequest,
+    args: tuple[object, ...],
+    kwargs: dict[str, object],
+    asynchronous: Literal[False],
+    host: _CompletedLifecycleHost,
+) -> object: ...
+@overload
+def _messages_lifecycle(
+    request: _MessagesLifecycleRequest,
+    args: tuple[object, ...],
+    kwargs: dict[str, object],
+    asynchronous: Literal[True],
+    host: _CompletedLifecycleHost,
+) -> Coroutine[object, object, object]: ...
+@overload
 def _chat_completions_lifecycle(
-    request: object, args: tuple[object, ...], kwargs: dict[str, object], asynchronous: bool, host: object
-) -> Never: ...
+    request: _ChatCompletionsLifecycleRequest,
+    args: tuple[object, ...],
+    kwargs: dict[str, object],
+    asynchronous: Literal[False],
+    host: _CompletedLifecycleHost,
+) -> object: ...
+@overload
+def _chat_completions_lifecycle(
+    request: _ChatCompletionsLifecycleRequest,
+    args: tuple[object, ...],
+    kwargs: dict[str, object],
+    asynchronous: Literal[True],
+    host: _CompletedLifecycleHost,
+) -> Coroutine[object, object, object]: ...
+@overload
 def _transcription_lifecycle(
-    request: object, args: tuple[object, ...], kwargs: dict[str, object], asynchronous: bool, host: object
-) -> Never: ...
+    request: _TranscriptionLifecycleRequest,
+    args: tuple[object, ...],
+    kwargs: dict[str, object],
+    asynchronous: Literal[False],
+    host: _CompletedLifecycleHost,
+) -> object: ...
+@overload
+def _transcription_lifecycle(
+    request: _TranscriptionLifecycleRequest,
+    args: tuple[object, ...],
+    kwargs: dict[str, object],
+    asynchronous: Literal[True],
+    host: _CompletedLifecycleHost,
+) -> Coroutine[object, object, object]: ...
 def _embeddings_lifecycle(
     request: object, args: tuple[object, ...], kwargs: dict[str, object], asynchronous: bool, host: object
 ) -> Never: ...
@@ -181,11 +261,6 @@ def gil_stats() -> dict[str, int]: ...
 
 __all__ = [
     "_OCR_MAX_FILE_BYTES",
-    "ResponsesWebSocketConnection",
-    "RustBridgeDeclined",
-    "RustBridgeUnavailable",
-    "RustHostCallbackError",
-    "RustUpstreamError",
     "_chat_completions_lifecycle",
     "_embeddings_lifecycle",
     "_image_edit_lifecycle",
@@ -200,6 +275,11 @@ __all__ = [
     "_responses_lifecycle",
     "_speech_lifecycle",
     "_transcription_lifecycle",
+    "ResponsesWebSocketConnection",
+    "RustBridgeDeclined",
+    "RustBridgeUnavailable",
+    "RustHostCallbackError",
+    "RustUpstreamError",
     "achat_completions",
     "amessages",
     "aocr",

litellm/rust_bridge/catalog.py

@@ -113,7 +113,7 @@ COMPONENTS: Final[Mapping[ComponentName, NativeComponent]] = MappingProxyType(
         ),
         ComponentName.CHAT_COMPLETIONS: _component(
             ComponentName.CHAT_COMPLETIONS,
-            _experimental_completed,
+            _experimental(),
             ("chat_completions", "achat_completions", "_chat_completions_lifecycle"),
         ),
         ComponentName.TRANSCRIPTION: _component(

tests/test_litellm/rust_bridge/test_chat_completions.py

@@ -113,13 +113,19 @@ def test_decline_has_no_logging_effect_and_runs_one_fallback(monkeypatch: pytest
     assert events == []
 
 
-def test_streaming_uses_python_without_loading_native() -> None:
-    native_call: Final = _RecordingCall()
+def test_streaming_decline_comes_from_the_native_call(monkeypatch: pytest.MonkeyPatch) -> None:
+    _fake_native_bridge(monkeypatch)
+    native_call: Final = _RecordingCall(error=_FakeDeclined("streaming"))
     bridge.set_rust_chat_completions(chat_completions=native_call)
     kwargs: Final = _call_kwargs(ModelResponse())
     kwargs["stream"] = True
     assert bridge.chat_completions(**kwargs) == "python"
-    assert native_call.calls == []
+    assert len(native_call.calls) == 1
+    assert native_call.calls[0]["host_facts"] == {
+        "stream": True,
+        "anthropic_user_id": False,
+        "bedrock_metadata_owned": False,
+    }
 
 
 def test_host_facts_reach_the_single_native_call(monkeypatch: pytest.MonkeyPatch) -> None:

tests/test_litellm_rust/test_route_foundation.py

@@ -7,26 +7,20 @@ import pytest
 from litellm.rust_bridge import _native
 from litellm.rust_bridge.bindings import NativeBinding
 from litellm.rust_bridge.catalog import NATIVE_EXPORTS
-from litellm.rust_bridge.chat_completions.lifecycle import LIFECYCLE as CHAT_COMPLETIONS
-from litellm.rust_bridge.configuration import ExecutionDecision, ComponentName
+from litellm.rust_bridge.configuration import ComponentName, ExecutionDecision
 from litellm.rust_bridge.embeddings.lifecycle import LIFECYCLE as EMBEDDINGS
 from litellm.rust_bridge.image_edit.lifecycle import LIFECYCLE as IMAGE_EDIT
 from litellm.rust_bridge.image_generation.lifecycle import LIFECYCLE as IMAGE_GENERATION
-from litellm.rust_bridge.messages.lifecycle import LIFECYCLE as MESSAGES
 from litellm.rust_bridge.moderation.lifecycle import LIFECYCLE as MODERATION
 from litellm.rust_bridge.rerank.lifecycle import LIFECYCLE as RERANK
 from litellm.rust_bridge.responses.lifecycle import LIFECYCLE as RESPONSES
 from litellm.rust_bridge.route import ComponentExecution, NativeLifecycle
 from litellm.rust_bridge.runtime import BridgeErrorContext, invoke
 from litellm.rust_bridge.speech.lifecycle import LIFECYCLE as SPEECH
-from litellm.rust_bridge.transcription.lifecycle import LIFECYCLE as TRANSCRIPTION
 
 pytestmark = pytest.mark.requires_rust_extension
 
 UNIMPLEMENTED: Final[dict[ComponentName, NativeBinding[NativeLifecycle[object, object]]]] = {
-    ComponentName.MESSAGES: MESSAGES,
-    ComponentName.CHAT_COMPLETIONS: CHAT_COMPLETIONS,
-    ComponentName.TRANSCRIPTION: TRANSCRIPTION,
     ComponentName.EMBEDDINGS: EMBEDDINGS,
     ComponentName.RERANK: RERANK,
     ComponentName.IMAGE_GENERATION: IMAGE_GENERATION,
@@ -126,6 +120,16 @@ def test_transcription_declines_audio_format_before_credentials() -> None:
         _native.transcription("model", {"format": "unsupported", "data": "YQ=="}, custom_llm_provider="bedrock")
 
 
+def test_transcription_lifecycle_declines_audio_format_before_host_work() -> None:
+    request: Final = {
+        "model": "model",
+        "audio": {"format": "unsupported", "data": "YQ=="},
+        "custom_llm_provider": "bedrock",
+    }
+    with pytest.raises(_native.RustBridgeDeclined, match="audio format"):
+        _native._transcription_lifecycle(request, (), {}, False, UntouchedInput())
+
+
 def test_websocket_declines_before_parsing_or_dialing_url() -> None:
     with pytest.raises(_native.RustBridgeDeclined):
         _native.ResponsesWebSocketConnection.connect("not a URL", custom_llm_provider="azure")