fix: honor ContextCacheConfig on the LiteLLM path

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 970091291

src/google/adk/models/_prompt_cache.py

@@ -0,0 +1,82 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Shared reading of ContextCacheConfig for models that cache a marked prefix.
+
+Gemini caches by creating a server-side resource, which
+``GeminiContextCacheManager`` owns. Claude instead caches whatever prefix the
+request marks, and a model reached through LiteLLM inherits whichever of the
+two its provider implements. The parts of the configuration that mean the same
+thing for every prefix-marking model live here, so those callers cannot drift
+apart on what one configuration means.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+  from ..agents.context_cache_config import ContextCacheConfig
+  from .llm_request import LlmRequest
+
+logger = logging.getLogger("google_adk." + __name__)
+
+# The longest prefix cache a prefix-marking model offers is an hour, and it
+# costs more to write than the short-lived default. Only a configured lifetime
+# of at least an hour is worth that price.
+_ONE_HOUR_TTL_SECONDS = 3600
+
+
+def resolve_cache_config(
+    llm_request: LlmRequest,
+) -> ContextCacheConfig | None:
+  """Returns the cache config governing this request, or None to not cache.
+
+  Args:
+    llm_request: Request whose cache configuration is being resolved.
+
+  Returns:
+    The cache config to honor, or None when the request should not be cached.
+  """
+  cache_config = llm_request.cache_config
+  if cache_config is None:
+    return None
+
+  # ``min_tokens`` gates on the previous turn's measured prompt size, the same
+  # signal the Gemini path uses. That size is unknown on the first turn, where
+  # marking a prefix costs nothing beyond writing the cache.
+  previous_prompt_tokens = llm_request.cacheable_contents_token_count
+  if (
+      previous_prompt_tokens is not None
+      and previous_prompt_tokens < cache_config.min_tokens
+  ):
+    logger.debug(
+        "Skipping cache breakpoints: the previous prompt of %d tokens is below"
+        " the configured minimum of %d.",
+        previous_prompt_tokens,
+        cache_config.min_tokens,
+    )
+    return None
+
+  return cache_config
+
+
+def use_one_hour_ttl(cache_config: ContextCacheConfig) -> bool:
+  """Reports whether to ask for the hour-long cache instead of the default.
+
+  An hour is the longest a prefix cache is kept, so a configured lifetime
+  beyond that gets an hour rather than what it asked for.
+  """
+  return cache_config.ttl_seconds >= _ONE_HOUR_TTL_SECONDS

src/google/adk/models/anthropic_llm.py

@@ -48,6 +48,7 @@ from pydantic import Field
 from pydantic import model_validator
 from typing_extensions import override
 
+from . import _prompt_cache
 from ..utils import _json_utils
 from ..utils._google_client_headers import get_tracking_headers
 from .base_llm import BaseLlm
@@ -101,11 +102,6 @@ _RATE_LIMIT_POSSIBLE_FIX_MESSAGE = (
     "https://docs.anthropic.com/en/api/errors#http-errors"
 )
 
-# Claude offers exactly two cache lifetimes, five minutes and one hour, and
-# charges a higher write price for the longer one. Only a lifetime of at least
-# an hour is worth that price.
-_ONE_HOUR_CACHE_TTL_SECONDS = 3600
-
 # Claude rejects a cache breakpoint on a reasoning block.
 _UNCACHEABLE_BLOCK_TYPES = frozenset({"thinking", "redacted_thinking"})
 
@@ -814,49 +810,11 @@ def function_declaration_to_tool_param(
   )
 
 
-def _resolve_cache_config(
-    llm_request: LlmRequest,
-) -> ContextCacheConfig | None:
-  """Returns the cache config governing this request, or None to not cache.
-
-  Args:
-    llm_request: Request whose cache configuration is being resolved.
-
-  Returns:
-    The cache config to honor, or None when the request should not be cached.
-  """
-  cache_config = llm_request.cache_config
-  if cache_config is None:
-    return None
-
-  # ``min_tokens`` gates on the previous turn's measured prompt size, the same
-  # signal the Gemini path uses. That size is unknown on the first turn, where
-  # a breakpoint costs nothing beyond writing the cache.
-  previous_prompt_tokens = llm_request.cacheable_contents_token_count
-  if (
-      previous_prompt_tokens is not None
-      and previous_prompt_tokens < cache_config.min_tokens
-  ):
-    logger.debug(
-        "Skipping cache breakpoints: the previous prompt of %d tokens is below"
-        " the configured minimum of %d.",
-        previous_prompt_tokens,
-        cache_config.min_tokens,
-    )
-    return None
-
-  return cache_config
-
-
 def _to_cache_control(
     cache_config: ContextCacheConfig,
 ) -> anthropic_types.CacheControlEphemeralParam:
-  """Maps the configured cache lifetime onto one Claude actually offers.
-
-  An hour is the longest Claude keeps a cached prefix, so a longer configured
-  lifetime gets an hour rather than what it asked for.
-  """
-  if cache_config.ttl_seconds >= _ONE_HOUR_CACHE_TTL_SECONDS:
+  """Maps the configured cache lifetime onto one Claude actually offers."""
+  if _prompt_cache.use_one_hour_ttl(cache_config):
     return anthropic_types.CacheControlEphemeralParam(
         type="ephemeral", ttl="1h"
     )
@@ -1015,7 +973,7 @@ class AnthropicLlm(BaseLlm):
         system = system_str
 
     system_param: str | list[anthropic_types.TextBlockParam] | NotGiven = system
-    cache_config = _resolve_cache_config(llm_request)
+    cache_config = _prompt_cache.resolve_cache_config(llm_request)
     if cache_config is not None:
       system_param = _apply_cache_breakpoints(
           cache_config=cache_config,

src/google/adk/models/lite_llm.py

@@ -57,6 +57,7 @@ from typing_extensions import NotRequired
 from typing_extensions import override
 from typing_extensions import Required
 
+from . import _prompt_cache
 from ..utils._google_client_headers import merge_tracking_headers
 from ._capabilities import LlmCapabilities
 from .base_llm import BaseLlm
@@ -80,6 +81,8 @@ if TYPE_CHECKING:
   from litellm import ModelResponseStream
   from litellm import OpenAIMessageContent
   from litellm.types.utils import Delta
+
+  from ..agents.context_cache_config import ContextCacheConfig
 else:
   litellm = None
   acompletion = None
@@ -1075,6 +1078,38 @@ def _extract_cache_creation_tokens(usage: Any) -> Optional[int]:
   return None
 
 
+def _cache_control_injection_points(
+    cache_config: ContextCacheConfig,
+) -> List[Dict[str, Any]]:
+  """Describes the prefix LiteLLM should mark as cacheable.
+
+  LiteLLM applies these itself and then lets each provider decide what to do
+  with them, so the same two points are correct whatever the model turns out
+  to be: a provider that caches by marked prefix, such as Claude, honors them,
+  and a provider that caches automatically or not at all has them dropped
+  before the request leaves.
+
+  The system instruction is one point because it is the stable head of the
+  prompt. The final message is the other, which caches the conversation so far
+  and moves forward on its own as the conversation grows. Tool definitions get
+  no point of their own, because LiteLLM's only tool-level location is
+  specific to one provider.
+
+  Args:
+    cache_config: Cache configuration for the request.
+
+  Returns:
+    Injection points to hand to LiteLLM.
+  """
+  control: Dict[str, Any] = {"type": "ephemeral"}
+  if _prompt_cache.use_one_hour_ttl(cache_config):
+    control["ttl"] = "1h"
+  return [
+      {"location": "message", "role": "system", "control": control},
+      {"location": "message", "index": -1, "control": control},
+  ]
+
+
 def _decode_thought_signature(value: Any) -> Optional[bytes]:
   """Safely decodes a thought_signature value to bytes.
 
@@ -3082,6 +3117,18 @@ class LiteLlm(BaseLlm):
     }
     completion_args.update(self._additional_args)
 
+    # A caller who named their own injection points at construction has said
+    # more about their provider than the app-level config can, so leave those
+    # alone.
+    cache_config = _prompt_cache.resolve_cache_config(llm_request)
+    if (
+        cache_config is not None
+        and "cache_control_injection_points" not in completion_args
+    ):
+      completion_args["cache_control_injection_points"] = (
+          _cache_control_injection_points(cache_config)
+      )
+
     # merge headers
     if _is_litellm_vertex_model(effective_model) or _is_litellm_gemini_model(
         effective_model

tests/unittests/models/test_litellm.py

@@ -27,6 +27,7 @@ from unittest.mock import Mock
 from unittest.mock import patch
 import warnings
 
+from google.adk.agents.context_cache_config import ContextCacheConfig
 from google.adk.models.lite_llm import _aggregate_streaming_thought_parts
 from google.adk.models.lite_llm import _append_fallback_user_content_if_missing
 from google.adk.models.lite_llm import _BraceDepthTracker
@@ -7212,3 +7213,129 @@ async def test_generate_content_async_omits_tool_choice_when_functions_override(
   _, kwargs = mock_acompletion.call_args
   assert kwargs.get("tools") is None
   assert "tool_choice" not in kwargs
+
+
+def _cache_llm_request(cache_config=None):
+  return LlmRequest(
+      contents=[
+          types.Content(
+              role="user", parts=[types.Part.from_text(text="Cache this")]
+          )
+      ],
+      config=types.GenerateContentConfig(
+          system_instruction="You are a helpful assistant",
+      ),
+      cache_config=cache_config,
+  )
+
+
+def _injection_points(mock_acompletion):
+  _, kwargs = mock_acompletion.call_args
+  return kwargs.get("cache_control_injection_points")
+
+
+@pytest.mark.asyncio
+async def test_no_cache_config_sends_no_injection_points(
+    lite_llm_instance, mock_acompletion
+):
+  """Caching stays off unless the app configured it."""
+  async for _ in lite_llm_instance.generate_content_async(_cache_llm_request()):
+    pass
+
+  assert _injection_points(mock_acompletion) is None
+
+
+@pytest.mark.asyncio
+async def test_cache_config_marks_system_and_last_message(
+    lite_llm_instance, mock_acompletion
+):
+  """The stable head of the prompt and the conversation each get a point."""
+  async for _ in lite_llm_instance.generate_content_async(
+      _cache_llm_request(ContextCacheConfig())
+  ):
+    pass
+
+  assert _injection_points(mock_acompletion) == [
+      {
+          "location": "message",
+          "role": "system",
+          "control": {"type": "ephemeral"},
+      },
+      {
+          "location": "message",
+          "index": -1,
+          "control": {"type": "ephemeral"},
+      },
+  ]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "ttl_seconds,expected_control",
+    [
+        (300, {"type": "ephemeral"}),
+        (1800, {"type": "ephemeral"}),
+        (3599, {"type": "ephemeral"}),
+        (3600, {"type": "ephemeral", "ttl": "1h"}),
+        (86400, {"type": "ephemeral", "ttl": "1h"}),
+    ],
+)
+async def test_cache_ttl_maps_onto_an_offered_lifetime(
+    lite_llm_instance, mock_acompletion, ttl_seconds, expected_control
+):
+  """Five minutes or an hour; a shorter ask gets five minutes."""
+  async for _ in lite_llm_instance.generate_content_async(
+      _cache_llm_request(ContextCacheConfig(ttl_seconds=ttl_seconds))
+  ):
+    pass
+
+  points = _injection_points(mock_acompletion)
+  assert [point["control"] for point in points] == [expected_control] * 2
+
+
+@pytest.mark.asyncio
+async def test_cache_config_below_min_tokens_sends_no_injection_points(
+    lite_llm_instance, mock_acompletion
+):
+  """A prompt the app called too small to cache is sent unmarked."""
+  llm_request = _cache_llm_request(ContextCacheConfig(min_tokens=5000))
+  llm_request.cacheable_contents_token_count = 4999
+
+  async for _ in lite_llm_instance.generate_content_async(llm_request):
+    pass
+
+  assert _injection_points(mock_acompletion) is None
+
+
+@pytest.mark.asyncio
+async def test_cache_config_at_min_tokens_sends_injection_points(
+    lite_llm_instance, mock_acompletion
+):
+  """Reaching the configured minimum is enough to start caching."""
+  llm_request = _cache_llm_request(ContextCacheConfig(min_tokens=5000))
+  llm_request.cacheable_contents_token_count = 5000
+
+  async for _ in lite_llm_instance.generate_content_async(llm_request):
+    pass
+
+  assert len(_injection_points(mock_acompletion)) == 2
+
+
+@pytest.mark.asyncio
+async def test_injection_points_given_at_construction_are_kept(
+    mock_client, mock_acompletion
+):
+  """A caller who named their own points knows their provider better."""
+  chosen = [{"location": "tool_config"}]
+  lite_llm_instance = LiteLlm(
+      model="test_model",
+      llm_client=mock_client,
+      cache_control_injection_points=chosen,
+  )
+
+  async for _ in lite_llm_instance.generate_content_async(
+      _cache_llm_request(ContextCacheConfig())
+  ):
+    pass
+
+  assert _injection_points(mock_acompletion) == chosen

tests/unittests/models/test_prompt_cache.py

@@ -0,0 +1,77 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for the ContextCacheConfig reading shared by prefix-marking models."""
+
+from google.adk.agents.context_cache_config import ContextCacheConfig
+from google.adk.models._prompt_cache import resolve_cache_config
+from google.adk.models._prompt_cache import use_one_hour_ttl
+from google.adk.models.llm_request import LlmRequest
+import pytest
+
+
+def _request(cache_config=None, previous_prompt_tokens=None):
+  return LlmRequest(
+      model="test-model",
+      cache_config=cache_config,
+      cacheable_contents_token_count=previous_prompt_tokens,
+  )
+
+
+def test_no_cache_config_resolves_to_none():
+  assert resolve_cache_config(_request()) is None
+
+
+def test_no_cache_config_resolves_to_none_with_a_known_prompt_size():
+  """The size is read only after a config is known to exist."""
+  assert resolve_cache_config(_request(previous_prompt_tokens=10_000)) is None
+
+
+def test_cache_config_resolves_before_the_first_prompt_size_is_known():
+  cache_config = ContextCacheConfig(min_tokens=5000)
+
+  assert resolve_cache_config(_request(cache_config)) is cache_config
+
+
+@pytest.mark.parametrize(
+    "previous_prompt_tokens,expected",
+    [(0, None), (4999, None), (5000, "config"), (5001, "config")],
+)
+def test_min_tokens_gates_on_the_previous_prompt_size(
+    previous_prompt_tokens, expected
+):
+  cache_config = ContextCacheConfig(min_tokens=5000)
+
+  resolved = resolve_cache_config(
+      _request(cache_config, previous_prompt_tokens)
+  )
+
+  assert resolved is (cache_config if expected else None)
+
+
+def test_a_prompt_size_of_zero_is_a_size_not_an_absent_one():
+  """Zero is below any positive minimum, and is not "not measured yet"."""
+  cache_config = ContextCacheConfig(min_tokens=1)
+
+  assert resolve_cache_config(_request(cache_config, 0)) is None
+
+
+@pytest.mark.parametrize(
+    "ttl_seconds,expected",
+    [(1, False), (300, False), (3599, False), (3600, True), (86400, True)],
+)
+def test_only_an_hour_or_more_asks_for_the_long_cache(ttl_seconds, expected):
+  assert (
+      use_one_hour_ttl(ContextCacheConfig(ttl_seconds=ttl_seconds)) is expected
+  )