From 88b307957d9e5dbff67d9cd87b6cad752bfa2789 Mon Sep 17 00:00:00 2001 From: Roo Code Date: Wed, 3 Sep 2025 23:12:13 +0000 Subject: [PATCH] feat: add minimum token requirement for intelligent condensing - Add minimumCondenseTokens configuration option to global settings - Implement iterative refinement in summarizeConversation to meet minimum token requirements - Add expandSummaryToMeetMinimum function for multi-request expansion - Update sliding-window and Task modules to pass minimum token configuration - Add comprehensive tests for the new functionality - Ensure backward compatibility for existing use cases Fixes #7644 --- packages/types/src/global-settings.ts | 1 + src/core/condense/__tests__/index.spec.ts | 382 ++++++++++++++++++++++ src/core/condense/index.ts | 187 +++++++++++ src/core/sliding-window/index.ts | 3 + src/core/task/Task.ts | 4 + src/shared/ExtensionMessage.ts | 1 + 6 files changed, 578 insertions(+) diff --git a/packages/types/src/global-settings.ts b/packages/types/src/global-settings.ts index 81c6ae6dfe..a623aefcca 100644 --- a/packages/types/src/global-settings.ts +++ b/packages/types/src/global-settings.ts @@ -75,6 +75,7 @@ export const globalSettingsSchema = z.object({ allowedMaxCost: z.number().nullish(), autoCondenseContext: z.boolean().optional(), autoCondenseContextPercent: z.number().optional(), + minimumCondenseTokens: z.number().optional(), // Minimum tokens for condensed output maxConcurrentFileReads: z.number().optional(), /** diff --git a/src/core/condense/__tests__/index.spec.ts b/src/core/condense/__tests__/index.spec.ts index 11a25a0693..10903e2402 100644 --- a/src/core/condense/__tests__/index.spec.ts +++ b/src/core/condense/__tests__/index.spec.ts @@ -792,3 +792,385 @@ describe("summarizeConversation with custom settings", () => { ) }) }) + +describe("summarizeConversation with minimum token requirements", () => { + // Mock ApiHandler + let mockApiHandler: ApiHandler + let mockCondensingApiHandler: ApiHandler + const defaultSystemPrompt = "You are a helpful assistant." + const taskId = "test-task-id" + + // Sample messages for testing + const sampleMessages: 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 }, + { role: "assistant", content: "Not much", ts: 6 }, + { role: "user", content: "Tell me more", ts: 7 }, + ] + + beforeEach(() => { + // Reset mocks + vi.clearAllMocks() + + // Setup mock API handler + mockApiHandler = { + createMessage: vi.fn(), + countTokens: vi.fn(), + getModel: vi.fn().mockReturnValue({ + id: "test-model", + info: { + contextWindow: 8000, + supportsImages: true, + supportsComputerUse: true, + supportsVision: true, + maxTokens: 4000, + supportsPromptCache: true, + maxCachePoints: 10, + minTokensPerCachePoint: 100, + cachableFields: ["system", "messages"], + }, + }), + } as unknown as ApiHandler + + mockCondensingApiHandler = { + createMessage: vi.fn(), + countTokens: vi.fn(), + getModel: vi.fn().mockReturnValue({ + id: "condensing-model", + info: { + contextWindow: 4000, + supportsImages: true, + supportsComputerUse: false, + supportsVision: false, + maxTokens: 2000, + supportsPromptCache: false, + maxCachePoints: 0, + minTokensPerCachePoint: 0, + cachableFields: [], + }, + }), + } as unknown as ApiHandler + }) + + it("should not expand summary when minimum tokens is not specified", async () => { + // Setup initial summary stream + const initialStream = (async function* () { + yield { type: "text" as const, text: "Short summary" } + yield { type: "usage" as const, totalCost: 0.02, outputTokens: 50 } + })() + + mockApiHandler.createMessage = vi.fn().mockReturnValueOnce(initialStream) as any + mockApiHandler.countTokens = vi.fn().mockResolvedValue(100) as any + + const result = await summarizeConversation( + sampleMessages, + mockApiHandler, + defaultSystemPrompt, + taskId, + 1000, // prevContextTokens + false, + undefined, + undefined, + undefined, // No minimum tokens specified + ) + + // Should only call createMessage once (no expansion) + expect(mockApiHandler.createMessage).toHaveBeenCalledTimes(1) + expect(result.summary).toBe("Short summary") + expect(result.newContextTokens).toBe(150) // 50 output + 100 counted + expect(result.error).toBeUndefined() + }) + + it("should not expand summary when current tokens already meet minimum requirement", async () => { + // Setup initial summary stream with enough tokens + const initialStream = (async function* () { + yield { type: "text" as const, text: "This is a longer summary with more content" } + yield { type: "usage" as const, totalCost: 0.05, outputTokens: 300 } + })() + + mockApiHandler.createMessage = vi.fn().mockReturnValueOnce(initialStream) as any + mockApiHandler.countTokens = vi.fn().mockResolvedValue(250) as any + + const result = await summarizeConversation( + sampleMessages, + mockApiHandler, + defaultSystemPrompt, + taskId, + 1000, // prevContextTokens + false, + undefined, + undefined, + 500, // minimumCondenseTokens - already met by 550 total + ) + + // Should only call createMessage once (no expansion needed) + expect(mockApiHandler.createMessage).toHaveBeenCalledTimes(1) + expect(result.summary).toBe("This is a longer summary with more content") + expect(result.newContextTokens).toBe(550) // 300 output + 250 counted + expect(result.error).toBeUndefined() + }) + + it("should expand summary when below minimum token requirement", async () => { + // Setup initial summary stream with too few tokens + const initialStream = (async function* () { + yield { type: "text" as const, text: "Short summary" } + yield { type: "usage" as const, totalCost: 0.02, outputTokens: 50 } + })() + + // Setup expansion stream + const expansionStream = (async function* () { + yield { + type: "text" as const, + text: "This is a much more detailed and expanded summary with lots of additional context and information", + } + yield { type: "usage" as const, totalCost: 0.08, outputTokens: 400 } + })() + + mockApiHandler.createMessage = vi + .fn() + .mockReturnValueOnce(initialStream) + .mockReturnValueOnce(expansionStream) as any + + mockApiHandler.countTokens = vi + .fn() + .mockResolvedValueOnce(100) // First count after initial summary + .mockResolvedValueOnce(150) // Count after first expansion attempt + .mockResolvedValueOnce(150) as any // Count after second expansion attempt (final) + + const result = await summarizeConversation( + sampleMessages, + mockApiHandler, + defaultSystemPrompt, + taskId, + 1000, // prevContextTokens + false, + undefined, + undefined, + 500, // minimumCondenseTokens - requires expansion + ) + + // Should call createMessage three times (initial + 2 expansions due to mock setup) + expect(mockApiHandler.createMessage).toHaveBeenCalledTimes(3) + + // Check the expansion request includes the expansion prompt + const secondCall = (mockApiHandler.createMessage as Mock).mock.calls[1] + const expansionMessages = secondCall[1] + const lastMessage = expansionMessages[expansionMessages.length - 1] + expect(lastMessage.content).toContain("The current summary has") + expect(lastMessage.content).toContain("tokens, but we need at least") + + expect(result.summary).toBe( + "This is a much more detailed and expanded summary with lots of additional context and information", + ) + expect(result.newContextTokens).toBe(150) // Final count from mock + expect(result.cost).toBe(0.1) // 0.02 + 0.08 + expect(result.error).toBeUndefined() + }) + + it("should use condensing API handler for expansion when provided", async () => { + // Setup initial summary stream with too few tokens + const initialStream = (async function* () { + yield { type: "text" as const, text: "Short summary" } + yield { type: "usage" as const, totalCost: 0.02, outputTokens: 50 } + })() + + // Setup expansion stream from condensing handler + const expansionStream = (async function* () { + yield { type: "text" as const, text: "Expanded summary from condensing handler" } + yield { type: "usage" as const, totalCost: 0.06, outputTokens: 350 } + })() + + mockCondensingApiHandler.createMessage = vi + .fn() + .mockReturnValueOnce(initialStream) + .mockReturnValueOnce(expansionStream) as any + + mockApiHandler.countTokens = vi + .fn() + .mockResolvedValueOnce(100) // First count + .mockResolvedValueOnce(200) // After first expansion + .mockResolvedValueOnce(200) as any // After second expansion (final) + + const result = await summarizeConversation( + sampleMessages, + mockApiHandler, + defaultSystemPrompt, + taskId, + 1000, // prevContextTokens + false, + "Custom prompt", + mockCondensingApiHandler, + 500, // minimumCondenseTokens + ) + + // Should use condensing handler for all calls (initial + expansions) + expect(mockCondensingApiHandler.createMessage).toHaveBeenCalledTimes(3) + expect(mockApiHandler.createMessage).not.toHaveBeenCalled() + + expect(result.summary).toBe("Expanded summary from condensing handler") + expect(result.newContextTokens).toBe(200) // Final count from mock + expect(result.cost).toBe(0.08) // 0.02 + 0.06 + expect(result.error).toBeUndefined() + }) + + it("should stop expansion after MAX_ITERATIONS to prevent infinite loops", async () => { + // Setup streams that always return insufficient tokens + const createSmallStream = () => + (async function* () { + yield { type: "text" as const, text: "Still too short" } + yield { type: "usage" as const, totalCost: 0.01, outputTokens: 30 } + })() + + let callCount = 0 + mockApiHandler.createMessage = vi.fn().mockImplementation(() => { + callCount++ + return createSmallStream() + }) as any + + // Always return low token count + mockApiHandler.countTokens = vi.fn().mockResolvedValue(50) as any + + const result = await summarizeConversation( + sampleMessages, + mockApiHandler, + defaultSystemPrompt, + taskId, + 2000, // prevContextTokens + false, + undefined, + undefined, + 1000, // minimumCondenseTokens - impossible to reach + ) + + // Should stop after MAX_ITERATIONS (5) + 1 initial call = 6 total + expect(mockApiHandler.createMessage).toHaveBeenCalledTimes(6) + expect(result.summary).toBe("Still too short") + expect(result.error).toBeUndefined() + }) + + it("should revert to previous summary if expansion exceeds context limit", async () => { + // Setup initial summary stream + const initialStream = (async function* () { + yield { type: "text" as const, text: "Initial summary" } + yield { type: "usage" as const, totalCost: 0.02, outputTokens: 100 } + })() + + // Setup expansion stream that's too large + const expansionStream = (async function* () { + yield { type: "text" as const, text: "Extremely long expanded summary that exceeds context" } + yield { type: "usage" as const, totalCost: 0.1, outputTokens: 800 } + })() + + mockApiHandler.createMessage = vi + .fn() + .mockReturnValueOnce(initialStream) + .mockReturnValueOnce(expansionStream) as any + + mockApiHandler.countTokens = vi + .fn() + .mockResolvedValueOnce(150) // First count + .mockResolvedValueOnce(500) as any // After expansion - too large! + + const result = await summarizeConversation( + sampleMessages, + mockApiHandler, + defaultSystemPrompt, + taskId, + 400, // prevContextTokens - will be exceeded + false, + undefined, + undefined, + 300, // minimumCondenseTokens + ) + + // Should revert to initial summary + expect(result.summary).toBe("Initial summary") + expect(result.newContextTokens).toBe(250) // 100 output + 150 counted (initial) + expect(result.cost).toBeCloseTo(0.02, 5) // Only initial cost, expansion cost excluded + expect(result.error).toBeUndefined() + }) + + it("should handle empty expansion response gracefully", async () => { + // Setup initial summary stream + const initialStream = (async function* () { + yield { type: "text" as const, text: "Initial summary" } + yield { type: "usage" as const, totalCost: 0.02, outputTokens: 100 } + })() + + // Setup empty expansion stream + const emptyExpansionStream = (async function* () { + yield { type: "text" as const, text: "" } + yield { type: "usage" as const, totalCost: 0.01, outputTokens: 0 } + })() + + mockApiHandler.createMessage = vi + .fn() + .mockReturnValueOnce(initialStream) + .mockReturnValueOnce(emptyExpansionStream) as any + + mockApiHandler.countTokens = vi.fn().mockResolvedValue(150) as any + + const result = await summarizeConversation( + sampleMessages, + mockApiHandler, + defaultSystemPrompt, + taskId, + 1000, // prevContextTokens + false, + undefined, + undefined, + 500, // minimumCondenseTokens - requires expansion but gets empty response + ) + + // Should keep initial summary when expansion fails + expect(result.summary).toBe("Initial summary") + expect(result.newContextTokens).toBe(250) // 100 output + 150 counted + expect(result.cost).toBe(0.02) // Only initial cost since expansion produced empty result + expect(result.error).toBeUndefined() + }) + + it("should use custom prompt for expansion when provided", async () => { + const customPrompt = "Custom summarization instructions" + + // Setup initial summary stream with too few tokens + const initialStream = (async function* () { + yield { type: "text" as const, text: "Short summary" } + yield { type: "usage" as const, totalCost: 0.02, outputTokens: 50 } + })() + + // Setup expansion stream + const expansionStream = (async function* () { + yield { type: "text" as const, text: "Expanded with custom prompt" } + yield { type: "usage" as const, totalCost: 0.05, outputTokens: 300 } + })() + + mockApiHandler.createMessage = vi + .fn() + .mockReturnValueOnce(initialStream) + .mockReturnValueOnce(expansionStream) as any + + mockApiHandler.countTokens = vi.fn().mockResolvedValueOnce(100).mockResolvedValueOnce(200) as any + + const result = await summarizeConversation( + sampleMessages, + mockApiHandler, + defaultSystemPrompt, + taskId, + 1000, + false, + customPrompt, // Custom prompt provided + undefined, + 400, // minimumCondenseTokens + ) + + // Check that custom prompt was used in expansion + const expansionCall = (mockApiHandler.createMessage as Mock).mock.calls[1] + expect(expansionCall[0]).toBe(customPrompt) + + expect(result.summary).toBe("Expanded with custom prompt") + expect(result.error).toBeUndefined() + }) +}) diff --git a/src/core/condense/index.ts b/src/core/condense/index.ts index 3b73b1915c..d0492707f6 100644 --- a/src/core/condense/index.ts +++ b/src/core/condense/index.ts @@ -80,6 +80,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 {number} minimumCondenseTokens - Optional minimum token requirement for condensed output * @returns {SummarizeResponse} - The result of the summarization operation (see above) */ export async function summarizeConversation( @@ -91,6 +92,7 @@ export async function summarizeConversation( isAutomaticTrigger?: boolean, customCondensingPrompt?: string, condensingApiHandler?: ApiHandler, + minimumCondenseTokens?: number, ): Promise { TelemetryService.instance.captureContextCondensed( taskId, @@ -203,6 +205,27 @@ export async function summarizeConversation( const error = t("common:errors.condense_context_grew") return { ...response, cost, error } } + + // Check if minimum token requirement is met + if (minimumCondenseTokens && minimumCondenseTokens > 0 && newContextTokens < minimumCondenseTokens) { + // Need to make additional API requests to expand the summary + const expandedResult = await expandSummaryToMeetMinimum( + messages, + summary, + keepMessages, + apiHandler, + systemPrompt, + taskId, + prevContextTokens, + newContextTokens, + minimumCondenseTokens, + cost, + customCondensingPrompt, + condensingApiHandler || apiHandler, + ) + return expandedResult + } + return { messages: newMessages, summary, cost, newContextTokens } } @@ -226,3 +249,167 @@ export function getMessagesSinceLastSummary(messages: ApiMessage[]): ApiMessage[ } return [userMessage, ...messagesSinceSummary] } + +/** + * Expands a summary by making additional API requests until minimum token requirement is met + * + * @param {ApiMessage[]} originalMessages - The original conversation messages + * @param {string} currentSummary - The current summary that needs expansion + * @param {ApiMessage[]} keepMessages - Messages to keep from the end + * @param {ApiHandler} apiHandler - The API handler to use + * @param {string} systemPrompt - The system prompt + * @param {string} taskId - The task ID for telemetry + * @param {number} prevContextTokens - Previous context token count + * @param {number} currentTokens - Current token count after initial summary + * @param {number} minimumTokens - Minimum required tokens + * @param {number} totalCost - Accumulated cost + * @param {string} customPrompt - Optional custom prompt + * @param {ApiHandler} condensingHandler - Handler for condensing operations + * @returns {SummarizeResponse} - The expanded summary response + */ +async function expandSummaryToMeetMinimum( + originalMessages: ApiMessage[], + currentSummary: string, + keepMessages: ApiMessage[], + apiHandler: ApiHandler, + systemPrompt: string, + taskId: string, + prevContextTokens: number, + currentTokens: number, + minimumTokens: number, + totalCost: number, + customPrompt?: string, + condensingHandler?: ApiHandler, +): Promise { + let expandedSummary = currentSummary + let accumulatedCost = totalCost + let tokenCount = currentTokens + let iterationCount = 0 + const MAX_ITERATIONS = 5 // Prevent infinite loops + + // Messages to summarize (excluding the ones we're keeping) + const messagesToSummarize = getMessagesSinceLastSummary(originalMessages.slice(0, -N_MESSAGES_TO_KEEP)) + + while (tokenCount < minimumTokens && iterationCount < MAX_ITERATIONS) { + iterationCount++ + + // Create an expansion prompt that requests more detail + const expansionPrompt = ` +The current summary has ${tokenCount} tokens, but we need at least ${minimumTokens} tokens to maintain sufficient context. +Please expand the following summary by adding more relevant details from the conversation history. +Focus on: +1. Technical implementation details and code patterns +2. Specific file changes and modifications +3. Problem-solving approaches and decisions made +4. Any error messages or debugging steps +5. Important context that would help continue the task + +Current Summary: +${expandedSummary} + +Please provide an expanded version with approximately ${Math.ceil(minimumTokens * 1.1)} tokens (10% buffer). +Include all the original information plus additional relevant details. +` + + // Prepare messages for expansion request + const expansionMessages: Anthropic.MessageParam[] = [ + ...messagesToSummarize.map(({ role, content }) => ({ role, content })), + { + role: "assistant", + content: expandedSummary, + }, + { + role: "user", + content: expansionPrompt, + }, + ] + + // Use the condensing handler if available + const handler = condensingHandler || apiHandler + const stream = handler.createMessage(customPrompt?.trim() || SUMMARY_PROMPT, expansionMessages) + + let additionalSummary = "" + let iterationCost = 0 + let outputTokens = 0 + + if (!stream) { + // Failed to create stream, break out of loop + break + } + + for await (const chunk of stream) { + if (chunk.type === "text") { + additionalSummary += chunk.text + } else if (chunk.type === "usage") { + iterationCost = chunk.totalCost ?? 0 + outputTokens = chunk.outputTokens ?? 0 + } + } + + additionalSummary = additionalSummary.trim() + + if (additionalSummary.length === 0) { + // Failed to get additional content, return what we have + break + } + + // Update the expanded summary + expandedSummary = additionalSummary + accumulatedCost += iterationCost + + // Recalculate token count + const summaryMessage: ApiMessage = { + role: "assistant", + content: expandedSummary, + ts: keepMessages[0]?.ts || Date.now(), + isSummary: true, + } + + const systemPromptMessage: ApiMessage = { role: "user", content: systemPrompt } + const contextMessages = [systemPromptMessage, summaryMessage, ...keepMessages] + const contextBlocks = contextMessages.flatMap((message) => + typeof message.content === "string" ? [{ text: message.content, type: "text" as const }] : message.content, + ) + + tokenCount = await apiHandler.countTokens(contextBlocks) + + // Check if we've exceeded the previous context (safety check) + if (tokenCount >= prevContextTokens) { + // We've grown too much, revert to previous iteration + // Don't include the cost of this failed iteration + const prevSummaryMessage: ApiMessage = { + role: "assistant", + content: currentSummary, + ts: keepMessages[0]?.ts || Date.now(), + isSummary: true, + } + const newMessages = [...originalMessages.slice(0, -N_MESSAGES_TO_KEEP), prevSummaryMessage, ...keepMessages] + return { + messages: newMessages, + summary: currentSummary, + cost: accumulatedCost - iterationCost, + newContextTokens: currentTokens, + } + } + + currentSummary = expandedSummary + currentTokens = tokenCount + } + + // Create final message structure + const finalSummaryMessage: ApiMessage = { + role: "assistant", + content: expandedSummary, + ts: keepMessages[0]?.ts || Date.now(), + isSummary: true, + } + + const newMessages = [...originalMessages.slice(0, -N_MESSAGES_TO_KEEP), finalSummaryMessage, ...keepMessages] + + return { + messages: newMessages, + summary: expandedSummary, + cost: accumulatedCost, + newContextTokens: tokenCount, + } +} diff --git a/src/core/sliding-window/index.ts b/src/core/sliding-window/index.ts index 1e518c9a56..5a6f5f5f7b 100644 --- a/src/core/sliding-window/index.ts +++ b/src/core/sliding-window/index.ts @@ -77,6 +77,7 @@ type TruncateOptions = { condensingApiHandler?: ApiHandler profileThresholds: Record currentProfileId: string + minimumCondenseTokens?: number } type TruncateResponse = SummarizeResponse & { prevContextTokens: number } @@ -102,6 +103,7 @@ export async function truncateConversationIfNeeded({ condensingApiHandler, profileThresholds, currentProfileId, + minimumCondenseTokens, }: TruncateOptions): Promise { let error: string | undefined let cost = 0 @@ -155,6 +157,7 @@ export async function truncateConversationIfNeeded({ true, // automatic trigger customCondensingPrompt, condensingApiHandler, + minimumCondenseTokens, ) if (result.error) { error = result.error diff --git a/src/core/task/Task.ts b/src/core/task/Task.ts index be932b098f..d49adc2de6 100644 --- a/src/core/task/Task.ts +++ b/src/core/task/Task.ts @@ -950,6 +950,7 @@ export class Task extends EventEmitter implements TaskLike { const customCondensingPrompt = state?.customCondensingPrompt const condensingApiConfigId = state?.condensingApiConfigId const listApiConfigMeta = state?.listApiConfigMeta + const minimumCondenseTokens = state?.minimumCondenseTokens // Determine API handler to use let condensingApiHandler: ApiHandler | undefined @@ -984,6 +985,7 @@ export class Task extends EventEmitter implements TaskLike { false, // manual trigger customCondensingPrompt, // User's custom prompt condensingApiHandler, // Specific handler for condensing + minimumCondenseTokens, // Minimum token requirement ) if (error) { this.say( @@ -2453,6 +2455,7 @@ export class Task extends EventEmitter implements TaskLike { autoCondenseContext = true, autoCondenseContextPercent = 100, profileThresholds = {}, + minimumCondenseTokens, } = state ?? {} // Get condensing configuration for automatic triggers. @@ -2536,6 +2539,7 @@ export class Task extends EventEmitter implements TaskLike { condensingApiHandler, profileThresholds, currentProfileId, + minimumCondenseTokens, }) if (truncateResult.messages !== this.apiConversationHistory) { await this.overwriteApiConversationHistory(truncateResult.messages) diff --git a/src/shared/ExtensionMessage.ts b/src/shared/ExtensionMessage.ts index d4caf2f674..74d35381a6 100644 --- a/src/shared/ExtensionMessage.ts +++ b/src/shared/ExtensionMessage.ts @@ -272,6 +272,7 @@ export type ExtensionState = Pick< | "enhancementApiConfigId" | "condensingApiConfigId" | "customCondensingPrompt" + | "minimumCondenseTokens" | "codebaseIndexConfig" | "codebaseIndexModels" | "profileThresholds"