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>

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)
     )

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,