-
Notifications
You must be signed in to change notification settings - Fork 2.4k
Fix: Preserve first message during conversation condensing #7910
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,260 @@ | ||
| // npx vitest src/core/condense/__tests__/condense.spec.ts | ||
|
|
||
| import { Anthropic } from "@anthropic-ai/sdk" | ||
| import type { ModelInfo } from "@roo-code/types" | ||
| import { TelemetryService } from "@roo-code/telemetry" | ||
|
|
||
| import { BaseProvider } from "../../../api/providers/base-provider" | ||
| import { ApiMessage } from "../../task-persistence/apiMessages" | ||
| import { summarizeConversation, getMessagesSinceLastSummary, N_MESSAGES_TO_KEEP } from "../index" | ||
|
|
||
| // Create a mock ApiHandler for testing | ||
| class MockApiHandler extends BaseProvider { | ||
| createMessage(): any { | ||
| // Mock implementation for testing - returns an async iterable stream | ||
| const mockStream = { | ||
| async *[Symbol.asyncIterator]() { | ||
| yield { type: "text", text: "Mock summary of the conversation" } | ||
| yield { type: "usage", inputTokens: 100, outputTokens: 50, totalCost: 0.01 } | ||
| }, | ||
| } | ||
| return mockStream | ||
| } | ||
|
|
||
| getModel(): { id: string; info: ModelInfo } { | ||
| return { | ||
| id: "test-model", | ||
| info: { | ||
| contextWindow: 100000, | ||
| maxTokens: 50000, | ||
| supportsPromptCache: true, | ||
| supportsImages: false, | ||
| inputPrice: 0, | ||
| outputPrice: 0, | ||
| description: "Test model", | ||
| }, | ||
| } | ||
| } | ||
|
|
||
| override async countTokens(content: Array<Anthropic.Messages.ContentBlockParam>): Promise<number> { | ||
| // Simple token counting for testing | ||
| let tokens = 0 | ||
| for (const block of content) { | ||
| if (block.type === "text") { | ||
| tokens += Math.ceil(block.text.length / 4) // Rough approximation | ||
| } | ||
| } | ||
| return tokens | ||
| } | ||
| } | ||
|
|
||
| const mockApiHandler = new MockApiHandler() | ||
| const taskId = "test-task-id" | ||
|
|
||
| describe("Condense", () => { | ||
| beforeEach(() => { | ||
| if (!TelemetryService.hasInstance()) { | ||
| TelemetryService.createInstance([]) | ||
| } | ||
| }) | ||
|
|
||
| describe("summarizeConversation", () => { | ||
| it("should preserve the first message when summarizing", async () => { | ||
| const messages: ApiMessage[] = [ | ||
| { role: "user", content: "First message with /prr command content" }, | ||
| { role: "assistant", content: "Second message" }, | ||
| { role: "user", content: "Third message" }, | ||
| { role: "assistant", content: "Fourth message" }, | ||
| { role: "user", content: "Fifth message" }, | ||
| { role: "assistant", content: "Sixth message" }, | ||
| { role: "user", content: "Seventh message" }, | ||
| { role: "assistant", content: "Eighth message" }, | ||
| { role: "user", content: "Ninth message" }, | ||
| ] | ||
|
|
||
| const result = await summarizeConversation(messages, mockApiHandler, "System prompt", taskId, 5000, false) | ||
|
|
||
| // Verify the first message is preserved | ||
| expect(result.messages[0]).toEqual(messages[0]) | ||
| expect(result.messages[0].content).toBe("First message with /prr command content") | ||
|
|
||
| // Verify we have a summary message | ||
| const summaryMessage = result.messages.find((msg) => msg.isSummary) | ||
| expect(summaryMessage).toBeTruthy() | ||
| expect(summaryMessage?.content).toBe("Mock summary of the conversation") | ||
|
|
||
| // Verify we have the expected number of messages | ||
| // [first message, summary, last N messages] | ||
| expect(result.messages.length).toBe(1 + 1 + N_MESSAGES_TO_KEEP) | ||
|
|
||
| // Verify the last N messages are preserved | ||
| const lastMessages = result.messages.slice(-N_MESSAGES_TO_KEEP) | ||
| expect(lastMessages).toEqual(messages.slice(-N_MESSAGES_TO_KEEP)) | ||
| }) | ||
|
|
||
| it("should preserve slash command content in the first message", async () => { | ||
| const slashCommandContent = "/prr #123 - Fix authentication bug" | ||
| const messages: ApiMessage[] = [ | ||
| { role: "user", content: slashCommandContent }, | ||
| { role: "assistant", content: "I'll help you fix that authentication bug" }, | ||
| { role: "user", content: "The issue is with JWT tokens" }, | ||
| { role: "assistant", content: "Let me examine the JWT implementation" }, | ||
| { role: "user", content: "It's failing on refresh" }, | ||
| { role: "assistant", content: "I found the issue" }, | ||
| { role: "user", content: "Great, can you fix it?" }, | ||
| { role: "assistant", content: "Here's the fix" }, | ||
| { role: "user", content: "Thanks!" }, | ||
| ] | ||
|
|
||
| const result = await summarizeConversation(messages, mockApiHandler, "System prompt", taskId, 5000, false) | ||
|
|
||
| // The first message with slash command should be intact | ||
| expect(result.messages[0].content).toBe(slashCommandContent) | ||
| expect(result.messages[0]).toEqual(messages[0]) | ||
| }) | ||
|
|
||
| it("should handle complex first message content", async () => { | ||
| const complexContent: Anthropic.Messages.ContentBlockParam[] = [ | ||
| { type: "text", text: "/mode code" }, | ||
| { type: "text", text: "Additional context from the user" }, | ||
| ] | ||
|
|
||
| const messages: ApiMessage[] = [ | ||
| { role: "user", content: complexContent }, | ||
| { role: "assistant", content: "Switching to code mode" }, | ||
| { role: "user", content: "Write a function" }, | ||
| { role: "assistant", content: "Here's the function" }, | ||
| { role: "user", content: "Add error handling" }, | ||
| { role: "assistant", content: "Added error handling" }, | ||
| { role: "user", content: "Add tests" }, | ||
| { role: "assistant", content: "Tests added" }, | ||
| { role: "user", content: "Perfect!" }, | ||
| ] | ||
|
|
||
| const result = await summarizeConversation(messages, mockApiHandler, "System prompt", taskId, 5000, false) | ||
|
|
||
| // The first message with complex content should be preserved | ||
| expect(result.messages[0].content).toEqual(complexContent) | ||
| expect(result.messages[0]).toEqual(messages[0]) | ||
| }) | ||
|
|
||
| it("should return error when not enough messages to summarize", async () => { | ||
| const messages: ApiMessage[] = [ | ||
| { role: "user", content: "First message with /command" }, | ||
| { role: "assistant", content: "Second message" }, | ||
| { role: "user", content: "Third message" }, | ||
| { role: "assistant", content: "Fourth message" }, | ||
| ] | ||
|
|
||
| const result = await summarizeConversation(messages, mockApiHandler, "System prompt", taskId, 5000, false) | ||
|
|
||
| // Should return an error since we have only 4 messages (first + 3 to keep) | ||
| expect(result.error).toBeDefined() | ||
| expect(result.messages).toEqual(messages) // Original messages unchanged | ||
| expect(result.summary).toBe("") | ||
| }) | ||
|
|
||
| it("should not summarize messages that already contain a recent summary", async () => { | ||
| const messages: ApiMessage[] = [ | ||
| { role: "user", content: "First message with /command" }, | ||
| { role: "assistant", content: "Old message" }, | ||
| { role: "user", content: "Message before summary" }, | ||
| { role: "assistant", content: "Response" }, | ||
| { role: "user", content: "Another message" }, | ||
| { role: "assistant", content: "Previous summary", isSummary: true }, // Summary in last N messages | ||
| { role: "user", content: "Final message" }, | ||
| ] | ||
|
|
||
| const result = await summarizeConversation(messages, mockApiHandler, "System prompt", taskId, 5000, false) | ||
|
|
||
| // Should return an error due to recent summary in last N messages | ||
| expect(result.error).toBeDefined() | ||
| expect(result.messages).toEqual(messages) | ||
| expect(result.summary).toBe("") | ||
| }) | ||
|
|
||
| it("should handle empty summary from API gracefully", async () => { | ||
| // Mock handler that returns empty summary | ||
| class EmptyMockApiHandler extends MockApiHandler { | ||
| override createMessage(): any { | ||
| const mockStream = { | ||
| async *[Symbol.asyncIterator]() { | ||
| yield { type: "text", text: "" } | ||
| yield { type: "usage", inputTokens: 100, outputTokens: 0, totalCost: 0.01 } | ||
| }, | ||
| } | ||
| return mockStream | ||
| } | ||
| } | ||
|
|
||
| const emptyHandler = new EmptyMockApiHandler() | ||
| const messages: ApiMessage[] = [ | ||
| { role: "user", content: "First message" }, | ||
| { role: "assistant", content: "Second" }, | ||
| { role: "user", content: "Third" }, | ||
| { role: "assistant", content: "Fourth" }, | ||
| { role: "user", content: "Fifth" }, | ||
| { role: "assistant", content: "Sixth" }, | ||
| { role: "user", content: "Seventh" }, | ||
| ] | ||
|
|
||
| const result = await summarizeConversation(messages, emptyHandler, "System prompt", taskId, 5000, false) | ||
|
|
||
| expect(result.error).toBeDefined() | ||
| expect(result.messages).toEqual(messages) | ||
| expect(result.cost).toBeGreaterThan(0) | ||
| }) | ||
| }) | ||
|
|
||
| describe("getMessagesSinceLastSummary", () => { | ||
| it("should return all messages when no summary exists", () => { | ||
| const messages: ApiMessage[] = [ | ||
| { role: "user", content: "First message" }, | ||
| { role: "assistant", content: "Second message" }, | ||
| { role: "user", content: "Third message" }, | ||
| ] | ||
|
|
||
| const result = getMessagesSinceLastSummary(messages) | ||
| expect(result).toEqual(messages) | ||
| }) | ||
|
|
||
| it("should return messages since last summary including the summary", () => { | ||
| const messages: ApiMessage[] = [ | ||
| { role: "user", content: "First message" }, | ||
| { role: "assistant", content: "Second message" }, | ||
| { role: "assistant", content: "Summary content", isSummary: true }, | ||
| { role: "user", content: "Message after summary" }, | ||
| { role: "assistant", content: "Final message" }, | ||
| ] | ||
|
|
||
| const result = getMessagesSinceLastSummary(messages) | ||
|
|
||
| // Should include a user message prefix for Bedrock compatibility, the summary, and messages after | ||
| expect(result[0].role).toBe("user") | ||
| expect(result[0].content).toBe("Please continue from the following summary:") | ||
| expect(result[1]).toEqual(messages[2]) // The summary | ||
| expect(result[2]).toEqual(messages[3]) | ||
| expect(result[3]).toEqual(messages[4]) | ||
| }) | ||
|
|
||
| it("should handle multiple summaries and return from the last one", () => { | ||
| const messages: ApiMessage[] = [ | ||
| { role: "user", content: "First message" }, | ||
| { role: "assistant", content: "First summary", isSummary: true }, | ||
| { role: "user", content: "Middle message" }, | ||
| { role: "assistant", content: "Second summary", isSummary: true }, | ||
| { role: "user", content: "Recent message" }, | ||
| { role: "assistant", content: "Final message" }, | ||
| ] | ||
|
|
||
| const result = getMessagesSinceLastSummary(messages) | ||
|
|
||
| // Should only include from the last summary | ||
| expect(result[0].role).toBe("user") | ||
| expect(result[0].content).toBe("Please continue from the following summary:") | ||
| expect(result[1]).toEqual(messages[3]) // Second summary | ||
| expect(result[2]).toEqual(messages[4]) | ||
| expect(result[3]).toEqual(messages[5]) | ||
| }) | ||
| }) | ||
| }) |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -100,7 +100,11 @@ export async function summarizeConversation( | |
| ) | ||
|
|
||
| const response: SummarizeResponse = { messages, cost: 0, summary: "" } | ||
| const messagesToSummarize = getMessagesSinceLastSummary(messages.slice(0, -N_MESSAGES_TO_KEEP)) | ||
|
|
||
| // Always preserve the first message (which may contain slash command content) | ||
| const firstMessage = messages[0] | ||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Could we add a safety check here to handle the edge case where the messages array might be empty? While unlikely in practice, it would make the code more robust. Consider checking if messages.length === 0 before accessing messages[0]. |
||
| // Get messages to summarize, excluding the first message and last N messages | ||
| const messagesToSummarize = getMessagesSinceLastSummary(messages.slice(1, -N_MESSAGES_TO_KEEP)) | ||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The logic here creates an implicit minimum of 4 messages (first + 3 to keep). Would it be clearer to extract this as a constant like MIN_MESSAGES_FOR_CONDENSING = N_MESSAGES_TO_KEEP + 1? Then the condition could be more explicit about the requirement. |
||
|
|
||
| if (messagesToSummarize.length <= 1) { | ||
| const error = | ||
|
|
@@ -184,7 +188,8 @@ export async function summarizeConversation( | |
| isSummary: true, | ||
| } | ||
|
|
||
| const newMessages = [...messages.slice(0, -N_MESSAGES_TO_KEEP), summaryMessage, ...keepMessages] | ||
| // Reconstruct messages: [first message, summary, last N messages] | ||
| const newMessages = [firstMessage, summaryMessage, ...keepMessages] | ||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Nice work maintaining consistency with the structure. The reconstruction clearly shows the preserved first message, followed by the summary, and then the recent messages. |
||
|
|
||
| // Count the tokens in the context for the next API request | ||
| // We only estimate the tokens in summaryMesage if outputTokens is 0, otherwise we use outputTokens | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Consider enhancing this comment to explicitly mention why we preserve the first message: 'Always preserve the first message (which may contain slash command content, initial user request, and important context)'. This would help future maintainers understand the critical importance of this preservation.