From a3ee9bf3892bde7105e1092aee967661e3a2cf8c Mon Sep 17 00:00:00 2001 From: Hannes Rudolph Date: Thu, 4 Dec 2025 14:18:40 -0700 Subject: [PATCH 1/8] feat: add MessageManager layer for centralized history coordination Implements the proposal from docs/proposals/message-manager-layer.md This introduces a MessageManager class that centralizes all conversation rewind operations, ensuring consistent handling of both UI (clineMessages) and API (apiConversationHistory) message histories. Key changes: - Add MessageManager class with rewindToTimestamp() and rewindToIndex() methods - Add 22 comprehensive unit tests for MessageManager - Refactor webviewMessageHandler.ts to use MessageManager (removes 79-line inline function) - Fix checkpoint restore to properly handle context-management events - Add messageManager getter to Task class with lazy initialization Bug fix: Checkpoint restore now properly: - Collects condenseIds from removed condense_context messages - Collects truncationIds from removed sliding_window_truncation messages - Removes orphaned Summary messages from API history - Removes orphaned truncation markers from API history - Calls cleanupAfterTruncation() to clear orphaned parent tags All 4508 tests pass with no regressions. --- src/core/checkpoints/index.ts | 16 +- src/core/message-manager/index.spec.ts | 731 ++++++++++++++++++ src/core/message-manager/index.ts | 184 +++++ src/core/task/Task.ts | 22 + .../webviewMessageHandler.delete.spec.ts | 17 +- src/core/webview/webviewMessageHandler.ts | 94 +-- 6 files changed, 965 insertions(+), 99 deletions(-) create mode 100644 src/core/message-manager/index.spec.ts create mode 100644 src/core/message-manager/index.ts diff --git a/src/core/checkpoints/index.ts b/src/core/checkpoints/index.ts index 3efdb466e66..8adb04a7bb6 100644 --- a/src/core/checkpoints/index.ts +++ b/src/core/checkpoints/index.ts @@ -4,6 +4,7 @@ import * as vscode from "vscode" import { TelemetryService } from "@roo-code/telemetry" import { Task } from "../task/Task" +import { MessageManager } from "../message-manager" import { getWorkspacePath } from "../../utils/path" import { checkGitInstalled } from "../../utils/git" @@ -258,20 +259,21 @@ export async function checkpointRestore( await provider?.postMessageToWebview({ type: "currentCheckpointUpdated", text: commitHash }) if (mode === "restore") { - await task.overwriteApiConversationHistory(task.apiConversationHistory.filter((m) => !m.ts || m.ts < ts)) - + // Calculate metrics from messages that will be deleted (must be done before rewind) const deletedMessages = task.clineMessages.slice(index + 1) const { totalTokensIn, totalTokensOut, totalCacheWrites, totalCacheReads, totalCost } = getApiMetrics( task.combineMessages(deletedMessages), ) - // For delete operations, exclude the checkpoint message itself - // For edit operations, include the checkpoint message (to be edited) - const endIndex = operation === "edit" ? index + 1 : index - await task.overwriteClineMessages(task.clineMessages.slice(0, endIndex)) + // Use MessageManager to properly handle context-management events + // This ensures orphaned Summary messages and truncation markers are cleaned up + const manager = new MessageManager(task) + await manager.rewindToTimestamp(ts, { + includeTargetMessage: operation === "edit", + }) - // TODO: Verify that this is working as expected. + // Report the deleted API request metrics await task.say( "api_req_deleted", JSON.stringify({ diff --git a/src/core/message-manager/index.spec.ts b/src/core/message-manager/index.spec.ts new file mode 100644 index 00000000000..84ee10e8956 --- /dev/null +++ b/src/core/message-manager/index.spec.ts @@ -0,0 +1,731 @@ +import { MessageManager } from "./index" +import * as condenseModule from "../condense" + +describe("MessageManager", () => { + let mockTask: any + let manager: MessageManager + let cleanupAfterTruncationSpy: any + + beforeEach(() => { + mockTask = { + clineMessages: [], + apiConversationHistory: [], + overwriteClineMessages: vi.fn(), + overwriteApiConversationHistory: vi.fn(), + } + manager = new MessageManager(mockTask) + + // Mock cleanupAfterTruncation to track calls and return input by default + cleanupAfterTruncationSpy = vi.spyOn(condenseModule, "cleanupAfterTruncation") + cleanupAfterTruncationSpy.mockImplementation((messages: any[]) => messages) + }) + + afterEach(() => { + cleanupAfterTruncationSpy.mockRestore() + }) + + describe("Basic rewind operations", () => { + it("should remove messages at and after the target timestamp", async () => { + mockTask.clineMessages = [ + { ts: 100, say: "user", text: "First" }, + { ts: 200, say: "assistant", text: "Response" }, + { ts: 300, say: "user", text: "Second" }, + { ts: 400, say: "assistant", text: "Response 2" }, + ] + + mockTask.apiConversationHistory = [ + { ts: 100, role: "user", content: [{ type: "text", text: "First" }] }, + { ts: 200, role: "assistant", content: [{ type: "text", text: "Response" }] }, + { ts: 300, role: "user", content: [{ type: "text", text: "Second" }] }, + { ts: 400, role: "assistant", content: [{ type: "text", text: "Response 2" }] }, + ] + + await manager.rewindToTimestamp(300) + + // Should keep messages before ts=300 + expect(mockTask.overwriteClineMessages).toHaveBeenCalledWith([ + { ts: 100, say: "user", text: "First" }, + { ts: 200, say: "assistant", text: "Response" }, + ]) + + // Should keep API messages before ts=300 + const apiCall = mockTask.overwriteApiConversationHistory.mock.calls[0][0] + expect(apiCall).toHaveLength(2) + expect(apiCall[0].ts).toBe(100) + expect(apiCall[1].ts).toBe(200) + }) + + it("should keep target message when includeTargetMessage is true", async () => { + mockTask.clineMessages = [ + { ts: 100, say: "user", text: "First" }, + { ts: 200, say: "assistant", text: "Response" }, + { ts: 300, say: "user", text: "Second" }, + ] + + mockTask.apiConversationHistory = [ + { ts: 100, role: "user", content: [{ type: "text", text: "First" }] }, + { ts: 200, role: "assistant", content: [{ type: "text", text: "Response" }] }, + { ts: 300, role: "user", content: [{ type: "text", text: "Second" }] }, + ] + + await manager.rewindToTimestamp(300, { includeTargetMessage: true }) + + // Should keep messages up to and including ts=300 in clineMessages + expect(mockTask.overwriteClineMessages).toHaveBeenCalledWith([ + { ts: 100, say: "user", text: "First" }, + { ts: 200, say: "assistant", text: "Response" }, + { ts: 300, say: "user", text: "Second" }, + ]) + + // API history uses ts < cutoffTs, so excludes the message at ts=300 + // This is correct for edit scenarios - keep UI message but truncate API before it + const apiCall = mockTask.overwriteApiConversationHistory.mock.calls[0][0] + expect(apiCall).toHaveLength(2) + expect(apiCall[0].ts).toBe(100) + expect(apiCall[1].ts).toBe(200) + }) + + it("should throw error when timestamp not found", async () => { + mockTask.clineMessages = [ + { ts: 100, say: "user", text: "First" }, + { ts: 200, say: "assistant", text: "Response" }, + ] + + await expect(manager.rewindToTimestamp(999)).rejects.toThrow( + "Message with timestamp 999 not found in clineMessages", + ) + }) + + it("should remove messages at and after the target index", async () => { + mockTask.clineMessages = [ + { ts: 100, say: "user", text: "First" }, + { ts: 200, say: "assistant", text: "Response" }, + { ts: 300, say: "user", text: "Second" }, + { ts: 400, say: "assistant", text: "Response 2" }, + ] + + mockTask.apiConversationHistory = [ + { ts: 100, role: "user", content: [{ type: "text", text: "First" }] }, + { ts: 200, role: "assistant", content: [{ type: "text", text: "Response" }] }, + { ts: 300, role: "user", content: [{ type: "text", text: "Second" }] }, + { ts: 400, role: "assistant", content: [{ type: "text", text: "Response 2" }] }, + ] + + await manager.rewindToIndex(2) + + // Should keep messages [0, 2) - index 0 and 1 + expect(mockTask.overwriteClineMessages).toHaveBeenCalledWith([ + { ts: 100, say: "user", text: "First" }, + { ts: 200, say: "assistant", text: "Response" }, + ]) + + // Should keep API messages before ts=300 + const apiCall = mockTask.overwriteApiConversationHistory.mock.calls[0][0] + expect(apiCall).toHaveLength(2) + }) + }) + + describe("Condense handling", () => { + it("should preserve Summary when condense_context is preserved", async () => { + const condenseId = "summary-123" + + mockTask.clineMessages = [ + { ts: 100, say: "user", text: "First" }, + { ts: 200, say: "assistant", text: "Response" }, + { ts: 300, say: "condense_context", contextCondense: { condenseId, summary: "Summary" } }, + { ts: 400, say: "user", text: "After condense" }, + ] + + mockTask.apiConversationHistory = [ + { ts: 100, role: "user", content: [{ type: "text", text: "First" }] }, + { + ts: 200, + role: "assistant", + content: [{ type: "text", text: "Response" }], + condenseParent: condenseId, + }, + { + ts: 299, + role: "assistant", + content: [{ type: "text", text: "Summary" }], + isSummary: true, + condenseId, + }, + { ts: 400, role: "user", content: [{ type: "text", text: "After condense" }] }, + ] + + // Rewind to ts=400, which preserves condense_context at ts=300 + await manager.rewindToTimestamp(400) + + // Summary should still exist + const apiCall = mockTask.overwriteApiConversationHistory.mock.calls[0][0] + const hasSummary = apiCall.some((m: any) => m.isSummary && m.condenseId === condenseId) + expect(hasSummary).toBe(true) + }) + + it("should remove Summary when condense_context is removed", async () => { + const condenseId = "summary-123" + + mockTask.clineMessages = [ + { ts: 100, say: "user", text: "First" }, + { ts: 200, say: "assistant", text: "Response" }, + { ts: 300, say: "user", text: "Second" }, + { ts: 400, say: "condense_context", contextCondense: { condenseId, summary: "Summary" } }, + { ts: 500, say: "user", text: "Third" }, + ] + + mockTask.apiConversationHistory = [ + { ts: 100, role: "user", content: [{ type: "text", text: "First" }] }, + { + ts: 200, + role: "assistant", + content: [{ type: "text", text: "Response" }], + condenseParent: condenseId, + }, + { + ts: 299, + role: "assistant", + content: [{ type: "text", text: "Summary" }], + isSummary: true, + condenseId, + }, + { ts: 300, role: "user", content: [{ type: "text", text: "Second" }] }, + { ts: 500, role: "user", content: [{ type: "text", text: "Third" }] }, + ] + + // Rewind to ts=300, which removes condense_context at ts=400 + await manager.rewindToTimestamp(300) + + // Summary should be removed + const apiCall = mockTask.overwriteApiConversationHistory.mock.calls[0][0] + const hasSummary = apiCall.some((m: any) => m.isSummary) + expect(hasSummary).toBe(false) + }) + + it("should clear orphaned condenseParent tags via cleanup", async () => { + const condenseId = "summary-123" + + mockTask.clineMessages = [ + { ts: 100, say: "user", text: "First" }, + { ts: 200, say: "condense_context", contextCondense: { condenseId, summary: "Summary" } }, + ] + + mockTask.apiConversationHistory = [ + { ts: 100, role: "user", content: [{ type: "text", text: "First" }] }, + { + ts: 150, + role: "assistant", + content: [{ type: "text", text: "Response" }], + condenseParent: condenseId, + }, + { + ts: 199, + role: "assistant", + content: [{ type: "text", text: "Summary" }], + isSummary: true, + condenseId, + }, + ] + + // Rewind to ts=100, which removes condense_context + await manager.rewindToTimestamp(100) + + // cleanupAfterTruncation should be called to remove orphaned tags + expect(cleanupAfterTruncationSpy).toHaveBeenCalled() + }) + + it("should handle multiple condense_context removals", async () => { + const condenseId1 = "summary-1" + const condenseId2 = "summary-2" + + mockTask.clineMessages = [ + { ts: 100, say: "user", text: "First" }, + { + ts: 200, + say: "condense_context", + contextCondense: { condenseId: condenseId1, summary: "Summary 1" }, + }, + { ts: 300, say: "user", text: "Second" }, + { + ts: 400, + say: "condense_context", + contextCondense: { condenseId: condenseId2, summary: "Summary 2" }, + }, + { ts: 500, say: "user", text: "Third" }, + ] + + mockTask.apiConversationHistory = [ + { ts: 100, role: "user", content: [{ type: "text", text: "First" }] }, + { + ts: 199, + role: "assistant", + content: [{ type: "text", text: "Summary 1" }], + isSummary: true, + condenseId: condenseId1, + }, + { ts: 300, role: "user", content: [{ type: "text", text: "Second" }] }, + { + ts: 399, + role: "assistant", + content: [{ type: "text", text: "Summary 2" }], + isSummary: true, + condenseId: condenseId2, + }, + { ts: 500, role: "user", content: [{ type: "text", text: "Third" }] }, + ] + + // Rewind to ts=200, which removes both condense_context messages + await manager.rewindToTimestamp(200) + + // Both summaries should be removed + const apiCall = mockTask.overwriteApiConversationHistory.mock.calls[0][0] + const hasSummary1 = apiCall.some((m: any) => m.condenseId === condenseId1) + const hasSummary2 = apiCall.some((m: any) => m.condenseId === condenseId2) + expect(hasSummary1).toBe(false) + expect(hasSummary2).toBe(false) + }) + }) + + describe("Truncation handling", () => { + it("should preserve truncation marker when sliding_window_truncation is preserved", async () => { + const truncationId = "trunc-123" + + mockTask.clineMessages = [ + { ts: 100, say: "user", text: "First" }, + { ts: 200, say: "sliding_window_truncation", contextTruncation: { truncationId, reason: "window" } }, + { ts: 300, say: "user", text: "After truncation" }, + ] + + mockTask.apiConversationHistory = [ + { ts: 100, role: "user", content: [{ type: "text", text: "First" }], truncationParent: truncationId }, + { + ts: 199, + role: "assistant", + content: [{ type: "text", text: "..." }], + isTruncationMarker: true, + truncationId, + }, + { ts: 300, role: "user", content: [{ type: "text", text: "After truncation" }] }, + ] + + // Rewind to ts=300, which preserves sliding_window_truncation at ts=200 + await manager.rewindToTimestamp(300) + + // Truncation marker should still exist + const apiCall = mockTask.overwriteApiConversationHistory.mock.calls[0][0] + const hasMarker = apiCall.some((m: any) => m.isTruncationMarker && m.truncationId === truncationId) + expect(hasMarker).toBe(true) + }) + + it("should remove truncation marker when sliding_window_truncation is removed", async () => { + const truncationId = "trunc-123" + + mockTask.clineMessages = [ + { ts: 100, say: "user", text: "First" }, + { ts: 200, say: "user", text: "Second" }, + { ts: 300, say: "sliding_window_truncation", contextTruncation: { truncationId, reason: "window" } }, + { ts: 400, say: "user", text: "Third" }, + ] + + mockTask.apiConversationHistory = [ + { ts: 100, role: "user", content: [{ type: "text", text: "First" }], truncationParent: truncationId }, + { ts: 200, role: "user", content: [{ type: "text", text: "Second" }] }, + { + ts: 299, + role: "assistant", + content: [{ type: "text", text: "..." }], + isTruncationMarker: true, + truncationId, + }, + { ts: 400, role: "user", content: [{ type: "text", text: "Third" }] }, + ] + + // Rewind to ts=200, which removes sliding_window_truncation at ts=300 + await manager.rewindToTimestamp(200) + + // Truncation marker should be removed + const apiCall = mockTask.overwriteApiConversationHistory.mock.calls[0][0] + const hasMarker = apiCall.some((m: any) => m.isTruncationMarker) + expect(hasMarker).toBe(false) + }) + + it("should clear orphaned truncationParent tags via cleanup", async () => { + const truncationId = "trunc-123" + + mockTask.clineMessages = [ + { ts: 100, say: "user", text: "First" }, + { ts: 200, say: "sliding_window_truncation", contextTruncation: { truncationId, reason: "window" } }, + ] + + mockTask.apiConversationHistory = [ + { ts: 100, role: "user", content: [{ type: "text", text: "First" }], truncationParent: truncationId }, + { + ts: 199, + role: "assistant", + content: [{ type: "text", text: "..." }], + isTruncationMarker: true, + truncationId, + }, + ] + + // Rewind to ts=100, which removes sliding_window_truncation + await manager.rewindToTimestamp(100) + + // cleanupAfterTruncation should be called to remove orphaned tags + expect(cleanupAfterTruncationSpy).toHaveBeenCalled() + }) + + it("should handle multiple truncation removals", async () => { + const truncationId1 = "trunc-1" + const truncationId2 = "trunc-2" + + mockTask.clineMessages = [ + { ts: 100, say: "user", text: "First" }, + { + ts: 200, + say: "sliding_window_truncation", + contextTruncation: { truncationId: truncationId1, reason: "window" }, + }, + { ts: 300, say: "user", text: "Second" }, + { + ts: 400, + say: "sliding_window_truncation", + contextTruncation: { truncationId: truncationId2, reason: "window" }, + }, + { ts: 500, say: "user", text: "Third" }, + ] + + mockTask.apiConversationHistory = [ + { ts: 100, role: "user", content: [{ type: "text", text: "First" }] }, + { + ts: 199, + role: "assistant", + content: [{ type: "text", text: "..." }], + isTruncationMarker: true, + truncationId: truncationId1, + }, + { ts: 300, role: "user", content: [{ type: "text", text: "Second" }] }, + { + ts: 399, + role: "assistant", + content: [{ type: "text", text: "..." }], + isTruncationMarker: true, + truncationId: truncationId2, + }, + { ts: 500, role: "user", content: [{ type: "text", text: "Third" }] }, + ] + + // Rewind to ts=200, which removes both truncation messages + await manager.rewindToTimestamp(200) + + // Both markers should be removed + const apiCall = mockTask.overwriteApiConversationHistory.mock.calls[0][0] + const hasMarker1 = apiCall.some((m: any) => m.truncationId === truncationId1) + const hasMarker2 = apiCall.some((m: any) => m.truncationId === truncationId2) + expect(hasMarker1).toBe(false) + expect(hasMarker2).toBe(false) + }) + }) + + describe("Checkpoint scenarios", () => { + it("should preserve Summary when checkpoint restore is BEFORE condense", async () => { + const condenseId = "summary-abc" + + mockTask.clineMessages = [ + { ts: 100, say: "user", text: "Task" }, + { ts: 500, say: "condense_context", contextCondense: { condenseId, summary: "Summary" } }, + { ts: 600, say: "checkpoint_saved", text: "checkpoint-hash" }, + { ts: 700, say: "user", text: "After checkpoint" }, + ] + + mockTask.apiConversationHistory = [ + { ts: 100, role: "user", content: [{ type: "text", text: "Task" }] }, + { + ts: 200, + role: "assistant", + content: [{ type: "text", text: "Response 1" }], + condenseParent: condenseId, + }, + { + ts: 499, + role: "assistant", + content: [{ type: "text", text: "Summary" }], + isSummary: true, + condenseId, + }, + { ts: 700, role: "user", content: [{ type: "text", text: "After checkpoint" }] }, + ] + + // Restore checkpoint at ts=600 (like checkpoint restore does) + await manager.rewindToTimestamp(600, { includeTargetMessage: true }) + + // Since condense_context (ts=500) is BEFORE checkpoint, it should be preserved + const clineCall = mockTask.overwriteClineMessages.mock.calls[0][0] + const hasCondenseContext = clineCall.some((m: any) => m.say === "condense_context") + expect(hasCondenseContext).toBe(true) + + // And the Summary should still exist + const apiCall = mockTask.overwriteApiConversationHistory.mock.calls[0][0] + const hasSummary = apiCall.some((m: any) => m.isSummary) + expect(hasSummary).toBe(true) + }) + + it("should remove Summary when checkpoint restore is AFTER condense", async () => { + const condenseId = "summary-xyz" + + mockTask.clineMessages = [ + { ts: 100, say: "user", text: "Task" }, + { ts: 200, say: "checkpoint_saved", text: "checkpoint-hash" }, + { ts: 300, say: "condense_context", contextCondense: { condenseId, summary: "Summary" } }, + { ts: 400, say: "user", text: "After condense" }, + ] + + mockTask.apiConversationHistory = [ + { ts: 100, role: "user", content: [{ type: "text", text: "Task" }] }, + { + ts: 150, + role: "assistant", + content: [{ type: "text", text: "Response" }], + condenseParent: condenseId, + }, + { + ts: 299, + role: "assistant", + content: [{ type: "text", text: "Summary" }], + isSummary: true, + condenseId, + }, + { ts: 400, role: "user", content: [{ type: "text", text: "After condense" }] }, + ] + + // Restore checkpoint at ts=200 (before the condense happened) + await manager.rewindToTimestamp(200, { includeTargetMessage: true }) + + // condense_context (ts=300) is AFTER checkpoint, so it should be removed + const clineCall = mockTask.overwriteClineMessages.mock.calls[0][0] + const hasCondenseContext = clineCall.some((m: any) => m.say === "condense_context") + expect(hasCondenseContext).toBe(false) + + // And the Summary should be removed too + const apiCall = mockTask.overwriteApiConversationHistory.mock.calls[0][0] + const hasSummary = apiCall.some((m: any) => m.isSummary) + expect(hasSummary).toBe(false) + }) + + it("should preserve truncation marker when checkpoint restore is BEFORE truncation", async () => { + const truncationId = "trunc-abc" + + mockTask.clineMessages = [ + { ts: 100, say: "user", text: "Task" }, + { ts: 500, say: "sliding_window_truncation", contextTruncation: { truncationId, reason: "window" } }, + { ts: 600, say: "checkpoint_saved", text: "checkpoint-hash" }, + { ts: 700, say: "user", text: "After checkpoint" }, + ] + + mockTask.apiConversationHistory = [ + { ts: 100, role: "user", content: [{ type: "text", text: "Task" }], truncationParent: truncationId }, + { + ts: 499, + role: "assistant", + content: [{ type: "text", text: "..." }], + isTruncationMarker: true, + truncationId, + }, + { ts: 700, role: "user", content: [{ type: "text", text: "After checkpoint" }] }, + ] + + // Restore checkpoint at ts=600 + await manager.rewindToTimestamp(600, { includeTargetMessage: true }) + + // Truncation should be preserved + const clineCall = mockTask.overwriteClineMessages.mock.calls[0][0] + const hasTruncation = clineCall.some((m: any) => m.say === "sliding_window_truncation") + expect(hasTruncation).toBe(true) + + // Marker should still exist + const apiCall = mockTask.overwriteApiConversationHistory.mock.calls[0][0] + const hasMarker = apiCall.some((m: any) => m.isTruncationMarker) + expect(hasMarker).toBe(true) + }) + + it("should remove truncation marker when checkpoint restore is AFTER truncation", async () => { + const truncationId = "trunc-xyz" + + mockTask.clineMessages = [ + { ts: 100, say: "user", text: "Task" }, + { ts: 200, say: "checkpoint_saved", text: "checkpoint-hash" }, + { ts: 300, say: "sliding_window_truncation", contextTruncation: { truncationId, reason: "window" } }, + { ts: 400, say: "user", text: "After truncation" }, + ] + + mockTask.apiConversationHistory = [ + { ts: 100, role: "user", content: [{ type: "text", text: "Task" }], truncationParent: truncationId }, + { + ts: 299, + role: "assistant", + content: [{ type: "text", text: "..." }], + isTruncationMarker: true, + truncationId, + }, + { ts: 400, role: "user", content: [{ type: "text", text: "After truncation" }] }, + ] + + // Restore checkpoint at ts=200 (before truncation happened) + await manager.rewindToTimestamp(200, { includeTargetMessage: true }) + + // Truncation should be removed + const clineCall = mockTask.overwriteClineMessages.mock.calls[0][0] + const hasTruncation = clineCall.some((m: any) => m.say === "sliding_window_truncation") + expect(hasTruncation).toBe(false) + + // Marker should be removed + const apiCall = mockTask.overwriteApiConversationHistory.mock.calls[0][0] + const hasMarker = apiCall.some((m: any) => m.isTruncationMarker) + expect(hasMarker).toBe(false) + }) + }) + + describe("Skip cleanup option", () => { + it("should NOT call cleanupAfterTruncation when skipCleanup is true", async () => { + const condenseId = "summary-123" + + mockTask.clineMessages = [ + { ts: 100, say: "user", text: "First" }, + { ts: 200, say: "condense_context", contextCondense: { condenseId, summary: "Summary" } }, + ] + + mockTask.apiConversationHistory = [ + { ts: 100, role: "user", content: [{ type: "text", text: "First" }] }, + { + ts: 150, + role: "assistant", + content: [{ type: "text", text: "Response" }], + condenseParent: condenseId, + }, + { + ts: 199, + role: "assistant", + content: [{ type: "text", text: "Summary" }], + isSummary: true, + condenseId, + }, + ] + + // Rewind with skipCleanup + await manager.rewindToTimestamp(100, { skipCleanup: true }) + + // cleanupAfterTruncation should NOT be called + expect(cleanupAfterTruncationSpy).not.toHaveBeenCalled() + }) + + it("should call cleanupAfterTruncation by default", async () => { + mockTask.clineMessages = [ + { ts: 100, say: "user", text: "First" }, + { ts: 200, say: "user", text: "Second" }, + ] + + mockTask.apiConversationHistory = [ + { ts: 100, role: "user", content: [{ type: "text", text: "First" }] }, + { ts: 200, role: "user", content: [{ type: "text", text: "Second" }] }, + ] + + // Rewind without options (skipCleanup defaults to false) + await manager.rewindToTimestamp(100) + + // cleanupAfterTruncation should be called + expect(cleanupAfterTruncationSpy).toHaveBeenCalled() + }) + + it("should call cleanupAfterTruncation when skipCleanup is explicitly false", async () => { + mockTask.clineMessages = [ + { ts: 100, say: "user", text: "First" }, + { ts: 200, say: "user", text: "Second" }, + ] + + mockTask.apiConversationHistory = [ + { ts: 100, role: "user", content: [{ type: "text", text: "First" }] }, + { ts: 200, role: "user", content: [{ type: "text", text: "Second" }] }, + ] + + // Rewind with skipCleanup explicitly false + await manager.rewindToTimestamp(100, { skipCleanup: false }) + + // cleanupAfterTruncation should be called + expect(cleanupAfterTruncationSpy).toHaveBeenCalled() + }) + }) + + describe("Combined scenarios", () => { + it("should handle both condense and truncation removal in the same rewind", async () => { + const condenseId = "summary-123" + const truncationId = "trunc-456" + + mockTask.clineMessages = [ + { ts: 100, say: "user", text: "First" }, + { ts: 200, say: "condense_context", contextCondense: { condenseId, summary: "Summary" } }, + { ts: 300, say: "sliding_window_truncation", contextTruncation: { truncationId, reason: "window" } }, + { ts: 400, say: "user", text: "After both" }, + ] + + mockTask.apiConversationHistory = [ + { ts: 100, role: "user", content: [{ type: "text", text: "First" }] }, + { + ts: 199, + role: "assistant", + content: [{ type: "text", text: "Summary" }], + isSummary: true, + condenseId, + }, + { + ts: 299, + role: "assistant", + content: [{ type: "text", text: "..." }], + isTruncationMarker: true, + truncationId, + }, + { ts: 400, role: "user", content: [{ type: "text", text: "After both" }] }, + ] + + // Rewind to ts=100, which removes both + await manager.rewindToTimestamp(100) + + // Both Summary and marker should be removed + const apiCall = mockTask.overwriteApiConversationHistory.mock.calls[0][0] + const hasSummary = apiCall.some((m: any) => m.isSummary) + const hasMarker = apiCall.some((m: any) => m.isTruncationMarker) + expect(hasSummary).toBe(false) + expect(hasMarker).toBe(false) + }) + + it("should handle empty clineMessages array", async () => { + mockTask.clineMessages = [] + mockTask.apiConversationHistory = [] + + await manager.rewindToIndex(0) + + expect(mockTask.overwriteClineMessages).toHaveBeenCalledWith([]) + // API history write is skipped when nothing changed (optimization) + expect(mockTask.overwriteApiConversationHistory).not.toHaveBeenCalled() + }) + + it("should handle messages without timestamps in API history", async () => { + mockTask.clineMessages = [ + { ts: 100, say: "user", text: "First" }, + { ts: 200, say: "user", text: "Second" }, + ] + + mockTask.apiConversationHistory = [ + { role: "system", content: [{ type: "text", text: "System message" }] }, // No ts + { ts: 100, role: "user", content: [{ type: "text", text: "First" }] }, + { ts: 200, role: "user", content: [{ type: "text", text: "Second" }] }, + ] + + await manager.rewindToTimestamp(100) + + // Should keep system message (no ts) and message at ts=100 + const apiCall = mockTask.overwriteApiConversationHistory.mock.calls[0][0] + expect(apiCall).toHaveLength(1) + expect(apiCall[0].role).toBe("system") + }) + }) +}) diff --git a/src/core/message-manager/index.ts b/src/core/message-manager/index.ts new file mode 100644 index 00000000000..bb65fea3acd --- /dev/null +++ b/src/core/message-manager/index.ts @@ -0,0 +1,184 @@ +import { Task } from "../task/Task" +import { ClineMessage } from "@roo-code/types" +import { ApiMessage } from "../task-persistence/apiMessages" +import { cleanupAfterTruncation } from "../condense" + +export interface RewindOptions { + /** Whether to include the target message in deletion (edit=true, delete=false) */ + includeTargetMessage?: boolean + /** Skip cleanup for special cases (default: false) */ + skipCleanup?: boolean +} + +interface ContextEventIds { + condenseIds: Set + truncationIds: Set +} + +/** + * MessageManager provides centralized handling for all conversation rewind operations. + * + * This ensures that whenever UI chat history is rewound (delete, edit, checkpoint restore, etc.), + * the API conversation history is properly maintained, including: + * - Removing orphaned Summary messages when their condense_context is removed + * - Removing orphaned truncation markers when their sliding_window_truncation is removed + * - Cleaning up orphaned condenseParent/truncationParent tags + * + * Usage: + * ```typescript + * const manager = new MessageManager(task) + * await manager.rewindToTimestamp(messageTs, { includeTargetMessage: false }) + * ``` + */ +export class MessageManager { + constructor(private task: Task) {} + + /** + * Rewind conversation to a specific timestamp. + * This is the SINGLE entry point for all message deletion operations. + * + * @param ts - The timestamp to rewind to + * @param options - Rewind options + * @throws Error if timestamp not found in clineMessages + */ + async rewindToTimestamp(ts: number, options: RewindOptions = {}): Promise { + const { includeTargetMessage = false, skipCleanup = false } = options + + // Find the index in clineMessages + const clineIndex = this.task.clineMessages.findIndex((m) => m.ts === ts) + if (clineIndex === -1) { + throw new Error(`Message with timestamp ${ts} not found in clineMessages`) + } + + // Calculate the actual cutoff index + const cutoffIndex = includeTargetMessage ? clineIndex + 1 : clineIndex + + await this.performRewind(cutoffIndex, ts, { skipCleanup }) + } + + /** + * Rewind conversation to a specific index in clineMessages. + * Keeps messages [0, toIndex) and removes [toIndex, end]. + * + * @param toIndex - The index to rewind to (exclusive) + * @param options - Rewind options + */ + async rewindToIndex(toIndex: number, options: RewindOptions = {}): Promise { + const cutoffTs = this.task.clineMessages[toIndex]?.ts ?? Date.now() + await this.performRewind(toIndex, cutoffTs, options) + } + + /** + * Internal method that performs the actual rewind operation. + */ + private async performRewind(toIndex: number, cutoffTs: number, options: RewindOptions): Promise { + const { skipCleanup = false } = options + + // Step 1: Collect context event IDs from messages being removed + const removedIds = this.collectRemovedContextEventIds(toIndex) + + // Step 2: Truncate clineMessages + await this.truncateClineMessages(toIndex) + + // Step 3: Truncate and clean API history (combined with cleanup for efficiency) + await this.truncateApiHistoryWithCleanup(cutoffTs, removedIds, skipCleanup) + } + + /** + * Collect condenseIds and truncationIds from context-management events + * that will be removed during the rewind. + * + * This is critical for maintaining the linkage between: + * - condense_context (clineMessage) ↔ Summary (apiMessage) + * - sliding_window_truncation (clineMessage) ↔ Truncation marker (apiMessage) + */ + private collectRemovedContextEventIds(fromIndex: number): ContextEventIds { + const condenseIds = new Set() + const truncationIds = new Set() + + for (let i = fromIndex; i < this.task.clineMessages.length; i++) { + const msg = this.task.clineMessages[i] + + // Collect condenseIds from condense_context events + if (msg.say === "condense_context" && msg.contextCondense?.condenseId) { + condenseIds.add(msg.contextCondense.condenseId) + console.log(`[MessageManager] Found condense_context to remove: ${msg.contextCondense.condenseId}`) + } + + // Collect truncationIds from sliding_window_truncation events + if (msg.say === "sliding_window_truncation" && msg.contextTruncation?.truncationId) { + truncationIds.add(msg.contextTruncation.truncationId) + console.log( + `[MessageManager] Found sliding_window_truncation to remove: ${msg.contextTruncation.truncationId}`, + ) + } + } + + return { condenseIds, truncationIds } + } + + /** + * Truncate clineMessages to the specified index. + */ + private async truncateClineMessages(toIndex: number): Promise { + await this.task.overwriteClineMessages(this.task.clineMessages.slice(0, toIndex)) + } + + /** + * Truncate API history by timestamp, remove orphaned summaries/markers, + * and clean up orphaned tags - all in a single write operation. + * + * This combined approach: + * 1. Avoids multiple writes to API history + * 2. Only writes if the history actually changed + * 3. Handles both truncation and cleanup atomically + */ + private async truncateApiHistoryWithCleanup( + cutoffTs: number, + removedIds: ContextEventIds, + skipCleanup: boolean, + ): Promise { + const originalHistory = this.task.apiConversationHistory + let apiHistory = [...originalHistory] + + // Step 1: Filter by timestamp + apiHistory = apiHistory.filter((m) => !m.ts || m.ts < cutoffTs) + + // Step 2: Remove Summaries whose condense_context was removed + if (removedIds.condenseIds.size > 0) { + apiHistory = apiHistory.filter((msg) => { + if (msg.isSummary && msg.condenseId && removedIds.condenseIds.has(msg.condenseId)) { + console.log(`[MessageManager] Removing orphaned Summary with condenseId: ${msg.condenseId}`) + return false + } + return true + }) + } + + // Step 3: Remove truncation markers whose sliding_window_truncation was removed + if (removedIds.truncationIds.size > 0) { + apiHistory = apiHistory.filter((msg) => { + if (msg.isTruncationMarker && msg.truncationId && removedIds.truncationIds.has(msg.truncationId)) { + console.log( + `[MessageManager] Removing orphaned truncation marker with truncationId: ${msg.truncationId}`, + ) + return false + } + return true + }) + } + + // Step 4: Cleanup orphaned tags (unless skipped) + if (!skipCleanup) { + apiHistory = cleanupAfterTruncation(apiHistory) + } + + // Only write if the history actually changed + const historyChanged = + apiHistory.length !== originalHistory.length || apiHistory.some((msg, i) => msg !== originalHistory[i]) + + if (historyChanged) { + await this.task.overwriteApiConversationHistory(apiHistory) + } + } +} diff --git a/src/core/task/Task.ts b/src/core/task/Task.ts index 493263b9d41..71b1c3a5fd9 100644 --- a/src/core/task/Task.ts +++ b/src/core/task/Task.ts @@ -125,6 +125,7 @@ import { processUserContentMentions } from "../mentions/processUserContentMentio import { getMessagesSinceLastSummary, summarizeConversation, getEffectiveApiHistory } from "../condense" import { MessageQueueService } from "../message-queue/MessageQueueService" import { AutoApprovalHandler, checkAutoApproval } from "../auto-approval" +import { MessageManager } from "../message-manager" const MAX_EXPONENTIAL_BACKOFF_SECONDS = 600 // 10 minutes const DEFAULT_USAGE_COLLECTION_TIMEOUT_MS = 5000 // 5 seconds @@ -327,6 +328,9 @@ export class Task extends EventEmitter implements TaskLike { // Initial status for the task's history item (set at creation time to avoid race conditions) private readonly initialStatus?: "active" | "delegated" | "completed" + // MessageManager for high-level message operations (lazy initialized) + private _messageManager?: MessageManager + constructor({ provider, apiConfiguration, @@ -4032,6 +4036,24 @@ export class Task extends EventEmitter implements TaskLike { return this.workspacePath } + /** + * Provides convenient access to high-level message operations. + * Uses lazy initialization - the MessageManager is only created when first accessed. + * Subsequent accesses return the same cached instance. + * + * @example + * ```typescript + * // Rewind conversation to a specific timestamp + * await task.messageManager.rewindToTimestamp(ts) + * ``` + */ + get messageManager(): MessageManager { + if (!this._messageManager) { + this._messageManager = new MessageManager(this) + } + return this._messageManager + } + /** * Broadcast browser session updates to the browser panel (if open) */ diff --git a/src/core/webview/__tests__/webviewMessageHandler.delete.spec.ts b/src/core/webview/__tests__/webviewMessageHandler.delete.spec.ts index 8450ab5d7ab..1a373809c9c 100644 --- a/src/core/webview/__tests__/webviewMessageHandler.delete.spec.ts +++ b/src/core/webview/__tests__/webviewMessageHandler.delete.spec.ts @@ -387,8 +387,10 @@ describe("webviewMessageHandler delete functionality", () => { expect(summary1.condenseParent).toBe(condenseId2) // Still tagged }) - it("should prefer non-summary message when timestamps collide for deletion target", async () => { - // When multiple messages share the same timestamp, prefer non-summary for targeting + it("should use timestamp-based truncation when multiple messages share same timestamp", async () => { + // When multiple messages share the same timestamp, timestamp-based truncation + // removes ALL messages at or after that timestamp. This is different from + // index-based truncation which would preserve earlier array indices. const sharedTs = 1000 getCurrentTaskMock.clineMessages = [ @@ -405,7 +407,8 @@ describe("webviewMessageHandler delete functionality", () => { { ts: 1100, role: "assistant", content: "Response" }, ] - // Delete at shared timestamp - should target non-summary message (index 2) + // Delete at shared timestamp - MessageManager uses ts < cutoffTs, so ALL + // messages at ts=1000 are removed (including the Summary) await webviewMessageHandler(provider, { type: "deleteMessageConfirm", messageTs: sharedTs, @@ -414,12 +417,10 @@ describe("webviewMessageHandler delete functionality", () => { expect(getCurrentTaskMock.overwriteApiConversationHistory).toHaveBeenCalled() const result = getCurrentTaskMock.overwriteApiConversationHistory.mock.calls[0][0] - // Truncation at index 2 means we keep indices 0-1: previous message and summary - expect(result.length).toBe(2) + // Timestamp-based truncation keeps only messages with ts < 1000 + // Both the Summary (ts=1000) and non-summary (ts=1000) are removed + expect(result.length).toBe(1) expect(result[0].content).toBe("Previous message") - // The summary is kept since it's before truncation point - expect(result[1].content).toBe("Summary") - expect(result[1].isSummary).toBe(true) }) it("should remove Summary when its condense_context clineMessage is deleted", async () => { diff --git a/src/core/webview/webviewMessageHandler.ts b/src/core/webview/webviewMessageHandler.ts index c1c8e6aa204..9d865fec6ea 100644 --- a/src/core/webview/webviewMessageHandler.ts +++ b/src/core/webview/webviewMessageHandler.ts @@ -21,7 +21,6 @@ import { TelemetryService } from "@roo-code/telemetry" import { type ApiMessage } from "../task-persistence/apiMessages" import { saveTaskMessages } from "../task-persistence" -import { cleanupAfterTruncation } from "../condense" import { ClineProvider } from "./ClineProvider" import { BrowserSessionPanelManager } from "./BrowserSessionPanelManager" @@ -30,6 +29,7 @@ import { changeLanguage, t } from "../../i18n" import { Package } from "../../shared/package" import { type RouterName, type ModelRecord, toRouterName } from "../../shared/api" import { MessageEnhancer } from "./messageEnhancer" +import { MessageManager } from "../message-manager" import { type WebviewMessage, @@ -111,85 +111,6 @@ export const webviewMessageHandler = async ( ) } - /** - * Removes the target message and all subsequent messages. - * After truncation, cleans up orphaned condenseParent and truncationParent references for any - * summaries or truncation markers that were removed by the truncation. - * - * Design: Rewind/delete operations preserve earlier condense and truncation states. - * Only summaries and truncation markers that are removed by the truncation (i.e., were created - * after the rewind point) have their associated tags cleared. - * This allows nested condensing and multiple truncations to work correctly - rewinding past the - * second condense restores visibility of messages condensed by it, while keeping the first condense intact. - * Same applies to truncation markers. - */ - const removeMessagesThisAndSubsequent = async ( - currentCline: any, - messageIndex: number, - apiConversationHistoryIndex: number, - ) => { - // Step 1: Collect condenseIds from condense_context messages being removed. - // These IDs link clineMessages to their corresponding Summaries in apiConversationHistory. - const removedCondenseIds = new Set() - // Step 1b: Collect truncationIds from sliding_window_truncation messages being removed. - // These IDs link clineMessages to their corresponding truncation markers in apiConversationHistory. - const removedTruncationIds = new Set() - - for (let i = messageIndex; i < currentCline.clineMessages.length; i++) { - const msg = currentCline.clineMessages[i] - if (msg.say === "condense_context" && msg.contextCondense?.condenseId) { - removedCondenseIds.add(msg.contextCondense.condenseId) - } - if (msg.say === "sliding_window_truncation" && msg.contextTruncation?.truncationId) { - removedTruncationIds.add(msg.contextTruncation.truncationId) - } - } - - // Step 2: Delete this message and all that follow - await currentCline.overwriteClineMessages(currentCline.clineMessages.slice(0, messageIndex)) - - if (apiConversationHistoryIndex !== -1) { - // Step 3: Truncate API history by timestamp/index - let truncatedApiHistory = currentCline.apiConversationHistory.slice(0, apiConversationHistoryIndex) - - // Step 4: Remove Summaries whose condenseId was in a removed condense_context message. - // This handles the case where Summary.ts < truncation point but condense_context.ts > truncation point. - // Without this, the Summary would survive truncation but its corresponding UI event would be gone. - if (removedCondenseIds.size > 0) { - truncatedApiHistory = truncatedApiHistory.filter((msg: ApiMessage) => { - if (msg.isSummary && msg.condenseId && removedCondenseIds.has(msg.condenseId)) { - console.log( - `[removeMessagesThisAndSubsequent] Removing orphaned Summary with condenseId=${msg.condenseId}`, - ) - return false - } - return true - }) - } - - // Step 4b: Remove truncation markers whose truncationId was in a removed sliding_window_truncation message. - // Same logic as condense - without this, the marker would survive but its UI event would be gone. - if (removedTruncationIds.size > 0) { - truncatedApiHistory = truncatedApiHistory.filter((msg: ApiMessage) => { - if (msg.isTruncationMarker && msg.truncationId && removedTruncationIds.has(msg.truncationId)) { - console.log( - `[removeMessagesThisAndSubsequent] Removing orphaned truncation marker with truncationId=${msg.truncationId}`, - ) - return false - } - return true - }) - } - - // Step 5: Clean up orphaned condenseParent and truncationParent references for messages whose - // summary or truncation marker was removed by the truncation. Summaries, truncation markers, and messages - // from earlier condense/truncation operations are preserved. - const cleanedApiHistory = cleanupAfterTruncation(truncatedApiHistory) - - await currentCline.overwriteApiConversationHistory(cleanedApiHistory) - } - } - /** * Handles message deletion operations with user confirmation */ @@ -281,8 +202,9 @@ export const webviewMessageHandler = async ( } } - // Delete this message and all subsequent messages - await removeMessagesThisAndSubsequent(currentCline, messageIndex, apiIndexToUse) + // Delete this message and all subsequent messages using MessageManager + const manager = new MessageManager(currentCline) + await manager.rewindToTimestamp(targetMessage.ts!, { includeTargetMessage: false }) // Restore checkpoint associations for preserved messages for (const [ts, checkpoint] of preservedCheckpoints) { @@ -448,8 +370,12 @@ export const webviewMessageHandler = async ( } } - // Delete the original (user) message and all subsequent messages - await removeMessagesThisAndSubsequent(currentCline, deleteFromMessageIndex, deleteFromApiIndex) + // Delete the original (user) message and all subsequent messages using MessageManager + const manager = new MessageManager(currentCline) + const rewindTs = currentCline.clineMessages[deleteFromMessageIndex]?.ts + if (rewindTs) { + await manager.rewindToTimestamp(rewindTs, { includeTargetMessage: false }) + } // Restore checkpoint associations for preserved messages for (const [ts, checkpoint] of preservedCheckpoints) { From b8c372819aeef8798f29619851760563cbf948e6 Mon Sep 17 00:00:00 2001 From: Hannes Rudolph Date: Thu, 4 Dec 2025 14:34:25 -0700 Subject: [PATCH 2/8] refactor: consolidate MessageManager usage through Task.messageManager getter Address PR review feedback from roomote to route all MessageManager usages through Task.messageManager instead of creating new instances directly. This ensures a single coordination point and avoids multiple instances. Changes: - checkpoints/index.ts: use task.messageManager instead of new MessageManager(task) - webviewMessageHandler.ts: use currentCline.messageManager instead of new MessageManager(currentCline) - Updated tests to support the messageManager property on mock task objects --- src/core/checkpoints/index.ts | 4 +--- .../webview/__tests__/webviewMessageHandler.delete.spec.ts | 4 ++++ src/core/webview/webviewMessageHandler.ts | 7 ++----- 3 files changed, 7 insertions(+), 8 deletions(-) diff --git a/src/core/checkpoints/index.ts b/src/core/checkpoints/index.ts index 8adb04a7bb6..64a8ad1cfe5 100644 --- a/src/core/checkpoints/index.ts +++ b/src/core/checkpoints/index.ts @@ -4,7 +4,6 @@ import * as vscode from "vscode" import { TelemetryService } from "@roo-code/telemetry" import { Task } from "../task/Task" -import { MessageManager } from "../message-manager" import { getWorkspacePath } from "../../utils/path" import { checkGitInstalled } from "../../utils/git" @@ -268,8 +267,7 @@ export async function checkpointRestore( // Use MessageManager to properly handle context-management events // This ensures orphaned Summary messages and truncation markers are cleaned up - const manager = new MessageManager(task) - await manager.rewindToTimestamp(ts, { + await task.messageManager.rewindToTimestamp(ts, { includeTargetMessage: operation === "edit", }) diff --git a/src/core/webview/__tests__/webviewMessageHandler.delete.spec.ts b/src/core/webview/__tests__/webviewMessageHandler.delete.spec.ts index 1a373809c9c..4d57608e9c7 100644 --- a/src/core/webview/__tests__/webviewMessageHandler.delete.spec.ts +++ b/src/core/webview/__tests__/webviewMessageHandler.delete.spec.ts @@ -2,6 +2,7 @@ import { describe, it, expect, beforeEach, vi } from "vitest" import { webviewMessageHandler } from "../webviewMessageHandler" import * as vscode from "vscode" import { ClineProvider } from "../ClineProvider" +import { MessageManager } from "../../message-manager" // Mock the saveTaskMessages function vi.mock("../../task-persistence", () => ({ @@ -63,6 +64,9 @@ describe("webviewMessageHandler delete functionality", () => { overwriteApiConversationHistory: vi.fn(async () => {}), taskId: "test-task-id", } + // Add messageManager using a real MessageManager instance (must be added after object creation + // to avoid circular reference issues with 'this') + getCurrentTaskMock.messageManager = new MessageManager(getCurrentTaskMock as any) // Create mock provider provider = { diff --git a/src/core/webview/webviewMessageHandler.ts b/src/core/webview/webviewMessageHandler.ts index 9d865fec6ea..f9bc7a85cda 100644 --- a/src/core/webview/webviewMessageHandler.ts +++ b/src/core/webview/webviewMessageHandler.ts @@ -29,7 +29,6 @@ import { changeLanguage, t } from "../../i18n" import { Package } from "../../shared/package" import { type RouterName, type ModelRecord, toRouterName } from "../../shared/api" import { MessageEnhancer } from "./messageEnhancer" -import { MessageManager } from "../message-manager" import { type WebviewMessage, @@ -203,8 +202,7 @@ export const webviewMessageHandler = async ( } // Delete this message and all subsequent messages using MessageManager - const manager = new MessageManager(currentCline) - await manager.rewindToTimestamp(targetMessage.ts!, { includeTargetMessage: false }) + await currentCline.messageManager.rewindToTimestamp(targetMessage.ts!, { includeTargetMessage: false }) // Restore checkpoint associations for preserved messages for (const [ts, checkpoint] of preservedCheckpoints) { @@ -371,10 +369,9 @@ export const webviewMessageHandler = async ( } // Delete the original (user) message and all subsequent messages using MessageManager - const manager = new MessageManager(currentCline) const rewindTs = currentCline.clineMessages[deleteFromMessageIndex]?.ts if (rewindTs) { - await manager.rewindToTimestamp(rewindTs, { includeTargetMessage: false }) + await currentCline.messageManager.rewindToTimestamp(rewindTs, { includeTargetMessage: false }) } // Restore checkpoint associations for preserved messages From 7b0c10840eacbb6e3171f66bc4b2117367a18be3 Mon Sep 17 00:00:00 2001 From: Hannes Rudolph Date: Thu, 4 Dec 2025 14:40:49 -0700 Subject: [PATCH 3/8] ci: retry tests From b2a8b403827f70e0ffa14b223b466bec97c9f4f0 Mon Sep 17 00:00:00 2001 From: Hannes Rudolph Date: Thu, 4 Dec 2025 15:25:42 -0700 Subject: [PATCH 4/8] docs: add architectural guidance for MessageManager getter usage Update Task.messageManager getter documentation to clarify that all MessageManager operations must go through this getter rather than instantiating new instances directly. This addresses the PR review feedback about updating class docs. --- src/core/task/Task.ts | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/src/core/task/Task.ts b/src/core/task/Task.ts index 71b1c3a5fd9..fc3184d8255 100644 --- a/src/core/task/Task.ts +++ b/src/core/task/Task.ts @@ -4041,10 +4041,21 @@ export class Task extends EventEmitter implements TaskLike { * Uses lazy initialization - the MessageManager is only created when first accessed. * Subsequent accesses return the same cached instance. * + * ## Important: Single Coordination Point + * + * **All MessageManager operations must go through this getter** rather than + * instantiating `new MessageManager(task)` directly. This ensures: + * - A single shared instance for consistent behavior + * - Centralized coordination of all rewind/message operations + * - Ability to add internal state or instrumentation in the future + * * @example * ```typescript - * // Rewind conversation to a specific timestamp + * // Correct: Use the getter * await task.messageManager.rewindToTimestamp(ts) + * + * // Incorrect: Do NOT create new instances directly + * // const manager = new MessageManager(task) // Don't do this! * ``` */ get messageManager(): MessageManager { From 956b3a24dc18a03025de299bf910549254943d61 Mon Sep 17 00:00:00 2001 From: Hannes Rudolph Date: Thu, 4 Dec 2025 16:36:24 -0700 Subject: [PATCH 5/8] fix(tests): update Task mocks to include messageManager for proper testing --- .../checkpoints/__tests__/checkpoint.test.ts | 2 + .../webview/__tests__/ClineProvider.spec.ts | 61 +++++++++++++------ .../webviewMessageHandler.checkpoint.spec.ts | 2 + .../webviewMessageHandler.edit.spec.ts | 2 + 4 files changed, 48 insertions(+), 19 deletions(-) diff --git a/src/core/checkpoints/__tests__/checkpoint.test.ts b/src/core/checkpoints/__tests__/checkpoint.test.ts index e9e565853f2..299ff823b87 100644 --- a/src/core/checkpoints/__tests__/checkpoint.test.ts +++ b/src/core/checkpoints/__tests__/checkpoint.test.ts @@ -2,6 +2,7 @@ import { describe, it, expect, vi, beforeEach, afterEach, Mock } from "vitest" import { Task } from "../../task/Task" import { ClineProvider } from "../../webview/ClineProvider" import { checkpointSave, checkpointRestore, checkpointDiff, getCheckpointService } from "../index" +import { MessageManager } from "../../message-manager" import * as vscode from "vscode" // Mock vscode @@ -102,6 +103,7 @@ describe("Checkpoint functionality", () => { overwriteApiConversationHistory: vi.fn(), combineMessages: vi.fn().mockReturnValue([]), } + mockTask.messageManager = new MessageManager(mockTask) // Update the mock to return our mockCheckpointService const checkpointsModule = await import("../../../services/checkpoints") diff --git a/src/core/webview/__tests__/ClineProvider.spec.ts b/src/core/webview/__tests__/ClineProvider.spec.ts index c3a515b9de2..97fbc43f43f 100644 --- a/src/core/webview/__tests__/ClineProvider.spec.ts +++ b/src/core/webview/__tests__/ClineProvider.spec.ts @@ -21,6 +21,7 @@ import { Task, TaskOptions } from "../../task/Task" import { safeWriteJson } from "../../../utils/safeWriteJson" import { ClineProvider } from "../ClineProvider" +import { MessageManager } from "../../message-manager" // Mock setup must come before imports. vi.mock("../../prompts/sections/custom-instructions") @@ -208,25 +209,21 @@ vi.mock("../../../integrations/workspace/WorkspaceTracker", () => { }) vi.mock("../../task/Task", () => ({ - Task: vi - .fn() - .mockImplementation( - (_provider, _apiConfiguration, _customInstructions, _diffEnabled, _fuzzyMatchThreshold, _task, taskId) => ({ - api: undefined, - abortTask: vi.fn(), - handleWebviewAskResponse: vi.fn(), - clineMessages: [], - apiConversationHistory: [], - overwriteClineMessages: vi.fn(), - overwriteApiConversationHistory: vi.fn(), - getTaskNumber: vi.fn().mockReturnValue(0), - setTaskNumber: vi.fn(), - setParentTask: vi.fn(), - setRootTask: vi.fn(), - taskId: taskId || "test-task-id", - emit: vi.fn(), - }), - ), + Task: vi.fn().mockImplementation((options: any) => ({ + api: undefined, + abortTask: vi.fn(), + handleWebviewAskResponse: vi.fn(), + clineMessages: [], + apiConversationHistory: [], + overwriteClineMessages: vi.fn(), + overwriteApiConversationHistory: vi.fn(), + getTaskNumber: vi.fn().mockReturnValue(0), + setTaskNumber: vi.fn(), + setParentTask: vi.fn(), + setRootTask: vi.fn(), + taskId: options?.historyItem?.id || "test-task-id", + emit: vi.fn(), + })), })) vi.mock("../../../integrations/misc/extract-text", () => ({ @@ -341,6 +338,32 @@ afterAll(() => { }) describe("ClineProvider", () => { + beforeAll(() => { + vi.mocked(Task).mockImplementation((options: any) => { + const task: any = { + api: undefined, + abortTask: vi.fn(), + handleWebviewAskResponse: vi.fn(), + clineMessages: [], + apiConversationHistory: [], + overwriteClineMessages: vi.fn(), + overwriteApiConversationHistory: vi.fn(), + getTaskNumber: vi.fn().mockReturnValue(0), + setTaskNumber: vi.fn(), + setParentTask: vi.fn(), + setRootTask: vi.fn(), + taskId: options?.historyItem?.id || "test-task-id", + emit: vi.fn(), + } + + Object.defineProperty(task, "messageManager", { + get: () => new MessageManager(task), + }) + + return task + }) + }) + let defaultTaskOptions: TaskOptions let provider: ClineProvider diff --git a/src/core/webview/__tests__/webviewMessageHandler.checkpoint.spec.ts b/src/core/webview/__tests__/webviewMessageHandler.checkpoint.spec.ts index 08e60f3cb1c..a0687d6cc19 100644 --- a/src/core/webview/__tests__/webviewMessageHandler.checkpoint.spec.ts +++ b/src/core/webview/__tests__/webviewMessageHandler.checkpoint.spec.ts @@ -2,6 +2,7 @@ import { describe, it, expect, vi, beforeEach } from "vitest" import { webviewMessageHandler } from "../webviewMessageHandler" import { saveTaskMessages } from "../../task-persistence" import { handleCheckpointRestoreOperation } from "../checkpointRestoreHandler" +import { MessageManager } from "../../message-manager" // Mock dependencies vi.mock("../../task-persistence") @@ -40,6 +41,7 @@ describe("webviewMessageHandler - checkpoint operations", () => { overwriteClineMessages: vi.fn(), overwriteApiConversationHistory: vi.fn(), } + mockCline.messageManager = new MessageManager(mockCline) // Setup mock provider mockProvider = { diff --git a/src/core/webview/__tests__/webviewMessageHandler.edit.spec.ts b/src/core/webview/__tests__/webviewMessageHandler.edit.spec.ts index d467f5cd92d..dfbce361e45 100644 --- a/src/core/webview/__tests__/webviewMessageHandler.edit.spec.ts +++ b/src/core/webview/__tests__/webviewMessageHandler.edit.spec.ts @@ -40,6 +40,7 @@ import { webviewMessageHandler } from "../webviewMessageHandler" import type { ClineProvider } from "../ClineProvider" import type { ClineMessage } from "@roo-code/types" import type { ApiMessage } from "../../task-persistence/apiMessages" +import { MessageManager } from "../../message-manager" describe("webviewMessageHandler - Edit Message with Timestamp Fallback", () => { let mockClineProvider: ClineProvider @@ -57,6 +58,7 @@ describe("webviewMessageHandler - Edit Message with Timestamp Fallback", () => { overwriteApiConversationHistory: vi.fn(), handleWebviewAskResponse: vi.fn(), } + mockCurrentTask.messageManager = new MessageManager(mockCurrentTask) // Create mock provider mockClineProvider = { From 39ca0baf21dc5bd6528a695c2f131890c94f1673 Mon Sep 17 00:00:00 2001 From: Hannes Rudolph Date: Thu, 4 Dec 2025 17:55:21 -0700 Subject: [PATCH 6/8] docs: update MessageManager usage example to use Task.messageManager getter --- src/core/message-manager/index.ts | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/core/message-manager/index.ts b/src/core/message-manager/index.ts index bb65fea3acd..b78aa95d681 100644 --- a/src/core/message-manager/index.ts +++ b/src/core/message-manager/index.ts @@ -24,11 +24,12 @@ interface ContextEventIds { * - Removing orphaned truncation markers when their sliding_window_truncation is removed * - Cleaning up orphaned condenseParent/truncationParent tags * - * Usage: + * Usage (always access via Task.messageManager getter): * ```typescript - * const manager = new MessageManager(task) - * await manager.rewindToTimestamp(messageTs, { includeTargetMessage: false }) + * await task.messageManager.rewindToTimestamp(messageTs, { includeTargetMessage: false }) * ``` + * + * @see Task.messageManager - The getter that provides lazy-initialized access to this manager */ export class MessageManager { constructor(private task: Task) {} From cef2492e699ebefd1101abd9670d980ff81ff3c0 Mon Sep 17 00:00:00 2001 From: Hannes Rudolph Date: Thu, 4 Dec 2025 18:16:28 -0700 Subject: [PATCH 7/8] test: stabilize mixed markdown list e2e --- apps/vscode-e2e/src/suite/markdown-lists.test.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/apps/vscode-e2e/src/suite/markdown-lists.test.ts b/apps/vscode-e2e/src/suite/markdown-lists.test.ts index 9b5c1bd8457..ec0cda50872 100644 --- a/apps/vscode-e2e/src/suite/markdown-lists.test.ts +++ b/apps/vscode-e2e/src/suite/markdown-lists.test.ts @@ -150,7 +150,7 @@ suite("Markdown List Rendering", function () { text: "Please create a list that has both numbered items and bullet points, mixing ordered and unordered lists", }) - await waitUntilCompleted({ api, taskId }) + await waitUntilCompleted({ api, taskId, timeout: 60_000 }) // Find a message that contains both types of lists const listMessage = messages.find( From 6aae0c4c9a9766df0d258e2b12d1f9577f79ff9c Mon Sep 17 00:00:00 2001 From: Hannes Rudolph Date: Thu, 4 Dec 2025 18:18:55 -0700 Subject: [PATCH 8/8] Discard changes to apps/vscode-e2e/src/suite/markdown-lists.test.ts --- apps/vscode-e2e/src/suite/markdown-lists.test.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/apps/vscode-e2e/src/suite/markdown-lists.test.ts b/apps/vscode-e2e/src/suite/markdown-lists.test.ts index ec0cda50872..9b5c1bd8457 100644 --- a/apps/vscode-e2e/src/suite/markdown-lists.test.ts +++ b/apps/vscode-e2e/src/suite/markdown-lists.test.ts @@ -150,7 +150,7 @@ suite("Markdown List Rendering", function () { text: "Please create a list that has both numbered items and bullet points, mixing ordered and unordered lists", }) - await waitUntilCompleted({ api, taskId, timeout: 60_000 }) + await waitUntilCompleted({ api, taskId }) // Find a message that contains both types of lists const listMessage = messages.find(