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

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

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

@@ -28,6 +28,11 @@ pub trait PythonBinding: PythonOwned {
     ) -> PyResult<Py<PyAny>>;
 
     /// What the stream carries at hand-off, as the caller's stream receives it.
+    /// The unary response after the deployment transform, before success hooks run.
+    fn deliver_response(&mut self, _py: Python<'_>, _response: &Bound<'_, PyAny>) -> PyResult<()> {
+        Ok(())
+    }
+
     fn encode_stream_head(
         &mut self,
         py: Python<'_>,

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

@@ -399,8 +399,16 @@ where
             Ok(step) => step,
             Err(error) => return self.failure(error, FailureOrigin::Host).await,
         };
-        match self.settle(step).await {
-            Ok(response) => self.succeeded(response).await,
+        let response = match self.settle(step).await {
+            Ok(response) => response,
+            Err(error) => return self.failure(error, FailureOrigin::Host).await,
+        };
+        let delivered = {
+            let mut parts = self.parts();
+            Python::attach(|py| parts.binding.deliver_response(py, response.bind(py)))
+        };
+        match delivered {
+            Ok(()) => self.succeeded(response).await,
             Err(error) => self.failure(error, FailureOrigin::Host).await,
         }
     }

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 deployment transform has returned the response the caller receives 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

@@ -244,6 +244,10 @@ impl PythonBinding for MessagesPythonHost {
         public_response(py, ROUTE_HOST_MODULE, &response)
     }
 
+    fn deliver_response(&mut self, py: Python<'_>, response: &Bound<'_, PyAny>) -> PyResult<()> {
+        self.cache.store_response(py, response)
+    }
+
     fn encode_stream_head(
         &mut self,
         py: Python<'_>,

litellm/rust_bridge/host/cache.py

@@ -1,12 +1,54 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from collections.abc import Mapping
+from typing import TYPE_CHECKING, Final
+
+from pydantic import TypeAdapter
+
+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"
+_HEADERS: Final = TypeAdapter(dict[str, object])
+
 
 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.validate_python(headers).items()
+                if name != _RUST_RESPONSE_HEADER
+            },
+        },
+    }
+
+
+def store_response(
+    cache: Cache,
+    response: Mapping[str, object],
+    /,
+    **kwargs: object,  # kwargs-ok: the cache call's own arguments
+) -> None:
+    from litellm.caching.caching_handler import create_cache_write_task
+
+    create_cache_write_task(
+        lambda: cache.async_add_cache(  # pyright: ignore[reportUnknownMemberType]  # Cache.async_add_cache leaves its result untyped
+            _python_entry(response),
+            **kwargs,  # pyright: ignore[reportArgumentType]  # the cache call's own keyword arguments, as Python passes them
+        )
+    )

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,
@@ -19,6 +23,8 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelReq
 from litellm.proxy.utils import InternalUsageCache
 from litellm.redis.stores.rate_limits import RateLimitScope
 from litellm.rust_bridge.catalog import Route
+from litellm.types.llms.anthropic import AnthropicResponseContentBlockText
+from litellm.types.utils import CallTypes, LLMResponseTypes
 from tests.test_litellm_rust.cache.support import cache_key, collect, invoke, payload
 from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_EVENTS, MESSAGES_MODEL, MESSAGES_RESPONSE
 from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging, request_body
@@ -251,9 +257,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
@@ -278,14 +286,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):
@@ -295,13 +308,82 @@ 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, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    from litellm.caching.in_memory_cache import InMemoryCache
+
+    storage: Final = InMemoryCache()
+    monkeypatch.setattr(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("native", (False, True), ids=("python", "rust"))
+async def test_messages_cache_stores_the_response_a_yielding_deployment_hook_returned(
+    recording_server: RecordingServer, native: bool, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    from litellm.caching.in_memory_cache import InMemoryCache
+    from litellm.integrations.custom_logger import CustomLogger
+
+    blocks: Final = TypeAdapter(list[AnthropicResponseContentBlockText])
+
+    class Redact(CustomLogger):
+        async def async_post_call_success_deployment_hook(
+            self, request_data: dict[str, object], response: LLMResponseTypes, call_type: CallTypes | None
+        ) -> LLMResponseTypes | None:
+            await asyncio.sleep(0.05)
+            assert isinstance(response, dict)
+            return {**response, "content": [{"type": "text", "text": "redacted"}]}  # pyright: ignore[reportReturnType]  # messages hooks see and return the plain dict the route serves
+
+    monkeypatch.setenv("LITELLM_RUST", "1" if native else "0")
+    storage: Final = InMemoryCache()
+    monkeypatch.setattr(litellm, "cache", Cache(_backend=storage))
+    monkeypatch.setattr(litellm, "callbacks", [Redact()])
+    first: Final = await invoke("messages", recording_server, {}, native=native)
+    assert isinstance(first, dict)
+    assert [block.text for block in blocks.validate_python(first["content"])] == ["redacted"]
+    await asyncio.gather(*tuple(_PENDING_CACHE_WRITES))
+    hit: Final = await invoke("messages", recording_server, {}, native=native)
+    assert isinstance(hit, dict)
+    assert [block.text for block in blocks.validate_python(hit["content"])] == ["redacted"]
+    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(

tests/unit/rust_bridge/host/test_cache.py

@@ -1,10 +1,12 @@
+import asyncio
 from typing import Final
 
 import pytest
 
 import litellm
 from litellm.caching.caching import Cache
-from litellm.rust_bridge.host.cache import configured_cache
+from litellm.caching.caching_handler import _PENDING_CACHE_WRITES
+from litellm.rust_bridge.host.cache import configured_cache, store_response
 from litellm.types.caching import LiteLLMCacheType
 
 
@@ -19,3 +21,33 @@ def test_configured_cache_preserves_identity_and_observes_runtime_replacement(mo
 
     monkeypatch.setattr(litellm, "cache", None)
     assert configured_cache() is None
+
+
+_REQUEST: Final = {"model": "claude-sonnet-4-5", "messages": [{"role": "user", "content": "hi"}]}
+
+
+async def _stored(response: dict[str, object]) -> object:
+    cache: Final = Cache(type=LiteLLMCacheType.LOCAL)
+    store_response(cache, response, **_REQUEST)
+    await asyncio.gather(*_PENDING_CACHE_WRITES)
+    return await cache.async_get_cache(**_REQUEST)
+
+
+@pytest.mark.asyncio
+async def test_store_response_drops_only_the_rust_marker_header_from_the_cached_entry() -> None:
+    stored: Final = await _stored(
+        {
+            "id": "msg_1",
+            "_hidden_params": {"additional_headers": {"x-litellm-rust": "true", "llm_provider-x": "1"}, "model": "m"},
+        }
+    )
+    assert stored == {
+        "id": "msg_1",
+        "_hidden_params": {"additional_headers": {"llm_provider-x": "1"}, "model": "m"},
+    }
+
+
+@pytest.mark.asyncio
+async def test_store_response_keeps_an_entry_without_the_rust_marker_unchanged() -> None:
+    response: Final = {"id": "msg_2", "_hidden_params": {"additional_headers": {"llm_provider-x": "1"}}}
+    assert await _stored(response) == response