Explorar o código

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

Daniel hai 1 mes
pai
achega
311940b085

+ 294 - 1
src/core/condense/__tests__/index.spec.ts

@@ -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", () => {

+ 83 - 3
src/core/condense/index.ts

@@ -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,
 	}

+ 2 - 0
src/core/context-management/__tests__/context-management.spec.ts

@@ -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

+ 3 - 0
src/core/context-management/index.ts

@@ -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

+ 17 - 0
src/core/task/Task.ts

@@ -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)