From 08f348bebd52d51deb96fb026d43d908621a1a98 Mon Sep 17 00:00:00 2001
From: nuernber <benjamin.nuernberger@jpl.nasa.gov>
Date: Fri, 11 Sep 2026 12:19:08 -0700
Subject: [PATCH] feat(prompt-caching): match affinity TTL to cache control TTL

Add logic to extract TTL from cache_control blocks and use 1 hour (3600s) for "1h" TTL or default 5 minutes (300s) otherwise. Apply the extracted TTL when storing model affinity in both sync and async add_model_id methods.
---
 litellm/router_utils/prompt_caching_cache.py  | 29 ++++++++++++-
 .../test_prompt_caching_deployment_check.py   | 41 +++++++++++++++++++
 2 files changed, 68 insertions(+), 2 deletions(-)

diff --git a/litellm/router_utils/prompt_caching_cache.py b/litellm/router_utils/prompt_caching_cache.py
index 39708e168f5..171a086de81 100644
--- a/litellm/router_utils/prompt_caching_cache.py
+++ b/litellm/router_utils/prompt_caching_cache.py
@@ -139,6 +139,27 @@ class PromptCachingCache:
 
         return cacheable_prefix
 
+    @staticmethod
+    def get_prompt_caching_ttl(messages: list[AllMessageValues] | None) -> int:
+        if messages is None:
+            return 300
+
+        cacheable_prefix: Final = PromptCachingCache.extract_cacheable_prefix(messages)
+        cache_control_ttls: Final = tuple(
+            cache_control.get("ttl")
+            for message in cacheable_prefix
+            for cache_control in (
+                message.get("cache_control"),
+                *(
+                    content_block.get("cache_control")
+                    for content_block in (message.get("content") if isinstance(message.get("content"), list) else ())
+                    if isinstance(content_block, dict)
+                ),
+            )
+            if isinstance(cache_control, dict) and cache_control.get("type") == "ephemeral"
+        )
+        return 3600 if "1h" in cache_control_ttls else 300
+
     @staticmethod
     def get_prompt_caching_cache_key(
         messages: list[AllMessageValues] | None,
@@ -189,7 +210,11 @@ class PromptCachingCache:
         if cache_key is None:
             return
 
-        self.cache.set_cache(cache_key, PromptCachingCacheValue(model_id=model_id), ttl=300)
+        self.cache.set_cache(
+            cache_key,
+            PromptCachingCacheValue(model_id=model_id),
+            ttl=PromptCachingCache.get_prompt_caching_ttl(messages),
+        )
         return
 
     async def async_add_model_id(
@@ -209,7 +234,7 @@ class PromptCachingCache:
         await self.cache.async_set_cache(
             cache_key,
             PromptCachingCacheValue(model_id=model_id),
-            ttl=300,  # store for 5 minutes
+            ttl=PromptCachingCache.get_prompt_caching_ttl(messages),
         )
         return
 
diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py
index 333e7b2ff31..fde779b3feb 100644
--- a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py
+++ b/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py
@@ -1,6 +1,7 @@
 import asyncio
 import copy
 from typing import List, cast
+from unittest.mock import AsyncMock
 
 import pytest
 
@@ -60,6 +61,46 @@ def _messages(word_count: int) -> List[AllMessageValues]:
     )
 
 
+@pytest.mark.parametrize(
+    ("ttl", "expected_affinity_ttl"),
+    ((None, 300), ("5m", 300), ("1h", 3600)),
+)
+def test_prompt_caching_affinity_ttl_matches_cache_control(ttl: str | None, expected_affinity_ttl: int):
+    cache_control: dict[str, str] = {"type": "ephemeral", **({"ttl": ttl} if ttl is not None else {})}
+    messages = cast(
+        List[AllMessageValues],
+        [{"role": "system", "content": [{"type": "text", "text": "cached", "cache_control": cache_control}]}],
+    )
+
+    assert PromptCachingCache.get_prompt_caching_ttl(messages) == expected_affinity_ttl
+
+
+@pytest.mark.asyncio
+async def test_async_add_model_id_uses_one_hour_affinity_ttl():
+    cache = DualCache()
+    async_set_cache = AsyncMock()
+    cache.async_set_cache = async_set_cache
+    messages = cast(
+        List[AllMessageValues],
+        [
+            {
+                "role": "system",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": "cached",
+                        "cache_control": {"type": "ephemeral", "ttl": "1h"},
+                    }
+                ],
+            }
+        ],
+    )
+
+    await PromptCachingCache(cache=cache).async_add_model_id("dep-1", messages, None)
+
+    assert async_set_cache.call_args.kwargs["ttl"] == 3600
+
+
 def test_get_min_token_count_for_deployments_takes_min_across_mixed_group():
     """
     A group may legally mix models whose real minimums differ, and one gate decides for every