diff --git a/src/core/sliding-window/__tests__/sliding-window.test.ts b/src/core/sliding-window/__tests__/sliding-window.test.ts index 7890b55ec8..fe3b71f4eb 100644 --- a/src/core/sliding-window/__tests__/sliding-window.test.ts +++ b/src/core/sliding-window/__tests__/sliding-window.test.ts @@ -11,6 +11,7 @@ import { truncateConversationIfNeeded, } from "../index" import { ApiMessage } from "../../task-persistence/apiMessages" +import * as condenseModule from "../../condense" // Create a mock ApiHandler for testing class MockApiHandler extends BaseProvider { @@ -248,7 +249,14 @@ describe("truncateConversationIfNeeded", () => { maxTokens: modelInfo.maxTokens, apiHandler: mockApiHandler, }) - expect(result).toEqual(messagesWithSmallContent) // No truncation occurs + + // Check the new return type + expect(result).toEqual({ + messages: messagesWithSmallContent, + summary: "", + cost: 0, + prevContextTokens: totalTokens, + }) }) it("should truncate if tokens are above max tokens threshold", async () => { @@ -260,7 +268,7 @@ describe("truncateConversationIfNeeded", () => { // When truncating, always uses 0.5 fraction // With 4 messages after the first, 0.5 fraction means remove 2 messages - const expectedResult = [messagesWithSmallContent[0], messagesWithSmallContent[3], messagesWithSmallContent[4]] + const expectedMessages = [messagesWithSmallContent[0], messagesWithSmallContent[3], messagesWithSmallContent[4]] const result = await truncateConversationIfNeeded({ messages: messagesWithSmallContent, @@ -269,7 +277,13 @@ describe("truncateConversationIfNeeded", () => { maxTokens: modelInfo.maxTokens, apiHandler: mockApiHandler, }) - expect(result).toEqual(expectedResult) + + expect(result).toEqual({ + messages: expectedMessages, + summary: "", + cost: 0, + prevContextTokens: totalTokens, + }) }) it("should work with non-prompt caching models the same as prompt caching models", async () => { @@ -298,7 +312,10 @@ describe("truncateConversationIfNeeded", () => { apiHandler: mockApiHandler, }) - expect(result1).toEqual(result2) + expect(result1.messages).toEqual(result2.messages) + expect(result1.summary).toEqual(result2.summary) + expect(result1.cost).toEqual(result2.cost) + expect(result1.prevContextTokens).toEqual(result2.prevContextTokens) // Test above threshold const aboveThreshold = 70001 @@ -318,7 +335,10 @@ describe("truncateConversationIfNeeded", () => { apiHandler: mockApiHandler, }) - expect(result3).toEqual(result4) + expect(result3.messages).toEqual(result4.messages) + expect(result3.summary).toEqual(result4.summary) + expect(result3.cost).toEqual(result4.cost) + expect(result3.prevContextTokens).toEqual(result4.prevContextTokens) }) it("should consider incoming content when deciding to truncate", async () => { @@ -344,7 +364,12 @@ describe("truncateConversationIfNeeded", () => { maxTokens, apiHandler: mockApiHandler, }) - expect(resultWithSmall).toEqual(messagesWithSmallContent) // No truncation + expect(resultWithSmall).toEqual({ + messages: messagesWithSmallContent, + summary: "", + cost: 0, + prevContextTokens: baseTokensForSmall + smallContentTokens, + }) // No truncation // Test case 2: Large content that will push us over the threshold const largeContent = [ @@ -368,7 +393,10 @@ describe("truncateConversationIfNeeded", () => { maxTokens, apiHandler: mockApiHandler, }) - expect(resultWithLarge).not.toEqual(messagesWithLargeContent) // Should truncate + expect(resultWithLarge.messages).not.toEqual(messagesWithLargeContent) // Should truncate + expect(resultWithLarge.summary).toBe("") + expect(resultWithLarge.cost).toBe(0) + expect(resultWithLarge.prevContextTokens).toBe(baseTokensForLarge + largeContentTokens) // Test case 3: Very large content that will definitely exceed threshold const veryLargeContent = [{ type: "text" as const, text: "X".repeat(1000) }] @@ -387,7 +415,10 @@ describe("truncateConversationIfNeeded", () => { maxTokens, apiHandler: mockApiHandler, }) - expect(resultWithVeryLarge).not.toEqual(messagesWithVeryLargeContent) // Should truncate + expect(resultWithVeryLarge.messages).not.toEqual(messagesWithVeryLargeContent) // Should truncate + expect(resultWithVeryLarge.summary).toBe("") + expect(resultWithVeryLarge.cost).toBe(0) + expect(resultWithVeryLarge.prevContextTokens).toBe(baseTokensForVeryLarge + veryLargeContentTokens) }) it("should truncate if tokens are within TOKEN_BUFFER_PERCENTAGE of the threshold", async () => { @@ -409,7 +440,140 @@ describe("truncateConversationIfNeeded", () => { maxTokens: modelInfo.maxTokens, apiHandler: mockApiHandler, }) - expect(result).toEqual(expectedResult) + expect(result).toEqual({ + messages: expectedResult, + summary: "", + cost: 0, + prevContextTokens: totalTokens, + }) + }) + + it("should use summarizeConversation when autoCondenseContext is true and tokens exceed threshold", async () => { + // Mock the summarizeConversation function + const mockSummary = "This is a summary of the conversation" + const mockCost = 0.05 + const mockSummarizeResponse: condenseModule.SummarizeResponse = { + messages: [ + { role: "user", content: "First message" }, + { role: "assistant", content: mockSummary, isSummary: true }, + { role: "user", content: "Last message" }, + ], + summary: mockSummary, + cost: mockCost, + newContextTokens: 100, + } + + const summarizeSpy = jest + .spyOn(condenseModule, "summarizeConversation") + .mockResolvedValue(mockSummarizeResponse) + + const modelInfo = createModelInfo(100000, 30000) + const totalTokens = 70001 // Above threshold + const messagesWithSmallContent = [...messages.slice(0, -1), { ...messages[messages.length - 1], content: "" }] + + const result = await truncateConversationIfNeeded({ + messages: messagesWithSmallContent, + totalTokens, + contextWindow: modelInfo.contextWindow, + maxTokens: modelInfo.maxTokens, + apiHandler: mockApiHandler, + autoCondenseContext: true, + systemPrompt: "System prompt", + }) + + // Verify summarizeConversation was called with the right parameters + expect(summarizeSpy).toHaveBeenCalledWith(messagesWithSmallContent, mockApiHandler, "System prompt") + + // Verify the result contains the summary information + expect(result).toMatchObject({ + messages: mockSummarizeResponse.messages, + summary: mockSummary, + cost: mockCost, + prevContextTokens: totalTokens, + }) + // newContextTokens might be present, but we don't need to verify its exact value + + // Clean up + summarizeSpy.mockRestore() + }) + + it("should fall back to truncateConversation when autoCondenseContext is true but summarization fails", async () => { + // Mock the summarizeConversation function to return empty summary + const mockSummarizeResponse: condenseModule.SummarizeResponse = { + messages: messages, // Original messages unchanged + summary: "", // Empty summary indicates failure + cost: 0.01, + } + + const summarizeSpy = jest + .spyOn(condenseModule, "summarizeConversation") + .mockResolvedValue(mockSummarizeResponse) + + const modelInfo = createModelInfo(100000, 30000) + const totalTokens = 70001 // Above threshold + const messagesWithSmallContent = [...messages.slice(0, -1), { ...messages[messages.length - 1], content: "" }] + + // When truncating, always uses 0.5 fraction + // With 4 messages after the first, 0.5 fraction means remove 2 messages + const expectedMessages = [messagesWithSmallContent[0], messagesWithSmallContent[3], messagesWithSmallContent[4]] + + const result = await truncateConversationIfNeeded({ + messages: messagesWithSmallContent, + totalTokens, + contextWindow: modelInfo.contextWindow, + maxTokens: modelInfo.maxTokens, + apiHandler: mockApiHandler, + autoCondenseContext: true, + }) + + // Verify summarizeConversation was called + expect(summarizeSpy).toHaveBeenCalled() + + // Verify it fell back to truncation + expect(result.messages).toEqual(expectedMessages) + expect(result.summary).toBe("") + expect(result.prevContextTokens).toBe(totalTokens) + // The cost might be different than expected, so we don't check it + + // Clean up + summarizeSpy.mockRestore() + }) + + it("should not call summarizeConversation when autoCondenseContext is false", async () => { + // Reset any previous mock calls + jest.clearAllMocks() + const summarizeSpy = jest.spyOn(condenseModule, "summarizeConversation") + + const modelInfo = createModelInfo(100000, 30000) + const totalTokens = 70001 // Above threshold + const messagesWithSmallContent = [...messages.slice(0, -1), { ...messages[messages.length - 1], content: "" }] + + // When truncating, always uses 0.5 fraction + // With 4 messages after the first, 0.5 fraction means remove 2 messages + const expectedMessages = [messagesWithSmallContent[0], messagesWithSmallContent[3], messagesWithSmallContent[4]] + + const result = await truncateConversationIfNeeded({ + messages: messagesWithSmallContent, + totalTokens, + contextWindow: modelInfo.contextWindow, + maxTokens: modelInfo.maxTokens, + apiHandler: mockApiHandler, + autoCondenseContext: false, + }) + + // Verify summarizeConversation was not called + expect(summarizeSpy).not.toHaveBeenCalled() + + // Verify it used truncation + expect(result).toEqual({ + messages: expectedMessages, + summary: "", + cost: 0, + prevContextTokens: totalTokens, + }) + + // Clean up + summarizeSpy.mockRestore() }) }) @@ -449,7 +613,12 @@ describe("getMaxTokens", () => { maxTokens: modelInfo.maxTokens, apiHandler: mockApiHandler, }) - expect(result1).toEqual(messagesWithSmallContent) + expect(result1).toEqual({ + messages: messagesWithSmallContent, + summary: "", + cost: 0, + prevContextTokens: 39999, + }) // Above max tokens - truncate const result2 = await truncateConversationIfNeeded({ @@ -459,8 +628,11 @@ describe("getMaxTokens", () => { maxTokens: modelInfo.maxTokens, apiHandler: mockApiHandler, }) - expect(result2).not.toEqual(messagesWithSmallContent) - expect(result2.length).toBe(3) // Truncated with 0.5 fraction + expect(result2.messages).not.toEqual(messagesWithSmallContent) + expect(result2.messages.length).toBe(3) // Truncated with 0.5 fraction + expect(result2.summary).toBe("") + expect(result2.cost).toBe(0) + expect(result2.prevContextTokens).toBe(50001) }) it("should use 20% of context window as buffer when maxTokens is undefined", async () => { @@ -479,7 +651,12 @@ describe("getMaxTokens", () => { maxTokens: modelInfo.maxTokens, apiHandler: mockApiHandler, }) - expect(result1).toEqual(messagesWithSmallContent) + expect(result1).toEqual({ + messages: messagesWithSmallContent, + summary: "", + cost: 0, + prevContextTokens: 69999, + }) // Above max tokens - truncate const result2 = await truncateConversationIfNeeded({ @@ -489,8 +666,11 @@ describe("getMaxTokens", () => { maxTokens: modelInfo.maxTokens, apiHandler: mockApiHandler, }) - expect(result2).not.toEqual(messagesWithSmallContent) - expect(result2.length).toBe(3) // Truncated with 0.5 fraction + expect(result2.messages).not.toEqual(messagesWithSmallContent) + expect(result2.messages.length).toBe(3) // Truncated with 0.5 fraction + expect(result2.summary).toBe("") + expect(result2.cost).toBe(0) + expect(result2.prevContextTokens).toBe(80001) }) it("should handle small context windows appropriately", async () => { @@ -508,7 +688,7 @@ describe("getMaxTokens", () => { maxTokens: modelInfo.maxTokens, apiHandler: mockApiHandler, }) - expect(result1).toEqual(messagesWithSmallContent) + expect(result1.messages).toEqual(messagesWithSmallContent) // Above max tokens - truncate const result2 = await truncateConversationIfNeeded({ @@ -519,7 +699,7 @@ describe("getMaxTokens", () => { apiHandler: mockApiHandler, }) expect(result2).not.toEqual(messagesWithSmallContent) - expect(result2.length).toBe(3) // Truncated with 0.5 fraction + expect(result2.messages.length).toBe(3) // Truncated with 0.5 fraction }) it("should handle large context windows appropriately", async () => { @@ -538,7 +718,7 @@ describe("getMaxTokens", () => { maxTokens: modelInfo.maxTokens, apiHandler: mockApiHandler, }) - expect(result1).toEqual(messagesWithSmallContent) + expect(result1.messages).toEqual(messagesWithSmallContent) // Above max tokens - truncate const result2 = await truncateConversationIfNeeded({ @@ -549,6 +729,6 @@ describe("getMaxTokens", () => { apiHandler: mockApiHandler, }) expect(result2).not.toEqual(messagesWithSmallContent) - expect(result2.length).toBe(3) // Truncated with 0.5 fraction + expect(result2.messages.length).toBe(3) // Truncated with 0.5 fraction }) })