@bermudi/pi-delegate 0.1.0 → 0.1.1

package/pool.ts

@@ -97,10 +97,25 @@ let poolAbortTimeoutMs = DEFAULT_POOL_ABORT_TIMEOUT_MS;
  *  delegate calls with the same sessionId queue instead of interleaving. */
 const sessionLocks = new Map<string, Promise<void>>();
 
+type PoolShutdownState = "open" | "closing" | "closed";
+let poolState: PoolShutdownState = "open";
+let closePromise: Promise<void> | null = null;
+
 function now(): number {
   return Date.now();
 }
 
+function assertPoolOpenForNormalWork(): void {
+  if (poolState !== "open") {
+    throw new Error("Session pool is not accepting new work.");
+  }
+}
+
+function waitForActiveSessionLocks(): Promise<void> {
+  const locks = [...sessionLocks.values()];
+  return Promise.all(locks).then(() => undefined);
+}
+
 // ── Read + validate ───────────────────────────────────────────────────────
 
 /** Look up a pooled session and validate a reuse request against its frozen
@@ -115,8 +130,8 @@ function now(): number {
  *  - mismatch → pooled but its immutable configuration conflicts with this
  *               request. Caller formats the structured diff into an error.
  *
- *  lastUsed is bumped by commit() on a successful run, not here — so a checkout
- *  that bails (e.g. a resumeFrom conflict at the caller) does not affect stats. */
+ *  lastUsed is bumped by recordUse() after a completed pool hit, not here — so
+ *  a checkout that bails (e.g. a resumeFrom conflict) does not affect stats. */
 export function checkout(
   sessionId: string,
   candidate: ConfigCandidate,
@@ -187,28 +202,26 @@ export function checkout(
 
 // ── The sole mutator (besides close) ──────────────────────────────────────
 
-/** Record the outcome of a run against a sessionId. Decides insert-vs-recordUse
- *  internally by map presence:
- *   - present (pool hit)  → bump lastUsed, totalTokens, promptCount.
- *   - absent  (fresh/resume success) → insert with the frozen config.
+/**
+ * Insert the first successful prompt for a fresh/resumed session into the pool.
  *
- *  MUST be called inside withSessionLock(sessionId, …) — the map-presence
- *  decision is sound only because the lock serializes same-sessionId tasks, so
- *  no concurrent commit can race the insert. MUST only be called on run success
- *  (insert-only-on-success is caller-gated). Returns true when the session is
- *  now pool-owned, false when a fresh session could not be inserted because it
- *  lacks the manager/file required by a pooled entry. */
+ *  MUST be called inside withSessionLock(sessionId, …) and only when this run
+ *  should transfer ownership from lifecycle to the pool. Returns true when the
+ *  session is now pool-owned, false when insertion is blocked (shutdown) or
+ *  impossible (missing manager/file).
+ */
 export function commit(sessionId: string, payload: CommitPayload): boolean {
-  const existing = agentPool.get(sessionId);
-  if (existing) {
-    // Pool hit: session already pooled, just bump stats.
-    existing.lastUsed = now();
-    existing.totalTokens += payload.tokens;
-    existing.promptCount++;
-    return true;
+  // Shutdown-aware policy: once shutdown has started, inserting a new pool entry
+  // is unsafe. The lifecycle handles the still-owned session (abort/cleanup) and
+  // should dispose it instead of inserting it after the barrier begins.
+  if (poolState !== "open") {
+    return false;
   }
+
   // Miss → success: insert. A pooled entry needs a concrete file/manager.
   if (!payload.sessionManager || !payload.sessionFile) return false;
+
+  if (agentPool.has(sessionId)) return false;
   agentPool.set(sessionId, {
     session: payload.session,
     sessionManager: payload.sessionManager,
@@ -222,6 +235,20 @@ export function commit(sessionId: string, payload: CommitPayload): boolean {
   return true;
 }
 
+/** Record a completed run against an existing pooled session. This is separate
+ *  from commit() so lifecycle can distinguish ownership transfer from hit
+ *  accounting. Returns true only when the sessionId is currently pooled.
+ */
+export function recordUse(sessionId: string, tokens: number): boolean {
+  const existing = agentPool.get(sessionId);
+  if (!existing) return false;
+
+  existing.lastUsed = now();
+  existing.totalTokens += tokens;
+  existing.promptCount++;
+  return true;
+}
+
 // ── Read-only defaults (for task-resolution) ──────────────────────────────
 
 /** Frozen config for a pooled session, or undefined if not pooled. Lock-free —
@@ -242,10 +269,26 @@ export function configFor(
  *  (lifecycle) can bracket the ENTIRE acquire/run/commit flow; checkout/commit
  *  do NOT lock internally because their only caller is already inside this
  *  bracket. Close is also invoked under this lock, so it never disposes an
- *  in-flight prompt. */
+ *  in-flight prompt.
+ *
+ *  External callers must go through this exported path only while the pool is
+ *  open. Internal callers that are already synchronized by a higher-level lock
+ *  can use `withSessionLockInternal`.
+ */
 export async function withSessionLock<T>(
   sessionId: string,
   fn: () => Promise<T>,
+): Promise<T> {
+  assertPoolOpenForNormalWork();
+  return withSessionLockInternal(sessionId, fn);
+}
+
+/** Internal lock variant that does not reject during shutdown. Use only from
+ *  lifecycle/control paths that already account for pool shutdown state.
+ */
+async function withSessionLockInternal<T>(
+  sessionId: string,
+  fn: () => Promise<T>,
 ): Promise<T> {
   const prev = sessionLocks.get(sessionId);
   let resolve!: () => void;
@@ -303,8 +346,8 @@ function beginAbort(session: AgentSession): Promise<AbortOutcome> {
 }
 
 /** Close and dispose one pooled session. A caller holds the per-session lock
- * while closing, so abort cannot race a reuse. All cleanup is attempted before
- * an error is surfaced; a removed session is never silently retained. */
+ *  while closing, so abort cannot race a reuse. All cleanup is attempted before
+ *  an error is surfaced; a removed session is never silently retained. */
 async function closePooledAgentAfterAbort(
   sessionId: string,
   abort: Promise<AbortOutcome>,
@@ -346,44 +389,117 @@ async function closePooledAgentAfterAbort(
   return true;
 }
 
-/** Abort, dispose, and remove one pooled session. Returns false when the id
- * is already absent; cleanup failures are aggregated after removal. */
-export async function closePooledAgent(sessionId: string): Promise<boolean> {
+/** Internal close helper for callers that already own (or are obtaining) the
+ *  session lock. */
+async function closePooledAgentAfterLock(sessionId: string): Promise<boolean> {
   const pooled = agentPool.get(sessionId);
   if (!pooled) return false;
   return closePooledAgentAfterAbort(sessionId, beginAbort(pooled.session));
 }
 
+/** @internal Close while the caller already owns the per-session lock. */
+export async function _closePooledAgentWithoutLock(
+  sessionId: string,
+): Promise<boolean> {
+  return closePooledAgentAfterLock(sessionId);
+}
+
+/** Abort, dispose, and remove one pooled session. Returns false when the id
+ * is already absent; cleanup failures are aggregated after removal.
+ * During shutdown or after close, this waits for shutdown to finish and
+ * returns `false` instead of throwing.
+ */
+export async function closePooledAgent(sessionId: string): Promise<boolean> {
+  if (poolState !== "open") {
+    const existing = closePromise;
+    if (existing) {
+      try {
+        await existing;
+      } catch {
+        // Keep public close idempotent during and after shutdown.
+      }
+    }
+    return false;
+  }
+
+  let awaitShutdown: Promise<void> | undefined;
+  let closed = false;
+
+  await withSessionLockInternal(sessionId, async () => {
+    // A shutdown can begin while this call waits for its lock. Avoid waiting
+    // for closePromise inside this lock: do that work after releasing it, so
+    // closeAll can acquire the lock and dispose the same session.
+    if (poolState !== "open") {
+      awaitShutdown = closePromise ?? undefined;
+      return;
+    }
+    closed = await closePooledAgentAfterLock(sessionId);
+  });
+
+  if (awaitShutdown) {
+    try {
+      await awaitShutdown;
+    } catch {
+      // Keep public close idempotent during and after shutdown.
+    }
+  }
+
+  return closed;
+}
+
 /** Dispose every live pooled session when the parent Pi session ends. First
- * request cancellation immediately, then acquire each session's lock before
- * disposal. The lock prevents an in-flight lifecycle from committing a live
- * session after shutdown has removed it. Attempts all cleanup before reporting
- * any failures. */
+ * request cancellation immediately, then wait for all in-flight session locks,
+ * then acquire each remaining session's lock before disposal. Attempts are
+ * executed for all sessions before reporting failures. Idempotent callers share
+ * the same completion promise, including failures.
+ */
 export async function closeAllPooledAgents(): Promise<void> {
-  const aborts = new Map<string, Promise<AbortOutcome>>(
-    [...agentPool].map(([sessionId, pooled]) => [
-      sessionId,
-      beginAbort(pooled.session),
-    ]),
-  );
-  const results = await Promise.allSettled(
-    [...aborts].map(([sessionId, abort]) =>
-      withSessionLock(sessionId, () =>
-        closePooledAgentAfterAbort(sessionId, abort),
+  if (closePromise) {
+    return closePromise;
+  }
+  if (poolState === "closed") {
+    const completed = Promise.resolve();
+    closePromise = completed;
+    return completed;
+  }
+
+  poolState = "closing";
+
+  const completion = (async () => {
+    const aborts = new Map<string, Promise<AbortOutcome>>(
+      [...agentPool].map(([sessionId, pooled]) => [
+        sessionId,
+        beginAbort(pooled.session),
+      ]),
+    );
+
+    await waitForActiveSessionLocks();
+
+    const results = await Promise.allSettled(
+      [...aborts].map(([sessionId, abort]) =>
+        withSessionLockInternal(sessionId, () =>
+          closePooledAgentAfterAbort(sessionId, abort),
+        ),
       ),
-    ),
-  );
-  const failures = results
-    .filter(
-      (result): result is PromiseRejectedResult => result.status === "rejected",
-    )
-    .map((result) => result.reason);
-  if (failures.length) {
-    throw new AggregateError(
-      failures,
-      "Failed to close one or more pooled sessions.",
     );
-  }
+    const failures = results
+      .filter(
+        (result): result is PromiseRejectedResult =>
+          result.status === "rejected",
+      )
+      .map((result) => result.reason);
+    if (failures.length) {
+      throw new AggregateError(
+        failures,
+        "Failed to close one or more pooled sessions.",
+      );
+    }
+  })();
+
+  closePromise = completion.finally(() => {
+    poolState = "closed";
+  });
+  return closePromise;
 }
 
 /** List live pooled agents. Sessions remain available until explicit close or
@@ -416,5 +532,7 @@ export function _setPoolAbortTimeoutForTesting(
 export function _resetPoolForTesting(): void {
   agentPool.clear();
   sessionLocks.clear();
+  poolState = "open";
+  closePromise = null;
   poolAbortTimeoutMs = DEFAULT_POOL_ABORT_TIMEOUT_MS;
 }

package/render-branches.ts

@@ -1,4 +1,4 @@
-import { Markdown } from "@earendil-works/pi-tui";
+import { Markdown, type MarkdownTheme } from "@earendil-works/pi-tui";
 import { getMarkdownTheme, type Theme } from "@earendil-works/pi-coding-agent";
 import {
   fmtDuration,
@@ -17,6 +17,25 @@ import { stripAnsi, resolveCarriageReturn } from "./utils.ts";
 import { getMaxConcurrent } from "./config.ts";
 import type { TaskProgress, TaskResult } from "./types.ts";
 
+/**
+ * Render markdown output when a compatible host theme hook exists. If the host
+ * does not expose `getMarkdownTheme`, return plain text lines as a safe fallback
+ * so a single missing host hook cannot crash the renderer.
+ */
+function renderOutputLines(raw: string, width: number): string[] {
+  const trimmed = raw.trim();
+  if (!trimmed) return [];
+  try {
+    if (typeof getMarkdownTheme !== "function") return trimmed.split("\n");
+    const theme = getMarkdownTheme();
+    if (typeof theme !== "object" || !theme) return trimmed.split("\n");
+    const md = new Markdown(trimmed, 0, 0, theme as MarkdownTheme);
+    return md.render(width);
+  } catch (_error) {
+    return trimmed.split("\n");
+  }
+}
+
 /** Renderer state — the live subset of Pi's `ToolRenderContext.state`. */
 export interface RenderState {
   startedAt?: number;
@@ -365,8 +384,7 @@ export function renderFinalBranch(ctx: BranchCtx, h: RenderHelpers): void {
       let mdLines: string[] | undefined = state[cacheKey] as
         string[] | undefined;
       if (!mdLines || state[`${cacheKey}_src`] !== r.output) {
-        const md = new Markdown(r.output.trim(), 0, 0, getMarkdownTheme());
-        mdLines = md.render(Math.max(20, w - ind.length));
+        mdLines = renderOutputLines(r.output, Math.max(20, w - ind.length));
         state[`${cacheKey}_src`] = r.output;
         state[cacheKey] = mdLines;
       }

package/runner.ts

@@ -296,7 +296,8 @@ export async function runAgentSession(
   const initialMessages = session.messages;
   const initialMessageCount = initialMessages.length;
   const initialMessageSnapshot = initialMessages.slice();
-  let transcriptMayHaveBeenReplaced = false;
+  let compactionInProgress = false;
+  let completedCompaction = false;
   const assistantMessagesForAttempt: AgentMessage[] = [];
   let partialAssistantMessage: AgentMessage | undefined;
   type AttemptCapture = {
@@ -426,11 +427,13 @@ export async function runAgentSession(
     // Keep an append-only fallback for providers/fakes that reject before
     // emitting message_end. It is deliberately used only when event capture
     // is empty; event capture is authoritative across compaction and retries.
-    // Requiring the original array and unchanged historical prefix prevents a
-    // replacement/compaction transcript from leaking old assistant output.
+    // An aborted compaction may leave the original append-only transcript
+    // untouched. Completed or indeterminate compaction remains fail-closed even
+    // if a host happened to preserve the array identity.
     const currentMessages = session.messages;
     if (
-      transcriptMayHaveBeenReplaced ||
+      compactionInProgress ||
+      completedCompaction ||
       currentMessages !== initialMessages ||
       currentMessages.length < initialMessageCount
     ) {
@@ -525,17 +528,23 @@ export async function runAgentSession(
         // delay before ordinary inactivity detection resumes.
         noteActivity("waiting to retry", event.delayMs);
         break;
-      case "auto_retry_end":
+      case "auto_retry_end": {
+        const autoRetry = event as {
+          success?: unknown;
+        };
+        if (autoRetry.success === false && pendingCompactionAttempt) {
+          pendingCompactionAttempt.omitFinalAssistant = false;
+        }
         noteActivity("waiting for model output");
         break;
+      }
       case "compaction_start":
-        // Even if a host mutates the transcript in place rather than replacing
-        // its array, historical messages are no longer a safe fallback source.
-        transcriptMayHaveBeenReplaced = true;
+        compactionInProgress = true;
         noteActivity("compacting context");
         break;
       case "compaction_end":
-        transcriptMayHaveBeenReplaced = true;
+        compactionInProgress = false;
+        if (!event.aborted) completedCompaction = true;
         // Context-overflow agent_end is intentionally emitted with
         // willRetry=false because Pi's retry decision belongs to compaction.
         // Only an overflow compaction that actually retries may retract that

package/schema.ts

@@ -26,7 +26,8 @@ export const delegateTaskSchema = Type.Object({
   ),
   agent: Type.Optional(
     Type.String({
-      description: "Named agent profile; omit for an ad-hoc subagent.",
+      description:
+        "Use `default`: parent model/thinking/native tools/base prompt. Omit=ad-hoc; unknown fails call.",
     }),
   ),
   cwd: Type.Optional(
@@ -54,19 +55,19 @@ export const delegateTaskSchema = Type.Object({
   tools: Type.Optional(
     Type.Array(Type.String(), {
       description:
-        "Omit to inherit; `*`=read/write/edit/bash (mutating); `ro`=read/grep/find/ls (read-only).",
+        "`default`=parent natives; ad-hoc=*; *=read/write/edit/bash (mutating); ro=read/grep/find/ls (read-only).",
     }),
   ),
   thinking: Type.Optional(
     StringEnum(VALID_THINKING_LEVELS, {
       description:
-        "Thinking: off/minimal/low/medium/high/xhigh/max; defaults to agent/off.",
+        "Thinking: off/minimal/low/medium/high/xhigh/max; default=parent; others=agent/off.",
     }),
   ),
   sessionId: Type.Optional(
     Type.String({
       description:
-        "Optional live pool key for multi-turn reuse; omit for one-shot tasks.",
+        "Live pool key for multi-turn reuse; omit for one-shot tasks.",
     }),
   ),
   action: Type.Optional(
@@ -90,13 +91,13 @@ export const delegateArgumentsSchema = Type.Object({
   action: Type.Optional(
     StringEnum(["poll", "cancel", "wait"], {
       description:
-        "Ticket control: poll, cancel, or wait. Prefer wait; do not cancel for time.",
+        "Ticket control: poll=snapshot; wait=block until settled; cancel=abort. Prefer wait; never cancel for time.",
     }),
   ),
   async: Type.Optional(
     Type.Boolean({
       description:
-        "Detach work and return a ticket; results auto-deliver. Wait only when blocked.",
+        "Detach work, return a ticket; applies to ALL tasks. Results auto-deliver. Wait only if blocked.",
       default: false,
     }),
   ),
@@ -116,7 +117,7 @@ export const delegateArgumentsSchema = Type.Object({
     Type.Number({
       minimum: 0,
       description:
-        "How long wait blocks (ms); timeout does not cancel the ticket.",
+        "Bounds wait (ms); omit to block until settled. Timeout never cancels the ticket.",
     }),
   ),
   tasks: Type.Optional(
@@ -128,6 +129,28 @@ export const delegateArgumentsSchema = Type.Object({
   ),
 });
 
+/** Fields that belong to a task entry. Models sometimes place these at the
+ * top level of the arguments; the shim folds them back into a single task. */
+const TASK_FIELD_NAMES = [
+  "prompt",
+  "agent",
+  "cwd",
+  "systemPrompt",
+  "context",
+  "model",
+  "tools",
+  "thinking",
+  "sessionId",
+  "resumeFrom",
+] as const;
+
+/** Every field a task entry may carry. Anything else is a model mistake —
+ * e.g. `async` placed inside a task — and must fail loudly with a
+ * corrective message instead of being silently ignored (observed in the
+ * wild: a task-level `async: true` the caller believed had backgrounded
+ * the work while the call in fact ran synchronously). */
+const VALID_TASK_KEYS = new Set<string>([...TASK_FIELD_NAMES, "action"]);
+
 /** Validate the three operation modes after compatibility reshaping. */
 export function validateDelegateOperation(
   params: DelegateArguments,
@@ -165,6 +188,20 @@ export function validateDelegateOperation(
   }
 
   for (const [index, task] of tasks.entries()) {
+    const unknownKeys = Object.keys(task).filter(
+      (key) => !VALID_TASK_KEYS.has(key),
+    );
+    if (unknownKeys.length) {
+      const asyncHint = unknownKeys.includes("async")
+        ? " 'async' is a top-level flag; move it out of the task entry."
+        : "";
+      return (
+        `task ${index + 1}: unknown field(s) ${unknownKeys
+          .map((key) => `'${key}'`)
+          .join(", ")}.${asyncHint} ` +
+        `Valid task fields: ${[...VALID_TASK_KEYS].join(", ")}.`
+      );
+    }
     if (task.action === "close") {
       if (!task.sessionId) {
         return `task ${index + 1}: action 'close' requires sessionId.`;
@@ -187,21 +224,6 @@ export function validateDelegateOperation(
   return undefined;
 }
 
-/** Fields that belong to a task entry. Models sometimes place these at the
- * top level of the arguments; the shim folds them back into a single task. */
-const TASK_FIELD_NAMES = [
-  "prompt",
-  "agent",
-  "cwd",
-  "systemPrompt",
-  "context",
-  "model",
-  "tools",
-  "thinking",
-  "sessionId",
-  "resumeFrom",
-] as const;
-
 /** Actions that are only valid at task level. The top-level `action` is
  * ticket-scoped (poll/cancel/wait), so a flat close/list/prompt belongs to
  * the wrapped task. */
@@ -233,7 +255,8 @@ function normalizeToolsField(value: string): unknown {
  * - `tasks` as a JSON string instead of an array;
  * - task fields (`prompt`, `systemPrompt`, `tools`, ...) placed at the top
  *   level instead of inside a `tasks` entry — wrapped into a single task;
- * - `tools` as a JSON string (or bare token) inside a task entry.
+ * - `tools` as a JSON string (or bare token) inside a task entry;
+ * - `agent: ""` inside a task entry — treated as omitted (ad-hoc).
  * Skipped when a ticket action is in play. All other invalid input is left
  * for normal schema validation to reject loudly.
  *
@@ -275,13 +298,22 @@ export function normalizeDelegateArguments(args: unknown): DelegateArguments {
     if (Object.keys(task).length > 0) record.tasks = [task];
   }
 
-  // Stringified (or bare-token) `tools` inside task entries → real arrays.
+  // Per-entry recovery: stringified (or bare-token) `tools` → real arrays,
+  // and `agent: ""` → omitted (models emit an empty string to mean "ad-hoc";
+  // without this it only works by truthiness accident downstream).
   if (Array.isArray(record.tasks)) {
     record.tasks = record.tasks.map((entry: unknown) => {
       if (!entry || typeof entry !== "object") return entry;
       const e = entry as Record<string, unknown>;
-      if (typeof e.tools !== "string") return entry;
-      return { ...e, tools: normalizeToolsField(e.tools) };
+      const rawTools = e.tools;
+      const fixAgent = e.agent === "";
+      if (typeof rawTools !== "string" && !fixAgent) return entry;
+      const out = { ...e };
+      if (typeof rawTools === "string") {
+        out.tools = normalizeToolsField(rawTools);
+      }
+      if (fixAgent) delete out.agent;
+      return out;
     });
   }