FIX + feat: add MessageManager layer for centralized history coordination (#9842)

This commit is contained in:
Hannes Rudolph 2025-12-04 18:48:31 -07:00 committed by GitHub
parent a1d392f5cc
commit 630c8bce90
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 1024 additions and 118 deletions

View file

@ -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")

View file

@ -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({

View file

@ -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")
})
})
})

View file

@ -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<string>
truncationIds: Set<string>
}
/**
* 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<void> {
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<void> {
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<void> {
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<string>()
const truncationIds = new Set<string>()
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<void> {
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<void> {
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)
}
}
}

View file

@ -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<TaskEvents> 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<TaskEvents> 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)
*/

View file

@ -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

View file

@ -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 = {

View file

@ -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 () => {

View file

@ -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 = {

View file

@ -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<string>()
// 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<string>()
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) {