@herbertgao/pi-extensions 2026.8.5 → 2026.8.7

package/node_modules/@herbertgao/pi-subagents/src/agent-manager.ts

@@ -24,19 +24,25 @@ import type {
   SubagentType,
   ThinkingLevel,
 } from "./types.js"
-import { addUsage } from "./usage.js"
-import { cleanupWorktree, createWorktree, pruneWorktrees } from "./worktree.js"
+import { addUsage, type LifetimeUsage } from "./usage.js"
+import {
+  cleanupWorktree,
+  createWorktree,
+  isWorktreeIsolationEnabled,
+  pruneWorktrees,
+} from "./worktree.js"
 
 export type OnAgentComplete = (record: AgentRecord) => void
 export type OnAgentStart = (record: AgentRecord) => void
 export type OnAgentCompact = (record: AgentRecord, info: CompactionInfo) => void
+export type OnAgentUsage = (record: AgentRecord, usage: LifetimeUsage) => void
 export type CompactionInfo = {
   reason: "manual" | "threshold" | "overflow"
   tokensBefore: number
 }
 
 /** Default max concurrent background agents. */
-const DEFAULT_MAX_CONCURRENT = 4
+const DEFAULT_MAX_CONCURRENT = 10
 
 /**
  * Validate a caller-supplied SpawnOptions.cwd. `undefined`/`null` mean "unset"
@@ -125,11 +131,7 @@ interface SpawnOptions {
   /** Called at the end of each agentic turn with the cumulative count. */
   onTurnEnd?: (turnCount: number) => void
   /** Called once per assistant message_end with that message's usage delta. */
-  onAssistantUsage?: (usage: {
-    input: number
-    output: number
-    cacheWrite: number
-  }) => void
+  onAssistantUsage?: (usage: LifetimeUsage) => void
   /** Called when the session successfully compacts. */
   onCompaction?: (info: CompactionInfo) => void
   /** Nesting depth: top-level subagent = 1. */
@@ -158,11 +160,7 @@ interface ResumeOptions {
   /** Called at the end of each resumed agentic turn with the cumulative count. */
   onTurnEnd?: (turnCount: number) => void
   /** Called once per assistant message_end with that message's usage delta. */
-  onAssistantUsage?: (usage: {
-    input: number
-    output: number
-    cacheWrite: number
-  }) => void
+  onAssistantUsage?: (usage: LifetimeUsage) => void
   /** Called when the session successfully compacts. */
   onCompaction?: (info: CompactionInfo) => void
   /**
@@ -176,12 +174,38 @@ interface ResumeOptions {
   onStarted?: () => void
 }
 
+const CHILD_SHUTDOWN_TIMEOUT_MS = 3_000
+
+async function shutdownChildSession(
+  session: AgentSession | undefined,
+): Promise<void> {
+  try {
+    const runner = session?.extensionRunner
+    if (runner?.hasHandlers?.("session_shutdown")) {
+      await Promise.race([
+        runner.emit({ type: "session_shutdown", reason: "quit" }),
+        new Promise<void>((resolve) =>
+          setTimeout(resolve, CHILD_SHUTDOWN_TIMEOUT_MS).unref(),
+        ),
+      ])
+    }
+  } catch {
+    /* best-effort lifecycle cleanup */
+  }
+  try {
+    session?.dispose?.()
+  } catch {
+    /* ignore */
+  }
+}
+
 export class AgentManager {
   private agents = new Map<string, AgentRecord>()
   private cleanupInterval: ReturnType<typeof setInterval>
   private onComplete?: OnAgentComplete
   private onStart?: OnAgentStart
   private onCompact?: OnAgentCompact
+  private onUsage?: OnAgentUsage
   private maxConcurrent: number
   /** Base repos worktrees were created from — so dispose() can prune them all,
    *  not just the parent repo (caller-supplied cwd can target other repos). */
@@ -197,10 +221,12 @@ export class AgentManager {
     maxConcurrent = DEFAULT_MAX_CONCURRENT,
     onStart?: OnAgentStart,
     onCompact?: OnAgentCompact,
+    onUsage?: OnAgentUsage,
   ) {
     this.onComplete = onComplete
     this.onStart = onStart
     this.onCompact = onCompact
+    this.onUsage = onUsage
     this.maxConcurrent = maxConcurrent
     // Cleanup completed agents after 10 minutes (but keep sessions for resume)
     this.cleanupInterval = setInterval(() => this.cleanup(), 60_000)
@@ -302,7 +328,7 @@ export class AgentManager {
     // fail loud if not possible (no silent fallback to main tree). Done
     // BEFORE state mutation so a throw doesn't leave the record half-running.
     let worktreeCwd: string | undefined
-    if (options.isolation === "worktree") {
+    if (options.isolation === "worktree" && isWorktreeIsolationEnabled()) {
       const wt = createWorktree(baseCwd, id)
       if (!wt) {
         throw new Error(
@@ -365,6 +391,7 @@ export class AgentManager {
       onTextDelta: options.onTextDelta,
       onAssistantUsage: (usage) => {
         addUsage(record.lifetimeUsage, usage)
+        this.onUsage?.(record, usage)
         options.onAssistantUsage?.(usage)
       },
       onCompaction: (info) => {
@@ -662,6 +689,7 @@ export class AgentManager {
         onTurnEnd: options?.onTurnEnd,
         onAssistantUsage: (usage) => {
           addUsage(record.lifetimeUsage, usage)
+          this.onUsage?.(record, usage)
           options?.onAssistantUsage?.(usage)
         },
         onCompaction: (info) => {
@@ -766,6 +794,7 @@ export class AgentManager {
       onTurnEnd: options.onTurnEnd,
       onAssistantUsage: (usage) => {
         addUsage(record.lifetimeUsage, usage)
+        this.onUsage?.(record, usage)
         options.onAssistantUsage?.(usage)
       },
       onCompaction: (info) => {
@@ -851,9 +880,10 @@ export class AgentManager {
 
   /** Dispose a record's session and remove it from the map. */
   private removeRecord(id: string, record: AgentRecord): void {
-    record.session?.dispose?.()
+    const session = record.session
     record.session = undefined
     this.agents.delete(id)
+    void shutdownChildSession(session)
   }
 
   private cleanup() {
@@ -926,14 +956,13 @@ export class AgentManager {
     }
   }
 
-  dispose() {
+  async dispose(): Promise<void> {
     clearInterval(this.cleanupInterval)
     // Clear queue
     this.queue = []
-    for (const record of this.agents.values()) {
-      record.session?.dispose()
-    }
+    const sessions = [...this.agents.values()].map((record) => record.session)
     this.agents.clear()
+    await Promise.all(sessions.map((session) => shutdownChildSession(session)))
     // Prune any orphaned git worktrees (crash recovery)
     try {
       pruneWorktrees(process.cwd())

package/node_modules/@herbertgao/pi-subagents/src/agent-runner.ts

@@ -41,6 +41,7 @@ import {
 import { buildAgentPrompt, type PromptExtras } from "./prompts.js"
 import { preloadSkills } from "./skill-loader.js"
 import type { SubagentType, ThinkingLevel } from "./types.js"
+import type { LifetimeUsage } from "./usage.js"
 
 /**
  * Tool names registered by THIS extension. Single source of truth so the
@@ -349,6 +350,15 @@ export function setDefaultMaxTurns(n: number | undefined): void {
   defaultMaxTurns = normalizeMaxTurns(n)
 }
 
+export function resolveEffectiveMaxTurns(
+  type: string,
+  explicit?: number,
+): number | undefined {
+  return normalizeMaxTurns(
+    explicit ?? getAgentConfig(type)?.maxTurns ?? defaultMaxTurns,
+  )
+}
+
 /** Additional turns allowed after the soft limit steer message. */
 let graceTurns = 5
 
@@ -442,11 +452,7 @@ export interface RunOptions {
    * Lets callers maintain a lifetime accumulator that survives compaction
    * (which replaces session.state.messages and resets stats-derived sums).
    */
-  onAssistantUsage?: (usage: {
-    input: number
-    output: number
-    cacheWrite: number
-  }) => void
+  onAssistantUsage?: (usage: LifetimeUsage) => void
   /**
    * Called when the session successfully compacts. `tokensBefore` is upstream's
    * pre-compaction context size estimate. Aborted compactions don't fire.
@@ -958,18 +964,21 @@ export async function runAgent(
   type SessionOptions = NonNullable<Parameters<typeof createAgentSession>[0]>
   const parentModelRuntime = (
     ctx.modelRegistry as unknown as {
-      runtime?: SessionOptions["modelRuntime"] | null
+      runtime?: unknown
     }
   ).runtime
   const sessionOpts: SessionOptions & {
     modelRegistry: ExtensionContext["modelRegistry"]
+    modelRuntime?: unknown
   } = {
     cwd: effectiveCwd,
     agentDir,
     sessionManager,
     settingsManager,
     modelRegistry: ctx.modelRegistry,
-    ...(parentModelRuntime != null && { modelRuntime: parentModelRuntime }),
+    ...(parentModelRuntime != null && {
+      modelRuntime: parentModelRuntime as never,
+    }),
     model,
     tools: sessionTools,
     customTools: nestedTools,
@@ -1027,9 +1036,7 @@ export async function runAgent(
 
   // Track turns for graceful max_turns enforcement
   let turnCount = 0
-  const maxTurns = normalizeMaxTurns(
-    options.maxTurns ?? agentConfig?.maxTurns ?? defaultMaxTurns,
-  )
+  const maxTurns = resolveEffectiveMaxTurns(type, options.maxTurns)
   let softLimitReached = false
   let aborted = false
 
@@ -1076,6 +1083,8 @@ export async function runAgent(
           input: u.input ?? 0,
           output: u.output ?? 0,
           cacheWrite: u.cacheWrite ?? 0,
+          cacheRead: u.cacheRead ?? 0,
+          cost: u.cost?.total ?? 0,
         })
     }
     if (event.type === "compaction_end" && !event.aborted && event.result) {
@@ -1130,11 +1139,7 @@ export async function resumeAgent(
     onToolActivity?: (activity: ToolActivity) => void
     /** Called at the end of each resumed agentic turn with the 1-based count. */
     onTurnEnd?: (turnCount: number) => void
-    onAssistantUsage?: (usage: {
-      input: number
-      output: number
-      cacheWrite: number
-    }) => void
+    onAssistantUsage?: (usage: LifetimeUsage) => void
     onCompaction?: (info: {
       reason: "manual" | "threshold" | "overflow"
       tokensBefore: number
@@ -1177,6 +1182,8 @@ export async function resumeAgent(
                 input: u.input ?? 0,
                 output: u.output ?? 0,
                 cacheWrite: u.cacheWrite ?? 0,
+                cacheRead: u.cacheRead ?? 0,
+                cost: u.cost?.total ?? 0,
               })
           }
           if (

package/node_modules/@herbertgao/pi-subagents/src/custom-agents.ts

@@ -6,7 +6,12 @@ import { existsSync, readdirSync, readFileSync } from "node:fs"
 import { basename, join } from "node:path"
 import { getAgentDir, parseFrontmatter } from "@earendil-works/pi-coding-agent"
 import { BUILTIN_TOOL_NAMES } from "./agent-types.js"
-import type { AgentConfig, MemoryScope, ThinkingLevel } from "./types.js"
+import type {
+  AgentConfig,
+  IsolationMode,
+  MemoryScope,
+  ThinkingLevel,
+} from "./types.js"
 
 interface WarningState {
   previous: Set<string>
@@ -124,7 +129,7 @@ function loadFromDir(
           : undefined,
       isolated: fm.isolated != null ? fm.isolated === true : undefined,
       memory: parseMemory(fm.memory),
-      isolation: fm.isolation === "worktree" ? "worktree" : undefined,
+      isolation: parseIsolation(fm.isolation),
       enabled: fm.enabled !== false, // default true; explicitly false disables
       source,
       sourcePath: path,
@@ -270,6 +275,25 @@ function parseMemory(val: unknown): MemoryScope | undefined {
   return undefined
 }
 
+/**
+ * Parse the `isolation` frontmatter field.
+ *
+ * `off` is kept as a value rather than folded into `undefined` because the two
+ * do not mean the same thing here: agent config outranks tool-call params, so
+ * `off` vetoes a caller's `worktree` while an absent field lets it through.
+ *
+ * pi's frontmatter parser is not YAML 1.1 — bare `off` and `no` arrive as
+ * strings and only `false` becomes a boolean — so all three spellings are
+ * accepted rather than leaving an author's intent silently dropped. Anything
+ * else stays `undefined`, as before.
+ */
+function parseIsolation(val: unknown): IsolationMode | undefined {
+  if (val === "worktree") return "worktree"
+  if (val === "off" || val === "none" || val === "no" || val === false)
+    return "off"
+  return undefined
+}
+
 /**
  * Parse an inherit field (extensions, skills).
  * omitted/true → true (inherit all); false/"none"/empty → false; csv → listed names.