From 13c1c2a364b2eb04c9d7aea7bcfe302a51495657 Mon Sep 17 00:00:00 2001 From: Nate Armstrong Date: Mon, 9 Nov 2026 16:24:21 -0800 Subject: [PATCH] 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 --- .../crates/host-python/src/binding.rs | 5 + litellm-rust/crates/host-python/src/driver.rs | 12 ++- .../crates/python-bridge/src/cache/mod.rs | 2 + .../python-bridge/src/cache/python/AGENTS.md | 2 +- .../python-bridge/src/cache/python/host.rs | 27 +++++- .../python-bridge/src/cache/python/service.rs | 8 +- .../python-bridge/src/cache/selection.rs | 3 +- .../python-bridge/src/routes/messages/host.rs | 4 + litellm/rust_bridge/host/cache.py | 44 ++++++++- .../cache/test_python_cache.py | 96 +++++++++++++++++-- tests/unit/rust_bridge/host/test_cache.py | 34 ++++++- 11 files changed, 220 insertions(+), 17 deletions(-) diff --git a/litellm-rust/crates/host-python/src/binding.rs b/litellm-rust/crates/host-python/src/binding.rs index e36b7d70c28..fe6173dd100 100644 --- a/litellm-rust/crates/host-python/src/binding.rs +++ b/litellm-rust/crates/host-python/src/binding.rs @@ -28,6 +28,11 @@ pub trait PythonBinding: PythonOwned { ) -> PyResult>; /// 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<'_>, diff --git a/litellm-rust/crates/host-python/src/driver.rs b/litellm-rust/crates/host-python/src/driver.rs index dbf843a5295..052156cbc5d 100644 --- a/litellm-rust/crates/host-python/src/driver.rs +++ b/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, } } diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index 4e3847951d8..fe9e06b3d23 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/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")] diff --git a/litellm-rust/crates/python-bridge/src/cache/python/AGENTS.md b/litellm-rust/crates/python-bridge/src/cache/python/AGENTS.md index 81a0997f02a..3b516251f43 100644 --- a/litellm-rust/crates/python-bridge/src/cache/python/AGENTS.md +++ b/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 diff --git a/litellm-rust/crates/python-bridge/src/cache/python/host.rs b/litellm-rust/crates/python-bridge/src/cache/python/host.rs index 82b7d881df6..a8cbd0ce607 100644 --- a/litellm-rust/crates/python-bridge/src/cache/python/host.rs +++ b/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, Error>>), @@ -21,6 +21,7 @@ pub(crate) struct PythonCache { arguments: Option>, pending: Option, 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; } diff --git a/litellm-rust/crates/python-bridge/src/cache/python/service.rs b/litellm-rust/crates/python-bridge/src/cache/python/service.rs index 38872b85d30..018c8647043 100644 --- a/litellm-rust/crates/python-bridge/src/cache/python/service.rs +++ b/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>, }, } @@ -101,8 +102,13 @@ where ) -> Pin> + 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)? }) diff --git a/litellm-rust/crates/python-bridge/src/cache/selection.rs b/litellm-rust/crates/python-bridge/src/cache/selection.rs index e96f079e5da..78c1606ea40 100644 --- a/litellm-rust/crates/python-bridge/src/cache/selection.rs +++ b/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

(std::marker::PhantomData

); diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs index e92daa67be4..2a5f79ca393 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs +++ b/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<'_>, diff --git a/litellm/rust_bridge/host/cache.py b/litellm/rust_bridge/host/cache.py index cac4b75677f..bf2655f9f62 100644 --- a/litellm/rust_bridge/host/cache.py +++ b/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 + ) + ) diff --git a/tests/test_litellm_rust/cache/test_python_cache.py b/tests/test_litellm_rust/cache/test_python_cache.py index d26fe47d69d..d88feec7d5e 100644 --- a/tests/test_litellm_rust/cache/test_python_cache.py +++ b/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( diff --git a/tests/unit/rust_bridge/host/test_cache.py b/tests/unit/rust_bridge/host/test_cache.py index 9d532631ef2..c2b87cca1dc 100644 --- a/tests/unit/rust_bridge/host/test_cache.py +++ b/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