refactor(rust): make Python host calls async fns

Co-Authored-By: Yujong Lee <yujong@berri.ai>

litellm-rust/Cargo.lock

@@ -4240,6 +4240,7 @@ dependencies = [
  "litellm-storage-clickhouse",
  "litellm-token-counter",
  "litellm-tracing",
+ "parking_lot",
  "prost",
  "pyo3",
  "pyo3-async-runtimes",

litellm-rust/crates/host-python/src/binding.rs

@@ -13,7 +13,7 @@ pub trait PythonBinding: PythonOwned {
 
     /// Decodes the keyword view returned by `prepare_arguments`, including preflight rewrites.
     fn decode_request(
-        &mut self,
+        &self,
         py: Python<'_>,
         arguments: &Bound<'_, PyDict>,
     ) -> Result<
@@ -22,21 +22,21 @@ pub trait PythonBinding: PythonOwned {
     >;
 
     fn encode_response(
-        &mut self,
+        &self,
         py: Python<'_>,
         response: <Self::Protocol as Protocol>::Response,
     ) -> PyResult<Py<PyAny>>;
 
     /// What the stream carries at hand-off, as the caller's stream receives it.
     fn encode_stream_head(
-        &mut self,
+        &self,
         py: Python<'_>,
         head: <Self::Protocol as Protocol>::StreamHead,
     ) -> PyResult<Py<PyAny>>;
 
     /// One streamed chunk as the caller receives it.
     fn encode_chunk(
-        &mut self,
+        &self,
         py: Python<'_>,
         chunk: <Self::Protocol as Protocol>::Chunk,
     ) -> PyResult<Py<PyAny>>;

litellm-rust/crates/host-python/src/driver.rs

@@ -91,15 +91,17 @@ impl Observation {
     }
 }
 
-/// Hooks lock their own state around Python work, so the handle can traverse them while a
-/// hook is suspended.
+/// What the running call shares with its handle: the binding and hooks the handle
+/// traverses and closes, and the Python values the call retains across a suspension.
+/// Binding and hooks lock their own state only around Python work, so the handle can
+/// traverse them while a hook or host call is suspended.
 struct Shared<H, L> {
+    binding: H,
     hooks: L,
-    parts: Mutex<Parts<H>>,
+    parts: Mutex<Parts>,
 }
 
-struct Parts<H> {
-    binding: H,
+struct Parts {
     arguments: Option<Py<PyDict>>,
     interrupted: Option<Py<PyBaseException>>,
     selected: Option<Selected>,
@@ -151,7 +153,7 @@ where
     M: Machine<Protocol = H::Protocol> + 'static,
     M::Complete: Into<HostedCompletion<ResponseOf<H>>>,
 {
-    fn parts(&self) -> MutexGuard<'_, Parts<H>> {
+    fn parts(&self) -> MutexGuard<'_, Parts> {
         self.shared.parts.lock()
     }
 
@@ -159,6 +161,10 @@ where
         &self.shared.hooks
     }
 
+    fn binding(&self) -> &H {
+        &self.shared.binding
+    }
+
     async fn run(mut self) -> PyResult<Py<PyAny>> {
         match self.begin().await {
             Ok(()) => self.execute().await,
@@ -198,12 +204,8 @@ where
         self.hooks()
             .arguments_prepared(&prepared)
             .map_err(Setup::Host)?;
-        let decoded = {
-            let mut parts = self.parts();
-            let decoded = Python::attach(|py| parts.binding.decode_request(py, prepared.bind(py)));
-            parts.arguments = Some(prepared);
-            decoded
-        };
+        let decoded = Python::attach(|py| self.binding().decode_request(py, prepared.bind(py)));
+        self.parts().arguments = Some(prepared);
         let request = decoded.map_err(Setup::Decode)?;
         let start = self
             .start
@@ -271,7 +273,9 @@ where
         op: HostRequest<ProtocolOf<H>>,
     ) -> Result<(), InvokeError<ErrorOf<H>>> {
         match op {
-            HostRequest::HostCall(call) => self.host_call(call).await,
+            HostRequest::HostCall(call) => {
+                self.binding().handle_host_call(&self.caller, call).await
+            }
             HostRequest::Intercept(InterceptRequest::BeforeProviderRequest {
                 wire,
                 context,
@@ -302,19 +306,13 @@ where
             }
             HostRequest::Stream(StreamDelivery::Open(head, reply)) => {
                 self.streaming = true;
-                let head = {
-                    let mut parts = self.parts();
-                    Python::attach(|py| parts.binding.encode_stream_head(py, head))?
-                };
+                let head = Python::attach(|py| self.binding().encode_stream_head(py, head))?;
                 self.hooks().on_stream_open(&head)?;
                 reply.send(self.caller.open_stream(head).await);
                 Ok(())
             }
             HostRequest::Stream(StreamDelivery::Chunk(chunk, reply)) => {
-                let chunk = {
-                    let mut parts = self.parts();
-                    Python::attach(|py| parts.binding.encode_chunk(py, chunk))?
-                };
+                let chunk = Python::attach(|py| self.binding().encode_chunk(py, chunk))?;
                 self.hooks().on_stream_chunk(&chunk)?;
                 reply.send(self.caller.yield_chunk(chunk).await);
                 Ok(())
@@ -322,24 +320,6 @@ where
         }
     }
 
-    async fn host_call(
-        &self,
-        call: <ProtocolOf<H> as Protocol>::HostCall,
-    ) -> Result<(), InvokeError<ErrorOf<H>>> {
-        let mut pending = {
-            let mut parts = self.parts();
-            Python::attach(|py| parts.binding.begin_host_call(py, call))?
-        };
-        while let Some(awaitable) = pending {
-            let result = self.caller.await_python(awaitable).await;
-            pending = {
-                let mut parts = self.parts();
-                Python::attach(|py| parts.binding.resume_host_call(py, result))?
-            };
-        }
-        Ok(())
-    }
-
     async fn notify(&self, event: PythonCallEvent<'_>) -> PyResult<()> {
         self.parts().observation.emit(&event);
         self.hooks().on_event(&self.caller, event).await
@@ -357,11 +337,7 @@ where
                 return self.succeeded(none).await;
             }
         };
-        let encoded = {
-            let mut parts = self.parts();
-            Python::attach(|py| parts.binding.encode_response(py, response))
-        };
-        let public = match encoded {
+        let public = match Python::attach(|py| self.binding().encode_response(py, response)) {
             Ok(public) => public,
             Err(error) => return self.failure(error, FailureOrigin::Call).await,
         };
@@ -385,7 +361,7 @@ where
             parts.observation.ended_at.get_or_insert_with(epoch_seconds);
             Python::attach(|py| match parts.interrupted.take() {
                 Some(retained) => PyErr::from_value(retained.into_bound(py).into_any()),
-                None => classified(py, &parts.binding, error),
+                None => classified(py, self.binding(), error),
             })
         };
         self.failure(error, FailureOrigin::Call).await
@@ -527,11 +503,11 @@ where
     }
 
     fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
+        self.shared.binding.traverse(visit)?;
         self.shared.hooks.traverse(visit)?;
         let Some(parts) = self.shared.parts.try_lock() else {
             return Ok(());
         };
-        parts.binding.traverse(visit)?;
         visit.call(&parts.arguments)?;
         visit.call(&parts.interrupted)?;
         match &parts.selected {
@@ -550,14 +526,15 @@ where
     fn drop(&mut self) {
         self.body.get_mut().cancel();
         self.reply = None;
-        let Some(Shared { hooks, parts }) = Arc::get_mut(&mut self.shared) else {
+        let Some(Shared {
+            binding,
+            hooks,
+            parts,
+        }) = Arc::get_mut(&mut self.shared)
+        else {
             return;
         };
-        let Parts {
-            binding,
-            observation,
-            ..
-        } = parts.get_mut();
+        let observation = &mut parts.get_mut().observation;
         Python::attach(|py| {
             if observation.observing {
                 let timing = observation.timing();
@@ -592,9 +569,9 @@ where
     M::Complete: Into<HostedCompletion<ResponseOf<H>>>,
 {
     let shared = Arc::new(Shared {
+        binding,
         hooks,
         parts: Mutex::new(Parts {
-            binding,
             arguments: Some(arguments),
             interrupted: None,
             selected: None,

litellm-rust/crates/host-python/src/services.rs

@@ -1,27 +1,12 @@
-use crate::{InvokeError, PythonOwned};
+use std::future::Future;
+
+use crate::{Awaiter, InvokeError, PythonOwned};
 use litellm_host::protocol::Protocol;
-use pyo3::prelude::*;
 
 pub trait PythonHostCalls<P: Protocol>: PythonOwned {
     fn handle_host_call(
-        &mut self,
-        py: Python<'_>,
+        &self,
+        task: &dyn Awaiter,
         call: P::HostCall,
-    ) -> Result<(), InvokeError<P::Error>>;
-
-    fn begin_host_call(
-        &mut self,
-        py: Python<'_>,
-        call: P::HostCall,
-    ) -> Result<Option<Py<PyAny>>, InvokeError<P::Error>> {
-        self.handle_host_call(py, call).map(|()| None)
-    }
-
-    fn resume_host_call(
-        &mut self,
-        _: Python<'_>,
-        result: PyResult<Py<PyAny>>,
-    ) -> Result<Option<Py<PyAny>>, InvokeError<P::Error>> {
-        result.map(|_| None).map_err(InvokeError::Python)
-    }
+    ) -> impl Future<Output = Result<(), InvokeError<P::Error>>> + Send;
 }

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

@@ -167,7 +167,6 @@ pub struct SyntheticBinding {
     log: Log,
     op: OpScript,
     classifier_fails: bool,
-    pending_reply: Option<Reply<String>>,
 }
 
 pub fn binding(log: &Log, op: OpScript) -> SyntheticBinding {
@@ -175,7 +174,6 @@ pub fn binding(log: &Log, op: OpScript) -> SyntheticBinding {
         log: log.clone(),
         op,
         classifier_fails: false,
-        pending_reply: None,
     }
 }
 
@@ -206,7 +204,7 @@ impl PythonBinding for SyntheticBinding {
     type Failure = Classified;
 
     fn decode_request(
-        &mut self,
+        &self,
         _: Python<'_>,
         arguments: &Bound<'_, PyDict>,
     ) -> Result<String, InvokeError<Error>> {
@@ -220,22 +218,18 @@ impl PythonBinding for SyntheticBinding {
     }
 
     fn encode_stream_head(
-        &mut self,
+        &self,
         _: Python<'_>,
         head: std::convert::Infallible,
     ) -> PyResult<Py<PyAny>> {
         match head {}
     }
 
-    fn encode_chunk(
-        &mut self,
-        _: Python<'_>,
-        chunk: std::convert::Infallible,
-    ) -> PyResult<Py<PyAny>> {
+    fn encode_chunk(&self, _: Python<'_>, chunk: std::convert::Infallible) -> PyResult<Py<PyAny>> {
         match chunk {}
     }
 
-    fn encode_response(&mut self, py: Python<'_>, response: String) -> PyResult<Py<PyAny>> {
+    fn encode_response(&self, py: Python<'_>, response: String) -> PyResult<Py<PyAny>> {
         self.log.push("complete");
         Ok(PyString::new(py, &response).into_any().unbind())
     }
@@ -254,49 +248,34 @@ impl PythonBinding for SyntheticBinding {
 }
 
 impl PythonHostCalls<Synthetic> for SyntheticBinding {
-    fn handle_host_call(
-        &mut self,
-        _: Python<'_>,
+    async fn handle_host_call(
+        &self,
+        task: &dyn Awaiter,
         (op, reply): (&'static str, Reply<String>),
     ) -> Result<(), InvokeError<Error>> {
         self.log.push(format!("op:{op}"));
-        self.answer(|| op.to_string())
-            .map(|answer| reply.send(answer))
-    }
-
-    fn begin_host_call(
-        &mut self,
-        py: Python<'_>,
-        (op, reply): (&'static str, Reply<String>),
-    ) -> Result<Option<Py<PyAny>>, InvokeError<Error>> {
         if !matches!(self.op, OpScript::AwaitAnswer | OpScript::AwaitFailure) {
-            return self.handle_host_call(py, (op, reply)).map(|()| None);
+            return self
+                .answer(|| op.to_string())
+                .map(|answer| reply.send(answer));
         }
-        self.pending_reply = Some(reply);
-        let module = PyModule::from_code(
-            py,
-            pyo3::ffi::c_str!(
-                "async def answer(fail):\n    if fail:\n        raise LookupError('async host failed')\n    return 'awaited'\n"
-            ),
-            pyo3::ffi::c_str!("host_op.py"),
-            pyo3::ffi::c_str!("host_op"),
-        )?;
-        Ok(Some(
-            module
+        let awaitable = Python::attach(|py| -> PyResult<Py<PyAny>> {
+            let module = PyModule::from_code(
+                py,
+                pyo3::ffi::c_str!(
+                    "async def answer(fail):\n    if fail:\n        raise LookupError('async host failed')\n    return 'awaited'\n"
+                ),
+                pyo3::ffi::c_str!("host_op.py"),
+                pyo3::ffi::c_str!("host_op"),
+            )?;
+            Ok(module
                 .getattr("answer")?
                 .call1((matches!(self.op, OpScript::AwaitFailure),))?
-                .unbind(),
-        ))
-    }
-
-    fn resume_host_call(
-        &mut self,
-        py: Python<'_>,
-        result: PyResult<Py<PyAny>>,
-    ) -> Result<Option<Py<PyAny>>, InvokeError<Error>> {
-        let answer = result?.extract::<String>(py)?;
-        self.pending_reply.take().unwrap().send(answer);
-        Ok(None)
+                .unbind())
+        })?;
+        let answer = task.awaited(awaitable).await?;
+        reply.send(Python::attach(|py| answer.extract::<String>(py))?);
+        Ok(())
     }
 }
 
@@ -561,7 +540,7 @@ impl PythonBinding for StreamingBinding {
     type Failure = Classified;
 
     fn decode_request(
-        &mut self,
+        &self,
         _: Python<'_>,
         _: &Bound<'_, PyDict>,
     ) -> Result<(), InvokeError<Error>> {
@@ -569,7 +548,7 @@ impl PythonBinding for StreamingBinding {
     }
 
     fn encode_stream_head(
-        &mut self,
+        &self,
         py: Python<'_>,
         head: Vec<(&'static str, &'static str)>,
     ) -> PyResult<Py<PyAny>> {
@@ -582,11 +561,11 @@ impl PythonBinding for StreamingBinding {
         Ok(hidden.into_any().unbind())
     }
 
-    fn encode_chunk(&mut self, py: Python<'_>, chunk: &'static str) -> PyResult<Py<PyAny>> {
+    fn encode_chunk(&self, py: Python<'_>, chunk: &'static str) -> PyResult<Py<PyAny>> {
         Ok(PyString::new(py, chunk).into_any().unbind())
     }
 
-    fn encode_response(&mut self, py: Python<'_>, (): ()) -> PyResult<Py<PyAny>> {
+    fn encode_response(&self, py: Python<'_>, (): ()) -> PyResult<Py<PyAny>> {
         Ok(py.None())
     }
 
@@ -600,9 +579,9 @@ impl PythonBinding for StreamingBinding {
 }
 
 impl PythonHostCalls<Streaming> for StreamingBinding {
-    fn handle_host_call(
-        &mut self,
-        _: Python<'_>,
+    async fn handle_host_call(
+        &self,
+        _: &dyn Awaiter,
         op: std::convert::Infallible,
     ) -> Result<(), InvokeError<Error>> {
         match op {}

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

@@ -26,6 +26,7 @@ litellm-storage-clickhouse.workspace = true
 litellm-host.workspace = true
 bytes.workspace = true
 futures-util.workspace = true
+parking_lot.workspace = true
 litellm-cache.workspace = true
 litellm-cache-memory.workspace = true
 litellm-cache-redis.workspace = true

litellm-rust/crates/python-bridge/src/cache/python/host.rs

@@ -1,7 +1,9 @@
 use litellm_cache::Error;
 use litellm_host::protocol::Reply;
-use litellm_host_python::{from_py, to_py};
+use litellm_host_python::{Awaiter, from_py, to_py};
+use parking_lot::Mutex;
 use pyo3::{
+    exceptions::{PyException, PyRuntimeError},
     gc::{PyTraverseError, PyVisit},
     prelude::*,
     types::PyDict,
@@ -15,127 +17,116 @@ enum Pending {
     Store(Reply<Result<(), Error>>),
 }
 
+struct Configured {
+    cache: Py<PyAny>,
+    arguments: Py<PyDict>,
+}
+
 pub(crate) struct PythonCache {
-    cache: Option<Py<PyAny>>,
-    arguments: Option<Py<PyDict>>,
-    pending: Option<Pending>,
+    configured: Mutex<Option<Configured>>,
     asynchronous: bool,
 }
 
 impl PythonCache {
     pub fn new(asynchronous: bool) -> Self {
         Self {
-            cache: None,
-            arguments: None,
-            pending: None,
+            configured: Mutex::new(None),
             asynchronous,
         }
     }
 
-    pub(in crate::cache) fn bind(
-        &mut self,
-        cache: Bound<'_, PyAny>,
-        arguments: &Bound<'_, PyDict>,
-    ) {
-        self.cache = Some(cache.unbind());
-        self.arguments = Some(arguments.clone().unbind());
+    pub(in crate::cache) fn bind(&self, cache: Bound<'_, PyAny>, arguments: &Bound<'_, PyDict>) {
+        *self.configured.lock() = Some(Configured {
+            cache: cache.unbind(),
+            arguments: arguments.clone().unbind(),
+        });
     }
 
-    pub fn begin(&mut self, py: Python<'_>, call: CacheCall) -> PyResult<Option<Py<PyAny>>> {
-        let Some(cache) = self.cache.as_ref() else {
-            return Err(pyo3::exceptions::PyRuntimeError::new_err(
+    /// Runs one cache operation against the configured Python cache, awaiting it on the
+    /// caller's task for an async call. A cache failure is an answer to the route, a
+    /// cancellation ends the call.
+    pub async fn call(&self, task: &dyn Awaiter, call: CacheCall) -> PyResult<()> {
+        let (pending, invoked) = Python::attach(|py| self.invoke(py, call))?;
+        let result = match invoked {
+            Ok(awaitable) if self.asynchronous => task.awaited(awaitable).await,
+            invoked => invoked,
+        };
+        Python::attach(|py| answer(py, pending, result))
+    }
+
+    fn invoke(&self, py: Python<'_>, call: CacheCall) -> PyResult<(Pending, PyResult<Py<PyAny>>)> {
+        let configured = self.configured.lock();
+        let Some(configured) = configured.as_ref() else {
+            return Err(PyRuntimeError::new_err(
                 "cache operation without configured cache",
             ));
         };
-        let arguments = self
-            .arguments
-            .as_ref()
-            .ok_or_else(|| {
-                pyo3::exceptions::PyRuntimeError::new_err("cache arguments unavailable")
-            })?
-            .bind(py)
-            .copy()?;
-        let (method, result) = match call {
-            CacheCall::Lookup { reply } => {
-                self.pending = Some(Pending::Lookup(reply));
-                (
-                    if self.asynchronous {
+        let arguments = configured.arguments.bind(py).copy()?;
+        let cache = configured.cache.bind(py);
+        let (pending, invoked) = match (call, self.asynchronous) {
+            (CacheCall::Lookup { reply }, asynchronous) => (
+                Pending::Lookup(reply),
+                cache.call_method(
+                    if asynchronous {
                         "async_get_cache"
                     } else {
                         "get_cache"
                     },
-                    None,
-                )
-            }
-            CacheCall::Store { value, reply } => {
-                self.pending = Some(Pending::Store(reply));
-                (
-                    if self.asynchronous {
+                    (),
+                    Some(&arguments),
+                ),
+            ),
+            (CacheCall::Store { value, reply }, asynchronous) => (
+                Pending::Store(reply),
+                cache.call_method(
+                    if asynchronous {
                         "async_add_cache"
                     } else {
                         "add_cache"
                     },
-                    Some(to_py(py, &value)?),
-                )
-            }
+                    (to_py(py, &value)?,),
+                    Some(&arguments),
+                ),
+            ),
         };
-        let result = match result {
-            Some(value) => cache
-                .bind(py)
-                .call_method(method, (value,), Some(&arguments)),
-            None => cache.bind(py).call_method(method, (), Some(&arguments)),
-        }
-        .map(Bound::unbind);
-        if self.asynchronous && result.is_ok() {
-            return result.map(Some);
-        }
-        self.resume(py, result)
-    }
-
-    pub fn resume(
-        &mut self,
-        py: Python<'_>,
-        result: PyResult<Py<PyAny>>,
-    ) -> PyResult<Option<Py<PyAny>>> {
-        if let Err(error) = &result
-            && !error.is_instance_of::<pyo3::exceptions::PyException>(py)
-        {
-            self.pending = None;
-            return Err(result.err().unwrap());
-        }
-        match self.pending.take() {
-            Some(Pending::Lookup(reply)) => {
-                let value = result.map_err(|_| Error::Unavailable).and_then(|value| {
-                    if value.bind(py).is_none() {
-                        Ok(None)
-                    } else {
-                        from_py(value.bind(py))
-                            .map(Some)
-                            .map_err(|_| Error::InvalidEntry)
-                    }
-                });
-                reply.send(value);
-            }
-            Some(Pending::Store(reply)) => {
-                reply.send(result.map(|_| ()).map_err(|_| Error::Unavailable));
-            }
-            None => {
-                return Err(pyo3::exceptions::PyRuntimeError::new_err(
-                    "cache reply without pending operation",
-                ));
-            }
-        }
-        Ok(None)
+        Ok((pending, invoked.map(Bound::unbind)))
     }
 
     pub fn close(&mut self) {
-        self.pending = None;
-        self.cache = None;
-        self.arguments = None;
+        *self.configured.get_mut() = None;
     }
 
     pub fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
-        visit.call(&self.cache)?;
-        visit.call(&self.arguments)
+        let Some(configured) = self.configured.try_lock() else {
+            return Ok(());
+        };
+        let Some(configured) = configured.as_ref() else {
+            return Ok(());
+        };
+        visit.call(&configured.cache)?;
+        visit.call(&configured.arguments)
     }
 }
+
+fn answer(py: Python<'_>, pending: Pending, result: PyResult<Py<PyAny>>) -> PyResult<()> {
+    if let Err(error) = &result
+        && !error.is_instance_of::<PyException>(py)
+    {
+        return result.map(drop);
+    }
+    match pending {
+        Pending::Lookup(reply) => {
+            reply.send(result.map_err(|_| Error::Unavailable).and_then(|value| {
+                if value.bind(py).is_none() {
+                    Ok(None)
+                } else {
+                    from_py(value.bind(py))
+                        .map(Some)
+                        .map_err(|_| Error::InvalidEntry)
+                }
+            }))
+        }
+        Pending::Store(reply) => reply.send(result.map(drop).map_err(|_| Error::Unavailable)),
+    }
+    Ok(())
+}

litellm-rust/crates/python-bridge/src/cache/selection.rs

@@ -125,7 +125,7 @@ fn native_configuration(
 }
 
 pub(crate) fn configure(
-    python: &mut python::PythonCache,
+    python: &python::PythonCache,
     py: Python<'_>,
     arguments: &Bound<'_, PyDict>,
     call_type: &str,

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

@@ -1,7 +1,7 @@
 use std::convert::Infallible;
 
 use crate::routes::inference::InferenceHost;
-use litellm_host_python::{InvokeError, PythonBinding, PythonHostCalls, PythonOwned};
+use litellm_host_python::{Awaiter, InvokeError, PythonBinding, PythonHostCalls, PythonOwned};
 use litellm_inference_chat::{Error, route::ChatCompletions, types::ChatCompletionsCall};
 use pyo3::{
     gc::{PyTraverseError, PyVisit},
@@ -12,7 +12,7 @@ use pyo3::{
 pub(super) struct ChatCompletionsPythonHost(pub InferenceHost);
 
 pub(super) fn project(
-    host: &mut InferenceHost,
+    host: &InferenceHost,
     py: Python<'_>,
     arguments: &Bound<'_, PyDict>,
 ) -> PyResult<ChatCompletionsCall> {
@@ -34,11 +34,11 @@ impl PythonBinding for ChatCompletionsPythonHost {
     type Failure = PyErr;
 
     fn decode_request(
-        &mut self,
+        &self,
         py: Python<'_>,
         arguments: &Bound<'_, PyDict>,
     ) -> Result<ChatCompletionsCall, InvokeError<Error>> {
-        let call = project(&mut self.0, py, arguments).map_err(InvokeError::Python)?;
+        let call = project(&self.0, py, arguments).map_err(InvokeError::Python)?;
         if call
             .optional_params
             .get("stream")
@@ -52,7 +52,7 @@ impl PythonBinding for ChatCompletionsPythonHost {
     }
 
     fn encode_response(
-        &mut self,
+        &self,
         py: Python<'_>,
         response: <ChatCompletions as litellm_host::protocol::Protocol>::Response,
     ) -> PyResult<Py<PyAny>> {
@@ -60,7 +60,7 @@ impl PythonBinding for ChatCompletionsPythonHost {
     }
 
     fn encode_stream_head(
-        &mut self,
+        &self,
         _py: Python<'_>,
         head: <ChatCompletions as litellm_host::protocol::Protocol>::StreamHead,
     ) -> PyResult<Py<PyAny>> {
@@ -68,7 +68,7 @@ impl PythonBinding for ChatCompletionsPythonHost {
     }
 
     fn encode_chunk(
-        &mut self,
+        &self,
         _py: Python<'_>,
         chunk: <ChatCompletions as litellm_host::protocol::Protocol>::Chunk,
     ) -> PyResult<Py<PyAny>> {
@@ -84,9 +84,9 @@ impl PythonBinding for ChatCompletionsPythonHost {
 }
 
 impl PythonHostCalls<ChatCompletions> for ChatCompletionsPythonHost {
-    fn handle_host_call(
-        &mut self,
-        _: Python<'_>,
+    async fn handle_host_call(
+        &self,
+        _: &dyn Awaiter,
         op: Infallible,
     ) -> Result<(), InvokeError<Error>> {
         match op {}

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

@@ -44,7 +44,7 @@ impl InferenceHost {
     }
 
     pub fn project(
-        &mut self,
+        &self,
         py: Python<'_>,
         arguments: &Bound<'_, PyDict>,
         input: &str,

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

@@ -1,5 +1,5 @@
 use crate::cache::{CacheCall, Cached, PythonCache, Selection};
-use litellm_host_python::{PythonHostCalls, PythonOwned};
+use litellm_host_python::{Awaiter, PythonHostCalls, PythonOwned};
 
 use bytes::Bytes;
 use litellm_host_python::{InvokeError, PythonBinding, from_py, to_py};
@@ -261,16 +261,15 @@ impl PythonBinding for MessagesPythonHost {
     type Failure = PyErr;
 
     fn decode_request(
-        &mut self,
+        &self,
         py: Python<'_>,
         arguments: &Bound<'_, PyDict>,
     ) -> Result<(MessagesCall, Selection), InvokeError<Error>> {
         self.view
             .resolve(py, arguments)
             .map_err(InvokeError::Python)?;
-        let selection =
-            crate::cache::configure(&mut self.cache, py, arguments, "anthropic_messages")
-                .map_err(InvokeError::Python)?;
+        let selection = crate::cache::configure(&self.cache, py, arguments, "anthropic_messages")
+            .map_err(InvokeError::Python)?;
         self.projection(py, arguments)
             .map_err(|error| InvokeError::Python(self.map_failure(py, error)))?
             .map_err(InvokeError::Native)
@@ -278,25 +277,21 @@ impl PythonBinding for MessagesPythonHost {
     }
 
     fn encode_response(
-        &mut self,
+        &self,
         py: Python<'_>,
         response: Box<litellm_llms_types::formats::messages::MessagesResponse>,
     ) -> PyResult<Py<PyAny>> {
         public_response(py, ROUTE_HOST_MODULE, response.as_ref())
     }
 
-    fn encode_stream_head(
-        &mut self,
-        py: Python<'_>,
-        head: MessagesStreamHead,
-    ) -> PyResult<Py<PyAny>> {
+    fn encode_stream_head(&self, py: Python<'_>, head: MessagesStreamHead) -> PyResult<Py<PyAny>> {
         py.import(ROUTE_HOST_MODULE)?
             .getattr("stream_hidden_params")?
             .call1((to_py(py, &head.headers)?,))
             .map(Bound::unbind)
     }
 
-    fn encode_chunk(&mut self, py: Python<'_>, chunk: Bytes) -> PyResult<Py<PyAny>> {
+    fn encode_chunk(&self, py: Python<'_>, chunk: Bytes) -> PyResult<Py<PyAny>> {
         Ok(PyBytes::new(py, &chunk).into_any().unbind())
     }
 
@@ -315,31 +310,12 @@ impl PythonBinding for MessagesPythonHost {
 }
 
 impl PythonHostCalls<Cached<Messages>> for MessagesPythonHost {
-    fn handle_host_call(
-        &mut self,
-        py: Python<'_>,
+    async fn handle_host_call(
+        &self,
+        task: &dyn Awaiter,
         op: CacheCall,
     ) -> Result<(), InvokeError<Error>> {
-        self.cache
-            .begin(py, op)
-            .map(|_| ())
-            .map_err(InvokeError::Python)
-    }
-
-    fn begin_host_call(
-        &mut self,
-        py: Python<'_>,
-        op: CacheCall,
-    ) -> Result<Option<Py<PyAny>>, InvokeError<Error>> {
-        self.cache.begin(py, op).map_err(InvokeError::Python)
-    }
-
-    fn resume_host_call(
-        &mut self,
-        py: Python<'_>,
-        result: PyResult<Py<PyAny>>,
-    ) -> Result<Option<Py<PyAny>>, InvokeError<Error>> {
-        self.cache.resume(py, result).map_err(InvokeError::Python)
+        self.cache.call(task, op).await.map_err(InvokeError::Python)
     }
 }
 

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

@@ -15,6 +15,7 @@ use litellm_host_python::{Error as HostError, present};
 use litellm_host_python::{
     CallOptions, HookChain, PythonBinding, PythonCallHooks, PythonHostCalls, effective_py_args,
 };
+use parking_lot::Mutex;
 use pyo3::{
     gc::{PyTraverseError, PyVisit},
     prelude::*,
@@ -48,7 +49,7 @@ impl<'py> NativeCall<'py> {
 /// The keyword view a route host reads: the signature base under the caller's keywords,
 /// laid again once the hooks have rewritten them so a keyword a hook deleted stays deleted
 /// instead of falling back to the caller's value.
-pub(crate) struct RequestView(Option<RequestDicts>);
+pub(crate) struct RequestView(Mutex<Option<RequestDicts>>);
 
 struct RequestDicts {
     base: Py<PyDict>,
@@ -57,30 +58,25 @@ struct RequestDicts {
 
 impl RequestView {
     pub(crate) fn new(base: &Bound<'_, PyDict>, kwargs: &Bound<'_, PyDict>) -> PyResult<Self> {
-        Ok(Self(Some(RequestDicts {
+        Ok(Self(Mutex::new(Some(RequestDicts {
             base: base.clone().unbind(),
             request: effective_py_args(base, kwargs)?.unbind(),
-        })))
+        }))))
     }
 
-    fn dicts(&self) -> PyResult<&RequestDicts> {
-        self.0
-            .as_ref()
-            .ok_or_else(|| HostError::MissingState.into())
-    }
-
-    pub(crate) fn resolve(
-        &mut self,
-        py: Python<'_>,
-        arguments: &Bound<'_, PyDict>,
-    ) -> PyResult<()> {
-        let dicts = self.0.as_mut().ok_or(HostError::MissingState)?;
+    pub(crate) fn resolve(&self, py: Python<'_>, arguments: &Bound<'_, PyDict>) -> PyResult<()> {
+        let mut dicts = self.0.lock();
+        let dicts = dicts.as_mut().ok_or(HostError::MissingState)?;
         dicts.request = effective_py_args(dicts.base.bind(py), arguments)?.unbind();
         Ok(())
     }
 
-    pub(crate) fn request<'py>(&self, py: Python<'py>) -> PyResult<&Bound<'py, PyDict>> {
-        Ok(self.dicts()?.request.bind(py))
+    pub(crate) fn request<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyDict>> {
+        self.0
+            .lock()
+            .as_ref()
+            .map(|dicts| dicts.request.bind(py).clone())
+            .ok_or_else(|| HostError::MissingState.into())
     }
 
     pub(crate) fn argument<'py>(
@@ -89,15 +85,18 @@ impl RequestView {
         arguments: &Bound<'py, PyDict>,
         name: &str,
     ) -> PyResult<Option<Bound<'py, PyAny>>> {
-        present(arguments, self.request(py)?, name)
+        present(arguments, &self.request(py)?, name)
     }
 
     pub(crate) fn close(&mut self) {
-        self.0 = None;
+        *self.0.get_mut() = None;
     }
 
     pub(crate) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
-        let Some(dicts) = &self.0 else {
+        let Some(dicts) = self.0.try_lock() else {
+            return Ok(());
+        };
+        let Some(dicts) = dicts.as_ref() else {
             return Ok(());
         };
         visit.call(&dicts.base)?;
@@ -226,7 +225,7 @@ mod tests {
     ) {
         Python::initialize();
         Python::attach(|py| {
-            let mut view = super::RequestView::new(&dict(py, base), &dict(py, kwargs)).unwrap();
+            let view = super::RequestView::new(&dict(py, base), &dict(py, kwargs)).unwrap();
             let prepared = dict(py, prepared);
             view.resolve(py, &prepared).unwrap();
             let value = view

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

@@ -1,9 +1,10 @@
 use litellm_auth::ResolvedCredential;
-use litellm_host_python::{Error as HostError, InvokeError, PythonBinding};
+use litellm_host_python::{Awaiter, Error as HostError, InvokeError, PythonBinding};
 use litellm_host_python::{PythonHostCalls, PythonOwned};
 use litellm_inference_ocr::route::{Ocr, OcrCall, OcrOp};
 use litellm_llms::base_llm::ocr::error::Error;
 use litellm_llms_types::formats::ocr::LiteLLMOcrResponse;
+use parking_lot::Mutex;
 use pyo3::{
     exceptions::{PyBaseException, PyException},
     gc::{PyTraverseError, PyVisit},
@@ -28,40 +29,37 @@ enum OcrHostData {
 /// exception.
 pub(super) struct OcrPythonHost {
     view: RequestView,
-    data: OcrHostData,
+    data: Mutex<OcrHostData>,
 }
 
 impl OcrPythonHost {
     pub(super) fn new(view: RequestView) -> Self {
         Self {
             view,
-            data: OcrHostData::Unprojected,
-        }
-    }
-
-    fn handles(&self) -> PyResult<&OcrHostHandles> {
-        match &self.data {
-            OcrHostData::Projected(handles) => Ok(handles),
-            _ => Err(HostError::MissingState.into()),
+            data: Mutex::new(OcrHostData::Unprojected),
         }
     }
 
     fn acquire_azure_ad_token(&self, py: Python<'_>) -> PyResult<ResolvedCredential> {
-        self.handles()?
+        let OcrHostData::Projected(handles) = &*self.data.lock() else {
+            return Err(HostError::MissingState.into());
+        };
+        handles
             .azure_ad_token_provider
             .as_ref()
             .ok_or(HostError::MissingState)?
             .acquire(py)
     }
 
-    fn projection(&mut self, py: Python<'_>, arguments: &Bound<'_, PyDict>) -> PyResult<OcrCall> {
-        let OcrHostData::Unprojected = self.data else {
+    fn projection(&self, py: Python<'_>, arguments: &Bound<'_, PyDict>) -> PyResult<OcrCall> {
+        let mut data = self.data.lock();
+        let OcrHostData::Unprojected = *data else {
             return Err(HostError::MissingState.into());
         };
         self.view.resolve(py, arguments)?;
-        let (request, handles) = project_request(self.view.request(py)?, arguments)?;
+        let (request, handles) = project_request(&self.view.request(py)?, arguments)?;
         let caller_token = handles.azure_ad_token_provider.is_some();
-        self.data = OcrHostData::Projected(Box::new(handles));
+        *data = OcrHostData::Projected(Box::new(handles));
         Ok(OcrCall {
             request,
             caller_token,
@@ -72,7 +70,7 @@ impl OcrPythonHost {
         if !error.is_instance_of::<PyException>(py) || is_settings_error(py, &error) {
             return error;
         }
-        let provider = match &self.data {
+        let provider = match &*self.data.lock() {
             OcrHostData::Projected(handles) => handles.provider,
             _ => "",
         };
@@ -102,7 +100,7 @@ impl PythonBinding for OcrPythonHost {
     type Failure = PyErr;
 
     fn decode_request(
-        &mut self,
+        &self,
         py: Python<'_>,
         arguments: &Bound<'_, PyDict>,
     ) -> Result<OcrCall, InvokeError<Error>> {
@@ -110,27 +108,19 @@ impl PythonBinding for OcrPythonHost {
             .map_err(|error| InvokeError::Python(self.map_failure(py, error)))
     }
 
-    fn encode_response(
-        &mut self,
-        py: Python<'_>,
-        response: LiteLLMOcrResponse,
-    ) -> PyResult<Py<PyAny>> {
+    fn encode_response(&self, py: Python<'_>, response: LiteLLMOcrResponse) -> PyResult<Py<PyAny>> {
         public_response(py, "litellm.rust_bridge.ocr.route_host", &response)
     }
 
     fn encode_stream_head(
-        &mut self,
+        &self,
         _: Python<'_>,
         head: std::convert::Infallible,
     ) -> PyResult<Py<PyAny>> {
         match head {}
     }
 
-    fn encode_chunk(
-        &mut self,
-        _: Python<'_>,
-        chunk: std::convert::Infallible,
-    ) -> PyResult<Py<PyAny>> {
+    fn encode_chunk(&self, _: Python<'_>, chunk: std::convert::Infallible) -> PyResult<Py<PyAny>> {
         match chunk {}
     }
 
@@ -148,8 +138,8 @@ impl PythonBinding for OcrPythonHost {
     }
 }
 
-impl PythonHostCalls<Ocr> for OcrPythonHost {
-    fn handle_host_call(&mut self, py: Python<'_>, op: OcrOp) -> Result<(), InvokeError<Error>> {
+impl OcrPythonHost {
+    fn answer(&self, py: Python<'_>, op: OcrOp) -> Result<(), InvokeError<Error>> {
         match op {
             OcrOp::AcquireAzureAdToken(reply) => self
                 .acquire_azure_ad_token(py)
@@ -159,14 +149,23 @@ impl PythonHostCalls<Ocr> for OcrPythonHost {
     }
 }
 
+impl PythonHostCalls<Ocr> for OcrPythonHost {
+    async fn handle_host_call(&self, _: &dyn Awaiter, op: OcrOp) -> Result<(), InvokeError<Error>> {
+        Python::attach(|py| self.answer(py, op))
+    }
+}
+
 impl PythonOwned for OcrPythonHost {
     fn close(&mut self, _: Python<'_>) {
         self.view.close();
-        self.data = OcrHostData::Released;
+        *self.data.get_mut() = OcrHostData::Released;
     }
     fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
         self.view.traverse(visit)?;
-        if let OcrHostData::Projected(handles) = &self.data
+        let Some(data) = self.data.try_lock() else {
+            return Ok(());
+        };
+        if let OcrHostData::Projected(handles) = &*data
             && let Some(provider) = &handles.azure_ad_token_provider
         {
             provider.traverse(visit)?;
@@ -227,8 +226,7 @@ del provider
             drop(kwargs);
             let (reply, _) = litellm_host::protocol::reply();
             assert_eq!(
-                host.handle_host_call(py, OcrOp::AcquireAzureAdToken(reply))
-                    .is_ok(),
+                host.answer(py, OcrOp::AcquireAzureAdToken(reply)).is_ok(),
                 succeeds
             );
             let alive = || {

litellm-rust/crates/python-bridge/src/routes/responses/host.rs

@@ -1,7 +1,7 @@
 use std::convert::Infallible;
 
 use crate::routes::inference::InferenceHost;
-use litellm_host_python::{InvokeError, PythonBinding, PythonHostCalls, PythonOwned};
+use litellm_host_python::{Awaiter, InvokeError, PythonBinding, PythonHostCalls, PythonOwned};
 use litellm_inference_responses::{Error, route::Responses, types::ResponsesCall};
 use pyo3::{
     gc::{PyTraverseError, PyVisit},
@@ -12,7 +12,7 @@ use pyo3::{
 pub(super) struct ResponsesPythonHost(pub InferenceHost);
 
 pub(super) fn project(
-    host: &mut InferenceHost,
+    host: &InferenceHost,
     py: Python<'_>,
     arguments: &Bound<'_, PyDict>,
 ) -> PyResult<ResponsesCall> {
@@ -55,11 +55,11 @@ impl PythonBinding for ResponsesPythonHost {
     type Failure = PyErr;
 
     fn decode_request(
-        &mut self,
+        &self,
         py: Python<'_>,
         arguments: &Bound<'_, PyDict>,
     ) -> Result<ResponsesCall, InvokeError<Error>> {
-        let call = project(&mut self.0, py, arguments).map_err(InvokeError::Python)?;
+        let call = project(&self.0, py, arguments).map_err(InvokeError::Python)?;
         if call
             .optional_params
             .get("stream")
@@ -73,7 +73,7 @@ impl PythonBinding for ResponsesPythonHost {
     }
 
     fn encode_response(
-        &mut self,
+        &self,
         py: Python<'_>,
         response: <Responses as litellm_host::protocol::Protocol>::Response,
     ) -> PyResult<Py<PyAny>> {
@@ -81,7 +81,7 @@ impl PythonBinding for ResponsesPythonHost {
     }
 
     fn encode_stream_head(
-        &mut self,
+        &self,
         _py: Python<'_>,
         head: <Responses as litellm_host::protocol::Protocol>::StreamHead,
     ) -> PyResult<Py<PyAny>> {
@@ -92,7 +92,7 @@ impl PythonBinding for ResponsesPythonHost {
     }
 
     fn encode_chunk(
-        &mut self,
+        &self,
         _py: Python<'_>,
         chunk: <Responses as litellm_host::protocol::Protocol>::Chunk,
     ) -> PyResult<Py<PyAny>> {
@@ -111,9 +111,9 @@ impl PythonBinding for ResponsesPythonHost {
 }
 
 impl PythonHostCalls<Responses> for ResponsesPythonHost {
-    fn handle_host_call(
-        &mut self,
-        _: Python<'_>,
+    async fn handle_host_call(
+        &self,
+        _: &dyn Awaiter,
         op: Infallible,
     ) -> Result<(), InvokeError<Error>> {
         match op {}