fix(memory): preserve client cache boundaries across tool rounds

litellm/litellm_core_utils/prompt_templates/server_tools.py

@@ -151,6 +151,22 @@ def trailing_system_messages(data: Mapping[str, object], route: ServerToolRoute)
     )
 
 
+def uncached_system_directive(message: Mapping[str, object]) -> Mapping[str, object]:
+    return {  # mutable-ok: Provider wire format requires native JSON containers.
+        **{key: value for key, value in message.items() if key != "cache_control"},
+        **(
+            {  # mutable-ok: Provider wire format requires native JSON containers.
+                "content": [  # mutable-ok: Provider wire format requires native JSON containers.
+                    {key: value for key, value in _OBJECT.validate_python(block).items() if key != "cache_control"}
+                    for block in _items(message.get("content"))
+                ],
+            }
+            if isinstance(message.get("content"), (list, tuple))
+            else {}  # mutable-ok: Provider wire format requires native JSON containers.
+        ),
+    }
+
+
 def append_server_reference(data: Mapping[str, object], route: ServerToolRoute, reference: str) -> Mapping[str, object]:
     field: Final = "input" if route == "aresponses" else "messages"
     messages: Final = _items(data.get(field))

litellm/proxy/memory/continuation.py

@@ -177,7 +177,20 @@ class MemoryContinuations:
             if patch.replaces > index + 1 or patch.replaces < 1:
                 raise HTTPException(status_code=409, detail="Invalid memory continuation")
             prefix: Final = result[: -(patch.replaces - 1)] if patch.replaces > 1 else result
-            return _append_items(prefix, patch.replacement, self.route)
+            # Clients move cache breakpoints between turns. Reuse their current
+            # directives rather than restoring an obsolete cached copy.
+            directives: Final = MappingProxyType(
+                {
+                    prefix_hashes((item,), self.route)[0]: item
+                    for item in items[index + 1 - patch.replaces : index + 1]
+                    if item.get("role") == "system"
+                }
+            )
+            replacement: Final = tuple(
+                directives.get(prefix_hashes((item,), self.route)[0], item) if item.get("role") == "system" else item
+                for item in patch.replacement
+            )
+            return _append_items(prefix, replacement, self.route)
 
         return reduce(apply, range(len(items)), ())
 

litellm/proxy/memory/gateway.py

@@ -24,6 +24,7 @@ from litellm.litellm_core_utils.prompt_templates.server_tools import (
     continue_server_tools,
     inject_server_tools,
     trailing_system_messages,
+    uncached_system_directive,
 )
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.common_utils.sse_keepalive import wrap_passthrough_sse_bytes_with_keepalive_pings
@@ -147,7 +148,10 @@ class GatewayMemoryLoop:
             **self.data,
             **(
                 {  # mutable-ok: Native provider JSON containers.
-                    "messages": [*messages, *directives],  # mutable-ok: Provider request JSON.
+                    "messages": [  # mutable-ok: Provider request JSON.
+                        *messages,
+                        *(uncached_system_directive(directive) for directive in directives),
+                    ],
                 }
                 if directives and messages[-len(directives) :] != directives
                 else {}  # mutable-ok: Native provider JSON containers.

tests/test_litellm/proxy/memory/test_memory_v2_boundaries.py

@@ -524,6 +524,7 @@ async def test_trailing_system_messages_survive_client_tool_continuation(prisma_
         body = await incoming.json()
         observed.append(body)
         messages = body["messages"]
+        assert json.dumps(body).count('"cache_control"') == 4
         assert all(
             message["role"] != "system" or messages[index + 1]["role"] == "assistant"
             for index, message in enumerate(messages[:-1])
@@ -540,8 +541,16 @@ async def test_trailing_system_messages_survive_client_tool_continuation(prisma_
         "role": "user",
         "content": [{"type": "text", "text": "Read README.md", "cache_control": {"type": "ephemeral"}}],
     }
-    directive = {"role": "system", "content": "Use concise answers"}
+    directive = {
+        "role": "system",
+        "content": [{"type": "text", "text": "Use concise answers", "cache_control": {"type": "ephemeral"}}],
+    }
+    earlier_directive = {"role": "system", "content": [{"type": "text", "text": "Use concise answers"}]}
     original = {
+        "system": [
+            {"type": "text", "text": "Cached prefix " + str(i), "cache_control": {"type": "ephemeral"}}
+            for i in range(2)
+        ],
         "messages": [prefix, directive],
         "tools": [{"name": "Read", "input_schema": {"type": "object"}}],
         "tool_choice": {"type": "tool", "name": "Read"},
@@ -558,7 +567,7 @@ async def test_trailing_system_messages_survive_client_tool_continuation(prisma_
         "tool_choice": {"type": "none"},
         "messages": [
             prefix,
-            directive,
+            earlier_directive,
             {"role": "assistant", "content": [client_call]},
             {
                 "role": "user",
@@ -572,7 +581,8 @@ async def test_trailing_system_messages_survive_client_tool_continuation(prisma_
         pass
     assert len(observed) == 2
     assert observed[0]["messages"][0] == observed[1]["messages"][0] == prefix
-    assert observed[1]["messages"].count(directive) == 2
+    assert observed[1]["messages"].count(directive) == 1
+    assert observed[1]["messages"].count(earlier_directive) == 1
     assert observed[1]["messages"].count({"role": "assistant", "content": [client_call]}) == 1
     assert observed[1]["messages"][-1] == directive
     assert observed[1]["messages"][-3]["content"][0]["tool_use_id"] == "client_read"
@@ -580,16 +590,32 @@ async def test_trailing_system_messages_survive_client_tool_continuation(prisma_
 
 
 @pytest.mark.asyncio
-async def test_claude_output_directives_reach_search_answer_and_reflection_rounds(prisma_edge: MagicMock) -> None:
+@pytest.mark.parametrize("cached_directive", [False, True])
+async def test_claude_output_directives_reach_search_answer_and_reflection_rounds(
+    prisma_edge: MagicMock, cached_directive: bool
+) -> None:
     provider = FastAPI()
     observed = []
-    directive = {"role": "system", "content": [], "output_config": {"effort": "low"}}
+    directive = {
+        "role": "system",
+        "content": [{"type": "text", "text": "Reply briefly", "cache_control": {"type": "ephemeral"}}]
+        if cached_directive
+        else [],
+        "output_config": {"effort": "low"},
+    }
 
     @provider.post("/v1/messages")
     async def model(incoming: Request):
         body = await incoming.json()
         observed.append(body)
-        assert body["messages"][-1] == directive
+        assert json.dumps(body).count('"cache_control"') == 3 + int(cached_directive)
+        last = body["messages"][-1]
+        assert last["role"] == "system" and last["output_config"] == {"effort": "low"}
+        assert last["content"] == (
+            [{"type": "text", "text": "Reply briefly"}]
+            if cached_directive and len(observed) > 1
+            else directive["content"]
+        )
         if len(observed) == 1:
             content = [
                 {"type": "tool_use", "id": "search", "name": "litellm_memory_search", "input": {"query": "demo"}}
@@ -606,7 +632,19 @@ async def test_claude_output_directives_reach_search_answer_and_reflection_round
             "content": content,
         }
 
-    original = {"messages": [{"role": "user", "content": "My demo port?"}, directive]}
+    original = {
+        "system": [
+            {"type": "text", "text": "Cached prefix " + str(i), "cache_control": {"type": "ephemeral"}}
+            for i in range(2)
+        ],
+        "messages": [
+            {
+                "role": "user",
+                "content": [{"type": "text", "text": "My demo port?", "cache_control": {"type": "ephemeral"}}],
+            },
+            directive,
+        ],
+    }
     loop = GatewayMemoryLoop(provider, request(), original, "anthropic_messages", store(prisma_edge))
     async for _ in loop.run():
         pass