diff --git a/src/core/context-tracking/FileContextTracker.ts b/src/core/context-tracking/FileContextTracker.ts index 7753e93567..2c54807ab3 100644 --- a/src/core/context-tracking/FileContextTracker.ts +++ b/src/core/context-tracking/FileContextTracker.ts @@ -61,6 +61,56 @@ export class FileContextTracker { return this.currentMessageTs } + /** + * Rewinds the file context metadata to a specific timestamp. + * This should be called when the conversation is rewound (messages deleted) + * to ensure file context tracking stays in sync with the conversation state. + * + * - Removes all entries where containingMessageTs >= cutoffTs + * - Restores the newest "stale" entry to "active" for files that lost their active entry + * + * @param cutoffTs - The timestamp cutoff. Entries with containingMessageTs >= this value are removed. + */ + async rewindToTimestamp(cutoffTs: number): Promise { + try { + const metadata = await this.getTaskMetadata(this.taskId) + + // Track which files had entries removed + const affectedFiles = new Set() + + // Filter out entries where containingMessageTs >= cutoffTs + const filteredEntries = metadata.files_in_context.filter((entry) => { + if (entry.containingMessageTs && entry.containingMessageTs >= cutoffTs) { + affectedFiles.add(entry.path) + return false + } + return true + }) + + // For affected files, check if we need to restore a stale entry to active + for (const filePath of affectedFiles) { + const fileEntries = filteredEntries.filter((e) => e.path === filePath) + const hasActiveEntry = fileEntries.some((e) => e.record_state === "active") + + if (!hasActiveEntry && fileEntries.length > 0) { + // Find the newest stale entry (by roo_read_date) and restore it to active + const staleEntries = fileEntries + .filter((e) => e.record_state === "stale" && e.roo_read_date) + .sort((a, b) => (b.roo_read_date ?? 0) - (a.roo_read_date ?? 0)) + + if (staleEntries.length > 0) { + staleEntries[0].record_state = "active" + } + } + } + + metadata.files_in_context = filteredEntries + await this.saveTaskMetadata(this.taskId, metadata) + } catch (error) { + console.error("Failed to rewind file context metadata:", error) + } + } + // Gets the current working directory or returns undefined if it cannot be determined private getCwd(): string | undefined { const cwd = vscode.workspace.workspaceFolders?.map((folder) => folder.uri.fsPath).at(0) diff --git a/src/core/context-tracking/__tests__/FileContextTracker.rewind.test.ts b/src/core/context-tracking/__tests__/FileContextTracker.rewind.test.ts new file mode 100644 index 0000000000..5a9e6799e9 --- /dev/null +++ b/src/core/context-tracking/__tests__/FileContextTracker.rewind.test.ts @@ -0,0 +1,404 @@ +import { describe, it, expect, vi, beforeEach, afterEach } from "vitest" +import type { TaskMetadata, FileMetadataEntry } from "../FileContextTrackerTypes" + +// Mock dependencies before importing FileContextTracker +const mockGetTaskMetadata = vi.fn() +const mockSaveTaskMetadata = vi.fn() + +vi.mock("vscode", () => ({ + workspace: { + workspaceFolders: [{ uri: { fsPath: "/test/workspace" } }], + createFileSystemWatcher: vi.fn(() => ({ + onDidChange: vi.fn(), + dispose: vi.fn(), + })), + }, + Uri: { + file: vi.fn((path: string) => ({ fsPath: path })), + }, + RelativePattern: vi.fn(), +})) + +vi.mock("../../../utils/storage", () => ({ + getTaskDirectoryPath: vi.fn().mockResolvedValue("/test/task/dir"), +})) + +vi.mock("../../../utils/fs", () => ({ + fileExistsAtPath: vi.fn().mockResolvedValue(true), +})) + +vi.mock("../../../utils/safeWriteJson", () => ({ + safeWriteJson: vi.fn().mockResolvedValue(undefined), +})) + +vi.mock("fs/promises", () => ({ + default: { + readFile: vi.fn(), + }, +})) + +// Import after mocks +import { FileContextTracker } from "../FileContextTracker" + +describe("FileContextTracker.rewindToTimestamp", () => { + let tracker: FileContextTracker + let mockProvider: any + + beforeEach(() => { + vi.clearAllMocks() + + mockProvider = { + contextProxy: { + globalStorageUri: { fsPath: "/test/global/storage" }, + }, + } + + tracker = new FileContextTracker(mockProvider, "test-task-id") + + // Override getTaskMetadata and saveTaskMetadata for testing + tracker.getTaskMetadata = mockGetTaskMetadata + tracker.saveTaskMetadata = mockSaveTaskMetadata + }) + + afterEach(() => { + tracker.dispose() + }) + + describe("basic rewind functionality", () => { + it("should remove entries where containingMessageTs >= cutoffTs", async () => { + const metadata: TaskMetadata = { + files_in_context: [ + { + path: "file1.ts", + record_state: "active", + record_source: "read_tool", + roo_read_date: 1000, + roo_edit_date: null, + containingMessageTs: 1000, // Before cutoff - should keep + }, + { + path: "file2.ts", + record_state: "active", + record_source: "read_tool", + roo_read_date: 2000, + roo_edit_date: null, + containingMessageTs: 2000, // At cutoff - should remove + }, + { + path: "file3.ts", + record_state: "active", + record_source: "read_tool", + roo_read_date: 3000, + roo_edit_date: null, + containingMessageTs: 3000, // After cutoff - should remove + }, + ], + } + + mockGetTaskMetadata.mockResolvedValue(metadata) + mockSaveTaskMetadata.mockResolvedValue(undefined) + + await tracker.rewindToTimestamp(2000) + + expect(mockSaveTaskMetadata).toHaveBeenCalledWith("test-task-id", { + files_in_context: [ + expect.objectContaining({ + path: "file1.ts", + containingMessageTs: 1000, + }), + ], + }) + }) + + it("should not remove entries without containingMessageTs", async () => { + const metadata: TaskMetadata = { + files_in_context: [ + { + path: "legacy-file.ts", + record_state: "active", + record_source: "read_tool", + roo_read_date: 3000, + roo_edit_date: null, + // No containingMessageTs - legacy entry + }, + { + path: "new-file.ts", + record_state: "active", + record_source: "read_tool", + roo_read_date: 3000, + roo_edit_date: null, + containingMessageTs: 3000, // After cutoff - should remove + }, + ], + } + + mockGetTaskMetadata.mockResolvedValue(metadata) + mockSaveTaskMetadata.mockResolvedValue(undefined) + + await tracker.rewindToTimestamp(2000) + + expect(mockSaveTaskMetadata).toHaveBeenCalledWith("test-task-id", { + files_in_context: [ + expect.objectContaining({ + path: "legacy-file.ts", + record_state: "active", + }), + ], + }) + }) + }) + + describe("stale entry restoration", () => { + it("should restore newest stale entry to active when active entry is removed", async () => { + const metadata: TaskMetadata = { + files_in_context: [ + { + path: "file.ts", + record_state: "stale", + record_source: "read_tool", + roo_read_date: 1000, + roo_edit_date: null, + containingMessageTs: 1000, + }, + { + path: "file.ts", + record_state: "stale", + record_source: "read_tool", + roo_read_date: 1500, // Newer stale entry + roo_edit_date: null, + containingMessageTs: 1500, + }, + { + path: "file.ts", + record_state: "active", + record_source: "read_tool", + roo_read_date: 2000, + roo_edit_date: null, + containingMessageTs: 2000, // Will be removed + }, + ], + } + + mockGetTaskMetadata.mockResolvedValue(metadata) + mockSaveTaskMetadata.mockResolvedValue(undefined) + + await tracker.rewindToTimestamp(2000) + + // Verify the saved metadata + expect(mockSaveTaskMetadata).toHaveBeenCalledWith("test-task-id", { + files_in_context: expect.arrayContaining([ + expect.objectContaining({ + path: "file.ts", + roo_read_date: 1000, + record_state: "stale", // First entry stays stale + }), + expect.objectContaining({ + path: "file.ts", + roo_read_date: 1500, + record_state: "active", // Newest stale entry restored to active + }), + ]), + }) + + // Verify the removed entry is gone + const savedMetadata = mockSaveTaskMetadata.mock.calls[0][1] as TaskMetadata + expect(savedMetadata.files_in_context).toHaveLength(2) + expect(savedMetadata.files_in_context.find((e) => e.roo_read_date === 2000)).toBeUndefined() + }) + + it("should not restore stale entry if active entry still exists after rewind", async () => { + const metadata: TaskMetadata = { + files_in_context: [ + { + path: "file.ts", + record_state: "stale", + record_source: "read_tool", + roo_read_date: 500, + roo_edit_date: null, + containingMessageTs: 500, + }, + { + path: "file.ts", + record_state: "active", + record_source: "read_tool", + roo_read_date: 1000, + roo_edit_date: null, + containingMessageTs: 1000, // Before cutoff - stays active + }, + { + path: "file.ts", + record_state: "stale", // Already stale + record_source: "read_tool", + roo_read_date: 1500, + roo_edit_date: null, + containingMessageTs: 3000, // After cutoff - will be removed + }, + ], + } + + mockGetTaskMetadata.mockResolvedValue(metadata) + mockSaveTaskMetadata.mockResolvedValue(undefined) + + await tracker.rewindToTimestamp(2000) + + const savedMetadata = mockSaveTaskMetadata.mock.calls[0][1] as TaskMetadata + + // The original active entry should still be active + const activeEntry = savedMetadata.files_in_context.find((e) => e.roo_read_date === 1000) + expect(activeEntry?.record_state).toBe("active") + + // The remaining stale entry should still be stale (not promoted) + const staleEntry = savedMetadata.files_in_context.find((e) => e.roo_read_date === 500) + expect(staleEntry?.record_state).toBe("stale") + }) + + it("should handle multiple files with different rewind scenarios", async () => { + const metadata: TaskMetadata = { + files_in_context: [ + // File A: active entry removed, has stale to restore + { + path: "fileA.ts", + record_state: "stale", + record_source: "read_tool", + roo_read_date: 1000, + roo_edit_date: null, + containingMessageTs: 1000, + }, + { + path: "fileA.ts", + record_state: "active", + record_source: "read_tool", + roo_read_date: 3000, + roo_edit_date: null, + containingMessageTs: 3000, // Removed + }, + // File B: active entry stays (before cutoff) + { + path: "fileB.ts", + record_state: "active", + record_source: "read_tool", + roo_read_date: 1500, + roo_edit_date: null, + containingMessageTs: 1500, + }, + // File C: all entries removed + { + path: "fileC.ts", + record_state: "active", + record_source: "read_tool", + roo_read_date: 3000, + roo_edit_date: null, + containingMessageTs: 3000, // Removed + }, + ], + } + + mockGetTaskMetadata.mockResolvedValue(metadata) + mockSaveTaskMetadata.mockResolvedValue(undefined) + + await tracker.rewindToTimestamp(2000) + + const savedMetadata = mockSaveTaskMetadata.mock.calls[0][1] as TaskMetadata + + // FileA: stale entry should be restored to active + const fileAEntry = savedMetadata.files_in_context.find((e) => e.path === "fileA.ts") + expect(fileAEntry?.record_state).toBe("active") + expect(fileAEntry?.roo_read_date).toBe(1000) + + // FileB: should stay as-is + const fileBEntry = savedMetadata.files_in_context.find((e) => e.path === "fileB.ts") + expect(fileBEntry?.record_state).toBe("active") + expect(fileBEntry?.roo_read_date).toBe(1500) + + // FileC: should be completely removed + const fileCEntry = savedMetadata.files_in_context.find((e) => e.path === "fileC.ts") + expect(fileCEntry).toBeUndefined() + }) + }) + + describe("edge cases", () => { + it("should handle empty metadata", async () => { + const metadata: TaskMetadata = { + files_in_context: [], + } + + mockGetTaskMetadata.mockResolvedValue(metadata) + mockSaveTaskMetadata.mockResolvedValue(undefined) + + await tracker.rewindToTimestamp(2000) + + expect(mockSaveTaskMetadata).toHaveBeenCalledWith("test-task-id", { + files_in_context: [], + }) + }) + + it("should handle entries with null containingMessageTs", async () => { + const metadata: TaskMetadata = { + files_in_context: [ + { + path: "file.ts", + record_state: "active", + record_source: "read_tool", + roo_read_date: 3000, + roo_edit_date: null, + containingMessageTs: null, // Explicitly null + }, + ], + } + + mockGetTaskMetadata.mockResolvedValue(metadata) + mockSaveTaskMetadata.mockResolvedValue(undefined) + + await tracker.rewindToTimestamp(2000) + + // Entry with null containingMessageTs should be preserved + const savedMetadata = mockSaveTaskMetadata.mock.calls[0][1] as TaskMetadata + expect(savedMetadata.files_in_context).toHaveLength(1) + }) + + it("should handle stale entries without roo_read_date", async () => { + const metadata: TaskMetadata = { + files_in_context: [ + { + path: "file.ts", + record_state: "stale", + record_source: "user_edited", + roo_read_date: null, // No roo_read_date + roo_edit_date: null, + user_edit_date: 1000, + containingMessageTs: 1000, + }, + { + path: "file.ts", + record_state: "active", + record_source: "read_tool", + roo_read_date: 2000, + roo_edit_date: null, + containingMessageTs: 2000, // Will be removed + }, + ], + } + + mockGetTaskMetadata.mockResolvedValue(metadata) + mockSaveTaskMetadata.mockResolvedValue(undefined) + + await tracker.rewindToTimestamp(2000) + + // Stale entry without roo_read_date should NOT be restored + // (only entries with roo_read_date are considered for restoration) + const savedMetadata = mockSaveTaskMetadata.mock.calls[0][1] as TaskMetadata + const remainingEntry = savedMetadata.files_in_context.find((e) => e.path === "file.ts") + expect(remainingEntry?.record_state).toBe("stale") + }) + + it("should handle errors gracefully", async () => { + mockGetTaskMetadata.mockRejectedValue(new Error("Test error")) + + // Should not throw + await expect(tracker.rewindToTimestamp(2000)).resolves.not.toThrow() + + // saveTaskMetadata should not be called if getTaskMetadata fails + expect(mockSaveTaskMetadata).not.toHaveBeenCalled() + }) + }) +}) diff --git a/src/core/message-manager/index.spec.ts b/src/core/message-manager/index.spec.ts index e2c11db3b7..a79bcf54fa 100644 --- a/src/core/message-manager/index.spec.ts +++ b/src/core/message-manager/index.spec.ts @@ -12,6 +12,9 @@ describe("MessageManager", () => { apiConversationHistory: [], overwriteClineMessages: vi.fn(), overwriteApiConversationHistory: vi.fn(), + fileContextTracker: { + rewindToTimestamp: vi.fn().mockResolvedValue(undefined), + }, } manager = new MessageManager(mockTask) diff --git a/src/core/message-manager/index.ts b/src/core/message-manager/index.ts index e35f290c39..88036f1b8c 100644 --- a/src/core/message-manager/index.ts +++ b/src/core/message-manager/index.ts @@ -83,6 +83,10 @@ export class MessageManager { // Step 3: Truncate and clean API history (combined with cleanup for efficiency) await this.truncateApiHistoryWithCleanup(cutoffTs, removedIds, skipCleanup) + + // Step 4: Rewind file context metadata to match the conversation state + // This ensures smart read tracking stays in sync when messages are deleted + await this.task.fileContextTracker.rewindToTimestamp(cutoffTs) } /**