fix(messages): write the native unary cache entry after the router stamps it like python

Python stores a unary messages response from a background cache write task that runs after the router stamps _hidden_params, so the entry carries model_id and the router headers and a later hit returns x-litellm-cache-key. The native path stored the entry inside the native call, before the router saw the response. The async unary store is now deferred until the route encodes its public response and is handed to Python's own cache write task.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013FDwtkaNTifWAP4W7JZavE

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

@@ -6,6 +6,8 @@ pub(crate) use selection::{Cached, Selection, configure, ensure_cache_disabled};
 
 use pyo3::{exceptions::PyRuntimeError, prelude::*};
 
+const CACHE_HOST_MODULE: &str = "litellm.rust_bridge.host.cache";
+
 #[derive(Debug, thiserror::Error)]
 pub(crate) enum HandleError {
     #[error("this inference route does not support configured caches")]

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

@@ -4,6 +4,6 @@ This directory lets Rust inference use a selected Python cache. `service.rs` imp
 
 Keep the shared inference protocol wrapper and adapter selection in the parent module. Receive the configured cache and prepared arguments from the parent module. Do not discover global configuration, choose native backends, or move inference to Python. Core remains independent of Python objects and cache implementation details
 
-Await asynchronous cache operations through the existing host driver in the caller's task. Do not create another asyncio task or event loop. Cancellation must prevent subsequent provider requests and cache writes. Preserve ordinary cache failure handling without swallowing cancellation or other Python base exceptions
+Await asynchronous cache operations through the existing host driver in the caller's task. Do not create another asyncio task or event loop. The one exception is the asynchronous unary messages store: Python's `_async_set_cache` writes the returned response from `create_cache_write_task` after the router has stamped its `_hidden_params`, so `host.rs` defers that store until the route encodes its public response and hands that response to `litellm.rust_bridge.host.cache.store_response`, which schedules the same write task. Lookups, sync stores and stream stores stay in the caller's task. Cancellation must prevent subsequent provider requests and cache writes. Preserve ordinary cache failure handling without swallowing cancellation or other Python base exceptions
 
 Traverse retained Python references for GC and release pending replies when the call closes. Regression tests must require native inference with Python fallback disabled and assert observable hits, provider request counts, cache-key headers, task identity and cancellation

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

@@ -9,7 +9,7 @@ use pyo3::{
 use serde_json::Value;
 
 use super::service::CacheCall;
-use crate::cache::HandleError;
+use crate::cache::{CACHE_HOST_MODULE, HandleError};
 
 enum Pending {
     Lookup(Reply<Result<Option<Value>, Error>>),
@@ -21,6 +21,7 @@ pub(crate) struct PythonCache {
     arguments: Option<Py<PyDict>>,
     pending: Option<Pending>,
     asynchronous: bool,
+    deferred: bool,
 }
 
 impl PythonCache {
@@ -30,6 +31,7 @@ impl PythonCache {
             arguments: None,
             pending: None,
             asynchronous,
+            deferred: false,
         }
     }
 
@@ -64,7 +66,14 @@ impl PythonCache {
                     None,
                 )
             }
-            CacheCall::Store { value, reply } => {
+            CacheCall::Store {
+                unary: true, reply, ..
+            } if self.asynchronous => {
+                self.deferred = true;
+                reply.send(Ok(()));
+                return Ok(None);
+            }
+            CacheCall::Store { value, reply, .. } => {
                 self.pending = Some(Pending::Store(reply));
                 (
                     if self.asynchronous {
@@ -89,6 +98,19 @@ impl PythonCache {
         self.resume(py, result)
     }
 
+    pub fn store_response(&mut self, py: Python<'_>, response: &Bound<'_, PyAny>) -> PyResult<()> {
+        if !std::mem::take(&mut self.deferred) {
+            return Ok(());
+        }
+        let (Some(cache), Some(arguments)) = (self.cache.as_ref(), self.arguments.as_ref()) else {
+            return Err(HandleError::Unconfigured.into());
+        };
+        py.import(CACHE_HOST_MODULE)?
+            .getattr("store_response")?
+            .call((cache, response), Some(&arguments.bind(py).copy()?))?;
+        Ok(())
+    }
+
     pub fn resume(
         &mut self,
         py: Python<'_>,
@@ -124,6 +146,7 @@ impl PythonCache {
 
     pub fn close(&mut self) {
         self.pending = None;
+        self.deferred = false;
         self.cache = None;
         self.arguments = None;
     }

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

@@ -45,6 +45,7 @@ pub(crate) enum CacheCall {
     },
     Store {
         value: Value,
+        unary: bool,
         reply: Reply<Result<(), Error>>,
     },
 }
@@ -101,8 +102,13 @@ where
     ) -> Pin<Box<dyn Future<Output = Result<(), Error>> + Send + 'a>> {
         Box::pin(async move {
             let value = to_python(value)?;
+            let unary = value.get(STREAM_EVENTS_KEY).is_none();
             self.services
-                .call(|reply| CacheCall::Store { value, reply })
+                .call(|reply| CacheCall::Store {
+                    value,
+                    unary,
+                    reply,
+                })
                 .await
                 .map_err(Error::unavailable)?
         })

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

@@ -1,11 +1,10 @@
-use super::{HandleError, python};
+use super::{CACHE_HOST_MODULE, HandleError, python};
 use litellm_cache_response::{CacheOptions, CachePolicy, CacheScope, ScopedCache};
 use litellm_host::{
     machine::{HostServices, MachineFault},
     protocol::Protocol,
 };
 use pyo3::{prelude::*, types::PyDict};
-const CACHE_HOST_MODULE: &str = "litellm.rust_bridge.host.cache";
 
 pub(crate) struct Cached<P>(std::marker::PhantomData<P>);
 

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

@@ -241,7 +241,9 @@ impl PythonBinding for MessagesPythonHost {
         py: Python<'_>,
         response: Box<litellm_llms_types::formats::messages::MessagesResponse>,
     ) -> PyResult<Py<PyAny>> {
-        public_response(py, ROUTE_HOST_MODULE, response.as_ref())
+        let response = public_response(py, ROUTE_HOST_MODULE, response.as_ref())?;
+        self.cache.store_response(py, response.bind(py))?;
+        Ok(response)
     }
 
     fn encode_stream_head(

litellm/rust_bridge/host/cache.py

@@ -1,12 +1,39 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from collections.abc import Mapping
+from typing import TYPE_CHECKING, Final
+
+from litellm.litellm_core_utils.add_retry_fallback_headers import get_hidden_params_dict
 
 if TYPE_CHECKING:
     from litellm.caching.caching import Cache
 
+_RUST_RESPONSE_HEADER: Final = "x-litellm-rust"
+
 
 def configured_cache() -> Cache | None:
     import litellm
 
     return litellm.cache
+
+
+def _python_entry(response: Mapping[str, object]) -> Mapping[str, object]:
+    hidden: Final = get_hidden_params_dict(response)
+    headers: Final = hidden.get("additional_headers")
+    if not isinstance(headers, Mapping) or _RUST_RESPONSE_HEADER not in headers:
+        return response
+    return {
+        **response,
+        "_hidden_params": {
+            **hidden,
+            "additional_headers": {name: value for name, value in headers.items() if name != _RUST_RESPONSE_HEADER},
+        },
+    }
+
+
+def store_response(
+    cache: Cache, response: Mapping[str, object], /, **kwargs: object
+) -> None:  # kwargs-ok: the cache call's own arguments
+    from litellm.caching.caching_handler import create_cache_write_task
+
+    create_cache_write_task(lambda: cache.async_add_cache(_python_entry(response), **kwargs))

tests/test_litellm_rust/cache/test_python_cache.py

@@ -9,7 +9,11 @@ from litellm.caching.caching import Cache, CacheMode
 from litellm.caching.caching_handler import (
     _PENDING_CACHE_WRITES,  # pyright: ignore[reportPrivateUsage]  # await the existing background cache writer before the next request
 )
-from litellm.litellm_core_utils.add_retry_fallback_headers import get_hidden_params_dict
+from litellm.litellm_core_utils.add_retry_fallback_headers import (
+    _add_headers_to_response,  # pyright: ignore[reportPrivateUsage]  # stamp the response the way the router does
+    add_retry_headers_to_response,
+    get_hidden_params_dict,
+)
 from litellm.proxy._types import Litellm_EntityType
 from litellm.proxy.hooks.model_max_budget_limiter import (
     PROXY_VirtualKeyModelMaxBudgetLimiter,
@@ -250,9 +254,11 @@ async def test_rust_responses_rejects_a_configured_python_cache(
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("cancel_lookup", (False, True))
-async def test_python_cache_operations_stay_in_the_rust_callers_task(
+@pytest.mark.parametrize("stream", (False, True), ids=("response", "stream"))
+async def test_python_cache_operations_stay_in_the_rust_callers_task_except_the_unary_store(
     recording_server: RecordingServer,
     cancel_lookup: bool,
+    stream: bool,
 ) -> None:
     from litellm.caching.base_cache import BaseCache
     from litellm.caching.in_memory_cache import InMemoryCache
@@ -277,14 +283,19 @@ async def test_python_cache_operations_stay_in_the_rust_callers_task(
             return storage.get_cache(key, **kwargs)
 
         async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None:
-            assert asyncio.current_task() is caller
+            writer: Final = asyncio.current_task()
+            if stream:
+                assert writer is caller
+            else:
+                assert writer is not caller
+                assert writer in _PENDING_CACHE_WRITES
             await asyncio.sleep(0)
             storage.set_cache(key, value, **kwargs)
 
     litellm.cache = Cache(_backend=CallerCache())
     if cancel_lookup:
         recording_server.expected_requests = 0
-        task: Final = asyncio.create_task(invoke("messages", recording_server, {}))
+        task: Final = asyncio.create_task(invoke("messages", recording_server, {"stream": stream}))
         await asyncio.wait_for(entered.wait(), timeout=5)
         task.cancel()
         with pytest.raises(asyncio.CancelledError):
@@ -294,13 +305,50 @@ async def test_python_cache_operations_stay_in_the_rust_callers_task(
         assert len(recording_server.requests) == 0
         assert storage.cache_dict == {}
         return
-    first: Final = await invoke("messages", recording_server, {})
-    second: Final = await invoke("messages", recording_server, {})
-    assert payload(first) == payload(second)
+    first: Final = await invoke("messages", recording_server, {"stream": stream})
+    first_payload: Final = await collect(first) if stream else payload(first)
+    await asyncio.gather(*tuple(_PENDING_CACHE_WRITES))
+    second: Final = await invoke("messages", recording_server, {"stream": stream})
+    second_payload: Final = await collect(second) if stream else payload(second)
+    assert first_payload == second_payload
     assert cache_key(second)
     assert len(recording_server.requests) == 1
 
 
+@pytest.mark.asyncio
+async def test_rust_unary_messages_store_keeps_the_hidden_params_stamped_after_the_call(
+    recording_server: RecordingServer,
+) -> None:
+    from litellm.caching.in_memory_cache import InMemoryCache
+
+    storage: Final = InMemoryCache()
+    litellm.cache = Cache(_backend=storage)
+    first: Final = await invoke("messages", recording_server, {})
+    assert isinstance(first, dict)
+    assert storage.cache_dict == {}
+    add_retry_headers_to_response(
+        _add_headers_to_response(first, {"x-litellm-model-group": "group"}), attempted_retries=0
+    )
+    get_hidden_params_dict(first)["model_id"] = "deployment"
+    await asyncio.gather(*tuple(_PENDING_CACHE_WRITES))
+    (entry,) = storage.cache_dict.values()
+    stored: Final = TypeAdapter(dict[str, object]).validate_python(
+        TypeAdapter(dict[str, object]).validate_python(entry)["response"]
+    )
+    assert stored["_hidden_params"] == {
+        "additional_headers": {"x-litellm-model-group": "group", "x-litellm-attempted-retries": 0},
+        "model_id": "deployment",
+    }
+    assert get_hidden_params_dict(first)["additional_headers"] == {
+        "x-litellm-rust": "true",
+        "x-litellm-model-group": "group",
+        "x-litellm-attempted-retries": 0,
+    }
+    hit: Final = await invoke("messages", recording_server, {})
+    assert get_hidden_params_dict(hit)["cache_key"] == cache_key(hit)
+    assert len(recording_server.requests) == 1
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("namespace_source", ("cache", "metadata"))
 async def test_rust_messages_legacy_cache_honors_request_namespaces(