ContextManager Tests (RooCodeInc#2415)

* add tests around truncating context

* remove comments

Author: celestial-vault

src/core/context-management/ContextManager.test.ts

@@ -0,0 +1,155 @@
+import { ContextManager } from "./ContextManager"
+import { Anthropic } from "@anthropic-ai/sdk"
+import { expect } from "chai"
+
+describe("ContextManager", () => {
+	function createMessages(count: number): Anthropic.Messages.MessageParam[] {
+		const messages: Anthropic.Messages.MessageParam[] = []
+
+		messages.push({
+			role: "user",
+			content: "Initial task message",
+		})
+
+		let role: "user" | "assistant" = "assistant"
+		for (let i = 1; i < count; i++) {
+			messages.push({
+				role,
+				content: `Message ${i}`,
+			})
+			role = role === "user" ? "assistant" : "user"
+		}
+
+		return messages
+	}
+
+	describe("getNextTruncationRange", () => {
+		let contextManager: ContextManager
+
+		beforeEach(() => {
+			contextManager = new ContextManager()
+		})
+
+		it("first truncation with half keep", () => {
+			const messages = createMessages(11)
+			const result = contextManager.getNextTruncationRange(messages)
+
+			expect(result).to.deep.equal([1, 4])
+		})
+
+		it("first truncation with quarter keep", () => {
+			const messages = createMessages(11)
+			const result = contextManager.getNextTruncationRange(messages, undefined, "quarter")
+
+			expect(result).to.deep.equal([1, 6])
+		})
+
+		it("sequential truncation with half keep", () => {
+			const messages = createMessages(21)
+			const firstRange = contextManager.getNextTruncationRange(messages)
+			expect(firstRange).to.deep.equal([1, 10])
+
+			// Pass the previous range for sequential truncation
+			const secondRange = contextManager.getNextTruncationRange(messages, firstRange)
+			expect(secondRange).to.deep.equal([1, 14])
+		})
+
+		it("sequential truncation with quarter keep", () => {
+			const messages = createMessages(41)
+			const firstRange = contextManager.getNextTruncationRange(messages, undefined, "quarter")
+
+			const secondRange = contextManager.getNextTruncationRange(messages, firstRange, "quarter")
+
+			expect(secondRange[0]).to.equal(1)
+			expect(secondRange[1]).to.be.greaterThan(firstRange[1])
+		})
+
+		it("ensures the last message in range is a user message", () => {
+			const messages = createMessages(14)
+			const result = contextManager.getNextTruncationRange(messages)
+
+			// Check if the message at the end of range is a user message
+			const lastRemovedMessage = messages[result[1]]
+			expect(lastRemovedMessage.role).to.equal("user")
+
+			// Check if the next message after the range is an assistant message
+			const nextMessage = messages[result[1] + 1]
+			expect(nextMessage.role).to.equal("assistant")
+		})
+
+		it("handles small message arrays", () => {
+			const messages = createMessages(3)
+			const result = contextManager.getNextTruncationRange(messages)
+
+			expect(result).to.deep.equal([1, 0])
+		})
+
+		it("preserves the message structure when truncating", () => {
+			const messages = createMessages(20)
+			const result = contextManager.getNextTruncationRange(messages)
+
+			// Get messages after removing the range
+			const effectiveMessages = [...messages.slice(0, result[0]), ...messages.slice(result[1] + 1)]
+
+			// Check first message and alternating pattern
+			expect(effectiveMessages[0].role).to.equal("user")
+			for (let i = 1; i < effectiveMessages.length; i++) {
+				const expectedRole = i % 2 === 1 ? "assistant" : "user"
+				expect(effectiveMessages[i].role).to.equal(expectedRole)
+			}
+		})
+	})
+
+	describe("getTruncatedMessages", () => {
+		let contextManager: ContextManager
+
+		beforeEach(() => {
+			contextManager = new ContextManager()
+		})
+
+		it("returns original messages when no range is provided", () => {
+			const messages = createMessages(3)
+
+			const result = contextManager.getTruncatedMessages(messages, undefined)
+			expect(result).to.deep.equal(messages)
+		})
+
+		it("correctly removes messages in the specified range", () => {
+			const messages = createMessages(5)
+
+			const range: [number, number] = [1, 3]
+			const result = contextManager.getTruncatedMessages(messages, range)
+
+			expect(result).to.have.lengthOf(2)
+			expect(result[0]).to.deep.equal(messages[0])
+			expect(result[1]).to.deep.equal(messages[4])
+		})
+
+		it("works with a range that starts at the first message after task", () => {
+			const messages = createMessages(4)
+
+			const range: [number, number] = [1, 2]
+			const result = contextManager.getTruncatedMessages(messages, range)
+
+			expect(result).to.have.lengthOf(2)
+			expect(result[0]).to.deep.equal(messages[0])
+			expect(result[1]).to.deep.equal(messages[3])
+		})
+
+		it("correctly handles removing a range while preserving alternation pattern", () => {
+			const messages = createMessages(5)
+
+			const range: [number, number] = [1, 2]
+			const result = contextManager.getTruncatedMessages(messages, range)
+
+			expect(result).to.have.lengthOf(3)
+			expect(result[0]).to.deep.equal(messages[0])
+			expect(result[1]).to.deep.equal(messages[3])
+			expect(result[2]).to.deep.equal(messages[4])
+
+			expect(result[0].role).to.equal("user")
+			expect(result[1].role).to.equal("assistant")
+			expect(result[2].role).to.equal("user")
+		})
+	})
+})