@bermudi/pi-delegate 0.1.2 → 0.1.3

package/extension.ts

@@ -30,10 +30,87 @@ import {
   notifyActiveTicketsOnSettled,
   syncDelegateStatus,
 } from "./status.ts";
+import {
+  beginCall,
+  closeTelemetry,
+  getTelemetryGeneration,
+  prepareTelemetryForSession,
+  sealTelemetryWrites,
+} from "./telemetry.ts";
 import type { DelegateArguments } from "./types.ts";
 
+const DEFAULT_SHUTDOWN_DRAIN_TIMEOUT_MS = 10_000;
+let shutdownDrainTimeoutMs = DEFAULT_SHUTDOWN_DRAIN_TIMEOUT_MS;
+
+type ShutdownDrainResult =
+  { drained: true; failures: unknown[] } | { drained: false; failures: [] };
+
+/** Wait for shutdown workers without allowing a stuck provider/tool to hold
+ * Pi's reload or exit hostage. The allSettled promise is intentionally left
+ * attached after timeout so a late rejection cannot become unhandled. */
+export async function drainAsyncTickets(
+  completions: readonly Promise<void>[],
+  timeoutMs = shutdownDrainTimeoutMs,
+): Promise<ShutdownDrainResult> {
+  if (!completions.length) return { drained: true, failures: [] };
+
+  const drain = Promise.allSettled(completions).then((results) => ({
+    drained: true as const,
+    failures: results
+      .filter(
+        (result): result is PromiseRejectedResult =>
+          result.status === "rejected",
+      )
+      .map((result) => result.reason),
+  }));
+
+  let timer: ReturnType<typeof setTimeout> | undefined;
+  const timeout = new Promise<ShutdownDrainResult>((resolve) => {
+    timer = setTimeout(
+      () => resolve({ drained: false, failures: [] }),
+      timeoutMs,
+    );
+  });
+  const result = await Promise.race([drain, timeout]);
+  if (timer !== undefined) clearTimeout(timer);
+  return result;
+}
+
+/** Wait for another shutdown cleanup operation without allowing a stuck
+ * session lock to hold reload/quit forever. Rejection is consumed here; the
+ * caller reports it when it wins the race, and a late rejection stays handled.
+ */
+async function boundedShutdownCleanup(
+  cleanup: Promise<void>,
+  timeoutMs = shutdownDrainTimeoutMs,
+): Promise<"settled" | "timed-out" | "failed"> {
+  let timer: ReturnType<typeof setTimeout> | undefined;
+  const outcome = cleanup.then(
+    () => "settled" as const,
+    () => "failed" as const,
+  );
+  const timeout = new Promise<"timed-out">((resolve) => {
+    timer = setTimeout(() => resolve("timed-out"), timeoutMs);
+  });
+  const result = await Promise.race([outcome, timeout]);
+  if (timer !== undefined) clearTimeout(timer);
+  return result;
+}
+
+/** @internal Test-only timeout override. */
+export function _setShutdownDrainTimeoutForTesting(
+  timeoutMs: number | undefined,
+): void {
+  shutdownDrainTimeoutMs = timeoutMs ?? DEFAULT_SHUTDOWN_DRAIN_TIMEOUT_MS;
+}
+
 /** Register the delegate tool and clean up its parent-session resources. */
 export default function delegateExtension(pi: ExtensionAPI): void {
+  // A /reload can reuse this module instance after the previous runtime closed
+  // its SQLite handle. Permit the new runtime to open a fresh backend; stale
+  // workers from the old runtime remain blocked from reopening it.
+  prepareTelemetryForSession();
+
   pi.registerTool({
     name: "delegate",
     label: "Delegate to Subagents",
@@ -45,23 +122,70 @@ export default function delegateExtension(pi: ExtensionAPI): void {
     prepareArguments: normalizeDelegateArguments,
 
     async execute(_id, params: DelegateArguments, signal, onUpdate, ctx) {
+      const parentModelId = ctx.model?.id;
+      const tasks = params.tasks ?? [];
+      const parentSessionFile = (
+        ctx as { sessionManager?: { getSessionFile?(): string | undefined } }
+      ).sessionManager?.getSessionFile?.();
+
+      let mode: string;
+      if (params.ticketAction) {
+        mode = params.ticketAction;
+      } else if (tasks.length === 0) {
+        mode = "manual";
+      } else if (params.async) {
+        mode = "async";
+      } else {
+        mode = "sync";
+      }
+      const taskCount = params.ticketAction ? 0 : tasks.length;
+      const callSpan = beginCall({
+        parentModel: parentModelId,
+        mode,
+        taskCount,
+        parentSessionFile,
+      });
+
+      function failCall(): void {
+        callSpan.finish({
+          status: "failed",
+          totalTokens: 0,
+          totalCost: 0,
+          wallMs: Date.now() - callSpan.startedAt,
+        });
+      }
+
+      function succeedCall(): void {
+        callSpan.finish({
+          status: "success",
+          totalTokens: 0,
+          totalCost: 0,
+          wallMs: Date.now() - callSpan.startedAt,
+        });
+      }
+
       // Guard against pi dropping/renaming a symbol this extension imports
       // before any operation-specific validation or early return.
       const compatError = hostCompatError();
-      if (compatError) return compatError;
-
-      const parentModelId = ctx.model?.id;
-      const tasks = params.tasks ?? [];
+      if (compatError) {
+        failCall();
+        return compatError;
+      }
 
       const operationResult = validateDelegateOperationResult(
         params,
         parentModelId,
       );
-      if (operationResult) return operationResult;
+      if (operationResult) {
+        failCall();
+        return operationResult;
+      }
 
       // ── Poll action ───────────────────────────────────────────────────
       if (params.ticketAction === "poll") {
-        return handlePoll(params, ctx);
+        const result = handlePoll(params, ctx);
+        succeedCall();
+        return result;
       }
 
       // ── Cancel action ─────────────────────────────────────────────────
@@ -70,12 +194,20 @@ export default function delegateExtension(pi: ExtensionAPI): void {
         // A forced cancel flips the ticket to "cancelling" — keep the
         // footer status in step (deduped; the preview path is a no-op).
         syncDelegateStatus(ctx);
+        succeedCall();
         return result;
       }
 
       // ── Wait action ────────────────────────────────────────────────────
       if (params.ticketAction === "wait") {
-        return handleWait(params, signal, onUpdate, ctx);
+        try {
+          const result = await handleWait(params, signal, onUpdate, ctx);
+          succeedCall();
+          return result;
+        } catch (err) {
+          failCall();
+          throw err;
+        }
       }
 
       // Agent discovery is intentionally parent-cwd-scoped: agent profiles are a
@@ -87,6 +219,7 @@ export default function delegateExtension(pi: ExtensionAPI): void {
 
       // ── Help mode ─────────────────────────────────────────────────
       if (!tasks.length) {
+        succeedCall();
         return {
           content: [{ type: "text", text: getSubagentManualMarkdown(agents) }],
           details: {
@@ -108,25 +241,38 @@ export default function delegateExtension(pi: ExtensionAPI): void {
       // across dispatches: edits to auth/models/settings/context files must be
       // visible without restarting Pi.
       invalidateHostDepsCache();
-      return dispatchDelegate({
-        pi,
-        params,
-        ctx,
-        agents,
-        parentModelId,
-        parentDefaults: {
-          thinking: pi.getThinkingLevel(),
-          tools: pi.getActiveTools(),
-        },
-        signal,
-        onUpdate,
-      });
+      try {
+        return await dispatchDelegate({
+          pi,
+          params,
+          ctx,
+          agents,
+          parentModelId,
+          parentDefaults: {
+            thinking: pi.getThinkingLevel(),
+            tools: pi.getActiveTools(),
+          },
+          signal,
+          onUpdate,
+          callSpan,
+        });
+      } catch (err) {
+        failCall();
+        throw err;
+      }
     },
 
     renderCall: renderDelegateCall,
     renderResult: renderDelegateResult,
   });
 
+  // A reload may rebuild the extension runtime without re-invoking this
+  // module's default export. Re-open telemetry only after the prior shutdown
+  // handler has drained old workers and closed its connection.
+  pi.on("session_start", () => {
+    prepareTelemetryForSession();
+  });
+
   // ── Background-work visibility (see status.ts) ──────────────────────────
   // The turn settling with live tickets is the "looks idle but isn't" moment:
   // warn once per ticket. The footer status carries it from there.
@@ -155,52 +301,137 @@ export default function delegateExtension(pi: ExtensionAPI): void {
 
   // ── Session shutdown: abort tickets and dispose live pooled sessions ──
   pi.on("session_shutdown", async (event, ctx) => {
-    // Quit and /reload kill background work with no cancellable hook, so
-    // leave a trace. For quit the TUI is already stopped — stderr lands in
-    // the scrollback. For reload the TUI survives — warn in place. Switch
-    // and fork already passed the confirm guard above.
-    const active = activeTicketSummary();
-    if (active.tickets.length) {
-      if (event.reason === "quit") {
+    const shutdownGeneration = getTelemetryGeneration();
+    const shutdownDeadline = Date.now() + shutdownDrainTimeoutMs;
+    const ticketCompletions: Promise<void>[] = [];
+    let drainAttempted = false;
+    let telemetrySealed = false;
+
+    try {
+      // Quit and /reload kill background work with no cancellable hook, so
+      // leave a trace. For quit the TUI is already stopped — stderr lands in
+      // the scrollback. For reload the TUI survives — warn in place. Switch
+      // and fork already passed the confirm guard above.
+      const active = activeTicketSummary();
+      if (active.tickets.length) {
+        if (event.reason === "quit") {
+          console.error(
+            `[delegate] pi exited with ${describeActiveTickets(active)} — aborted.`,
+          );
+        } else if (event.reason === "reload") {
+          try {
+            ctx.ui.notify(
+              `[delegate] reload aborted ${describeActiveTickets(active)}`,
+              "warning",
+            );
+          } catch (error) {
+            console.error(
+              "[delegate] reload shutdown notification failed",
+              error,
+            );
+          }
+        }
+      }
+
+      for (const ticket of ticketRegistry.values()) {
+        if (ticket.status === "running" || ticket.status === "cancelling") {
+          cancelTicketForShutdown(ticket);
+        }
+        // Include already-cancelled tickets too: a repeated shutdown event can
+        // race the first handler while its workers are still unwinding.
+        if (ticket.completion) ticketCompletions.push(ticket.completion);
+      }
+      syncDelegateStatus(ctx);
+      // The runtime is invalidated right after this handler returns; aborted
+      // tickets keep unwinding asynchronously and must find no cached ctx (or
+      // captured pi) to touch. The cancelled completion path still writes one
+      // final aggregate after late task results arrive, but never delivers UI.
+      clearDelegateStatusContext();
+      // A replacement session starts on its own leaf; stale tracking would make
+      // every ticket look cross-leaf (or, worse, look same-leaf by accident).
+      resetLeafTracking();
+      // Do NOT clear the ticket registry here — completed tickets are retained
+      // until their TTL cleanup. Pooled AgentSessions, however, own listeners
+      // and must be disposed before the parent session exits.
+      //
+      // Start pooled cleanup before waiting for ticket completion. Its immediate
+      // abort requests can help a worker blocked on a pooled session unwind;
+      // SQLite still stays open until both cleanup paths finish.
+      let poolCleanup: Promise<void>;
+      try {
+        poolCleanup = closeAllPooledAgents();
+      } catch (error) {
+        console.error("[delegate] pooled-session shutdown start failed", error);
+        poolCleanup = Promise.resolve();
+      }
+      // Drain cooperatively, but never let a provider/tool hold reload or
+      // quit forever. A timed-out old runtime is sealed before a new one can
+      // reopen telemetry; its eventual task completion remains harmless.
+      drainAttempted = true;
+      const drain = await drainAsyncTickets(
+        ticketCompletions,
+        Math.max(0, shutdownDeadline - Date.now()),
+      );
+      if (!drain.drained) {
         console.error(
-          `[delegate] pi exited with ${describeActiveTickets(active)} — aborted.`,
+          `[delegate] async-ticket shutdown drain exceeded ${shutdownDrainTimeoutMs}ms; continuing without late results`,
         );
-      } else if (event.reason === "reload") {
-        ctx.ui.notify(
-          `[delegate] reload aborted ${describeActiveTickets(active)}`,
-          "warning",
+        // Seal immediately, before waiting on pooled-session cleanup. A worker
+        // that finishes during that second phase belongs to the old runtime
+        // and must not write into a freshly reopened backend.
+        telemetrySealed = sealTelemetryWrites(shutdownGeneration);
+      } else if (drain.failures.length) {
+        console.error(
+          "[delegate] async-ticket cleanup failed during shutdown:",
+          drain.failures,
         );
       }
-    }
-    for (const ticket of ticketRegistry.values()) {
-      cancelTicketForShutdown(ticket);
-    }
-    syncDelegateStatus(ctx);
-    // The runtime is invalidated right after this handler returns; aborted
-    // tickets keep unwinding asynchronously and must find no cached ctx (or
-    // captured pi) to touch. See the "cancelled"-at-entry guard in dispatch.
-    clearDelegateStatusContext();
-    // A replacement session starts on its own leaf; stale tracking would make
-    // every ticket look cross-leaf (or, worse, look same-leaf by accident).
-    resetLeafTracking();
-    // Do NOT clear the ticket registry here — completed tickets are retained
-    // until their TTL cleanup. Pooled AgentSessions, however, own listeners
-    // and must be disposed before the parent session exits.
-    //
-    // closeAllPooledAgents attempts every session (Promise.allSettled) before
-    // aggregating failures into an AggregateError, so swallowing here does not
-    // abandon remaining cleanup. Catch and log so a wedged session's failure
-    // stays observable instead of becoming an unhandled rejection — pi.on is
-    // EventEmitter-style and does not surface handler rejections — and so
-    // shutdown completes even when one pooled session failed to abort/dispose.
-    try {
-      await closeAllPooledAgents();
-    } catch (error) {
-      const failures = error instanceof AggregateError ? error.errors : [error];
-      console.error(
-        "[delegate] pooled-session cleanup failed during shutdown:",
-        failures,
+
+      // closeAllPooledAgents attempts every session (Promise.allSettled) before
+      // aggregating failures into an AggregateError, so swallowing here does
+      // not abandon remaining cleanup.
+      const poolResult = await boundedShutdownCleanup(
+        poolCleanup,
+        Math.max(0, shutdownDeadline - Date.now()),
       );
+      if (poolResult === "timed-out") {
+        console.error(
+          `[delegate] pooled-session shutdown cleanup exceeded ${shutdownDrainTimeoutMs}ms; continuing`,
+        );
+      } else if (poolResult === "failed") {
+        // closeAllPooledAgents aggregates every session failure before
+        // rejecting, so the rejection is already a complete diagnostic.
+        try {
+          await poolCleanup;
+        } catch (error) {
+          const failures =
+            error instanceof AggregateError ? error.errors : [error];
+          console.error(
+            "[delegate] pooled-session cleanup failed during shutdown:",
+            failures,
+          );
+        }
+      }
+    } catch (error) {
+      // pi.on is EventEmitter-style and does not surface handler rejections.
+      // Keep shutdown moving, but leave the failure visible.
+      console.error("[delegate] shutdown cleanup failed", error);
+    } finally {
+      if (!drainAttempted) {
+        // An unexpected earlier cleanup failure must not leave a worker with an
+        // unhandled rejection, but this fallback is bounded too.
+        await drainAsyncTickets(
+          ticketCompletions,
+          Math.max(0, shutdownDeadline - Date.now()),
+        );
+      }
+      // Once the bounded drain expires, old workers may still unwind. Seal
+      // their generation before closing SQLite so a later runtime cannot
+      // receive stale task/call writes.
+      if (!telemetrySealed) {
+        telemetrySealed = sealTelemetryWrites(shutdownGeneration);
+      }
+      closeTelemetry(shutdownGeneration);
     }
   });
 }

package/lifecycle.ts

@@ -27,6 +27,7 @@ import { resolveCwd, validateResumeFromPath } from "./utils.ts";
 import { getWholeTaskMaxRetries, getWholeTaskBaseDelayMs } from "./config.ts";
 import { addUsage, emptyUsage } from "./usage.ts";
 import { scheduleDeadline } from "./timer.ts";
+import { recordTask } from "./telemetry.ts";
 
 /** Internal seam for lifecycle-level tests without replacing session ownership. */
 type RunAgentSession = typeof runAgentSession;
@@ -187,8 +188,22 @@ function finishTask(
   env: TaskRunEnv,
   p: TaskProgress,
   r: TaskResult,
+  task: ResolvedTask,
+  retries = 0,
 ): TaskResult {
   updateProgressFromResult(p, r);
+  if (env.telemetryCallId) {
+    recordTask({
+      callId: env.telemetryCallId,
+      generation: env.telemetryGeneration,
+      async: env.async ?? false,
+      taskIndex: p.index,
+      task,
+      progress: p,
+      result: r,
+      retries,
+    });
+  }
   env.onStatusChange?.();
   return r;
 }
@@ -553,7 +568,7 @@ async function runResolvedTaskUnlocked(
   try {
     // ── Aborted before we started? ───────────────────────────────────
     if (env.signal?.aborted) {
-      return finishTask(env, p, failTask(task, "Aborted"));
+      return finishTask(env, p, failTask(task, "Aborted"), task);
     }
 
     // ── Session busy guard (defense-in-depth) ────────────────────────
@@ -563,7 +578,7 @@ async function runResolvedTaskUnlocked(
       const busyTicketId = isSessionBusy(task.sessionId);
       if (busyTicketId && busyTicketId !== env.ticketId) {
         const msg = `Session '${task.sessionId}' is already in use by ticket ${busyTicketId}. Each session can only handle one task at a time.`;
-        return finishTask(env, p, failTask(task, msg));
+        return finishTask(env, p, failTask(task, msg), task);
       }
     }
 
@@ -577,6 +592,7 @@ async function runResolvedTaskUnlocked(
           env,
           p,
           failTask(task, "sessionAction='close' requires sessionId."),
+          task,
         );
       }
       // The per-session lock for action-based operations is already held by the
@@ -593,6 +609,7 @@ async function runResolvedTaskUnlocked(
             : `Session '${task.sessionId}' not found.`,
           Date.now() - env.delegateStartedAt,
         ),
+        task,
       );
     }
 
@@ -605,6 +622,7 @@ async function runResolvedTaskUnlocked(
           `Active sessions:\n${pool.listPooledAgents().join("\n")}`,
           Date.now() - env.delegateStartedAt,
         ),
+        task,
       );
     }
 
@@ -821,16 +839,22 @@ async function runResolvedTaskUnlocked(
         task,
         err instanceof Error ? err.message : String(err),
       );
-      return finishTask(env, p, {
-        ...failure,
-        durationMs: Math.max(failure.durationMs, Date.now() - taskStartedAt),
-        tokens: accumulatedUsage.totalTokens,
-        usage: accumulatedUsage,
-      });
+      return finishTask(
+        env,
+        p,
+        {
+          ...failure,
+          durationMs: Math.max(failure.durationMs, Date.now() - taskStartedAt),
+          tokens: accumulatedUsage.totalTokens,
+          usage: accumulatedUsage,
+        },
+        task,
+      );
     }
 
     const maxRetries = resolvedWholeTaskMaxRetries();
     const baseDelayMs = resolvedWholeTaskBaseDelayMs();
+    let retriesExecuted = 0;
     for (
       let retry = 0;
       retry < maxRetries && canRetryWholeTask(task, result, hasBashExecution);
@@ -864,6 +888,7 @@ async function runResolvedTaskUnlocked(
       p.error = undefined;
       p.failureKind = undefined;
       env.onStatusChange?.();
+      retriesExecuted++;
       try {
         result = await runAttempt();
       } catch (err) {
@@ -881,12 +906,18 @@ async function runResolvedTaskUnlocked(
       }
     }
 
-    return finishTask(env, p, {
-      ...result,
-      durationMs: Math.max(result.durationMs, Date.now() - taskStartedAt),
-      tokens: Math.max(cumulativeTokens, accumulatedUsage.totalTokens),
-      usage: accumulatedUsage,
-    });
+    return finishTask(
+      env,
+      p,
+      {
+        ...result,
+        durationMs: Math.max(result.durationMs, Date.now() - taskStartedAt),
+        tokens: Math.max(cumulativeTokens, accumulatedUsage.totalTokens),
+        usage: accumulatedUsage,
+      },
+      task,
+      retriesExecuted,
+    );
   } catch (err) {
     // Any acquired session is released by runAttempt's finally before an
     // exception reaches this boundary. This outer catch handles unexpected
@@ -896,6 +927,7 @@ async function runResolvedTaskUnlocked(
       env,
       p,
       failTask(task, err instanceof Error ? err.message : String(err)),
+      task,
     );
   }
 }

package/manual.ts

@@ -185,8 +185,9 @@ export function getSubagentManualMarkdown(
     "```",
     "",
     '- `delegate({ ticketAction: "poll" })` \u2014 list all tickets',
-    '- `delegate({ ticketAction: "poll", ticket: "abc123" })` \u2014 check one ticket',
-    '- `delegate({ ticketAction: "wait", ticket: "abc123", timeoutMs: 600000 })` \u2014 block until finished or timeout',
+    '- `delegate({ ticketAction: "poll", ticket: "abc123" })` \u2014 take one progress snapshot',
+    '- `delegate({ ticketAction: "wait", ticket: "abc123" })` \u2014 block until finished; omit `timeoutMs` when the result is needed this turn',
+    '- `delegate({ ticketAction: "wait", ticket: "abc123", timeoutMs: 600000 })` \u2014 bounded wait; timeout includes the latest snapshot, so do not poll afterward',
     '- `delegate({ ticketAction: "cancel", ticket: "abc123" })` \u2014 preview activity and partial effects before cancelling',
     '- `delegate({ ticketAction: "cancel", ticket: "abc123", force: true })` \u2014 abort after review',
     "",

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@bermudi/pi-delegate",
-  "version": "0.1.2",
+  "version": "0.1.3",
   "description": "Delegate tool for the Pi coding agent.",
   "keywords": [
     "pi-package"

package/schema.ts

@@ -139,7 +139,7 @@ export const delegateArgumentsSchema = Type.Object({
     Type.Number({
       minimum: 0,
       description:
-        "Bounds wait (ms); omit to block until settled. Timeout never cancels the ticket.",
+        "Bounds wait (ms); omit to block until settled. Timeout returns a snapshot; do not poll afterward.",
     }),
   ),
   tasks: Type.Optional(