Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
260 changes: 260 additions & 0 deletions src/core/condense/__tests__/condense.spec.ts
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])
})
})
})
12 changes: 8 additions & 4 deletions src/core/condense/__tests__/index.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -188,11 +188,14 @@ describe("summarizeConversation", () => {
expect(maybeRemoveImageBlocks).toHaveBeenCalled()

// Verify the structure of the result
// The result should be: original messages (except last N) + summary + last N messages
expect(result.messages.length).toBe(messages.length + 1) // Original + summary
// The result should be: first message + summary + last N messages
expect(result.messages.length).toBe(1 + 1 + N_MESSAGES_TO_KEEP) // First + summary + last N

// Check that the first message is preserved
expect(result.messages[0]).toEqual(messages[0])

// Check that the summary message was inserted correctly
const summaryMessage = result.messages[result.messages.length - N_MESSAGES_TO_KEEP - 1]
const summaryMessage = result.messages[1]
expect(summaryMessage.role).toBe("assistant")
expect(summaryMessage.content).toBe("This is a summary")
expect(summaryMessage.isSummary).toBe(true)
Expand Down Expand Up @@ -395,7 +398,8 @@ describe("summarizeConversation", () => {
)

// Should successfully summarize
expect(result.messages.length).toBe(messages.length + 1) // Original + summary
// Result should be: first message + summary + last N messages
expect(result.messages.length).toBe(1 + 1 + N_MESSAGES_TO_KEEP) // First + summary + last N
expect(result.cost).toBe(0.03)
expect(result.summary).toBe("Concise summary")
expect(result.error).toBeUndefined()
Expand Down
9 changes: 7 additions & 2 deletions src/core/condense/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Copy link

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.

const firstMessage = messages[0]
Copy link

Choose a reason for hiding this comment

The 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))
Copy link

Choose a reason for hiding this comment

The 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 =
Expand Down Expand Up @@ -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]
Copy link

Choose a reason for hiding this comment

The 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
Expand Down
Loading