mirror of
https://github.com/RooVetGit/Roo-Code.git
synced 2026-08-28 05:27:24 +00:00
FIX + feat: add MessageManager layer for centralized history coordination (#9842)
This commit is contained in:
parent
a1d392f5cc
commit
630c8bce90
10 changed files with 1024 additions and 118 deletions
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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({
|
||||
|
|
|
|||
731
src/core/message-manager/index.spec.ts
Normal file
731
src/core/message-manager/index.spec.ts
Normal 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")
|
||||
})
|
||||
})
|
||||
})
|
||||
185
src/core/message-manager/index.ts
Normal file
185
src/core/message-manager/index.ts
Normal 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)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -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)
|
||||
*/
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 = {
|
||||
|
|
|
|||
|
|
@ -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 () => {
|
||||
|
|
|
|||
|
|
@ -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 = {
|
||||
|
|
|
|||
|
|
@ -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) {
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue