fix: preserve tool_use blocks in summary message during condensing with native tools (#9582)

This commit is contained in:
Daniel 2025-11-25 13:17:26 -05:00 committed by GitHub
parent 774b492eff
commit 311940b085
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 399 additions and 4 deletions

View file

@ -7,7 +7,12 @@ import { TelemetryService } from "@roo-code/telemetry"
import { ApiHandler } from "../../../api"
import { ApiMessage } from "../../task-persistence/apiMessages"
import { maybeRemoveImageBlocks } from "../../../api/transform/image-cleaning"
import { summarizeConversation, getMessagesSinceLastSummary, N_MESSAGES_TO_KEEP } from "../index"
import {
summarizeConversation,
getMessagesSinceLastSummary,
getKeepMessagesWithToolBlocks,
N_MESSAGES_TO_KEEP,
} from "../index"
vi.mock("../../../api/transform/image-cleaning", () => ({
maybeRemoveImageBlocks: vi.fn((messages: ApiMessage[], _apiHandler: ApiHandler) => [...messages]),
@ -24,6 +29,222 @@ vi.mock("@roo-code/telemetry", () => ({
const taskId = "test-task-id"
const DEFAULT_PREV_CONTEXT_TOKENS = 1000
describe("getKeepMessagesWithToolBlocks", () => {
it("should return keepMessages without tool blocks when no tool_result blocks in first kept message", () => {
const messages: ApiMessage[] = [
{ role: "user", content: "Hello", ts: 1 },
{ role: "assistant", content: "Hi there", ts: 2 },
{ role: "user", content: "How are you?", ts: 3 },
{ role: "assistant", content: "I'm good", ts: 4 },
{ role: "user", content: "What's new?", ts: 5 },
]
const result = getKeepMessagesWithToolBlocks(messages, 3)
expect(result.keepMessages).toHaveLength(3)
expect(result.toolUseBlocksToPreserve).toHaveLength(0)
})
it("should return all messages when messages.length <= keepCount", () => {
const messages: ApiMessage[] = [
{ role: "user", content: "Hello", ts: 1 },
{ role: "assistant", content: "Hi there", ts: 2 },
]
const result = getKeepMessagesWithToolBlocks(messages, 3)
expect(result.keepMessages).toEqual(messages)
expect(result.toolUseBlocksToPreserve).toHaveLength(0)
})
it("should preserve tool_use blocks when first kept message has tool_result blocks", () => {
const toolUseBlock = {
type: "tool_use" as const,
id: "toolu_123",
name: "read_file",
input: { path: "test.txt" },
}
const toolResultBlock = {
type: "tool_result" as const,
tool_use_id: "toolu_123",
content: "file contents",
}
const messages: ApiMessage[] = [
{ role: "user", content: "Hello", ts: 1 },
{ role: "assistant", content: "Let me read that file", ts: 2 },
{ role: "user", content: "Please continue", ts: 3 },
{
role: "assistant",
content: [{ type: "text" as const, text: "Reading file..." }, toolUseBlock],
ts: 4,
},
{
role: "user",
content: [toolResultBlock, { type: "text" as const, text: "Continue" }],
ts: 5,
},
{ role: "assistant", content: "Got it, the file says...", ts: 6 },
{ role: "user", content: "Thanks", ts: 7 },
]
const result = getKeepMessagesWithToolBlocks(messages, 3)
// keepMessages should be the last 3 messages
expect(result.keepMessages).toHaveLength(3)
expect(result.keepMessages[0].ts).toBe(5)
expect(result.keepMessages[1].ts).toBe(6)
expect(result.keepMessages[2].ts).toBe(7)
// Should preserve the tool_use block from the preceding assistant message
expect(result.toolUseBlocksToPreserve).toHaveLength(1)
expect(result.toolUseBlocksToPreserve[0]).toEqual(toolUseBlock)
})
it("should not preserve tool_use blocks when first kept message is assistant role", () => {
const toolUseBlock = {
type: "tool_use" as const,
id: "toolu_123",
name: "read_file",
input: { path: "test.txt" },
}
const messages: ApiMessage[] = [
{ role: "user", content: "Hello", ts: 1 },
{ role: "assistant", content: "Hi there", ts: 2 },
{ role: "user", content: "Please read", ts: 3 },
{
role: "assistant",
content: [{ type: "text" as const, text: "Reading..." }, toolUseBlock],
ts: 4,
},
{ role: "user", content: "Continue", ts: 5 },
{ role: "assistant", content: "Done", ts: 6 },
]
const result = getKeepMessagesWithToolBlocks(messages, 3)
// First kept message is assistant, not user with tool_result
expect(result.keepMessages).toHaveLength(3)
expect(result.keepMessages[0].role).toBe("assistant")
expect(result.toolUseBlocksToPreserve).toHaveLength(0)
})
it("should not preserve tool_use blocks when first kept user message has string content", () => {
const messages: ApiMessage[] = [
{ role: "user", content: "Hello", ts: 1 },
{ role: "assistant", content: "Hi there", ts: 2 },
{ role: "user", content: "How are you?", ts: 3 },
{ role: "assistant", content: "Good", ts: 4 },
{ role: "user", content: "Simple text message", ts: 5 }, // String content, not array
{ role: "assistant", content: "Response", ts: 6 },
{ role: "user", content: "More text", ts: 7 },
]
const result = getKeepMessagesWithToolBlocks(messages, 3)
expect(result.keepMessages).toHaveLength(3)
expect(result.toolUseBlocksToPreserve).toHaveLength(0)
})
it("should handle multiple tool_use blocks that need to be preserved", () => {
const toolUseBlock1 = {
type: "tool_use" as const,
id: "toolu_123",
name: "read_file",
input: { path: "file1.txt" },
}
const toolUseBlock2 = {
type: "tool_use" as const,
id: "toolu_456",
name: "read_file",
input: { path: "file2.txt" },
}
const toolResultBlock1 = {
type: "tool_result" as const,
tool_use_id: "toolu_123",
content: "contents 1",
}
const toolResultBlock2 = {
type: "tool_result" as const,
tool_use_id: "toolu_456",
content: "contents 2",
}
const messages: ApiMessage[] = [
{ role: "user", content: "Hello", ts: 1 },
{
role: "assistant",
content: [{ type: "text" as const, text: "Reading files..." }, toolUseBlock1, toolUseBlock2],
ts: 2,
},
{
role: "user",
content: [toolResultBlock1, toolResultBlock2],
ts: 3,
},
{ role: "assistant", content: "Got both files", ts: 4 },
{ role: "user", content: "Thanks", ts: 5 },
]
const result = getKeepMessagesWithToolBlocks(messages, 3)
// Should preserve both tool_use blocks
expect(result.toolUseBlocksToPreserve).toHaveLength(2)
expect(result.toolUseBlocksToPreserve).toContainEqual(toolUseBlock1)
expect(result.toolUseBlocksToPreserve).toContainEqual(toolUseBlock2)
})
it("should not preserve tool_use blocks when preceding message has no tool_use blocks", () => {
const toolResultBlock = {
type: "tool_result" as const,
tool_use_id: "toolu_123",
content: "file contents",
}
const messages: ApiMessage[] = [
{ role: "user", content: "Hello", ts: 1 },
{ role: "assistant", content: "Plain text response", ts: 2 }, // No tool_use blocks
{
role: "user",
content: [toolResultBlock], // Has tool_result but preceding message has no tool_use
ts: 3,
},
{ role: "assistant", content: "Response", ts: 4 },
{ role: "user", content: "Thanks", ts: 5 },
]
const result = getKeepMessagesWithToolBlocks(messages, 3)
expect(result.keepMessages).toHaveLength(3)
expect(result.toolUseBlocksToPreserve).toHaveLength(0)
})
it("should handle edge case when startIndex - 1 is negative", () => {
const toolResultBlock = {
type: "tool_result" as const,
tool_use_id: "toolu_123",
content: "file contents",
}
// Only 3 messages total, so startIndex = 0 and precedingIndex would be -1
const messages: ApiMessage[] = [
{
role: "user",
content: [toolResultBlock],
ts: 1,
},
{ role: "assistant", content: "Response", ts: 2 },
{ role: "user", content: "Thanks", ts: 3 },
]
const result = getKeepMessagesWithToolBlocks(messages, 3)
expect(result.keepMessages).toEqual(messages)
expect(result.toolUseBlocksToPreserve).toHaveLength(0)
})
})
describe("getMessagesSinceLastSummary", () => {
it("should return all messages when there is no summary", () => {
const messages: ApiMessage[] = [
@ -535,6 +756,78 @@ describe("summarizeConversation", () => {
// Restore console.error
console.error = originalError
})
it("should append tool_use blocks to summary message when first kept message has tool_result blocks", async () => {
const toolUseBlock = {
type: "tool_use" as const,
id: "toolu_123",
name: "read_file",
input: { path: "test.txt" },
}
const toolResultBlock = {
type: "tool_result" as const,
tool_use_id: "toolu_123",
content: "file contents",
}
const messages: ApiMessage[] = [
{ role: "user", content: "Hello", ts: 1 },
{ role: "assistant", content: "Let me read that file", ts: 2 },
{ role: "user", content: "Please continue", ts: 3 },
{
role: "assistant",
content: [{ type: "text" as const, text: "Reading file..." }, toolUseBlock],
ts: 4,
},
{
role: "user",
content: [toolResultBlock, { type: "text" as const, text: "Continue" }],
ts: 5,
},
{ role: "assistant", content: "Got it, the file says...", ts: 6 },
{ role: "user", content: "Thanks", ts: 7 },
]
// Create a stream with usage information
const streamWithUsage = (async function* () {
yield { type: "text" as const, text: "Summary of conversation" }
yield { type: "usage" as const, totalCost: 0.05, outputTokens: 100 }
})()
mockApiHandler.createMessage = vi.fn().mockReturnValue(streamWithUsage) as any
mockApiHandler.countTokens = vi.fn().mockImplementation(() => Promise.resolve(50)) as any
const result = await summarizeConversation(
messages,
mockApiHandler,
defaultSystemPrompt,
taskId,
DEFAULT_PREV_CONTEXT_TOKENS,
false, // isAutomaticTrigger
undefined, // customCondensingPrompt
undefined, // condensingApiHandler
true, // useNativeTools - required for tool_use block preservation
)
// Verify the summary message has content array with text and tool_use blocks
const summaryMessage = result.messages[1]
expect(summaryMessage.role).toBe("assistant")
expect(summaryMessage.isSummary).toBe(true)
expect(Array.isArray(summaryMessage.content)).toBe(true)
// Content should be [text block, tool_use block]
const content = summaryMessage.content as any[]
expect(content).toHaveLength(2)
expect(content[0].type).toBe("text")
expect(content[0].text).toBe("Summary of conversation")
expect(content[1].type).toBe("tool_use")
expect(content[1].id).toBe("toolu_123")
expect(content[1].name).toBe("read_file")
// Verify the keepMessages are preserved correctly
expect(result.messages).toHaveLength(1 + 1 + N_MESSAGES_TO_KEEP) // first + summary + last 3
expect(result.error).toBeUndefined()
})
})
describe("summarizeConversation with custom settings", () => {

View file

@ -7,6 +7,66 @@ import { ApiHandler } from "../../api"
import { ApiMessage } from "../task-persistence/apiMessages"
import { maybeRemoveImageBlocks } from "../../api/transform/image-cleaning"
/**
* Checks if a message contains tool_result blocks.
* For native tools protocol, user messages with tool_result blocks require
* corresponding tool_use blocks from the previous assistant turn.
*/
function hasToolResultBlocks(message: ApiMessage): boolean {
if (message.role !== "user" || typeof message.content === "string") {
return false
}
return message.content.some((block) => block.type === "tool_result")
}
/**
* Gets the tool_use blocks from a message.
*/
function getToolUseBlocks(message: ApiMessage): Anthropic.Messages.ToolUseBlock[] {
if (message.role !== "assistant" || typeof message.content === "string") {
return []
}
return message.content.filter((block) => block.type === "tool_use") as Anthropic.Messages.ToolUseBlock[]
}
/**
* Extracts tool_use blocks that need to be preserved to match tool_result blocks in keepMessages.
* When the first kept message is a user message with tool_result blocks,
* we need to find the corresponding tool_use blocks from the preceding assistant message.
* These tool_use blocks will be appended to the summary message to maintain proper pairing.
*
* @param messages - The full conversation messages
* @param keepCount - The number of messages to keep from the end
* @returns Object containing keepMessages and any tool_use blocks to preserve
*/
export function getKeepMessagesWithToolBlocks(
messages: ApiMessage[],
keepCount: number,
): { keepMessages: ApiMessage[]; toolUseBlocksToPreserve: Anthropic.Messages.ToolUseBlock[] } {
if (messages.length <= keepCount) {
return { keepMessages: messages, toolUseBlocksToPreserve: [] }
}
const startIndex = messages.length - keepCount
const keepMessages = messages.slice(startIndex)
// Check if the first kept message is a user message with tool_result blocks
if (keepMessages.length > 0 && hasToolResultBlocks(keepMessages[0])) {
// Look for the preceding assistant message with tool_use blocks
const precedingIndex = startIndex - 1
if (precedingIndex >= 0) {
const precedingMessage = messages[precedingIndex]
const toolUseBlocks = getToolUseBlocks(precedingMessage)
if (toolUseBlocks.length > 0) {
// Return the tool_use blocks to be merged into the summary message
return { keepMessages, toolUseBlocksToPreserve: toolUseBlocks }
}
}
}
return { keepMessages, toolUseBlocksToPreserve: [] }
}
export const N_MESSAGES_TO_KEEP = 3
export const MIN_CONDENSE_THRESHOLD = 5 // Minimum percentage of context window to trigger condensing
export const MAX_CONDENSE_THRESHOLD = 100 // Maximum percentage of context window to trigger condensing
@ -80,6 +140,7 @@ export type SummarizeResponse = {
* @param {boolean} isAutomaticTrigger - Whether the summarization is triggered automatically
* @param {string} customCondensingPrompt - Optional custom prompt to use for condensing
* @param {ApiHandler} condensingApiHandler - Optional specific API handler to use for condensing
* @param {boolean} useNativeTools - Whether native tools protocol is being used (requires tool_use/tool_result pairing)
* @returns {SummarizeResponse} - The result of the summarization operation (see above)
*/
export async function summarizeConversation(
@ -91,6 +152,7 @@ export async function summarizeConversation(
isAutomaticTrigger?: boolean,
customCondensingPrompt?: string,
condensingApiHandler?: ApiHandler,
useNativeTools?: boolean,
): Promise<SummarizeResponse> {
TelemetryService.instance.captureContextCondensed(
taskId,
@ -103,6 +165,13 @@ export async function summarizeConversation(
// Always preserve the first message (which may contain slash command content)
const firstMessage = messages[0]
// Get keepMessages and any tool_use blocks that need to be preserved for tool_result pairing
// Only preserve tool_use blocks when using native tools protocol (XML protocol doesn't need them)
const { keepMessages, toolUseBlocksToPreserve } = useNativeTools
? getKeepMessagesWithToolBlocks(messages, N_MESSAGES_TO_KEEP)
: { keepMessages: messages.slice(-N_MESSAGES_TO_KEEP), toolUseBlocksToPreserve: [] }
// Get messages to summarize, including the first message and excluding the last N messages
const messagesToSummarize = getMessagesSinceLastSummary(messages.slice(0, -N_MESSAGES_TO_KEEP))
@ -114,9 +183,8 @@ export async function summarizeConversation(
return { ...response, error }
}
const keepMessages = messages.slice(-N_MESSAGES_TO_KEEP)
// Check if there's a recent summary in the messages we're keeping
const recentSummaryExists = keepMessages.some((message) => message.isSummary)
const recentSummaryExists = keepMessages.some((message: ApiMessage) => message.isSummary)
if (recentSummaryExists) {
const error = t("common:errors.condensed_recently")
@ -181,9 +249,21 @@ export async function summarizeConversation(
return { ...response, cost, error }
}
// Build the summary message content
// If there are tool_use blocks to preserve (for tool_result pairing), append them to the summary
let summaryContent: string | Anthropic.Messages.ContentBlockParam[]
if (toolUseBlocksToPreserve.length > 0) {
// Create content array with text block followed by tool_use blocks
// Use TextBlockParam which doesn't require citations field
const textBlock: Anthropic.Messages.TextBlockParam = { type: "text", text: summary }
summaryContent = [textBlock, ...toolUseBlocksToPreserve]
} else {
summaryContent = summary
}
const summaryMessage: ApiMessage = {
role: "assistant",
content: summary,
content: summaryContent,
ts: keepMessages[0].ts,
isSummary: true,
}

View file

@ -589,6 +589,7 @@ describe("Context Management", () => {
true,
undefined, // customCondensingPrompt
undefined, // condensingApiHandler
undefined, // useNativeTools
)
// Verify the result contains the summary information
@ -760,6 +761,7 @@ describe("Context Management", () => {
true,
undefined, // customCondensingPrompt
undefined, // condensingApiHandler
undefined, // useNativeTools
)
// Verify the result contains the summary information

View file

@ -86,6 +86,7 @@ export type ContextManagementOptions = {
condensingApiHandler?: ApiHandler
profileThresholds: Record<string, number>
currentProfileId: string
useNativeTools?: boolean
}
export type ContextManagementResult = SummarizeResponse & { prevContextTokens: number }
@ -110,6 +111,7 @@ export async function manageContext({
condensingApiHandler,
profileThresholds,
currentProfileId,
useNativeTools,
}: ContextManagementOptions): Promise<ContextManagementResult> {
let error: string | undefined
let cost = 0
@ -163,6 +165,7 @@ export async function manageContext({
true, // automatic trigger
customCondensingPrompt,
condensingApiHandler,
useNativeTools,
)
if (result.error) {
error = result.error

View file

@ -1272,6 +1272,11 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
const { contextTokens: prevContextTokens } = this.getTokenUsage()
// Determine if we're using native tool protocol for proper message handling
const modelInfo = this.api.getModel().info
const protocol = resolveToolProtocol(this.apiConfiguration, modelInfo)
const useNativeTools = isNativeProtocol(protocol)
const {
messages,
summary,
@ -1287,6 +1292,7 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
false, // manual trigger
customCondensingPrompt, // User's custom prompt
condensingApiHandler, // Specific handler for condensing
useNativeTools, // Pass native tools flag for proper message handling
)
if (error) {
this.say(
@ -3235,6 +3241,10 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
`Forcing truncation to ${FORCED_CONTEXT_REDUCTION_PERCENT}% of current context.`,
)
// Determine if we're using native tool protocol for proper message handling
const protocol = resolveToolProtocol(this.apiConfiguration, modelInfo)
const useNativeTools = isNativeProtocol(protocol)
// Force aggressive truncation by keeping only 75% of the conversation history
const truncateResult = await manageContext({
messages: this.apiConversationHistory,
@ -3248,6 +3258,7 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
taskId: this.taskId,
profileThresholds,
currentProfileId,
useNativeTools,
})
if (truncateResult.messages !== this.apiConversationHistory) {
@ -3350,6 +3361,11 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
// Get the current profile ID using the helper method
const currentProfileId = this.getCurrentProfileId(state)
// Determine if we're using native tool protocol for proper message handling
const modelInfoForProtocol = this.api.getModel().info
const protocol = resolveToolProtocol(this.apiConfiguration, modelInfoForProtocol)
const useNativeTools = isNativeProtocol(protocol)
const truncateResult = await manageContext({
messages: this.apiConversationHistory,
totalTokens: contextTokens,
@ -3364,6 +3380,7 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
condensingApiHandler,
profileThresholds,
currentProfileId,
useNativeTools,
})
if (truncateResult.messages !== this.apiConversationHistory) {
await this.overwriteApiConversationHistory(truncateResult.messages)