diff --git a/packages/cli/src/running-task-adaptor.ts b/packages/cli/src/running-task-adaptor.ts index 282a1a3be1..f9bd100be4 100644 --- a/packages/cli/src/running-task-adaptor.ts +++ b/packages/cli/src/running-task-adaptor.ts @@ -4,7 +4,9 @@ import { pochiConfig } from "@getpochi/common/configuration"; import type { McpHub } from "@getpochi/common/mcp-utils"; import { FileStateCache, + isVirtualPath, maybePersistToolResult, + resolvePath, } from "@getpochi/common/tool-utils"; import { type ValidCustomAgentFile, @@ -150,6 +152,14 @@ export class CliRunningTaskAdaptor implements RunningTaskAdaptor { logger.warn({ taskId, error }, "Task execution failed"); } + hydrateFileReadHistory(taskId: string, paths: string[]) { + this.getFileStateCache(taskId).hydrateReadHistory( + paths + .filter((filePath) => !isVirtualPath(filePath)) + .map((filePath) => resolvePath(filePath, this.cwd)), + ); + } + clearFileStateCache(taskId: string) { this.fileStateCaches.get(taskId)?.markAllAsWritten(); } @@ -181,13 +191,23 @@ export class CliRunningTaskAdaptor implements RunningTaskAdaptor { (sourceTaskId === this.parentTaskId ? this.parentFileStateCache : undefined); - const target = new FileStateCache(); + const target = existingTarget ?? new FileStateCache(); if (source) { for (const [key, value] of source) { target.set(key, { ...value }); } + const history = source.getReadHistorySnapshot(); + if (history.hydrated) { + target.hydrateReadHistory(history.paths); + } else { + for (const filePath of history.paths) { + target.recordRead(filePath); + } + } + } + if (!existingTarget) { + this.fileStateCaches.set(targetTaskId, target); } - this.fileStateCaches.set(targetTaskId, target); } private getFileStateCache(taskId: string) { diff --git a/packages/cli/src/task-runner.ts b/packages/cli/src/task-runner.ts index ae9cfd24a7..0d1326c90b 100644 --- a/packages/cli/src/task-runner.ts +++ b/packages/cli/src/task-runner.ts @@ -18,7 +18,9 @@ import { } from "@getpochi/common/message-utils"; import { FileStateCache, + isVirtualPath, maybePersistToolResult, + resolvePath, } from "@getpochi/common/tool-utils"; import { type ValidCustomAgentFile, @@ -34,6 +36,7 @@ import { type LiveKitStore, type Message, type Task, + collectPreviouslyReadFilePaths, processContentOutput, } from "@getpochi/livekit"; import { LiveChatKit } from "@getpochi/livekit/node"; @@ -209,6 +212,7 @@ export class TaskRunner { private backgroundJobManager: BackgroundJobManager; private fileSystem: FileSystem; private customAgent?: CustomAgent; + private fileReadHistoryHydrated = false; private attemptCompletionHook?: string; private asyncWaitTimeoutInMs: number; @@ -750,6 +754,8 @@ export class TaskRunner { toolCall: ToolUIPart, envs: Record | undefined, ): Promise { + this.hydrateFileReadHistory(); + try { return await processContentOutput( this.blobStore, @@ -769,6 +775,17 @@ export class TaskRunner { } } + private hydrateFileReadHistory() { + if (this.fileReadHistoryHydrated) return; + + this.toolCallOptions.fileStateCache.hydrateReadHistory( + collectPreviouslyReadFilePaths(this.chat.messages) + .filter((filePath) => !isVirtualPath(filePath)) + .map((filePath) => resolvePath(filePath, this.cwd)), + ); + this.fileReadHistoryHydrated = true; + } + // Helper method to run the command private runAttemptCompletionHook( command: string, diff --git a/packages/common/src/tool-utils/__tests__/file-state-cache.test.ts b/packages/common/src/tool-utils/__tests__/file-state-cache.test.ts index 48e261a2b4..7551e139e3 100644 --- a/packages/common/src/tool-utils/__tests__/file-state-cache.test.ts +++ b/packages/common/src/tool-utils/__tests__/file-state-cache.test.ts @@ -82,7 +82,7 @@ describe("FileStateCache", () => { await expect( checkStaleness(cache, "/tmp/file.txt", async () => 1234, "editing"), - ).rejects.toThrow("File has not been read yet"); + ).rejects.toThrow("No current read snapshot is available"); }); it("throws when writing a file that was never read but exists on disk", async () => { @@ -90,7 +90,59 @@ describe("FileStateCache", () => { await expect( checkStaleness(cache, "/tmp/file.txt", async () => 1234, "writing"), - ).rejects.toThrow("File has not been read yet"); + ).rejects.toThrow("No current read snapshot is available"); + }); + + it("explains when the file was read earlier but its current snapshot is unavailable", async () => { + const cache = new FileStateCache(); + cache.recordRead("/tmp/file.txt"); + + await expect( + checkStaleness(cache, "/tmp/file.txt", async () => 1234, "writing"), + ).rejects.toThrow( + "File was read earlier in this task, but no current read snapshot is available", + ); + }); + + it("explains when the available task history has no successful read", async () => { + const cache = new FileStateCache(); + cache.hydrateReadHistory([]); + + await expect( + checkStaleness(cache, "/tmp/file.txt", async () => 1234, "writing"), + ).rejects.toThrow("File has not been read in the available task history"); + }); + + it("normalizes hydrated read history and clears it with the cache", () => { + const cache = new FileStateCache(); + cache.hydrateReadHistory(["/tmp/src/../file.txt"]); + + expect(cache.getReadHistoryState("/tmp/file.txt")).toBe("read"); + expect(cache.getReadHistoryState("/tmp/other.txt")).toBe("not-read"); + + cache.clear(); + + expect(cache.getReadHistoryState("/tmp/file.txt")).toBe("unknown"); + }); + + it("records successful reads even when there is no text snapshot to cache", async () => { + const cache = new FileStateCache(); + + await withReadFileCache({ + cache, + path: "binary.png", + cwd: "/tmp", + startLine: undefined, + endLine: undefined, + getMtime: async () => 1, + doRead: async () => ({ + result: { content: "image" }, + fileCacheContent: null, + }), + }); + + expect(cache.getReadHistoryState("/tmp/binary.png")).toBe("read"); + expect(cache.has("/tmp/binary.png")).toBe(false); }); it("allows writing a new file that does not exist on disk and was never read", async () => { @@ -195,7 +247,26 @@ describe("withFileStateCacheGuard", () => { fileCacheContent: "new content", }), }), - ).rejects.toThrow("File has not been read yet"); + ).rejects.toThrow("No current read snapshot is available"); + }); + + it("normalizes historical paths before checking the resolved target path", async () => { + const cache = new FileStateCache(); + cache.recordRead("/tmp/src/../file.txt"); + + await expect( + withFileStateCacheGuard({ + cache, + path: "src/../file.txt", + cwd: "/tmp", + getMtime: async () => 1000, + operation: "writing", + doWork: async () => ({ + result: { success: true as const }, + fileCacheContent: "new content", + }), + }), + ).rejects.toThrow("File was read earlier in this task"); }); it("allows writing a brand-new file that does not yet exist on disk", async () => { @@ -216,5 +287,6 @@ describe("withFileStateCacheGuard", () => { }), }), ).resolves.toEqual({ success: true }); + expect(cache.getReadHistoryState("/tmp/brand-new.txt")).toBe("read"); }); }); diff --git a/packages/common/src/tool-utils/file-state-cache.ts b/packages/common/src/tool-utils/file-state-cache.ts index a4b4f8cbe3..b37b22fcf7 100644 --- a/packages/common/src/tool-utils/file-state-cache.ts +++ b/packages/common/src/tool-utils/file-state-cache.ts @@ -31,7 +31,8 @@ const DEFAULT_MAX_ENTRIES = 100; const DEFAULT_MAX_SIZE_BYTES = 25 * 1024 * 1024; /** - * LRU cache that tracks what file content the model has "seen". + * Cache that tracks both the model's current file snapshots and which files + * have become known through a successful read or write during the task. * * this serves three purposes: * @@ -46,10 +47,15 @@ const DEFAULT_MAX_SIZE_BYTES = 25 * 1024 * 1024; * 3. **Post-write cache update**: After a successful edit or write, the cache * is updated with the new content and mtime so subsequent reads can dedup. * - * The cache uses an LRU eviction policy with both entry-count and byte-size limits. + * Current snapshots use an LRU eviction policy with both entry-count and + * byte-size limits. Read history is lightweight provenance and is retained + * independently, so evicting a snapshot does not erase the fact that a file + * was read earlier. */ export class FileStateCache { private readonly entries: Map = new Map(); + private readonly readHistory = new Set(); + private readHistoryHydrated = false; private currentSizeBytes = 0; /** @@ -112,9 +118,37 @@ export class FileStateCache { clear(): void { this.entries.clear(); + this.readHistory.clear(); + this.readHistoryHydrated = false; this.currentSizeBytes = 0; } + recordRead(key: string): void { + this.readHistory.add(this.normalizeKey(key)); + } + + /** Merge read paths recovered from persisted task history. */ + hydrateReadHistory(keys: readonly string[]): void { + for (const key of keys) { + this.recordRead(key); + } + this.readHistoryHydrated = true; + } + + getReadHistoryState(key: string): "read" | "not-read" | "unknown" { + if (this.readHistory.has(this.normalizeKey(key))) { + return "read"; + } + return this.readHistoryHydrated ? "not-read" : "unknown"; + } + + getReadHistorySnapshot(): { paths: string[]; hydrated: boolean } { + return { + paths: [...this.readHistory], + hydrated: this.readHistoryHydrated, + }; + } + /** * Downgrade every cached entry to "written" state (`fromWrite: true`). * @@ -122,7 +156,7 @@ export class FileStateCache { * leave the conversation — a compaction summary, or a retry that strips a * completed read. Keeping the entries preserves the edit/write staleness * guard (so a later edit of an already-read file is not falsely rejected - * with "File has not been read yet"), while `fromWrite: true` stops them + * for lacking a current read snapshot), while `fromWrite: true` stops them * from producing a "File unchanged" dedup stub that would dangle onto a * tool_result no longer present in the conversation. */ @@ -221,12 +255,17 @@ export async function checkStaleness( const cachedState = cache.get(resolvedPath); if (!cachedState) { const currentMtime = await getMtime(resolvedPath); - // If the file exists on disk but was never read, require a read first. + // If the file exists on disk but has no current snapshot, require a read first. // A missing mtime means the file doesn't exist yet, so creating it is fine. if (currentMtime !== undefined) { - throw new Error( - `File has not been read yet. Please read the file before ${operation} it.`, - ); + const readHistoryState = cache.getReadHistoryState(resolvedPath); + const message = + readHistoryState === "read" + ? `File was read earlier in this task, but no current read snapshot is available. Please read the file again before ${operation} it.` + : readHistoryState === "not-read" + ? `File has not been read in the available task history. Please read the file before ${operation} it.` + : `No current read snapshot is available for this file. Please read the file before ${operation} it.`; + throw new Error(message); } return; } @@ -308,14 +347,18 @@ export async function withFileStateCacheGuard(opts: { const { result, fileCacheContent } = await doWork(resolvedPath); - // --- Update cache with new content --- - if (!isVirtual && cache && fileCacheContent !== null) { - await updateCacheAfterWrite( - cache, - resolvedPath, - fileCacheContent, - getMtime, - ); + // A successful write leaves the model knowing the resulting file content, + // so treat it the same as a successful read for historical messaging. + if (!isVirtual && cache) { + cache.recordRead(resolvedPath); + if (fileCacheContent !== null) { + await updateCacheAfterWrite( + cache, + resolvedPath, + fileCacheContent, + getMtime, + ); + } } return result; @@ -406,6 +449,10 @@ export async function withReadFileCache(opts: { const { result, fileCacheContent, fileCacheIsTruncated } = await doRead(resolvedPath); + if (shouldCache) { + cache.recordRead(resolvedPath); + } + // --- Populate cache --- // Store what the model has "seen" so that future reads can dedup, // and edit/write tools can detect external modifications. diff --git a/packages/common/src/vscode-webui-bridge/webview-stub.ts b/packages/common/src/vscode-webui-bridge/webview-stub.ts index 206e341d5f..ad50d21202 100644 --- a/packages/common/src/vscode-webui-bridge/webview-stub.ts +++ b/packages/common/src/vscode-webui-bridge/webview-stub.ts @@ -139,6 +139,12 @@ const VSCodeHostStub = { clearFileStateCache: (_taskId: string): Promise => { return Promise.resolve(); }, + hydrateFileReadHistory: ( + _taskId: string, + _paths: string[], + ): Promise => { + return Promise.resolve(); + }, readRecentFilesForCompact: (_taskId: string) => { return Promise.resolve([]); }, diff --git a/packages/common/src/vscode-webui-bridge/webview.ts b/packages/common/src/vscode-webui-bridge/webview.ts index 379295f41a..72821c8d2d 100644 --- a/packages/common/src/vscode-webui-bridge/webview.ts +++ b/packages/common/src/vscode-webui-bridge/webview.ts @@ -169,6 +169,9 @@ export interface VSCodeHostApi { */ clearFileStateCache(taskId: string): Promise; + /** Restore successful readFile paths from the available task history. */ + hydrateFileReadHistory(taskId: string, paths: string[]): Promise; + /** * Read recent file state cache entries for the given task ID. * Used by compaction to keep recently read file contents visible after diff --git a/packages/livekit/src/background-task/task-executor/__tests__/task-executor.test.ts b/packages/livekit/src/background-task/task-executor/__tests__/task-executor.test.ts index 5efb9799f6..97005de3e6 100644 --- a/packages/livekit/src/background-task/task-executor/__tests__/task-executor.test.ts +++ b/packages/livekit/src/background-task/task-executor/__tests__/task-executor.test.ts @@ -123,6 +123,54 @@ describe("TaskExecutor", () => { await executor.dispose(); }); + it("hydrates successful historical reads once before executing tools", async () => { + const store = new FakeLiveKitStore([ + makeTask({ id: "task", status: "pending-tool" }), + ]); + store.setMessages("task", [ + makeAssistantMessage([ + { + type: "tool-readFile", + toolCallId: "read", + state: "output-available", + input: { path: "src/index.ts" }, + output: { content: "hello", isTruncated: false }, + }, + ]), + makeAssistantMessage([ + makeToolPart("writeToFile", "write-1", { + path: "src/index.ts", + content: "one", + }), + makeToolPart("writeToFile", "write-2", { + path: "src/other.ts", + content: "two", + }), + ]), + ]); + const hydrateFileReadHistory = vi.fn(); + const executeToolCall = vi.fn(async () => ({ ok: true })); + const adaptor = { + ...makeAdaptor({ executeToolCall }), + hydrateFileReadHistory, + } satisfies RunningTaskAdaptor; + const executor = makeExecutor(store, adaptor, { + tools: ["writeToFile"], + }); + + await executor.drain(); + + expect(hydrateFileReadHistory).toHaveBeenCalledTimes(1); + expect(hydrateFileReadHistory).toHaveBeenCalledWith("task", [ + "src/index.ts", + ]); + expect(executeToolCall).toHaveBeenCalledTimes(2); + expect(hydrateFileReadHistory.mock.invocationCallOrder[0]).toBeLessThan( + executeToolCall.mock.invocationCallOrder[0], + ); + await executor.dispose(); + }); + it("does not start duplicate running tasks for the same active task", async () => { const store = new FakeLiveKitStore([ makeTask({ id: "task", status: "pending-tool" }), diff --git a/packages/livekit/src/background-task/task-executor/task-executor.ts b/packages/livekit/src/background-task/task-executor/task-executor.ts index 4aba34fe98..cf0e73d50f 100644 --- a/packages/livekit/src/background-task/task-executor/task-executor.ts +++ b/packages/livekit/src/background-task/task-executor/task-executor.ts @@ -31,6 +31,7 @@ import { import type { BlobStore } from "../../blob-store"; import type { PrepareRequestGetters } from "../../chat/flexible-chat-transport"; import { defaultCatalog as catalog } from "../../livestore"; +import { collectPreviouslyReadFilePaths } from "../../task-utils"; import type { LiveKitStore, Message, Task } from "../../types"; const logger = getLogger("TaskExecutor"); @@ -57,6 +58,7 @@ export interface RunningTaskAdaptor { taskId: string; cwd: string | undefined; }): PrepareRequestGetters; + hydrateFileReadHistory?(taskId: string, paths: string[]): MaybePromise; executeToolCall(args: TaskExecutorToolCallExecution): Promise; onTaskError?(taskId: string, error: Error): MaybePromise; } @@ -284,6 +286,7 @@ class RunningTask { private chatKit: RunningTaskChatKit | undefined; private retryCount = 0; private toolRejectionCount = 0; + private fileReadHistoryHydration: Promise | undefined; private disposed = false; readonly done: Promise; @@ -520,6 +523,7 @@ class RunningTask { } try { + await this.hydrateFileReadHistory(); const result = await this.adaptor.executeToolCall({ taskId: this.taskId, parentTaskId: this.taskState.parentTaskId, @@ -557,6 +561,16 @@ class RunningTask { } } + private hydrateFileReadHistory(): Promise { + this.fileReadHistoryHydration ??= Promise.resolve( + this.adaptor.hydrateFileReadHistory?.( + this.taskId, + collectPreviouslyReadFilePaths(this.chat.messages), + ), + ); + return this.fileReadHistoryHydration; + } + private validateToolCall( toolName: string, input: unknown, diff --git a/packages/livekit/src/index.ts b/packages/livekit/src/index.ts index 8354f42c3b..00dee1e085 100644 --- a/packages/livekit/src/index.ts +++ b/packages/livekit/src/index.ts @@ -24,6 +24,7 @@ export type { BlobStore } from "./blob-store"; export { processContentOutput, fileToUri, findBlob } from "./store-blob"; export { + collectPreviouslyReadFilePaths, extractAttemptCompletionResult, extractTaskResult, formatFollowupQuestions, diff --git a/packages/livekit/src/task-utils.test.ts b/packages/livekit/src/task-utils.test.ts index d09a2e9a72..670a20e8e6 100644 --- a/packages/livekit/src/task-utils.test.ts +++ b/packages/livekit/src/task-utils.test.ts @@ -1,11 +1,88 @@ import { describe, expect, it } from "vitest"; import { z } from "zod"; import { + collectPreviouslyReadFilePaths, extractAttemptCompletionResult, extractTaskResult, formatFollowupQuestions, } from "./task-utils"; +describe("collectPreviouslyReadFilePaths", () => { + it("collects unique paths from successful file reads and writes", () => { + const messages = [ + { + id: "assistant-1", + role: "assistant", + parts: [ + { + type: "tool-readFile", + toolCallId: "read-1", + state: "output-available", + input: { path: "src/index.ts" }, + output: { content: "first", isTruncated: false }, + }, + { + type: "tool-readFile", + toolCallId: "read-2", + state: "output-available", + input: { path: "src/index.ts", startLine: 1, endLine: 10 }, + output: { content: "first", isTruncated: false }, + }, + { + type: "tool-readFile", + toolCallId: "read-3", + state: "output-available", + input: { path: "src/other.ts" }, + output: { error: "File not found" }, + }, + { + type: "tool-writeToFile", + toolCallId: "write-1", + state: "output-available", + input: { path: "src/generated.ts", content: "generated" }, + output: { success: true }, + }, + { + type: "tool-applyDiff", + toolCallId: "write-2", + state: "output-available", + input: { path: "src/index.ts", patch: "patch" }, + output: { success: true }, + }, + { + type: "tool-editNotebook", + toolCallId: "write-3", + state: "output-available", + input: { path: "notebooks/demo.ipynb", edits: [] }, + output: { success: true }, + }, + { + type: "tool-multiApplyDiff", + toolCallId: "write-legacy", + state: "output-available", + input: { path: "src/legacy.ts", edits: [] }, + output: { success: true }, + }, + { + type: "tool-writeToFile", + toolCallId: "write-failed", + state: "output-available", + input: { path: "src/failed.ts", content: "failed" }, + output: { error: "Write failed" }, + }, + ], + }, + ] as any; + + expect(collectPreviouslyReadFilePaths(messages)).toEqual([ + "src/index.ts", + "src/generated.ts", + "notebooks/demo.ipynb", + "src/legacy.ts", + ]); + }); +}); + describe("formatFollowupQuestions", () => { it("formats all questions from the new askFollowupQuestion payload", () => { expect( diff --git a/packages/livekit/src/task-utils.ts b/packages/livekit/src/task-utils.ts index 0b5c21d827..5056e23200 100644 --- a/packages/livekit/src/task-utils.ts +++ b/packages/livekit/src/task-utils.ts @@ -12,6 +12,45 @@ export type TaskStatusLike = export type BackgroundJobStatus = "idle" | "running" | "completed"; +/** + * Collect paths from successful file reads and writes that are still present + * in the task history. A successful write means the resulting content was + * known to the model, so it counts as a read for historical messaging. + * Historical knowledge never bypasses the current mtime-based guard. + */ +export function collectPreviouslyReadFilePaths( + messages: readonly Message[], +): string[] { + const paths = new Set(); + + for (const message of messages) { + for (const part of message.parts) { + if ( + (part.type !== "tool-readFile" && + part.type !== "tool-writeToFile" && + part.type !== "tool-applyDiff" && + part.type !== "tool-multiApplyDiff" && + part.type !== "tool-editNotebook") || + part.state !== "output-available" + ) { + continue; + } + + const output = part.output as unknown; + if (output && typeof output === "object" && "error" in output) { + continue; + } + + const input = part.input as { path?: unknown }; + if (typeof input.path === "string") { + paths.add(input.path); + } + } + } + + return [...paths]; +} + function formatQuestion({ question, header, options }: Question) { const title = header ? `[${header}] ${question}` : question; if (!options?.length) return title; diff --git a/packages/tools/src/types.ts b/packages/tools/src/types.ts index bd79245e72..4b67cb8c4a 100644 --- a/packages/tools/src/types.ts +++ b/packages/tools/src/types.ts @@ -40,6 +40,10 @@ export interface IFileStateCache { has(key: string): boolean; delete(key: string): boolean; clear(): void; + recordRead(key: string): void; + hydrateReadHistory(keys: readonly string[]): void; + getReadHistoryState(key: string): "read" | "not-read" | "unknown"; + getReadHistorySnapshot(): { paths: string[]; hydrated: boolean }; } export type CompiledToolPolicy = diff --git a/packages/vscode-webui/src/features/chat/page.tsx b/packages/vscode-webui/src/features/chat/page.tsx index 6970ba0af0..86cd1a138b 100644 --- a/packages/vscode-webui/src/features/chat/page.tsx +++ b/packages/vscode-webui/src/features/chat/page.tsx @@ -19,7 +19,12 @@ import { constants, formatters } from "@getpochi/common"; import type { UserInfo } from "@getpochi/common/configuration"; import { hasActiveTodos } from "@getpochi/common/message-utils"; import type { PochiTaskInfo } from "@getpochi/common/vscode-webui-bridge"; -import { type Message, type Task, catalog } from "@getpochi/livekit"; +import { + type Message, + type Task, + catalog, + collectPreviouslyReadFilePaths, +} from "@getpochi/livekit"; import { useLiveChatKit } from "@getpochi/livekit/react"; import { parseOutputSchema } from "@getpochi/tools"; import { useStoreRegistry } from "@livestore/react"; @@ -105,6 +110,17 @@ function Chat({ user, uid, info }: ChatProps) { const isSubTask = !!task?.parentId; const messageRows = store.useQuery(catalog.queries.makeMessagesQuery(uid)); + const historicalReadFilePaths = useMemo( + () => + collectPreviouslyReadFilePaths( + messageRows.map((row) => row.data as Message), + ), + [messageRows], + ); + + useEffect(() => { + void vscodeHost.hydrateFileReadHistory(uid, historicalReadFilePaths); + }, [uid, historicalReadFilePaths]); // inherit autoApproveSettings from parent task useEffect(() => { diff --git a/packages/vscode-webui/src/features/tools/hooks/__tests__/use-live-sub-task.test.tsx b/packages/vscode-webui/src/features/tools/hooks/__tests__/use-live-sub-task.test.tsx index 1df0e69d46..7096618fbf 100644 --- a/packages/vscode-webui/src/features/tools/hooks/__tests__/use-live-sub-task.test.tsx +++ b/packages/vscode-webui/src/features/tools/hooks/__tests__/use-live-sub-task.test.tsx @@ -8,11 +8,15 @@ import { useLiveSubTask } from "../use-live-sub-task"; const useLiveChatKitGettersMock = vi.hoisted(() => vi.fn(() => ({}))); const storeMock = vi.hoisted(() => ({ storeId: "store-1", - useQuery: vi.fn(() => ({ - id: "subtask-1", - parentId: "parent-1", - status: "pending-model", - })), + useQuery: vi.fn((query: { kind?: string }) => + query.kind === "messages" + ? [] + : { + id: "subtask-1", + parentId: "parent-1", + status: "pending-model", + }, + ), })); const retryErrorMock = vi.hoisted<{ current: Error | undefined }>(() => ({ current: undefined, @@ -73,6 +77,7 @@ vi.mock("@/lib/use-default-store", () => ({ vi.mock("@/lib/vscode", () => ({ vscodeHost: { clearFileStateCache: vi.fn(), + hydrateFileReadHistory: vi.fn(), }, })); @@ -92,8 +97,13 @@ vi.mock("@getpochi/livekit", () => ({ catalog: { queries: { makeTaskQuery: vi.fn((taskId: string) => ({ taskId })), + makeMessagesQuery: vi.fn((taskId: string) => ({ + kind: "messages", + taskId, + })), }, }, + collectPreviouslyReadFilePaths: vi.fn(() => []), })); vi.mock("@getpochi/livekit/react", () => ({ diff --git a/packages/vscode-webui/src/features/tools/hooks/use-live-sub-task.tsx b/packages/vscode-webui/src/features/tools/hooks/use-live-sub-task.tsx index f4216b03cf..a259a01286 100644 --- a/packages/vscode-webui/src/features/tools/hooks/use-live-sub-task.tsx +++ b/packages/vscode-webui/src/features/tools/hooks/use-live-sub-task.tsx @@ -21,7 +21,11 @@ import { vscodeHost } from "@/lib/vscode"; import { useChat } from "@ai-sdk/react"; import { constants } from "@getpochi/common"; import type { BuiltinSubAgentInfo } from "@getpochi/common/vscode-webui-bridge"; -import { catalog } from "@getpochi/livekit"; +import { + type Message, + catalog, + collectPreviouslyReadFilePaths, +} from "@getpochi/livekit"; import { useLiveChatKit } from "@getpochi/livekit/react"; import { type Todo, @@ -85,6 +89,17 @@ export function useLiveSubTask( const store = useDefaultStore(); const task = store.useQuery(catalog.queries.makeTaskQuery(uid)); + const messageRows = store.useQuery(catalog.queries.makeMessagesQuery(uid)); + const historicalReadFilePaths = useMemo( + () => + collectPreviouslyReadFilePaths( + messageRows.map((row) => row.data as Message), + ), + [messageRows], + ); + useEffect(() => { + void vscodeHost.hydrateFileReadHistory(uid, historicalReadFilePaths); + }, [uid, historicalReadFilePaths]); const todosRef = useRef(undefined); todosRef.current = tool.state !== "input-streaming" && diff --git a/packages/vscode-webui/src/lib/vscode-running-task-adaptor.ts b/packages/vscode-webui/src/lib/vscode-running-task-adaptor.ts index 069208e978..e0c812af5a 100644 --- a/packages/vscode-webui/src/lib/vscode-running-task-adaptor.ts +++ b/packages/vscode-webui/src/lib/vscode-running-task-adaptor.ts @@ -107,6 +107,10 @@ export class VscodeRunningTaskAdaptor implements RunningTaskAdaptor { return result; } + hydrateFileReadHistory(taskId: string, paths: string[]) { + return vscodeHost.hydrateFileReadHistory(taskId, paths); + } + onTaskError(taskId: string, error: Error) { logger.warn({ taskId, error }, "Task execution failed"); } diff --git a/packages/vscode-webui/src/lib/vscode.ts b/packages/vscode-webui/src/lib/vscode.ts index 967d958cc1..d08e43626f 100644 --- a/packages/vscode-webui/src/lib/vscode.ts +++ b/packages/vscode-webui/src/lib/vscode.ts @@ -145,6 +145,7 @@ function createVSCodeHost(): VSCodeHostApi { "readPochiTabs", "closePochiTabs", "clearFileStateCache", + "hydrateFileReadHistory", "readRecentFilesForCompact", "queryGithubIssues", "readGitBranches", diff --git a/packages/vscode/src/integrations/webview/vscode-host-impl.ts b/packages/vscode/src/integrations/webview/vscode-host-impl.ts index bd964d5904..781f048cb2 100644 --- a/packages/vscode/src/integrations/webview/vscode-host-impl.ts +++ b/packages/vscode/src/integrations/webview/vscode-host-impl.ts @@ -442,6 +442,14 @@ export class VSCodeHostImpl implements VSCodeHostApi, vscode.Disposable { this.fileStateCacheRegistry.markAllAsWritten(taskId); }; + hydrateFileReadHistory = async ( + taskId: string, + paths: string[], + ): Promise => { + if (!this.cwd) return; + this.fileStateCacheRegistry.hydrateReadHistory(taskId, paths, this.cwd); + }; + readRecentFilesForCompact = async (taskId: string) => { return this.fileStateCacheRegistry.getRecentFiles(taskId); }; diff --git a/packages/vscode/src/lib/__test__/file-state-cache-registry.test.ts b/packages/vscode/src/lib/__test__/file-state-cache-registry.test.ts index b74c1008bd..0b244c1d89 100644 --- a/packages/vscode/src/lib/__test__/file-state-cache-registry.test.ts +++ b/packages/vscode/src/lib/__test__/file-state-cache-registry.test.ts @@ -98,6 +98,44 @@ describe("FileStateCacheRegistry", () => { ); }); + it("merges parent snapshots into a target that only has hydrated history", () => { + const registry = new FileStateCacheRegistry(); + + const parent = registry.get("parent-task"); + parent.set("/tmp/parent.txt", { + content: "parent", + timestamp: 1, + startLine: 1, + endLine: 1, + }); + parent.recordRead("/tmp/parent.txt"); + registry.hydrateReadHistory("fork-task", ["child.txt"], "/tmp"); + + registry.copyIfAbsent("parent-task", "fork-task"); + + const fork = registry.get("fork-task"); + assert.strictEqual(fork.get("/tmp/parent.txt")?.content, "parent"); + assert.strictEqual(fork.getReadHistoryState("/tmp/parent.txt"), "read"); + assert.strictEqual(fork.getReadHistoryState("/tmp/child.txt"), "read"); + }); + + it("resolves hydrated paths and ignores virtual files", () => { + const registry = new FileStateCacheRegistry(); + + registry.hydrateReadHistory( + "task", + ["src/../file.txt", "pochi://virtual/file.txt"], + "/repo", + ); + + const cache = registry.get("task"); + assert.strictEqual(cache.getReadHistoryState("/repo/file.txt"), "read"); + assert.strictEqual( + cache.getReadHistoryState("pochi://virtual/file.txt"), + "not-read", + ); + }); + it("disposes all retained caches", () => { const registry = new FileStateCacheRegistry(); diff --git a/packages/vscode/src/lib/file-state-cache-registry.ts b/packages/vscode/src/lib/file-state-cache-registry.ts index c42e9b60e5..765cfe7937 100644 --- a/packages/vscode/src/lib/file-state-cache-registry.ts +++ b/packages/vscode/src/lib/file-state-cache-registry.ts @@ -1,6 +1,8 @@ import { FileStateCache, type RecentFileState, + isVirtualPath, + resolvePath, } from "@getpochi/common/tool-utils"; import { injectable, singleton } from "tsyringe"; import type * as vscode from "vscode"; @@ -26,19 +28,37 @@ export class FileStateCacheRegistry implements vscode.Disposable { } const source = this.caches.get(sourceTaskId); - const target = new FileStateCache(); + const target = existingTarget ?? new FileStateCache(); if (source) { for (const [key, value] of source) { target.set(key, { ...value }); } + const history = source.getReadHistorySnapshot(); + if (history.hydrated) { + target.hydrateReadHistory(history.paths); + } else { + for (const filePath of history.paths) { + target.recordRead(filePath); + } + } + } + if (!existingTarget) { + this.caches.set(targetTaskId, target); } - this.caches.set(targetTaskId, target); } markAllAsWritten(taskId: string): void { this.caches.get(taskId)?.markAllAsWritten(); } + hydrateReadHistory(taskId: string, paths: string[], cwd: string): void { + this.get(taskId).hydrateReadHistory( + paths + .filter((filePath) => !isVirtualPath(filePath)) + .map((filePath) => resolvePath(filePath, cwd)), + ); + } + getRecentFiles(taskId: string): RecentFileState[] { return this.caches.get(taskId)?.getRecentFiles() ?? []; }