Enable prompt caching on AnthropicAgent (#23627)

* Add cache breakpoints to anthropic_client

* Fix little bugs

* Move imports only used for type checking to if TYPE_CHECKING

Author: luisorofino

ddev/src/ddev/ai/agent/anthropic_client.py

@@ -2,19 +2,35 @@
 # All rights reserved
 # Licensed under a 3-clause BSD style license (see LICENSE)
 
-from typing import Final
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Final, overload
 
 import anthropic
-from anthropic.types import MessageParam, ToolParam, ToolResultBlockParam
+from anthropic.types import MessageParam
 
 from ddev.ai.agent.base import BaseAgent
 from ddev.ai.agent.exceptions import AgentAPIError, AgentConnectionError, AgentError, AgentRateLimitError
 from ddev.ai.agent.types import AgentResponse, ContextUsage, StopReason, TokenUsage, ToolCall, ToolResultMessage
 from ddev.ai.tools.registry import ToolRegistry
 
+if TYPE_CHECKING:
+    from anthropic.types import (
+        CacheControlEphemeralParam,
+        TextBlockParam,
+        ToolParam,
+        ToolResultBlockParam,
+    )
+
 DEFAULT_MODEL: Final[str] = "claude-sonnet-4-6"
 DEFAULT_MAX_TOKENS: Final[int] = 8192  # max tokens per response
 
+# 1h TTL for the static prefix (system + tools): paid once, read for the whole session.
+STATIC_CACHE_CONTROL: Final[CacheControlEphemeralParam] = {"type": "ephemeral", "ttl": "1h"}
+# Default TTL (currently 5 min, but Anthropic may change it) for the sliding breakpoint
+# on the last user message: re-written each turn, so a longer TTL would be wasted.
+SLIDING_CACHE_CONTROL: Final[CacheControlEphemeralParam] = {"type": "ephemeral"}
+
 
 class AnthropicAgent(BaseAgent[MessageParam]):
     """A wrapper around the Anthropic API that provides a simple interface for interacting with agents."""
@@ -102,6 +118,32 @@ def _to_tool_result_params(self, messages: list[ToolResultMessage]) -> list[Tool
             for msg in messages
         ]
 
+    @overload
+    @staticmethod
+    def _with_user_cache_breakpoint(content: str) -> list[TextBlockParam]: ...
+
+    @overload
+    @staticmethod
+    def _with_user_cache_breakpoint(content: list[ToolResultBlockParam]) -> list[ToolResultBlockParam]: ...
+
+    @staticmethod
+    def _with_user_cache_breakpoint(
+        content: str | list[ToolResultBlockParam],
+    ) -> list[TextBlockParam] | list[ToolResultBlockParam]:
+        """Return a block list with a sliding cache breakpoint on the last block."""
+        if isinstance(content, str):
+            return [{"type": "text", "text": content, "cache_control": SLIDING_CACHE_CONTROL}]
+        if not content:
+            return []
+        return [*content[:-1], {**content[-1], "cache_control": SLIDING_CACHE_CONTROL}]
+
+    @staticmethod
+    def _with_tools_cache_breakpoint(tool_defs: list[ToolParam]) -> list[ToolParam]:
+        """Return a tool list with a static cache breakpoint on the last tool."""
+        if not tool_defs:
+            return tool_defs
+        return [*tool_defs[:-1], {**tool_defs[-1], "cache_control": STATIC_CACHE_CONTROL}]
+
     async def send(
         self,
         content: str | list[ToolResultMessage],
@@ -114,19 +156,27 @@ async def send(
         Returns:
             An AgentResponse object containing the response from the agent.
         """
-        tool_defs = self._get_tool_definitions(allowed_tools)
+        tool_defs = self._with_tools_cache_breakpoint(self._get_tool_definitions(allowed_tools))
 
         api_content: str | list[ToolResultBlockParam] = (
             self._to_tool_result_params(content) if isinstance(content, list) else content
         )
-        user_msg: MessageParam = {"role": "user", "content": api_content}
-        messages = [*self._history, user_msg]
+        user_msg_for_history: MessageParam = {"role": "user", "content": api_content}
+        user_msg_for_request: MessageParam = {
+            "role": "user",
+            "content": self._with_user_cache_breakpoint(api_content),
+        }
+        messages = [*self._history, user_msg_for_request]
+
+        system_param: list[TextBlockParam] = [
+            {"type": "text", "text": self._system_prompt, "cache_control": STATIC_CACHE_CONTROL}
+        ]
 
         try:
             response = await self._client.messages.create(
                 model=self._model,
                 max_tokens=self._max_tokens,
-                system=self._system_prompt,
+                system=system_param,
                 messages=messages,
                 tools=tool_defs if tool_defs else anthropic.NOT_GIVEN,
             )
@@ -174,7 +224,9 @@ async def send(
             usage=usage,
         )
 
-        # Save to history only after a successful response.
-        self._history.extend([user_msg, {"role": "assistant", "content": response.content}])
+        # Save to history only after a successful response. Use the unmarked form so the
+        # cache_control breakpoint is only ever on the latest user message — this keeps the
+        # request below the 4-marker limit regardless of conversation length.
+        self._history.extend([user_msg_for_history, {"role": "assistant", "content": response.content}])
 
         return agent_response

ddev/tests/ai/agent/test_anthropic_client.py

@@ -485,3 +485,151 @@ async def test_error_mid_conversation_leaves_history_unchanged() -> None:
         await agent.send("Second message")
 
     assert agent.history == history_after_first
+
+
+# ---------------------------------------------------------------------------
+# Prompt caching: static breakpoints (system + last tool, 1h TTL)
+# ---------------------------------------------------------------------------
+
+
+async def test_system_prompt_sent_as_block_with_static_cache_control() -> None:
+    resp = make_response("end_turn", [make_text_block("ok")])
+    agent, create_mock = make_agent(mock_response=resp)
+
+    await agent.send("Hi")
+
+    assert create_mock.call_args.kwargs["system"] == [
+        {
+            "type": "text",
+            "text": "You are helpful.",
+            "cache_control": {"type": "ephemeral", "ttl": "1h"},
+        }
+    ]
+
+
+@pytest.mark.parametrize(
+    "tool_names",
+    [["only"], ["a", "b"], ["a", "b", "c", "d"]],
+    ids=["single_tool", "two_tools", "four_tools"],
+)
+async def test_only_last_tool_carries_static_cache_control(tool_names: list[str]) -> None:
+    registry = ToolRegistry([FakeTool(n) for n in tool_names])
+    resp = make_response("end_turn", [make_text_block("ok")])
+    agent, create_mock = make_agent(tools=registry, mock_response=resp)
+
+    await agent.send("Hi")
+
+    sent_tools = create_mock.call_args.kwargs["tools"]
+    assert all("cache_control" not in t for t in sent_tools[:-1])
+    assert sent_tools[-1]["cache_control"] == {"type": "ephemeral", "ttl": "1h"}
+
+
+async def test_allowed_tools_subset_places_cache_control_on_last_of_subset() -> None:
+    registry = ToolRegistry([FakeTool(n) for n in ["a", "b", "c"]])
+    resp = make_response("end_turn", [make_text_block("ok")])
+    agent, create_mock = make_agent(tools=registry, mock_response=resp)
+
+    await agent.send("Hi", allowed_tools=["a", "b"])
+
+    sent_tools = create_mock.call_args.kwargs["tools"]
+    assert [t["name"] for t in sent_tools] == ["a", "b"]
+    assert "cache_control" not in sent_tools[0]
+    assert sent_tools[-1]["cache_control"] == {"type": "ephemeral", "ttl": "1h"}
+
+
+# ---------------------------------------------------------------------------
+# Prompt caching: sliding breakpoint on the last user message block (default TTL)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "content",
+    [
+        pytest.param("Hi there", id="str"),
+        pytest.param(
+            [ToolResultMessage(tool_call_id="t1", result=ToolResult(success=True, data="r1"))],
+            id="single_tool_result",
+        ),
+        pytest.param(
+            [ToolResultMessage(tool_call_id=f"t{i}", result=ToolResult(success=True, data=f"r{i}")) for i in range(3)],
+            id="multiple_tool_results",
+        ),
+    ],
+)
+async def test_sliding_cache_control_on_last_user_block_only(
+    content: str | list[ToolResultMessage],
+) -> None:
+    resp = make_response("end_turn", [make_text_block("ok")])
+    agent, create_mock = make_agent(mock_response=resp)
+
+    await agent.send(content)
+
+    blocks = create_mock.call_args.kwargs["messages"][-1]["content"]
+    assert isinstance(blocks, list) and blocks
+    assert all("cache_control" not in b for b in blocks[:-1])
+    assert blocks[-1]["cache_control"] == {"type": "ephemeral"}
+    assert "ttl" not in blocks[-1]["cache_control"]
+
+
+async def test_empty_tool_results_produces_empty_content_block() -> None:
+    resp = make_response("end_turn", [make_text_block("ok")])
+    agent, create_mock = make_agent(mock_response=resp)
+
+    await agent.send([])
+
+    blocks = create_mock.call_args.kwargs["messages"][-1]["content"]
+    assert blocks == []
+
+
+# ---------------------------------------------------------------------------
+# Prompt caching: history must not retain cache_control markers
+# (otherwise multi-turn requests would exceed the 4-marker limit)
+# ---------------------------------------------------------------------------
+
+
+async def test_history_str_message_keeps_raw_str_form() -> None:
+    resp = make_response("end_turn", [make_text_block("ok")])
+    agent, _ = make_agent(mock_response=resp)
+
+    await agent.send("Hi")
+
+    assert agent.history[0] == {"role": "user", "content": "Hi"}
+
+
+async def test_history_tool_result_blocks_have_no_cache_control() -> None:
+    resp = make_response("end_turn", [make_text_block("ok")])
+    agent, _ = make_agent(mock_response=resp)
+
+    tool_results = [
+        ToolResultMessage(tool_call_id=f"t{i}", result=ToolResult(success=True, data=f"r{i}")) for i in range(2)
+    ]
+    await agent.send(tool_results)
+
+    blocks = agent.history[0]["content"]
+    assert all("cache_control" not in b for b in blocks)
+
+
+async def test_multi_turn_only_latest_user_message_in_request_has_cache_control() -> None:
+    first_resp = make_response("tool_use", [make_tool_use_block(id="t1")])
+    second_resp = make_response("end_turn", [make_text_block("done")])
+
+    client = MagicMock(spec=anthropic.AsyncAnthropic)
+    client.messages = MagicMock()
+    client.messages.create = AsyncMock(side_effect=[first_resp, second_resp])
+    client.models = MagicMock()
+    client.models.retrieve = AsyncMock(return_value=SimpleNamespace(max_input_tokens=FAKE_CONTEXT_WINDOW))
+    agent = AnthropicAgent(client=client, tools=ToolRegistry([]), system_prompt="sp", name="t")
+
+    await agent.send("First")
+    await agent.send([ToolResultMessage(tool_call_id="t1", result=ToolResult(success=True, data="r"))])
+
+    first_call_messages = client.messages.create.call_args_list[0].kwargs["messages"]
+    assert first_call_messages[-1]["content"] == [
+        {"type": "text", "text": "First", "cache_control": {"type": "ephemeral"}}
+    ]
+
+    second_call_messages = client.messages.create.call_args_list[1].kwargs["messages"]
+    assert second_call_messages[0] == {"role": "user", "content": "First"}
+    latest_blocks = second_call_messages[-1]["content"]
+    assert all("cache_control" not in b for b in latest_blocks[:-1])
+    assert latest_blocks[-1]["cache_control"] == {"type": "ephemeral"}