diff --git a/src/core/condense/__tests__/journal.test.ts b/src/core/condense/__tests__/journal.test.ts new file mode 100644 index 0000000000..7b08dbfe5b --- /dev/null +++ b/src/core/condense/__tests__/journal.test.ts @@ -0,0 +1,387 @@ +import { describe, it, expect, beforeEach, afterEach, vi } from "vitest" +import * as fs from "fs/promises" +import * as path from "path" +import * as os from "os" +import { + type CondenseJournalEntry, + type CondenseJournal, + readJournal, + writeJournal, + appendJournalEntry, + createJournalEntry, + restoreMessagesForTimestamp, + findRemovedMessages, + getJournalPath, +} from "../journal" +import { type ApiMessage } from "../../task-persistence" + +// Mock safeWriteJson +vi.mock("../../../utils/safeWriteJson", () => ({ + safeWriteJson: vi.fn(async (filePath: string, data: any) => { + await fs.writeFile(filePath, JSON.stringify(data, null, 2), "utf-8") + }), +})) + +describe("Condense Journal", () => { + let testDir: string + + beforeEach(async () => { + // Create a temporary test directory + testDir = path.join(os.tmpdir(), `journal-test-${Date.now()}`) + await fs.mkdir(testDir, { recursive: true }) + }) + + afterEach(async () => { + // Clean up test directory + try { + await fs.rm(testDir, { recursive: true, force: true }) + } catch (error) { + // Ignore cleanup errors + } + }) + + describe("getJournalPath", () => { + it("should return the correct journal path", () => { + const journalPath = getJournalPath(testDir) + expect(journalPath).toBe(path.join(testDir, "condense_journal.json")) + }) + }) + + describe("readJournal", () => { + it("should return null when journal doesn't exist", async () => { + const journal = await readJournal(testDir) + expect(journal).toBeNull() + }) + + it("should read existing journal", async () => { + const testJournal: CondenseJournal = { + version: 1, + entries: [ + { + removed: [{ role: "user", content: "test", ts: 1000 }], + boundary: { firstKeptTs: 900, lastKeptTs: 1100, summaryTs: 1050 }, + createdAt: Date.now(), + type: "manual", + }, + ], + } + + const journalPath = getJournalPath(testDir) + await fs.writeFile(journalPath, JSON.stringify(testJournal), "utf-8") + + const journal = await readJournal(testDir) + expect(journal).toEqual(testJournal) + }) + + it("should handle corrupted journal file gracefully", async () => { + const journalPath = getJournalPath(testDir) + await fs.writeFile(journalPath, "invalid json", "utf-8") + + const journal = await readJournal(testDir) + expect(journal).toBeNull() + }) + }) + + describe("writeJournal", () => { + it("should write journal to disk", async () => { + const testJournal: CondenseJournal = { + version: 1, + entries: [ + { + removed: [{ role: "assistant", content: "response", ts: 2000 }], + boundary: { firstKeptTs: 1900, lastKeptTs: 2100 }, + createdAt: Date.now(), + type: "auto", + }, + ], + } + + await writeJournal(testDir, testJournal) + + const journalPath = getJournalPath(testDir) + const content = await fs.readFile(journalPath, "utf-8") + const savedJournal = JSON.parse(content) + expect(savedJournal).toEqual(testJournal) + }) + }) + + describe("appendJournalEntry", () => { + it("should create new journal if none exists", async () => { + const entry: CondenseJournalEntry = { + removed: [{ role: "user", content: "test message", ts: 3000 }], + boundary: { firstKeptTs: 2900, summaryTs: 3050 }, + createdAt: Date.now(), + type: "manual", + } + + await appendJournalEntry(testDir, entry) + + const journal = await readJournal(testDir) + expect(journal).not.toBeNull() + expect(journal?.version).toBe(1) + expect(journal?.entries).toHaveLength(1) + expect(journal?.entries[0]).toEqual(entry) + }) + + it("should append to existing journal", async () => { + const existingEntry: CondenseJournalEntry = { + removed: [{ role: "user", content: "old message", ts: 1000 }], + boundary: { firstKeptTs: 900 }, + createdAt: Date.now() - 10000, + type: "manual", + } + + const newEntry: CondenseJournalEntry = { + removed: [{ role: "assistant", content: "new message", ts: 2000 }], + boundary: { lastKeptTs: 2100 }, + createdAt: Date.now(), + type: "auto", + } + + // Create initial journal + await writeJournal(testDir, { version: 1, entries: [existingEntry] }) + + // Append new entry + await appendJournalEntry(testDir, newEntry) + + const journal = await readJournal(testDir) + expect(journal?.entries).toHaveLength(2) + expect(journal?.entries[0]).toEqual(existingEntry) + expect(journal?.entries[1]).toEqual(newEntry) + }) + }) + + describe("createJournalEntry", () => { + it("should create journal entry with all fields", () => { + const removed: ApiMessage[] = [ + { role: "user", content: "message 1", ts: 1000 }, + { role: "assistant", content: "message 2", ts: 1100 }, + ] + const firstKept: ApiMessage = { role: "user", content: "first kept", ts: 900 } + const lastKept: ApiMessage = { role: "assistant", content: "last kept", ts: 1200 } + const summary: ApiMessage = { role: "assistant", content: "summary", ts: 1150, isSummary: true } + + const entry = createJournalEntry(removed, firstKept, lastKept, summary, "manual") + + expect(entry.removed).toEqual(removed) + expect(entry.boundary.firstKeptTs).toBe(900) + expect(entry.boundary.lastKeptTs).toBe(1200) + expect(entry.boundary.summaryTs).toBe(1150) + expect(entry.type).toBe("manual") + expect(entry.createdAt).toBeGreaterThan(0) + }) + + it("should handle undefined boundary messages", () => { + const removed: ApiMessage[] = [{ role: "user", content: "message", ts: 1000 }] + + const entry = createJournalEntry(removed, undefined, undefined, undefined, "auto") + + expect(entry.removed).toEqual(removed) + expect(entry.boundary.firstKeptTs).toBeUndefined() + expect(entry.boundary.lastKeptTs).toBeUndefined() + expect(entry.boundary.summaryTs).toBeUndefined() + expect(entry.type).toBe("auto") + }) + }) + + describe("findRemovedMessages", () => { + it("should identify removed messages correctly", () => { + const original: ApiMessage[] = [ + { role: "user", content: "msg1", ts: 1000 }, + { role: "assistant", content: "msg2", ts: 1100 }, + { role: "user", content: "msg3", ts: 1200 }, + { role: "assistant", content: "msg4", ts: 1300 }, + ] + + const condensed: ApiMessage[] = [ + { role: "user", content: "msg1", ts: 1000 }, + { role: "assistant", content: "summary", ts: 1150, isSummary: true }, + { role: "assistant", content: "msg4", ts: 1300 }, + ] + + const removed = findRemovedMessages(original, condensed) + + expect(removed).toHaveLength(2) + expect(removed[0]).toEqual({ role: "assistant", content: "msg2", ts: 1100 }) + expect(removed[1]).toEqual({ role: "user", content: "msg3", ts: 1200 }) + }) + + it("should handle messages without timestamps", () => { + const original: ApiMessage[] = [ + { role: "user", content: "msg1", ts: 1000 }, + { role: "assistant", content: "no timestamp" }, // No ts field + { role: "user", content: "msg3", ts: 1200 }, + ] + + const condensed: ApiMessage[] = [ + { role: "user", content: "msg1", ts: 1000 }, + { role: "user", content: "msg3", ts: 1200 }, + ] + + const removed = findRemovedMessages(original, condensed) + + expect(removed).toHaveLength(0) // Message without timestamp is not included + }) + }) + + describe("restoreMessagesForTimestamp", () => { + it("should return null if target timestamp already exists", async () => { + const currentMessages: ApiMessage[] = [ + { role: "user", content: "msg1", ts: 1000 }, + { role: "assistant", content: "msg2", ts: 1100 }, + ] + + const result = await restoreMessagesForTimestamp(testDir, currentMessages, 1100) + expect(result).toBeNull() + }) + + it("should return null if no journal exists", async () => { + const currentMessages: ApiMessage[] = [{ role: "user", content: "msg1", ts: 1000 }] + + const result = await restoreMessagesForTimestamp(testDir, currentMessages, 2000) + expect(result).toBeNull() + }) + + it("should restore messages from single journal entry", async () => { + const currentMessages: ApiMessage[] = [ + { role: "user", content: "msg1", ts: 1000 }, + { role: "assistant", content: "summary", ts: 1500, isSummary: true }, + { role: "user", content: "msg5", ts: 1600 }, + ] + + const journal: CondenseJournal = { + version: 1, + entries: [ + { + removed: [ + { role: "assistant", content: "msg2", ts: 1100 }, + { role: "user", content: "msg3", ts: 1200 }, + { role: "assistant", content: "msg4", ts: 1300 }, + ], + boundary: { firstKeptTs: 1000, lastKeptTs: 1600, summaryTs: 1500 }, + createdAt: Date.now(), + type: "manual", + }, + ], + } + + await writeJournal(testDir, journal) + + const result = await restoreMessagesForTimestamp(testDir, currentMessages, 1200) + + expect(result).not.toBeNull() + expect(result).toHaveLength(6) // 3 current + 3 restored + expect(result?.find((m) => m.ts === 1200)).toBeDefined() + expect(result?.find((m) => m.ts === 1100)).toBeDefined() + expect(result?.find((m) => m.ts === 1300)).toBeDefined() + // Should be sorted by timestamp + expect(result?.[0].ts).toBe(1000) + expect(result?.[1].ts).toBe(1100) + expect(result?.[2].ts).toBe(1200) + expect(result?.[3].ts).toBe(1300) + expect(result?.[4].ts).toBe(1500) + expect(result?.[5].ts).toBe(1600) + }) + + it("should handle nested condenses correctly", async () => { + const currentMessages: ApiMessage[] = [ + { role: "user", content: "msg1", ts: 1000 }, + { role: "assistant", content: "summary2", ts: 2500, isSummary: true }, + { role: "user", content: "msg9", ts: 2600 }, + ] + + const journal: CondenseJournal = { + version: 1, + entries: [ + // First condense + { + removed: [ + { role: "assistant", content: "msg2", ts: 1100 }, + { role: "user", content: "msg3", ts: 1200 }, + ], + boundary: { firstKeptTs: 1000, summaryTs: 1500 }, + createdAt: Date.now() - 10000, + type: "manual", + }, + // Second condense (nested) - this wouldn't contain msg2 and msg3 again since they were already condensed + { + removed: [ + { role: "assistant", content: "summary1", ts: 1500, isSummary: true }, + { role: "user", content: "msg5", ts: 1600 }, + { role: "assistant", content: "msg6", ts: 1700 }, + ], + boundary: { firstKeptTs: 1000, summaryTs: 2500 }, + createdAt: Date.now(), + type: "manual", + }, + ], + } + + await writeJournal(testDir, journal) + + // Try to restore a message from the first condensed range + const result = await restoreMessagesForTimestamp(testDir, currentMessages, 1100) + + expect(result).not.toBeNull() + // Should restore messages that contain the target timestamp + expect(result?.find((m) => m.ts === 1100)).toBeDefined() + expect(result?.find((m) => m.ts === 1200)).toBeDefined() + // The restoration logic only restores messages needed to reach the target timestamp + // It doesn't necessarily restore all messages from all entries + }) + + it("should not restore messages that are already in current messages", async () => { + const currentMessages: ApiMessage[] = [ + { role: "user", content: "msg1", ts: 1000 }, + { role: "assistant", content: "msg2", ts: 1100 }, // Already present + { role: "assistant", content: "summary", ts: 1500, isSummary: true }, + ] + + const journal: CondenseJournal = { + version: 1, + entries: [ + { + removed: [ + { role: "assistant", content: "msg2", ts: 1100 }, // Duplicate + { role: "user", content: "msg3", ts: 1200 }, + ], + boundary: {}, + createdAt: Date.now(), + type: "manual", + }, + ], + } + + await writeJournal(testDir, journal) + + const result = await restoreMessagesForTimestamp(testDir, currentMessages, 1200) + + expect(result).not.toBeNull() + expect(result).toHaveLength(4) // 3 current + 1 restored (msg3) + // Should only have one msg2 + expect(result?.filter((m) => m.ts === 1100)).toHaveLength(1) + }) + + it("should return null if target timestamp not found in journal", async () => { + const currentMessages: ApiMessage[] = [{ role: "user", content: "msg1", ts: 1000 }] + + const journal: CondenseJournal = { + version: 1, + entries: [ + { + removed: [{ role: "assistant", content: "msg2", ts: 1100 }], + boundary: {}, + createdAt: Date.now(), + type: "manual", + }, + ], + } + + await writeJournal(testDir, journal) + + // Try to restore a timestamp that doesn't exist in journal + const result = await restoreMessagesForTimestamp(testDir, currentMessages, 9999) + expect(result).toBeNull() + }) + }) +}) diff --git a/src/core/condense/journal.ts b/src/core/condense/journal.ts new file mode 100644 index 0000000000..f8c0e7fbae --- /dev/null +++ b/src/core/condense/journal.ts @@ -0,0 +1,209 @@ +import * as fs from "fs/promises" +import * as path from "path" +import { type ApiMessage } from "../task-persistence" + +/** + * Represents a single entry in the condense journal + */ +export interface CondenseJournalEntry { + /** Messages that were removed during this condense operation */ + removed: ApiMessage[] + /** Boundary timestamps for the condense operation */ + boundary: { + /** Timestamp of the first message kept after condensing */ + firstKeptTs?: number + /** Timestamp of the last message kept before summary */ + lastKeptTs?: number + /** Timestamp of the summary message created */ + summaryTs?: number + } + /** When this journal entry was created */ + createdAt: number + /** Type of condense operation */ + type: "manual" | "auto" +} + +/** + * The complete journal containing all condense entries + */ +export interface CondenseJournal { + version: number + entries: CondenseJournalEntry[] +} + +const JOURNAL_FILENAME = "condense_journal.json" +const JOURNAL_VERSION = 1 + +/** + * Get the path to the condense journal file for a task + */ +export function getJournalPath(taskDirPath: string): string { + return path.join(taskDirPath, JOURNAL_FILENAME) +} + +/** + * Read the condense journal from disk + */ +export async function readJournal(taskDirPath: string): Promise { + const journalPath = getJournalPath(taskDirPath) + try { + const content = await fs.readFile(journalPath, "utf-8") + const journal = JSON.parse(content) as CondenseJournal + + // Validate version + if (journal.version !== JOURNAL_VERSION) { + console.warn(`Condense journal version mismatch: expected ${JOURNAL_VERSION}, got ${journal.version}`) + } + + return journal + } catch (error) { + // Journal doesn't exist yet, which is fine + if ((error as NodeJS.ErrnoException).code === "ENOENT") { + return null + } + console.error("Error reading condense journal:", error) + return null + } +} + +/** + * Write the condense journal to disk + */ +export async function writeJournal(taskDirPath: string, journal: CondenseJournal): Promise { + const journalPath = getJournalPath(taskDirPath) + const { safeWriteJson } = await import("../../utils/safeWriteJson") + await safeWriteJson(journalPath, journal) +} + +/** + * Append a new entry to the condense journal + */ +export async function appendJournalEntry(taskDirPath: string, entry: CondenseJournalEntry): Promise { + // Read existing journal or create new one + let journal = await readJournal(taskDirPath) + if (!journal) { + journal = { + version: JOURNAL_VERSION, + entries: [], + } + } + + // Append the new entry + journal.entries.push(entry) + + // Write back to disk + await writeJournal(taskDirPath, journal) +} + +/** + * Create a journal entry from the condense operation + */ +export function createJournalEntry( + removedMessages: ApiMessage[], + firstKeptMessage: ApiMessage | undefined, + lastKeptMessage: ApiMessage | undefined, + summaryMessage: ApiMessage | undefined, + type: "manual" | "auto" = "manual", +): CondenseJournalEntry { + return { + removed: removedMessages, + boundary: { + firstKeptTs: firstKeptMessage?.ts, + lastKeptTs: lastKeptMessage?.ts, + summaryTs: summaryMessage?.ts, + }, + createdAt: Date.now(), + type, + } +} + +/** + * Restore messages from the journal to make a target timestamp available + * @param taskDirPath Path to the task directory + * @param currentMessages Current API messages in memory + * @param targetTs Target timestamp we need to make available + * @returns Updated messages array with restored messages, or null if restoration wasn't needed + */ +export async function restoreMessagesForTimestamp( + taskDirPath: string, + currentMessages: ApiMessage[], + targetTs: number, +): Promise { + // Check if target timestamp already exists in current messages + const targetExists = currentMessages.some((msg) => msg.ts === targetTs) + if (targetExists) { + return null // No restoration needed + } + + // Read the journal + const journal = await readJournal(taskDirPath) + if (!journal || journal.entries.length === 0) { + return null // No journal or no entries to restore from + } + + // Create a map of current message timestamps for quick lookup + const currentTsSet = new Set(currentMessages.map((msg) => msg.ts).filter((ts) => ts !== undefined)) + + // Collect messages to restore + const messagesToRestore: ApiMessage[] = [] + + // Walk journal entries from newest to oldest + for (let i = journal.entries.length - 1; i >= 0; i--) { + const entry = journal.entries[i] + + // Check if this entry contains our target timestamp + const hasTarget = entry.removed.some((msg) => msg.ts === targetTs) + if (!hasTarget) { + continue + } + + // Add all removed messages from this entry that aren't already in current messages + for (const msg of entry.removed) { + if (msg.ts && !currentTsSet.has(msg.ts)) { + messagesToRestore.push(msg) + currentTsSet.add(msg.ts) + } + } + + // Check if we now have the target + if (messagesToRestore.some((msg) => msg.ts === targetTs)) { + break + } + } + + // If we didn't find the target, return null + if (!messagesToRestore.some((msg) => msg.ts === targetTs)) { + return null + } + + // Merge restored messages with current messages and sort by timestamp + const mergedMessages = [...currentMessages, ...messagesToRestore] + mergedMessages.sort((a, b) => { + const tsA = a.ts ?? 0 + const tsB = b.ts ?? 0 + return tsA - tsB + }) + + return mergedMessages +} + +/** + * Find messages that will be removed during a condense operation + * @param originalMessages Original messages before condensing + * @param condensedMessages Messages after condensing + * @returns Array of removed messages + */ +export function findRemovedMessages(originalMessages: ApiMessage[], condensedMessages: ApiMessage[]): ApiMessage[] { + // Create a set of timestamps from condensed messages for efficient lookup + const condensedTsSet = new Set(condensedMessages.map((msg) => msg.ts).filter((ts) => ts !== undefined)) + + // Find messages that exist in original but not in condensed + return originalMessages.filter((msg) => { + // Keep messages without timestamps (shouldn't happen but be safe) + if (!msg.ts) { + return false + } + // Message is removed if its timestamp is not in the condensed set + return !condensedTsSet.has(msg.ts) + }) +} diff --git a/src/core/task/Task.ts b/src/core/task/Task.ts index 851df91e6c..ba9b55da06 100644 --- a/src/core/task/Task.ts +++ b/src/core/task/Task.ts @@ -110,6 +110,7 @@ import { } from "../checkpoints" import { processUserContentMentions } from "../mentions/processUserContentMentions" import { getMessagesSinceLastSummary, summarizeConversation } from "../condense" +import { appendJournalEntry, createJournalEntry, findRemovedMessages } from "../condense/journal" import { Gpt5Metadata, ClineMessageWithMetadata } from "./types" import { MessageQueueService } from "../message-queue/MessageQueueService" @@ -1007,6 +1008,9 @@ export class Task extends EventEmitter implements TaskLike { const { contextTokens: prevContextTokens } = this.getTokenUsage() + // Store original messages before condensing for journal + const originalMessages = [...this.apiConversationHistory] + const { messages, summary, @@ -1035,6 +1039,38 @@ export class Task extends EventEmitter implements TaskLike { ) return } + + // Write journal entry before overwriting history + try { + // Find which messages were removed + const removedMessages = findRemovedMessages(originalMessages, messages) + + if (removedMessages.length > 0) { + // Find boundary messages for the journal + const firstKeptMessage = messages.find((msg) => msg.ts && !msg.isSummary) + const summaryMessage = messages.find((msg) => msg.isSummary) + const lastKeptBeforeSummary = summaryMessage + ? messages[messages.indexOf(summaryMessage) - 1] + : undefined + + // Create and append journal entry + const journalEntry = createJournalEntry( + removedMessages, + firstKeptMessage, + lastKeptBeforeSummary, + summaryMessage, + "manual", + ) + + // Get task directory path and write journal + const taskDirPath = path.join(this.globalStoragePath, this.taskId) + await appendJournalEntry(taskDirPath, journalEntry) + } + } catch (journalError) { + // Log but don't fail the condense operation if journal writing fails + console.error("Failed to write condense journal entry:", journalError) + } + await this.overwriteApiConversationHistory(messages) // Set flag to skip previous_response_id on the next API call after manual condense diff --git a/src/core/webview/webviewMessageHandler.ts b/src/core/webview/webviewMessageHandler.ts index af5f9925c3..facf27122e 100644 --- a/src/core/webview/webviewMessageHandler.ts +++ b/src/core/webview/webviewMessageHandler.ts @@ -18,6 +18,7 @@ import { TelemetryService } from "@roo-code/telemetry" import { type ApiMessage } from "../task-persistence/apiMessages" import { saveTaskMessages } from "../task-persistence" +import { restoreMessagesForTimestamp } from "../condense/journal" import { ClineProvider } from "./ClineProvider" import { handleCheckpointRestoreOperation } from "./checkpointRestoreHandler" @@ -155,7 +156,30 @@ export const webviewMessageHandler = async ( return } - const { messageIndex, apiConversationHistoryIndex } = findMessageIndices(messageTs, currentCline) + let { messageIndex, apiConversationHistoryIndex } = findMessageIndices(messageTs, currentCline) + + // If the message is not found in API history, try to restore from journal + if (apiConversationHistoryIndex === -1 && messageIndex !== -1) { + const tsThreshold = currentCline.clineMessages[messageIndex]?.ts + if (typeof tsThreshold === "number") { + // Try to restore messages from journal + const taskDirPath = path.join(provider.contextProxy.globalStorageUri.fsPath, currentCline.taskId) + const restoredMessages = await restoreMessagesForTimestamp( + taskDirPath, + currentCline.apiConversationHistory, + tsThreshold, + ) + + if (restoredMessages) { + // Update the API conversation history with restored messages + await currentCline.overwriteApiConversationHistory(restoredMessages) + // Re-find the indices with restored messages + const updatedIndices = findMessageIndices(messageTs, currentCline) + apiConversationHistoryIndex = updatedIndices.apiConversationHistoryIndex + } + } + } + // Determine API truncation index with timestamp fallback if exact match not found let apiIndexToUse = apiConversationHistoryIndex const tsThreshold = currentCline.clineMessages[messageIndex]?.ts @@ -285,7 +309,29 @@ export const webviewMessageHandler = async ( } // Use findMessageIndices to find messages based on timestamp - const { messageIndex, apiConversationHistoryIndex } = findMessageIndices(messageTs, currentCline) + let { messageIndex, apiConversationHistoryIndex } = findMessageIndices(messageTs, currentCline) + + // If the message is not found in API history, try to restore from journal + if (apiConversationHistoryIndex === -1 && messageIndex !== -1) { + const tsThreshold = currentCline.clineMessages[messageIndex]?.ts + if (typeof tsThreshold === "number") { + // Try to restore messages from journal + const taskDirPath = path.join(provider.contextProxy.globalStorageUri.fsPath, currentCline.taskId) + const restoredMessages = await restoreMessagesForTimestamp( + taskDirPath, + currentCline.apiConversationHistory, + tsThreshold, + ) + + if (restoredMessages) { + // Update the API conversation history with restored messages + await currentCline.overwriteApiConversationHistory(restoredMessages) + // Re-find the indices with restored messages + const updatedIndices = findMessageIndices(messageTs, currentCline) + apiConversationHistoryIndex = updatedIndices.apiConversationHistoryIndex + } + } + } if (messageIndex === -1) { const errorMessage = t("common:errors.message.message_not_found", { messageTs })