@esso0428/pi-subagents 0.15.5 → 0.15.7

package/src/agent-manager.ts

@@ -23,13 +23,101 @@ export type CompactionInfo = { reason: "manual" | "threshold" | "overflow"; toke
 
 /** Default max concurrent background agents. */
 const DEFAULT_MAX_CONCURRENT = 4;
-const TERMINAL_STATUSES = new Set<AgentRecord["status"]>([
-  "completed",
-  "steered",
-  "aborted",
-  "stopped",
-  "error",
-]);
+const RESTORABLE_STATUSES = new Set(["completed", "steered", "stopped", "aborted", "error"] as const);
+const THINKING_LEVELS = new Set(["minimal", "low", "medium", "high", "xhigh", "max", "off"]);
+
+type RestorableAgentStatus = "completed" | "steered" | "stopped" | "aborted" | "error";
+
+/** Narrow persisted strings before putting them into runtime/UI state. */
+function isSafePersistedString(value: unknown, maxLength: number): value is string {
+  return typeof value === "string" && value.length > 0 && value.length <= maxLength && !/[\0\r\n]/.test(value);
+}
+
+function isFiniteTimestamp(value: unknown): value is number {
+  return typeof value === "number" && Number.isFinite(value) && Number.isInteger(value) && value >= 0;
+}
+
+function isValidUsage(value: unknown): value is { input: number; output: number; cacheWrite: number } {
+  if (!value || typeof value !== "object") return false;
+  const usage = value as Record<string, unknown>;
+  return ["input", "output", "cacheWrite"].every((key) => {
+    const n = usage[key];
+    return typeof n === "number" && Number.isFinite(n) && n >= 0;
+  });
+}
+
+function isSafeTranscriptLocator(value: unknown): value is string {
+  return typeof value === "string"
+    && /^\.pi-subagents\/agent-transcripts\/[^/]+\.jsonl$/.test(value)
+    && !value.includes("..")
+    && !value.includes("\\")
+    && !value.includes("\0");
+}
+
+function isSafeInvocation(value: unknown): value is AgentInvocation {
+  if (!value || typeof value !== "object") return false;
+  const invocation = value as Record<string, unknown>;
+  for (const key of ["modelName", "effectiveModelName"]) {
+    if (invocation[key] !== undefined && !isSafePersistedString(invocation[key], 512)) return false;
+  }
+  for (const key of ["thinking", "effectiveThinking"]) {
+    if (invocation[key] !== undefined && (typeof invocation[key] !== "string" || !THINKING_LEVELS.has(invocation[key]))) return false;
+  }
+  if (invocation.maxTurns !== undefined && (!Number.isInteger(invocation.maxTurns) || (invocation.maxTurns as number) < 0)) return false;
+  for (const key of ["isolated", "inheritContext", "runInBackground"]) {
+    if (invocation[key] !== undefined && typeof invocation[key] !== "boolean") return false;
+  }
+  if (invocation.isolation !== undefined && invocation.isolation !== "worktree") return false;
+  return true;
+}
+
+function cloneInvocation(value: AgentInvocation | undefined): AgentInvocation | undefined {
+  return value ? {
+    modelName: value.modelName,
+    effectiveModelName: value.effectiveModelName,
+    thinking: value.thinking,
+    effectiveThinking: value.effectiveThinking,
+    maxTurns: value.maxTurns,
+    isolated: value.isolated,
+    inheritContext: value.inheritContext,
+    runInBackground: value.runInBackground,
+    isolation: value.isolation,
+  } : undefined;
+}
+
+/** Validate one persisted terminal record without constructing runtime handles. */
+export function isRestorableAgentRecord(value: unknown): value is {
+  id: string;
+  type: string;
+  description: string;
+  status: RestorableAgentStatus;
+  startedAt: number;
+  completedAt: number;
+  result?: string;
+  error?: string;
+  toolUses?: number;
+  lifetimeUsage?: { input: number; output: number; cacheWrite: number };
+  transcriptPath?: string;
+  invocation?: AgentInvocation;
+} {
+  if (!value || typeof value !== "object") return false;
+  const record = value as Record<string, unknown>;
+  if (!isSafePersistedString(record.id, 256)
+    || !isSafePersistedString(record.type, 256)
+    || !isSafePersistedString(record.description, 4096)
+    || typeof record.status !== "string"
+    || !RESTORABLE_STATUSES.has(record.status as RestorableAgentStatus)
+    || !isFiniteTimestamp(record.startedAt)
+    || !isFiniteTimestamp(record.completedAt)
+    || record.completedAt < record.startedAt) return false;
+  if (record.result !== undefined && !isSafePersistedString(record.result, 2_000_000)) return false;
+  if (record.error !== undefined && !isSafePersistedString(record.error, 64_000)) return false;
+  if (record.toolUses !== undefined && (!Number.isInteger(record.toolUses) || (record.toolUses as number) < 0)) return false;
+  if (record.lifetimeUsage !== undefined && !isValidUsage(record.lifetimeUsage)) return false;
+  if (record.transcriptPath !== undefined && !isSafeTranscriptLocator(record.transcriptPath)) return false;
+  if (record.invocation !== undefined && !isSafeInvocation(record.invocation)) return false;
+  return true;
+}
 
 /**
  * Validate a caller-supplied SpawnOptions.cwd. `undefined`/`null` mean "unset"
@@ -104,45 +192,6 @@ interface SpawnOptions {
   onCompaction?: (info: CompactionInfo) => void;
 }
 
-const RESTORABLE_STATUSES = new Set<AgentRecord["status"]>([
-  "completed",
-  "steered",
-  "aborted",
-  "stopped",
-  "error",
-]);
-
-type PersistedAgentRecord = Pick<
-  AgentRecord,
-  "id" | "type" | "description" | "status" | "result" | "error" | "startedAt" | "completedAt" | "transcriptPath"
->;
-
-function isRestorableRecord(value: unknown): value is PersistedAgentRecord {
-  if (!value || typeof value !== "object") return false;
-  const record = value as Record<string, unknown>;
-  const transcriptPath = record.transcriptPath;
-  const validTranscriptPath =
-    transcriptPath === undefined ||
-    (typeof transcriptPath === "string" &&
-      transcriptPath.startsWith(".pi-subagents/agent-transcripts/") &&
-      !transcriptPath.includes("..") &&
-      !transcriptPath.includes("\\"));
-  return (
-    typeof record.id === "string" &&
-    record.id.length > 0 &&
-    typeof record.type === "string" &&
-    typeof record.description === "string" &&
-    RESTORABLE_STATUSES.has(record.status as AgentRecord["status"]) &&
-    typeof record.startedAt === "number" &&
-    Number.isFinite(record.startedAt) &&
-    typeof record.completedAt === "number" &&
-    Number.isFinite(record.completedAt) &&
-    (record.result === undefined || typeof record.result === "string") &&
-    (record.error === undefined || typeof record.error === "string") &&
-    validTranscriptPath
-  );
-}
-
 export class AgentManager {
   private agents = new Map<string, AgentRecord>();
   private cleanupInterval: ReturnType<typeof setInterval>;
@@ -323,6 +372,12 @@ export class AgentManager {
       },
       onSessionCreated: (session) => {
         record.session = session;
+        const model = session.model;
+        record.invocation = {
+          ...(record.invocation ?? {}),
+          ...(model && { effectiveModelName: model.name ?? model.id }),
+          effectiveThinking: session.thinkingLevel,
+        };
         // Flush any steers that arrived before the session was ready
         if (record.pendingSteers?.length) {
           for (const msg of record.pendingSteers) {
@@ -501,6 +556,12 @@ export class AgentManager {
     record.completedAt = undefined;
     record.result = undefined;
     record.error = undefined;
+    const resumedModel = record.session.model;
+    record.invocation = {
+      ...(record.invocation ?? {}),
+      ...(resumedModel && { effectiveModelName: resumedModel.name ?? resumedModel.id }),
+      effectiveThinking: record.session.thinkingLevel,
+    };
 
     try {
       const { text, failure } = await resumeAgent(record.session, prompt, {
@@ -556,47 +617,61 @@ export class AgentManager {
     return this.agents.get(id);
   }
 
-  /**
-   * Restore terminal records persisted in a parent session.
-   *
-   * Restored records deliberately have no live session, promise, or abort
-   * controller. Invalid data is ignored because session entries are persisted
-   * extension data and may have been written by an older version.
-   */
-  restoreCompleted(records: readonly unknown[]): void {
-    const restoredIds = new Set<string>();
-    // getBranch() is chronological; newest persisted state wins on duplicate IDs.
-    for (const value of [...records].reverse()) {
-      if (!isRestorableRecord(value)) continue;
-      if (restoredIds.has(value.id) || this.agents.has(value.id)) continue;
-      restoredIds.add(value.id);
-
-      this.agents.set(value.id, {
-        id: value.id,
-        type: value.type,
-        description: value.description,
-        status: value.status,
-        result: value.result,
-        error: value.error,
-        transcriptPath: value.transcriptPath,
-        toolUses: 0,
-        startedAt: value.startedAt,
-        completedAt: value.completedAt,
-        lifetimeUsage: { input: 0, output: 0, cacheWrite: 0 },
-        compactionCount: 0,
-        // Historical records have no inline tool surface and should remain
-        // visible in the background widget.
-        isBackground: true,
-      });
-    }
-  }
-
   listAgents(): AgentRecord[] {
     return [...this.agents.values()].sort(
       (a, b) => b.startedAt - a.startedAt,
     );
   }
 
+  /** Restore terminal records persisted by a parent branch without runtime handles. */
+  restoreCompleted(records: readonly unknown[]): void {
+    const latest = new Map<string, ReturnType<typeof this.createRestoredRecord>>();
+    for (const value of records) {
+      if (isRestorableAgentRecord(value)) {
+        latest.set(value.id, this.createRestoredRecord(value));
+      }
+    }
+
+    for (const [id, restored] of latest) {
+      const existing = this.agents.get(id);
+      if (existing?.status === "running" || existing?.status === "queued") continue;
+      this.agents.set(id, restored);
+    }
+  }
+
+  private createRestoredRecord(record: {
+    id: string;
+    type: string;
+    description: string;
+    status: RestorableAgentStatus;
+    startedAt: number;
+    completedAt: number;
+    result?: string;
+    error?: string;
+    toolUses?: number;
+    lifetimeUsage?: { input: number; output: number; cacheWrite: number };
+    transcriptPath?: string;
+    invocation?: AgentInvocation;
+  }): AgentRecord {
+    return {
+      id: record.id,
+      type: record.type,
+      description: record.description,
+      status: record.status,
+      result: record.result,
+      error: record.error,
+      toolUses: record.toolUses ?? 0,
+      startedAt: record.startedAt,
+      completedAt: record.completedAt,
+      transcriptPath: record.transcriptPath,
+      invocation: cloneInvocation(record.invocation),
+      lifetimeUsage: record.lifetimeUsage
+        ? { ...record.lifetimeUsage }
+        : { input: 0, output: 0, cacheWrite: 0 },
+      compactionCount: 0,
+    };
+  }
+
   abort(id: string): boolean {
     const record = this.agents.get(id);
     if (!record) return false;
@@ -628,6 +703,19 @@ export class AgentManager {
     for (const [id, record] of this.agents) {
       if (record.status === "running" || record.status === "queued") continue;
       if ((record.completedAt ?? 0) >= cutoff) continue;
+      // A durable transcript is the source of truth for history. Release the
+      // live session after the TTL, but retain a lightweight record so opening
+      // history again in this session does not silently lose its identity or
+      // locator. Records without durable storage remain eligible for eviction.
+      if (record.transcriptPath) {
+        try { record.session?.dispose?.(); } catch { /* ignore cleanup failures */ }
+        record.session = undefined;
+        try { record.outputCleanup?.(); } catch { /* ignore cleanup failures */ }
+        record.outputCleanup = undefined;
+        record.outputFile = undefined;
+        record.historyFile = undefined;
+        continue;
+      }
       this.removeRecord(id, record);
     }
   }

package/src/agent-runner.ts

@@ -6,7 +6,7 @@ import { readFileSync } from "node:fs";
 import { homedir } from "node:os";
 import { basename, dirname, isAbsolute, join, resolve } from "node:path";
 import type { Model } from "@earendil-works/pi-ai";
-import type { ExtensionContext, LoadExtensionsResult } from "@earendil-works/pi-coding-agent";
+import type { ExtensionContext, LoadExtensionsResult, ModelRuntime } from "@earendil-works/pi-coding-agent";
 import {
   type AgentSession,
   type AgentSessionEvent,
@@ -25,6 +25,7 @@ import { buildMemoryBlock, buildReadOnlyMemoryBlock } from "./memory.js";
 import { buildAgentPrompt, type PromptExtras } from "./prompts.js";
 import { preloadSkills } from "./skill-loader.js";
 import type { SubagentType, ThinkingLevel } from "./types.js";
+import { createTrackedWriteTool } from "./write-execution.js";
 
 /**
  * Tool names registered by THIS extension. Single source of truth so the
@@ -781,32 +782,30 @@ export async function runAgent(
   const sessionManager = agentConfig?.persistSession
     ? SessionManager.create(effectiveCwd, configuredSessionDir ?? defaultSessionDir)
     : SessionManager.inMemory(effectiveCwd);
+  // Keep write tracking session-local. The custom definition shadows only this
+  // child session's built-in write entry; the parent registry is untouched.
+  const trackedWriteTool = toolNames.includes("write")
+    ? createTrackedWriteTool(effectiveCwd)
+    : undefined;
 
   // Pi 0.80.8 replaced createAgentSession's modelRegistry option with
   // modelRuntime, but ExtensionContext still exposes only the registry facade.
   // Pass both so the full supported Pi range retains the parent's providers.
-  const parentModelRuntime = (ctx.modelRegistry as unknown as { runtime?: unknown }).runtime;
-  type CreateAgentSessionOptions = NonNullable<Parameters<typeof createAgentSession>[0]>;
-  // `modelRuntime` was added in Pi 0.80.8. Infer it when available, while
-  // falling back to unknown for older Pi declarations.
-  type CompatibleModelRuntime = CreateAgentSessionOptions extends {
-    modelRuntime?: infer Runtime;
-  } ? Runtime : unknown;
-  const sessionOpts: CreateAgentSessionOptions & {
+  const parentModelRuntime = (ctx.modelRegistry as unknown as { runtime?: ModelRuntime | null }).runtime ?? undefined;
+  const sessionOpts: Parameters<typeof createAgentSession>[0] & {
     modelRegistry: ExtensionContext["modelRegistry"];
-    modelRuntime?: CompatibleModelRuntime;
+    modelRuntime?: ModelRuntime;
   } = {
     cwd: effectiveCwd,
     agentDir,
     sessionManager,
     settingsManager,
     modelRegistry: ctx.modelRegistry,
-    ...(parentModelRuntime != null && {
-      modelRuntime: parentModelRuntime as CompatibleModelRuntime,
-    }),
+    ...(parentModelRuntime !== undefined && { modelRuntime: parentModelRuntime }),
     model,
     tools: sessionTools,
     resourceLoader: loader,
+    ...(trackedWriteTool && { customTools: [trackedWriteTool as any] }),
   };
   if (sessionExcludeTools) {
     sessionOpts.excludeTools = sessionExcludeTools;