feat(session): rebuild session state to support restoring snippet from history.

Author: qorzj

src/common/state.ts

@@ -1,4 +1,6 @@
+import * as fs from "fs";
 import * as path from "path";
+import { readTextFileWithMetadata } from "./file-utils";
 import { posixPathToWindowsPath } from "./shell-utils";
 
 export type FileLineEnding = "LF" | "CRLF";
@@ -25,6 +27,11 @@ export type FileSnippet = {
   scopeType: "snippet" | "full";
 };
 
+export type SessionStateHistoryMessage = {
+  role?: unknown;
+  content?: unknown;
+};
+
 const fileStatesBySession = new Map<string, Map<string, FileState>>();
 const snippetsBySession = new Map<string, Map<string, FileSnippet>>();
 const snippetCountersBySession = new Map<string, number>();
@@ -39,9 +46,24 @@ export function clearSessionState(sessionId: string): void {
   fileStatesBySession.delete(sessionId);
   snippetsBySession.delete(sessionId);
   snippetCountersBySession.delete(sessionId);
+  fullFileSnippetCountersBySession.delete(sessionId);
   fileVersionsBySession.delete(sessionId);
 }
 
+export function hasSessionState(sessionId: string): boolean {
+  if (!sessionId) {
+    return false;
+  }
+
+  return Boolean(
+    fileStatesBySession.get(sessionId)?.size ||
+    snippetsBySession.get(sessionId)?.size ||
+    snippetCountersBySession.has(sessionId) ||
+    fullFileSnippetCountersBySession.has(sessionId) ||
+    fileVersionsBySession.get(sessionId)?.size
+  );
+}
+
 export function normalizeFilePath(filePath: string, platform: NodeJS.Platform = process.platform): string {
   const nativePath = normalizeNativeFilePath(filePath, platform);
   return platform === "win32" ? path.win32.normalize(nativePath) : path.normalize(nativePath);
@@ -178,6 +200,32 @@ export function createFullFileSnippet(
   return createSnippetWithId(sessionId, filePath, startLine, endLine, preview, `full_file_${nextCounter}`, "full");
 }
 
+export function restoreSnippet(
+  sessionId: string,
+  snippet: {
+    id: string;
+    filePath: string;
+    startLine: number;
+    endLine: number;
+    preview?: string;
+    scopeType?: FileSnippet["scopeType"];
+  }
+): FileSnippet | null {
+  const restored = createSnippetWithId(
+    sessionId,
+    snippet.filePath,
+    snippet.startLine,
+    snippet.endLine,
+    snippet.preview ?? "",
+    snippet.id,
+    snippet.scopeType ?? inferSnippetScopeType(snippet.id)
+  );
+  if (restored) {
+    updateSnippetCounters(sessionId, snippet.id);
+  }
+  return restored;
+}
+
 function createSnippetWithId(
   sessionId: string,
   filePath: string,
@@ -210,6 +258,27 @@ function createSnippetWithId(
   return snippet;
 }
 
+function inferSnippetScopeType(id: string): FileSnippet["scopeType"] {
+  return id.startsWith("full_file_") ? "full" : "snippet";
+}
+
+function updateSnippetCounters(sessionId: string, id: string): void {
+  const fullFileMatch = /^full_file_(\d+)$/.exec(id);
+  if (fullFileMatch) {
+    const nextCounter = Number(fullFileMatch[1]) + 1;
+    const current = fullFileSnippetCountersBySession.get(sessionId) ?? 0;
+    fullFileSnippetCountersBySession.set(sessionId, Math.max(current, nextCounter));
+    return;
+  }
+
+  const snippetMatch = /^snippet_(\d+)$/.exec(id);
+  if (snippetMatch) {
+    const currentCounter = Number(snippetMatch[1]);
+    const current = snippetCountersBySession.get(sessionId) ?? 0;
+    snippetCountersBySession.set(sessionId, Math.max(current, currentCounter));
+  }
+}
+
 export function getSnippet(sessionId: string, snippetId: string): FileSnippet | null {
   if (!sessionId || !snippetId) {
     return null;
@@ -220,3 +289,247 @@ export function getSnippet(sessionId: string, snippetId: string): FileSnippet |
 export function hasSnippetOutdatedFileVersion(sessionId: string, snippet: FileSnippet): boolean {
   return getFileVersion(sessionId, snippet.filePath) > snippet.fileVersion;
 }
+
+export function rebuildSessionStateFromHistory(
+  sessionId: string,
+  messages: Iterable<SessionStateHistoryMessage>
+): void {
+  if (!sessionId || hasSessionState(sessionId)) {
+    return;
+  }
+
+  for (const message of messages) {
+    if (message.role !== "tool" || typeof message.content !== "string") {
+      continue;
+    }
+
+    const result = parsePersistedToolResult(message.content);
+    if (!result || result.ok !== true) {
+      continue;
+    }
+
+    const metadata = asRecord(result.metadata);
+    if (!metadata) {
+      continue;
+    }
+
+    if (result.name === "read") {
+      rebuildReadResult(sessionId, result, metadata);
+    } else if (result.name === "edit") {
+      rebuildEditResult(sessionId, metadata);
+    } else if (result.name === "write") {
+      rebuildWriteResult(sessionId, metadata);
+    }
+  }
+}
+
+function rebuildReadResult(
+  sessionId: string,
+  result: Record<string, unknown>,
+  metadata: Record<string, unknown>
+): void {
+  const snippet = asRecord(metadata.snippet);
+  if (!snippet) {
+    return;
+  }
+
+  const restored = restoreSnippetFromRecord(sessionId, snippet, {
+    idKey: "id",
+    filePathKey: "filePath",
+    startLineKey: "startLine",
+    endLineKey: "endLine",
+    preview: typeof result.output === "string" ? result.output : "",
+  });
+  if (!restored) {
+    return;
+  }
+
+  refreshRebuiltFileState(sessionId, restored.filePath, {
+    scopeType: restored.scopeType,
+    startLine: restored.startLine,
+    endLine: restored.endLine,
+    incrementVersion: false,
+  });
+}
+
+function rebuildEditResult(sessionId: string, metadata: Record<string, unknown>): void {
+  const scope = asRecord(metadata.scope);
+  if (scope) {
+    restoreSnippetFromRecord(sessionId, scope, {
+      idKey: "snippet_id",
+      filePathKey: "file_path",
+      startLineKey: "start_line",
+      endLineKey: "end_line",
+      scopeType: metadata.read_scope_type === "full" ? "full" : undefined,
+    });
+  }
+
+  const scopeFilePath = typeof scope?.file_path === "string" ? scope.file_path : undefined;
+  rebuildCandidateSnippets(sessionId, metadata, scopeFilePath);
+
+  const filePath = typeof metadata.file_path === "string" ? metadata.file_path : scopeFilePath;
+  if (filePath && metadata.cache_refreshed === true) {
+    refreshRebuiltFileState(sessionId, filePath, { incrementVersion: true });
+  }
+}
+
+function rebuildWriteResult(sessionId: string, metadata: Record<string, unknown>): void {
+  if (metadata.cache_refreshed !== true || typeof metadata.file_path !== "string") {
+    return;
+  }
+
+  refreshRebuiltFileState(sessionId, metadata.file_path, { incrementVersion: true });
+}
+
+function rebuildCandidateSnippets(
+  sessionId: string,
+  metadata: Record<string, unknown>,
+  filePath: string | undefined
+): void {
+  if (!filePath) {
+    return;
+  }
+
+  const candidates = Array.isArray(metadata.candidates) ? metadata.candidates : [];
+  for (const candidate of candidates) {
+    const record = asRecord(candidate);
+    if (!record) {
+      continue;
+    }
+    restoreSnippetFromRecord(
+      sessionId,
+      { ...record, file_path: filePath },
+      {
+        idKey: "snippet_id",
+        filePathKey: "file_path",
+        startLineKey: "start_line",
+        endLineKey: "end_line",
+        scopeType: "snippet",
+        preview: typeof record.preview === "string" ? record.preview : "",
+      }
+    );
+  }
+
+  const closestMatch = asRecord(metadata.closest_match);
+  if (closestMatch) {
+    restoreSnippetFromRecord(
+      sessionId,
+      { ...closestMatch, file_path: filePath },
+      {
+        idKey: "snippet_id",
+        filePathKey: "file_path",
+        startLineKey: "start_line",
+        endLineKey: "end_line",
+        scopeType: "snippet",
+        preview: typeof closestMatch.preview === "string" ? closestMatch.preview : "",
+      }
+    );
+  }
+}
+
+function restoreSnippetFromRecord(
+  sessionId: string,
+  record: Record<string, unknown>,
+  options: {
+    idKey: string;
+    filePathKey: string;
+    startLineKey: string;
+    endLineKey: string;
+    preview?: string;
+    scopeType?: FileSnippet["scopeType"];
+  }
+): FileSnippet | null {
+  const rawId = record[options.idKey];
+  const rawFilePath = record[options.filePathKey];
+  const id = typeof rawId === "string" ? rawId.trim() : "";
+  const filePath = typeof rawFilePath === "string" ? normalizeFilePath(rawFilePath) : "";
+  const startLine = toPositiveInteger(record[options.startLineKey]);
+  const endLine = toPositiveInteger(record[options.endLineKey]);
+  if (!id || !filePath || startLine === null || endLine === null) {
+    return null;
+  }
+
+  return restoreSnippet(sessionId, {
+    id,
+    filePath,
+    startLine,
+    endLine,
+    preview: options.preview,
+    scopeType: options.scopeType,
+  });
+}
+
+function refreshRebuiltFileState(
+  sessionId: string,
+  rawFilePath: string,
+  options: {
+    scopeType?: FileSnippet["scopeType"];
+    startLine?: number;
+    endLine?: number;
+    incrementVersion?: boolean;
+  } = {}
+): void {
+  const filePath = normalizeFilePath(rawFilePath);
+  if (!filePath || !fs.existsSync(filePath)) {
+    return;
+  }
+
+  try {
+    const stat = fs.statSync(filePath);
+    if (stat.isDirectory()) {
+      return;
+    }
+
+    const metadata = readTextFileWithMetadata(filePath);
+    const isPartialView = options.scopeType === "snippet";
+    const content = isPartialView
+      ? metadata.content
+          .split("\n")
+          .slice((options.startLine ?? 1) - 1, options.endLine)
+          .join("\n")
+      : metadata.content;
+
+    recordFileState(
+      sessionId,
+      {
+        filePath,
+        content,
+        timestamp: metadata.timestamp,
+        offset: isPartialView ? options.startLine : undefined,
+        limit:
+          isPartialView && options.startLine !== undefined && options.endLine !== undefined
+            ? Math.max(1, options.endLine - options.startLine + 1)
+            : undefined,
+        isPartialView,
+        encoding: metadata.encoding,
+        lineEndings: metadata.lineEndings,
+      },
+      { incrementVersion: options.incrementVersion }
+    );
+  } catch {
+    // Best-effort restore: later tool execution will return the precise filesystem error.
+  }
+}
+
+function parsePersistedToolResult(content: string): Record<string, unknown> | null {
+  try {
+    return asRecord(JSON.parse(content));
+  } catch {
+    return null;
+  }
+}
+
+function asRecord(value: unknown): Record<string, unknown> | null {
+  if (value === null || typeof value !== "object" || Array.isArray(value)) {
+    return null;
+  }
+  return value as Record<string, unknown>;
+}
+
+function toPositiveInteger(value: unknown): number | null {
+  const numberValue = typeof value === "number" ? value : typeof value === "string" ? Number(value) : NaN;
+  if (!Number.isInteger(numberValue) || numberValue < 1) {
+    return null;
+  }
+  return numberValue;
+}

src/session.ts

@@ -31,7 +31,7 @@ import { logApiError } from "./common/error-logger";
 import { logOpenAIChatCompletionDebug, normalizeDebugError } from "./common/debug-logger";
 import { killProcessTree } from "./common/process-tree";
 import { GitFileHistory } from "./common/file-history";
-import { getSnippet } from "./common/state";
+import { clearSessionState, getSnippet, rebuildSessionStateFromHistory } from "./common/state";
 import {
   appendProjectPermissionAllows,
   buildPermissionToolExecution,
@@ -45,7 +45,6 @@ import {
   type UserToolPermission,
 } from "./common/permissions";
 import { clearSessionWorkingDir } from "./tools/bash-handler";
-import { clearSessionState } from "./common/state";
 
 export type { PermissionScope } from "./settings";
 export type {
@@ -1143,6 +1142,7 @@ ${skillMd}
     const { client, model, baseURL, thinkingEnabled, reasoningEffort, debugLogEnabled, notify, env } =
       this.createOpenAIClient();
     const now = new Date().toISOString();
+    rebuildSessionStateFromHistory(sessionId, this.listSessionMessages(sessionId));
 
     if (!client) {
       this.updateSessionEntry(sessionId, (entry) => ({