@meowlynxsea/koi 0.1.0

package/src/agent/session-fork.ts

@@ -0,0 +1,299 @@
+/**
+ * Session Fork Manager
+ *
+ * Manages fork metadata for session branching. Tracks:
+ * - Fork relationships between sessions
+ * - Task snapshots at fork time
+ * - Agent mode state at fork time
+ * - Pending plan content at fork time
+ *
+ * All fork metadata is persisted to ~/.config/koi/sessions/forks/<sessionId>.json
+ */
+
+import fs from "fs";
+import path from "path";
+import os from "os";
+import type { Task } from "./session-tasks.js";
+import type { AgentMode } from "./mode.js";
+
+const CONFIG_DIR = path.join(os.homedir(), ".config", "koi");
+const KOI_SESSIONS_DIR = path.join(CONFIG_DIR, "sessions");
+const FORK_METADATA_DIR = path.join(KOI_SESSIONS_DIR, "forks");
+
+export interface ForkMetadata {
+  /** ID of the forked session */
+  forkId: string;
+  /** ID of the source session this was forked from */
+  sourceSessionId: string;
+  /** Branch ID in the source session at the fork point */
+  sourceBranchId: string;
+  /** Entry ID where the fork occurred */
+  forkPoint: string;
+  /** Timestamp when the fork was created */
+  forkedAt: number;
+  /** Snapshot of tasks at fork time (before fork) */
+  tasksSnapshot: Task[];
+  /** Agent mode state at fork time */
+  agentMode: AgentMode;
+  /** Active tool names at fork time */
+  activeTools: string[];
+  /** Pending plan text at fork time (if in plan mode with pending content) */
+  pendingPlanText: string | null;
+}
+
+export interface ForkTreeNode {
+  sessionId: string;
+  forkId: string;
+  title: string;
+  forkedAt: number;
+  children: ForkTreeNode[];
+}
+
+/**
+ * File System Helpers
+ */
+
+function ensureDir(dir: string): void {
+  if (!fs.existsSync(dir)) {
+    fs.mkdirSync(dir, { recursive: true, mode: 0o700 });
+  }
+}
+
+function getForkMetadataPath(sessionId: string): string {
+  return path.join(FORK_METADATA_DIR, `${sessionId}.json`);
+}
+
+function safeReadFile<T>(filePath: string, parser: (raw: string) => T): T | null {
+  try {
+    if (!fs.existsSync(filePath)) return null;
+    const raw = fs.readFileSync(filePath, "utf-8");
+    return parser(raw);
+  } catch {
+    return null;
+  }
+}
+
+function safeWriteFile(filePath: string, data: string): void {
+  try {
+    ensureDir(path.dirname(filePath));
+    fs.writeFileSync(filePath, data, { mode: 0o600 });
+  } catch {
+    // Silently ignore write errors
+  }
+}
+
+function safeDeleteFile(filePath: string): void {
+  try {
+    if (fs.existsSync(filePath)) fs.unlinkSync(filePath);
+  } catch {
+    // ignore
+  }
+}
+
+/**
+ * ForkManager
+ *
+ * Singleton class managing all fork metadata for sessions.
+ * Provides methods to create, save, load, and query fork relationships.
+ */
+export class ForkManager {
+  private cache = new Map<string, ForkMetadata | null>();
+
+  /**
+   * Save fork metadata for a session.
+   * Also updates the parent session's childForks list.
+   */
+  saveForkMetadata(forkId: string, metadata: ForkMetadata): void {
+    this.cache.set(forkId, metadata);
+    safeWriteFile(
+      getForkMetadataPath(forkId),
+      JSON.stringify(metadata, null, 2) + "\n"
+    );
+
+    // Update parent session's child forks
+    const parentMeta = this.loadForkMetadata(metadata.sourceSessionId);
+    if (parentMeta) {
+      // Parent is a forked session, update its child forks tracking
+      const childForksPath = path.join(
+        KOI_SESSIONS_DIR,
+        metadata.sourceSessionId,
+        "child-forks.json"
+      );
+      const existingChildForks: string[] = safeReadFile<string[]>(childForksPath, (raw) => JSON.parse(raw) as string[]) ?? [];
+      if (!existingChildForks.includes(forkId)) {
+        safeWriteFile(childForksPath, JSON.stringify([...existingChildForks, forkId], null, 2) + "\n");
+      }
+    }
+  }
+
+  /**
+   * Load fork metadata for a session from disk.
+   */
+  loadForkMetadata(sessionId: string): ForkMetadata | null {
+    // Check cache first
+    if (this.cache.has(sessionId)) {
+      return this.cache.get(sessionId) ?? null;
+    }
+
+    const metadata = safeReadFile<ForkMetadata>(
+      getForkMetadataPath(sessionId),
+      (raw) => JSON.parse(raw) as ForkMetadata
+    );
+
+    this.cache.set(sessionId, metadata);
+    return metadata;
+  }
+
+  /**
+   * Get the fork metadata for a session, returning null if not a fork.
+   */
+  getForkMetadata(sessionId: string): ForkMetadata | null {
+    return this.loadForkMetadata(sessionId);
+  }
+
+  /**
+   * Check if a session is a fork.
+   */
+  isFork(sessionId: string): boolean {
+    const meta = this.loadForkMetadata(sessionId);
+    return meta !== null;
+  }
+
+  /**
+   * Get the source session ID for a fork.
+   */
+  getSourceSessionId(sessionId: string): string | null {
+    const meta = this.loadForkMetadata(sessionId);
+    return meta?.sourceSessionId ?? null;
+  }
+
+  /**
+   * Get all child fork IDs for a session.
+   */
+  getChildForks(sessionId: string): string[] {
+    const childForksPath = path.join(KOI_SESSIONS_DIR, sessionId, "child-forks.json");
+    return safeReadFile<string[]>(childForksPath, (raw) => JSON.parse(raw) as string[]) ?? [];
+  }
+
+  /**
+   * Calculate the fork depth of a session (0 for original, 1+ for forks).
+   */
+  getForkDepth(sessionId: string): number {
+    let depth = 0;
+    let currentId: string | null = sessionId;
+    const visited = new Set<string>();
+
+    while (currentId) {
+      if (visited.has(currentId)) break;
+      visited.add(currentId);
+
+      const meta = this.loadForkMetadata(currentId);
+      if (!meta) break;
+      if (meta.sourceSessionId === currentId) break; // self-referential
+
+      depth++;
+      currentId = meta.sourceSessionId;
+    }
+
+    return depth;
+  }
+
+  /**
+   * Build the full fork tree starting from a session.
+   */
+  getForkTree(sessionId: string, title: string = "Session", visited?: Set<string>): ForkTreeNode {
+    if (!visited) visited = new Set();
+    if (visited.has(sessionId)) {
+      return {
+        sessionId,
+        forkId: sessionId,
+        title: "(cycle)",
+        forkedAt: Date.now(),
+        children: [],
+      };
+    }
+    visited.add(sessionId);
+
+    const meta = this.loadForkMetadata(sessionId);
+    const forkedAt = meta?.forkedAt ?? Date.now();
+
+    const node: ForkTreeNode = {
+      sessionId,
+      forkId: meta?.forkId ?? sessionId,
+      title,
+      forkedAt,
+      children: [],
+    };
+
+    // Recursively build children
+    const childForks = this.getChildForks(sessionId);
+    for (const childId of childForks) {
+      const childMeta = this.loadForkMetadata(childId);
+      const childTitle = childMeta?.forkPoint
+        ? `Fork at ${childMeta.forkPoint.slice(0, 8)}...`
+        : "Fork";
+      node.children.push(this.getForkTree(childId, childTitle, visited));
+    }
+
+    return node;
+  }
+
+  /**
+   * Get all fork metadata for sessions that originated from a given session.
+   */
+  getForkHistory(sessionId: string, visited?: Set<string>): ForkMetadata[] {
+    if (!visited) visited = new Set();
+    if (visited.has(sessionId)) return [];
+    visited.add(sessionId);
+
+    const results: ForkMetadata[] = [];
+    const childForks = this.getChildForks(sessionId);
+
+    for (const childId of childForks) {
+      const meta = this.loadForkMetadata(childId);
+      if (meta) {
+        results.push(meta);
+        // Recursively get grandchildren
+        results.push(...this.getForkHistory(childId, visited));
+      }
+    }
+
+    return results;
+  }
+
+  /**
+   * Delete fork metadata for a session.
+   * Also removes the session ID from parent session's child-forks list.
+   */
+  deleteForkMetadata(sessionId: string): void {
+    const meta = this.loadForkMetadata(sessionId);
+    if (meta) {
+      // Remove from parent's child-forks
+      const parentChildForksPath = path.join(
+        KOI_SESSIONS_DIR,
+        meta.sourceSessionId,
+        "child-forks.json"
+      );
+      const existingChildForks: string[] = safeReadFile<string[]>(parentChildForksPath, (raw) => JSON.parse(raw) as string[]) ?? [];
+      const updatedChildForks = existingChildForks.filter(id => id !== sessionId);
+      if (updatedChildForks.length > 0) {
+        safeWriteFile(parentChildForksPath, JSON.stringify(updatedChildForks, null, 2) + "\n");
+      } else {
+        safeDeleteFile(parentChildForksPath);
+      }
+    }
+
+    this.cache.delete(sessionId);
+    safeDeleteFile(getForkMetadataPath(sessionId));
+  }
+
+  /**
+   * Clear the in-memory cache.
+   */
+  clearCache(): void {
+    this.cache.clear();
+  }
+}
+
+// Global singleton instance
+export const forkManager = new ForkManager();

package/src/agent/session-snapshots.ts

@@ -0,0 +1,216 @@
+/**
+ * Session Snapshots
+ *
+ * Records per-entry snapshots of Koi state (tasks, plan, agent mode) using
+ * Pi's CustomEntry mechanism. Each snapshot is stored as a `custom` entry
+ * with `customType: "koi:snapshot"` in the session tree.
+ *
+ * Because `appendCustomEntry` appends the snapshot as a child of the current
+ * leaf and then advances leaf, the snapshot sits **after** the turn it
+ * represents. When forking to entry X we must therefore look **forward**
+ * from X (in X's subtree) to find the snapshot that captured the state
+ * right after X's turn completed. When loading a session we look
+ * **backward** from the current leaf toward root.
+ */
+
+import type { AgentSession, SessionEntry } from "@mariozechner/pi-coding-agent";
+import type { Task } from "./session-tasks.js";
+import type { SessionTaskManager } from "./session-tasks.js";
+import type { AgentMode } from "./mode.js";
+import { setCurrentPlanText } from "./plan-ui.js";
+import fs from "fs";
+
+export interface KoiSnapshotData {
+  tasks: Task[];
+  planText: string | null;
+  agentMode: AgentMode;
+  activeTools: string[];
+}
+
+const SNAPSHOT_CUSTOM_TYPE = "koi:snapshot";
+
+/**
+ * Save a snapshot of the current Koi state as a CustomEntry in the Pi session.
+ * The snapshot is attached as a child of the current leaf entry, then leaf advances.
+ */
+export function saveSnapshot(
+  session: AgentSession,
+  data: KoiSnapshotData
+): string {
+  // Deep clone to prevent later mutations from corrupting persisted snapshot data.
+  // Pi's appendCustomEntry does not clone the payload, so object references
+  // (e.g. Task objects in the tasks array) would be mutated by subsequent
+  // task updates, making old snapshots reflect the latest state instead of
+  // the state at the time they were saved.
+  const cloned: KoiSnapshotData = {
+    tasks: data.tasks.map((t) => ({ ...t })),
+    planText: data.planText,
+    agentMode: data.agentMode,
+    activeTools: [...data.activeTools],
+  };
+  return session.sessionManager.appendCustomEntry(SNAPSHOT_CUSTOM_TYPE, cloned);
+}
+
+/**
+ * Save a snapshot only if the state has changed since the last snapshot.
+ * Returns the new entry id, or null if nothing changed.
+ */
+export function saveSnapshotIfChanged(
+  session: AgentSession,
+  data: KoiSnapshotData
+): string | null {
+  const leafId = session.sessionManager.getLeafId();
+  if (!leafId) return null;
+
+  const last = findSnapshotBeforeEntry(session, leafId);
+  fs.appendFileSync("/tmp/koi-snapshot-debug.log", `[snapshot] saveSnapshotIfChanged leafId=${leafId} last=${last ? last.entryId : "null"} tasks=${data.tasks.length}\n`);
+  if (last?.data) {
+    const s = last.data;
+    const sameTasks =
+      s.tasks.length === data.tasks.length &&
+      s.tasks.every(
+        (t, i) =>
+          t.id === data.tasks[i]!.id &&
+          t.content === data.tasks[i]!.content &&
+          t.status === data.tasks[i]!.status &&
+          t.priority === data.tasks[i]!.priority
+      );
+    const samePlan = s.planText === data.planText;
+    const sameMode = s.agentMode === data.agentMode;
+    const sameTools =
+      s.activeTools.length === data.activeTools.length &&
+      s.activeTools.every((t, i) => t === data.activeTools[i]);
+
+    fs.appendFileSync("/tmp/koi-snapshot-debug.log", `[snapshot] saveSnapshotIfChanged compare sameTasks=${sameTasks} samePlan=${samePlan} sameMode=${sameMode} sameTools=${sameTools}\n`);
+    if (sameTasks && samePlan && sameMode && sameTools) {
+      fs.appendFileSync("/tmp/koi-snapshot-debug.log", `[snapshot] saveSnapshotIfChanged -> skip saving (no change)\n`);
+      return null;
+    }
+  }
+
+  const id = saveSnapshot(session, data);
+  fs.appendFileSync("/tmp/koi-snapshot-debug.log", `[snapshot] saveSnapshotIfChanged -> saved new snapshot id=${id}\n`);
+  return id;
+}
+
+/**
+ * Walk the branch path from root to `targetEntryId` backwards and return the
+ * nearest custom entry with type `koi:snapshot`.
+ * Used when loading a session (leaf is known).
+ */
+export function findSnapshotBeforeEntry(
+  session: AgentSession,
+  targetEntryId: string
+): { entryId: string; data: KoiSnapshotData | undefined } | null {
+  const branch = session.sessionManager.getBranch(targetEntryId);
+  const snapshotIds: string[] = [];
+  for (let i = branch.length - 1; i >= 0; i--) {
+    const entry = branch[i]!;
+    if (
+      entry.type === "custom" &&
+      (entry as unknown as { customType?: string }).customType ===
+        SNAPSHOT_CUSTOM_TYPE
+    ) {
+      snapshotIds.push(entry.id);
+      fs.appendFileSync("/tmp/koi-snapshot-debug.log", `[snapshot] findSnapshotBeforeEntry target=${targetEntryId} found=${entry.id} tasks=${(entry as unknown as { data?: KoiSnapshotData }).data?.tasks?.length ?? 0} taskStatus=${(entry as unknown as { data?: KoiSnapshotData }).data?.tasks?.map(t => t.status).join(",") ?? "none"}\n`);
+      return {
+        entryId: entry.id,
+        data: (entry as unknown as { data?: KoiSnapshotData }).data,
+      };
+    }
+  }
+  fs.appendFileSync("/tmp/koi-snapshot-debug.log", `[snapshot] findSnapshotBeforeEntry target=${targetEntryId} no snapshot found on branch of ${branch.length} entries\n`);
+  return null;
+}
+
+/**
+ * Search forward from `targetEntryId` through its entire subtree (using all
+ * entries in the session file) to find the nearest snapshot.
+ * Used when forking: the snapshot lives after the selected entry, not on the
+ * branch path *to* the entry.
+ */
+export function findSnapshotAfterEntry(
+  session: AgentSession,
+  targetEntryId: string
+): { entryId: string; data: KoiSnapshotData | undefined } | null {
+  const allEntries = session.sessionManager.getEntries();
+  const byId = new Map<string, SessionEntry>();
+  const children = new Map<string, string[]>();
+
+  for (const entry of allEntries) {
+    byId.set(entry.id, entry);
+    const parentId = entry.parentId ?? "root";
+    const list = children.get(parentId) ?? [];
+    list.push(entry.id);
+    children.set(parentId, list);
+  }
+
+  const queue = [targetEntryId];
+  const visited = new Set<string>();
+  const visitedSnapshots: string[] = [];
+
+  while (queue.length > 0) {
+    const id = queue.shift()!;
+    if (visited.has(id)) continue;
+    visited.add(id);
+
+    const entry = byId.get(id);
+    if (
+      entry &&
+      entry.type === "custom" &&
+      (entry as unknown as { customType?: string }).customType ===
+        SNAPSHOT_CUSTOM_TYPE
+    ) {
+      visitedSnapshots.push(entry.id);
+      fs.appendFileSync("/tmp/koi-snapshot-debug.log", `[snapshot] findSnapshotAfterEntry target=${targetEntryId} found=${entry.id} tasks=${(entry as unknown as { data?: KoiSnapshotData }).data?.tasks?.length ?? 0} taskStatus=${(entry as unknown as { data?: KoiSnapshotData }).data?.tasks?.map(t => t.status).join(",") ?? "none"} visitedSnapshots=[${visitedSnapshots.join(",")}]\n`);
+      return {
+        entryId: entry.id,
+        data: (entry as unknown as { data?: KoiSnapshotData }).data,
+      };
+    }
+
+    const childIds = children.get(id) ?? [];
+    queue.push(...childIds);
+  }
+
+  fs.appendFileSync("/tmp/koi-snapshot-debug.log", `[snapshot] findSnapshotAfterEntry target=${targetEntryId} no snapshot found in subtree visited=${visited.size} entries\n`);
+  return null;
+}
+
+/**
+ * Restore tasks and plan text from the nearest snapshot related to
+ * `targetEntryId`.
+ *
+ * Strategy:
+ * 1. Try to find a snapshot **after** the entry (forward in the subtree).
+ *    This is the normal case when forking: the snapshot was appended as a
+ *    child of the turn that ended at/after the selected entry.
+ * 2. If no forward snapshot exists, fall back to a snapshot **before/at**
+ *    the entry (walking backward toward root). This handles the load-session
+ *    case where the leaf itself may be a snapshot or the snapshot sits on the
+ *    branch path before the leaf.
+ *
+ * Returns the snapshot data (including agentMode/activeTools) so the caller
+ * can restore mode state, or null if no snapshot was found.
+ */
+export function restoreSnapshot(
+  session: AgentSession,
+  targetEntryId: string,
+  taskManager: SessionTaskManager
+): KoiSnapshotData | null {
+  let snapshot = findSnapshotAfterEntry(session, targetEntryId);
+  fs.appendFileSync("/tmp/koi-snapshot-debug.log", `[snapshot] restoreSnapshot entryId=${targetEntryId} forward=${snapshot ? snapshot.entryId : "null"}\n`);
+  if (!snapshot?.data) {
+    snapshot = findSnapshotBeforeEntry(session, targetEntryId);
+    fs.appendFileSync("/tmp/koi-snapshot-debug.log", `[snapshot] restoreSnapshot fallback=${snapshot ? snapshot.entryId : "null"}\n`);
+  }
+  if (!snapshot?.data) return null;
+
+  const taskStatuses = snapshot.data.tasks.map(t => `${t.id}:${t.status}`).join(", ");
+  fs.appendFileSync("/tmp/koi-snapshot-debug.log", `[snapshot] restoreSnapshot -> restoring tasks=[${taskStatuses}] plan=${snapshot.data.planText?.slice(0, 20) ?? "null"} mode=${snapshot.data.agentMode}\n`);
+
+  taskManager.setTasks(snapshot.data.tasks);
+  setCurrentPlanText(snapshot.data.planText ?? "");
+
+  return snapshot.data;
+}