From 3da541aaf94d1119e9fd932cd10c82d40e6b6590 Mon Sep 17 00:00:00 2001
From: Aman Jaiswal <66757799+amanjaiswal73892@users.noreply.github.com>
Date: Fri, 25 Jul 2025 19:13:09 -0400
Subject: [PATCH] Fix: cache breakpoint handling in StructuredDiscussion class.

Anthropic requires maximum 4 cache breakpoints per request. This unsets cache breakpoints from earlier messages and only adds a single breakpoint in the flatten() method.
---
 src/agentlab/agents/tool_use_agent/tool_use_agent.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/agentlab/agents/tool_use_agent/tool_use_agent.py b/src/agentlab/agents/tool_use_agent/tool_use_agent.py
index b7494693..6ac61180 100644
--- a/src/agentlab/agents/tool_use_agent/tool_use_agent.py
+++ b/src/agentlab/agents/tool_use_agent/tool_use_agent.py
@@ -101,8 +101,10 @@ def flatten(self) -> list[MessageBuilder]:
                 messages.extend(group.messages)
             # Mark all summarized messages for caching
             if i == len(self.groups) - keep_last_n_obs:
-                if not isinstance(messages[i], ToolCalls):
-                    messages[i].mark_all_previous_msg_for_caching()
+                for msg in messages:  # unset previous cache breakpoints
+                    msg._cache_breakpoint = False
+                # set new cache breakpoint
+                messages[i].mark_all_previous_msg_for_caching()
         return messages
 
     def set_last_summary(self, summary: MessageBuilder):