From f70f01b8fa671a1a5b5f40f540dae7bf1ff4247c Mon Sep 17 00:00:00 2001
From: mateo <277851410+mateo-berri@users.noreply.github.com>
Date: Wed, 4 Nov 2026 16:28:54 +0000
Subject: [PATCH] fix(vertex_ai): cache the whole conversation prefix up to the
 cache_control block so Gemini sees turns in order

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 .../context_caching/transformation.py         |  9 +++--
 .../test_vertex_ai_context_caching.py         | 40 +++++++++++++++++++
 2 files changed, 45 insertions(+), 4 deletions(-)

diff --git a/litellm/llms/vertex_ai/context_caching/transformation.py b/litellm/llms/vertex_ai/context_caching/transformation.py
index a4601af944b..6d0b6d45436 100644
--- a/litellm/llms/vertex_ai/context_caching/transformation.py
+++ b/litellm/llms/vertex_ai/context_caching/transformation.py
@@ -110,16 +110,17 @@ def separate_cached_messages(
     messages: list[AllMessageValues],
 ) -> tuple[list[AllMessageValues], list[AllMessageValues]]:
     """
-    Returns the first continuous block of cached messages, plus every system message once
-    any message is cached, since generateContent rejects a system_instruction beside a
-    cachedContent, and the remaining messages.
+    Returns the conversation prefix ending at the first continuous block of cached messages,
+    plus every system message once any message is cached, since generateContent rejects a
+    system_instruction beside a cachedContent, and the remaining messages. Gemini reads
+    cachedContent as a prefix of contents, so the cache must hold everything before the block.
     """
     filtered_messages: Final = [(idx, message) for idx, message in enumerate(messages) if is_cached_message(message)]
     if not filtered_messages:
         return [], messages
 
     last_continuous_block_idx: Final = get_first_continuous_block_idx(filtered_messages)
-    cached_block: Final = range(filtered_messages[0][0], filtered_messages[last_continuous_block_idx][0] + 1)
+    cached_block: Final = range(filtered_messages[last_continuous_block_idx][0] + 1)
     goes_in_cache: Final = tuple(
         idx in cached_block or message.get("role") == "system" for idx, message in enumerate(messages)
     )
diff --git a/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py b/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py
index bc8421c80f8..aaa22a06766 100644
--- a/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py
+++ b/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py
@@ -1790,6 +1790,46 @@ class TestContextCachingEndpoints:
         self.mock_async_client.post.assert_not_called()
 
 
+def _marked(role: Literal["system", "user"], text: str) -> dict:
+    return {"role": role, "content": [{"type": "text", "text": text, "cache_control": {"type": "ephemeral"}}]}
+
+
+@pytest.mark.parametrize(
+    "messages",
+    [
+        pytest.param(
+            [
+                {"role": "system", "content": "sys"},
+                {"role": "user", "content": "Q1"},
+                {"role": "assistant", "content": "A1"},
+                _marked("user", "Q2"),
+            ],
+            id="marked-latest-turn",
+        ),
+        pytest.param(
+            [
+                {"role": "user", "content": "Q1"},
+                {"role": "assistant", "content": "A1"},
+                _marked("user", "Q2"),
+                {"role": "assistant", "content": "A2"},
+                {"role": "user", "content": "Q3"},
+            ],
+            id="marked-middle-turn",
+        ),
+        pytest.param(
+            [_marked("system", "sys"), {"role": "user", "content": "Q1"}],
+            id="marked-system",
+        ),
+    ],
+)
+def test_separate_cached_messages_keeps_conversation_order(messages):
+    from litellm.llms.vertex_ai.context_caching.transformation import separate_cached_messages
+
+    cached_messages, non_cached_messages = separate_cached_messages(messages)
+
+    assert cached_messages + non_cached_messages == messages
+
+
 def test_cached_messages_end_on_supported_turn():
     from litellm.llms.vertex_ai.context_caching.transformation import (
         cached_messages_end_on_supported_turn,