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/checkpoints/index.ts b/src/core/checkpoints/index.ts index 3efdb466e66..64a8ad1cfe5 100644 --- a/src/core/checkpoints/index.ts +++ b/src/core/checkpoints/index.ts @@ -258,20 +258,20 @@ 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 + await task.messageManager.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..b78aa95d681 --- /dev/null +++ b/src/core/message-manager/index.ts @@ -0,0 +1,185 @@ +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 (always access via Task.messageManager getter): + * ```typescript + * 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) {} + + /** + * 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..fc3184d8255 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,35 @@ 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. + * + * ## 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 + * // 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 { + 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__/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.delete.spec.ts b/src/core/webview/__tests__/webviewMessageHandler.delete.spec.ts index 8450ab5d7ab..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 = { @@ -387,8 +391,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 +411,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 +421,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/__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 = { diff --git a/src/core/webview/webviewMessageHandler.ts b/src/core/webview/webviewMessageHandler.ts index c1c8e6aa204..f9bc7a85cda 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" @@ -111,85 +110,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 +201,8 @@ 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 + await currentCline.messageManager.rewindToTimestamp(targetMessage.ts!, { includeTargetMessage: false }) // Restore checkpoint associations for preserved messages for (const [ts, checkpoint] of preservedCheckpoints) { @@ -448,8 +368,11 @@ 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 rewindTs = currentCline.clineMessages[deleteFromMessageIndex]?.ts + if (rewindTs) { + await currentCline.messageManager.rewindToTimestamp(rewindTs, { includeTargetMessage: false }) + } // Restore checkpoint associations for preserved messages for (const [ts, checkpoint] of preservedCheckpoints) {