@bermudi/pi-delegate 0.1.1 → 0.1.3

package/dispatch.ts

@@ -10,13 +10,21 @@ import {
   notifyWaiters,
 } from "./tickets.ts";
 import { getConcurrencyLimit, getMaxAsyncTickets } from "./config.ts";
+import { getCurrentLeafId } from "./leaf.ts";
 import { getModelKey, mapConcurrentByModel } from "./concurrency.ts";
-import { sumUsage } from "./usage.ts";
+import { aggregateTaskResults, sumUsage } from "./usage.ts";
 import { runResolvedTask, updateProgressFromRun } from "./lifecycle.ts";
-import { fmtDuration, formatCompletedTask, trunc } from "./format.ts";
+import {
+  fmtDuration,
+  formatCompletedTask,
+  trunc,
+  findTouchedOverlaps,
+  formatTouchedOverlapWarning,
+} from "./format.ts";
 import { validateDelegateOperation } from "./schema.ts";
-import { syncDelegateStatus } from "./status.ts";
+import { notifyCrossLeafDelivery, syncDelegateStatus } from "./status.ts";
 import { validateTasks, resolveTasks } from "./task-resolution.ts";
+import type { CallSpan } from "./telemetry.ts";
 import type {
   AgentConfig,
   AsyncTicket,
@@ -57,9 +65,10 @@ export function validateDelegateOperationResult(
 /** Build the initial per-task progress rows from resolved tasks. */
 export function initProgress(resolved: ResolvedTask[]): TaskProgress[] {
   return resolved.map((t, i) => ({
+    id: t.id,
     index: i,
     agent: t.agentName,
-    task: trunc(t.prompt || t.action || "", 50),
+    task: trunc(t.prompt || t.sessionAction || "", 50),
     status: "pending" as const,
     durationMs: 0,
     tokens: 0,
@@ -105,6 +114,7 @@ export interface AsyncDispatchInput {
   resolved: ResolvedTask[];
   progress: TaskProgress[];
   parentModelId: string | undefined;
+  callSpan?: CallSpan;
 }
 
 /** Inputs needed by the sync (blocking) dispatch path. */
@@ -116,6 +126,7 @@ export interface SyncDispatchInput {
   parentModelId: string | undefined;
   signal: AbortSignal | undefined;
   fire: () => void;
+  callSpan?: CallSpan;
 }
 
 /** Inputs for the normal task-validation, resolution, and dispatch path. */
@@ -128,6 +139,7 @@ export interface DelegateDispatchInput {
   parentDefaults: ParentAgentDefaults;
   signal: AbortSignal | undefined;
   onUpdate: AgentToolUpdateCallback<DelegateDetails> | undefined;
+  callSpan?: CallSpan;
 }
 
 /** Validate, resolve, and dispatch a non-short-circuit delegate operation. */
@@ -143,11 +155,20 @@ export async function dispatchDelegate(
     parentDefaults,
     signal,
     onUpdate,
+    callSpan,
   } = input;
   const tasks = params.tasks ?? [];
 
   const validationError = validateTasks(tasks, agents, parentModelId);
-  if (validationError) return validationError;
+  if (validationError) {
+    callSpan?.finish({
+      status: "failed",
+      totalTokens: 0,
+      totalCost: 0,
+      wallMs: Date.now() - callSpan.startedAt,
+    });
+    return validationError;
+  }
 
   const resolved = resolveTasks(tasks, ctx, agents, parentDefaults);
   const progress = initProgress(resolved);
@@ -168,6 +189,7 @@ export async function dispatchDelegate(
       resolved,
       progress,
       parentModelId,
+      callSpan,
     });
   }
 
@@ -179,20 +201,52 @@ export async function dispatchDelegate(
     parentModelId,
     signal,
     fire,
+    callSpan,
   });
 }
 
+/** Deliver a settled ticket and, when leaf affinity downgraded delivery to a
+ *  non-waking `nextTurn` message, tell the human — otherwise the completion is
+ *  silent apart from the footer clearing. */
+function finishTicketDelivery(pi: ExtensionAPI, ticket: AsyncTicket): void {
+  if (deliverTicketResults(pi, ticket) === "deferred") {
+    notifyCrossLeafDelivery(ticket);
+  }
+}
+
+function settleAsyncCall(
+  ticket: AsyncTicket,
+  callSpan: CallSpan | undefined,
+): void {
+  if (!callSpan) return;
+  const { totalTokens, totalCost } = aggregateTaskResults(ticket.results);
+  const wallMs = (ticket.completedAt ?? Date.now()) - callSpan.startedAt;
+  const status =
+    ticket.status === "done"
+      ? "done"
+      : ticket.status === "cancelled"
+        ? "cancelled"
+        : "failed";
+  callSpan.finish({ status, totalTokens, totalCost, wallMs });
+}
+
 /** Fire-and-forget background execution. Registers an `AsyncTicket`, kicks off
  *  the concurrent run, and returns the ticket acknowledgment immediately.
  *  Results are delivered via `deliverTicketResults` when all tasks settle. */
 export function dispatchAsync(input: AsyncDispatchInput): DelegateToolResult {
-  const { pi, ctx, tasks, resolved, progress, parentModelId } = input;
+  const { pi, ctx, tasks, resolved, progress, parentModelId, callSpan } = input;
 
   sweepTickets();
   const runningCount = [...ticketRegistry.values()].filter(
     (t) => t.status === "running" || t.status === "cancelling",
   ).length;
   if (runningCount >= getMaxAsyncTickets()) {
+    callSpan?.finish({
+      status: "failed",
+      totalTokens: 0,
+      totalCost: 0,
+      wallMs: Date.now() - (callSpan?.startedAt ?? Date.now()),
+    });
     return {
       content: [
         {
@@ -216,8 +270,16 @@ export function dispatchAsync(input: AsyncDispatchInput): DelegateToolResult {
     progress: [...progress],
     controller,
     parentModelId,
+    // Leaf affinity for delivery: a ticket that outlives a /tree navigation
+    // must not wake the agent on the branch the user moved to (issue #30).
+    spawnLeafId: getCurrentLeafId(),
+    callId: callSpan?.id,
+    callStartedAt: callSpan?.startedAt,
+    callRecord: callSpan ? { ...callSpan.baseRecord() } : undefined,
+    telemetryGeneration: callSpan?.generation,
   };
   ticketRegistry.set(ticketId, ticket);
+  callSpan?.spawn();
   // Footer visibility for the new background work (see status.ts). Uses the
   // ctx cached from the dispatch path in extension.ts — DelegateToolCtx is
   // the intentionally narrowed surface and does not carry `ui`.
@@ -234,6 +296,9 @@ export function dispatchAsync(input: AsyncDispatchInput): DelegateToolResult {
     parentSessionManager: ctx.sessionManager,
     ticketId,
     delegateStartedAt: ticket.created,
+    telemetryCallId: callSpan?.id,
+    telemetryGeneration: callSpan?.generation,
+    async: true,
     onProgress: (p, u) => {
       updateProgressFromRun(p, u);
       notifyWaiters(ticket);
@@ -251,7 +316,7 @@ export function dispatchAsync(input: AsyncDispatchInput): DelegateToolResult {
   // Worker must store the TaskResult back into ticket.results, since
   // formatCompletedTicket/handlePoll read from there. Without the write,
   // completed async tasks would be reported as PENDING.
-  mapConcurrentByModel(
+  const completion = mapConcurrentByModel(
     resolved,
     (t) => getModelKey(t.model),
     getConcurrencyLimit,
@@ -263,12 +328,14 @@ export function dispatchAsync(input: AsyncDispatchInput): DelegateToolResult {
     ticketSignal,
   )
     .then(() => {
-      // A ticket that is already terminally "cancelled" at this point was
-      // finalized by cancelTicketForShutdown (user cancels pass through
-      // "cancelling" first): the extension runtime is being torn down, the
-      // captured `pi` is stale or about to be, and a follow-up message has
-      // no live session to land in. Skip delivery entirely.
-      if (ticket.status === "cancelled") return;
+      // Shutdown marks the ticket terminal before cooperative worker aborts
+      // have finished. Still write one final aggregate after every result has
+      // landed; the immediate shutdown snapshot may have missed late usage.
+      // The runtime is being torn down, so never attempt UI delivery here.
+      if (ticket.status === "cancelled") {
+        settleAsyncCall(ticket, callSpan);
+        return;
+      }
       // All tasks settled — determine final ticket status.
       // Use progress (set by runResolvedTask) for settled-ness so the
       // status reflects work completion, not just result-array density.
@@ -289,12 +356,17 @@ export function dispatchAsync(input: AsyncDispatchInput): DelegateToolResult {
         syncTicketBusyIndex(ticket);
       }
       syncDelegateStatus();
-      deliverTicketResults(pi, ticket);
+      settleAsyncCall(ticket, callSpan);
+      finishTicketDelivery(pi, ticket);
     })
     .catch((err) => {
       // Defense-in-depth — should not happen if individual tasks catch properly.
-      // Same shutdown guard as the .then path above.
-      if (ticket.status === "cancelled") return;
+      // Even an unexpected worker rejection must leave the shutdown aggregate
+      // with every result that did settle, without touching the stale UI.
+      if (ticket.status === "cancelled") {
+        settleAsyncCall(ticket, callSpan);
+        return;
+      }
       if (ticket.status === "cancelling") {
         ticket.status = "cancelled";
       } else if (ticket.status === "running") {
@@ -304,8 +376,10 @@ export function dispatchAsync(input: AsyncDispatchInput): DelegateToolResult {
       ticket.completedAt = Date.now();
       syncTicketBusyIndex(ticket);
       syncDelegateStatus();
-      deliverTicketResults(pi, ticket);
+      settleAsyncCall(ticket, callSpan);
+      finishTicketDelivery(pi, ticket);
     });
+  ticket.completion = completion;
 
   return {
     content: [
@@ -315,9 +389,9 @@ export function dispatchAsync(input: AsyncDispatchInput): DelegateToolResult {
           `Async ticket: ${ticketId}`,
           `${resolved.length} task(s) dispatched · ${runningCount + 1}/${getMaxAsyncTickets()} async slots in use`,
           "",
-          "Completed task results are available via poll. Final results delivered automatically when all tasks complete.",
-          `Check progress: delegate({ action: "poll", ticket: "${ticketId}" }) — avoid polling in a tight loop`,
-          `Cancel if needed: delegate({ action: "cancel", ticket: "${ticketId}", force: true }) — first call without force is a preview`,
+          "Work is detached. Stop this turn to let final results auto-deliver.",
+          `If this turn must block for the result, call once: delegate({ ticketAction: "wait", ticket: "${ticketId}" }) — omit timeoutMs and do not poll`,
+          `Cancel if needed: delegate({ ticketAction: "cancel", ticket: "${ticketId}", force: true }) — first call without force is a preview`,
         ].join("\n"),
       },
     ],
@@ -337,7 +411,16 @@ export function dispatchAsync(input: AsyncDispatchInput): DelegateToolResult {
 export async function dispatchSync(
   input: SyncDispatchInput,
 ): Promise<DelegateToolResult> {
-  const { ctx, tasks, resolved, progress, parentModelId, signal, fire } = input;
+  const {
+    ctx,
+    tasks,
+    resolved,
+    progress,
+    parentModelId,
+    signal,
+    fire,
+    callSpan,
+  } = input;
 
   const startedAt = Date.now();
   const syncEnv: TaskRunEnv = {
@@ -346,6 +429,9 @@ export async function dispatchSync(
     parentSessionManager: ctx.sessionManager,
     ticketId: undefined,
     delegateStartedAt: startedAt,
+    telemetryCallId: callSpan?.id,
+    telemetryGeneration: callSpan?.generation,
+    async: false,
     onProgress: (p, u) => {
       updateProgressFromRun(p, u);
       fire();
@@ -376,6 +462,24 @@ export async function dispatchSync(
     parts.push(...formatCompletedTask(t, r));
   }
 
+  const overlapWarning = formatTouchedOverlapWarning(
+    findTouchedOverlaps(finalResults),
+  );
+  if (overlapWarning) parts.push("", overlapWarning);
+
+  const status = finalResults.some((r) => r.error) ? "failed" : "success";
+  const totalTokens = finalResults.reduce((sum, r) => sum + r.tokens, 0);
+  const totalCost = finalResults.reduce(
+    (sum, r) => sum + r.usage.cost.total,
+    0,
+  );
+  callSpan?.finish({
+    status,
+    totalTokens,
+    totalCost,
+    wallMs: Date.now() - callSpan.startedAt,
+  });
+
   return {
     content: [{ type: "text", text: parts.join("\n\n") }],
     details: {
@@ -383,6 +487,7 @@ export async function dispatchSync(
       results: finalResults,
       progress,
       parentModel: parentModelId,
+      overlapWarning: overlapWarning || undefined,
     },
     // Aggregate subagent spend so Pi folds it into the parent's
     // session/footer totals. Sync dispatch only — async results arrive via a