diff --git a/src/core/condense/__tests__/index.spec.ts b/src/core/condense/__tests__/index.spec.ts index 6f7423f9534..bf5560ff0b2 100644 --- a/src/core/condense/__tests__/index.spec.ts +++ b/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", () => { diff --git a/src/core/condense/index.ts b/src/core/condense/index.ts index 86cfa7ab1e5..29509788086 100644 --- a/src/core/condense/index.ts +++ b/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 { 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, } diff --git a/src/core/context-management/__tests__/context-management.spec.ts b/src/core/context-management/__tests__/context-management.spec.ts index a2c6daec728..712785b7cd8 100644 --- a/src/core/context-management/__tests__/context-management.spec.ts +++ b/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 diff --git a/src/core/context-management/index.ts b/src/core/context-management/index.ts index fa91fc0c9d3..40ad688ae9a 100644 --- a/src/core/context-management/index.ts +++ b/src/core/context-management/index.ts @@ -86,6 +86,7 @@ export type ContextManagementOptions = { condensingApiHandler?: ApiHandler profileThresholds: Record currentProfileId: string + useNativeTools?: boolean } export type ContextManagementResult = SummarizeResponse & { prevContextTokens: number } @@ -110,6 +111,7 @@ export async function manageContext({ condensingApiHandler, profileThresholds, currentProfileId, + useNativeTools, }: ContextManagementOptions): Promise { 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 diff --git a/src/core/task/Task.ts b/src/core/task/Task.ts index f7b0adc28ab..0db4019394e 100644 --- a/src/core/task/Task.ts +++ b/src/core/task/Task.ts @@ -1272,6 +1272,11 @@ export class Task extends EventEmitter 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 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( @@ -3182,6 +3188,10 @@ export class Task extends EventEmitter 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, @@ -3195,6 +3205,7 @@ export class Task extends EventEmitter implements TaskLike { taskId: this.taskId, profileThresholds, currentProfileId, + useNativeTools, }) if (truncateResult.messages !== this.apiConversationHistory) { @@ -3297,6 +3308,11 @@ export class Task extends EventEmitter 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, @@ -3311,6 +3327,7 @@ export class Task extends EventEmitter implements TaskLike { condensingApiHandler, profileThresholds, currentProfileId, + useNativeTools, }) if (truncateResult.messages !== this.apiConversationHistory) { await this.overwriteApiConversationHistory(truncateResult.messages)