@bermudi/pi-delegate 0.1.0 → 0.1.2

package/lifecycle.ts

@@ -11,20 +11,32 @@ import type {
   TaskProgress,
   TaskResult,
   TaskRunEnv,
+  ToolActivity,
 } from "./types.ts";
 import * as pool from "./pool.ts";
 import { isSessionBusy } from "./tickets.ts";
 import {
   createSubagentSessionManager,
-  setParentSession,
   persistSessionHeader,
+  setParentSession,
 } from "./sessions.ts";
-import { runAgentSession } from "./runner.ts";
+import { runAgentSession, formatDeadlineExceededError } from "./runner.ts";
 import { getGitChangedFiles } from "./file-tracking.ts";
 import { getHostDeps } from "./host.ts";
 import { resolveCwd, validateResumeFromPath } from "./utils.ts";
 import { getWholeTaskMaxRetries, getWholeTaskBaseDelayMs } from "./config.ts";
 import { addUsage, emptyUsage } from "./usage.ts";
+import { scheduleDeadline } from "./timer.ts";
+
+/** Internal seam for lifecycle-level tests without replacing session ownership. */
+type RunAgentSession = typeof runAgentSession;
+let runAgentSessionForTesting: RunAgentSession = runAgentSession;
+
+export function _setRunAgentSessionForTesting(
+  override: RunAgentSession | undefined,
+): void {
+  runAgentSessionForTesting = override ?? runAgentSession;
+}
 
 /**
  * Test-only overrides for whole-task retry settings. When set, these bypass
@@ -61,6 +73,7 @@ function failTask(
   sessionFile?: string,
 ): TaskResult {
   return {
+    id: task.id,
     agent: task.agentName,
     output: "",
     error,
@@ -69,6 +82,7 @@ function failTask(
     usage: emptyUsage(),
     sessionFile,
     touchedFiles: [],
+    attributedFiles: [],
   };
 }
 
@@ -80,6 +94,7 @@ function completeSessionAction(
   elapsedMs?: number,
 ): TaskResult {
   return {
+    id: task.id,
     agent: task.agentName,
     output,
     durationMs: elapsedMs ?? 0,
@@ -87,6 +102,7 @@ function completeSessionAction(
     usage: emptyUsage(),
     sessionFile: undefined,
     touchedFiles: [],
+    attributedFiles: [],
   };
 }
 
@@ -104,19 +120,57 @@ function disposeOwnedSession(acquired: AcquiredSession): void {
   }
 }
 
-/** Mirror a progress update from runAgent into a TaskProgress row. */
+/** Merge per-attempt activities into a live history list while preserving
+ * prior-attempt evidence. Activity IDs are used as a stable handle so in-flight
+ * updates can replace earlier skeletons for the same call. */
+function mergeToolActivities(
+  existing: TaskProgress["activities"],
+  incoming: ToolActivity[],
+): ToolActivity[] {
+  const byId = new Map<string, number>();
+  for (let i = 0; i < existing.length; i++) {
+    byId.set(existing[i]!.id, i);
+  }
+
+  const merged = [...existing];
+  for (const incomingActivity of incoming) {
+    const index = byId.get(incomingActivity.id);
+    if (index === undefined) {
+      byId.set(incomingActivity.id, merged.length);
+      merged.push(incomingActivity);
+    } else {
+      merged[index] = { ...merged[index], ...incomingActivity };
+    }
+  }
+  return merged;
+}
+
+/** Optional counters are used by retry-aware accounting to keep live progress
+ * monotonic across attempts. */
+interface RunUpdateOffset {
+  tokensOffset?: number;
+  toolUsesOffset?: number;
+}
+/** Mirror a progress update from runAgent into a TaskProgress row.
+ *
+ * Runner callbacks report attempt-local counters. Offset/merge them so a single
+ * TaskProgress row is monotonic across whole-task retries.
+ */
 export function updateProgressFromRun(
   p: TaskProgress,
   u: AgentProgressUpdate,
+  offsets: RunUpdateOffset = {},
 ): void {
-  p.tokens = u.tokens;
-  p.toolUses = u.toolUses;
-  p.durationMs = u.durationMs;
+  const cumulativeTokens = (offsets.tokensOffset ?? 0) + u.tokens;
+  const cumulativeTools = (offsets.toolUsesOffset ?? 0) + u.toolUses;
+  p.tokens = Math.max(p.tokens, cumulativeTokens);
+  p.toolUses = Math.max(p.toolUses, cumulativeTools);
+
+  p.durationMs = Math.max(p.durationMs, u.durationMs);
   p.lastActivityAt = u.lastActivityAt;
-  p.activities = u.activities;
+  p.activities = mergeToolActivities(p.activities, u.activities);
   p.failureKind = u.failureKind;
 }
-
 /** Mirror a completed TaskResult into a TaskProgress row (status/duration/error). */
 function updateProgressFromResult(p: TaskProgress, r: TaskResult): void {
   p.status = r.error ? "failed" : "done";
@@ -195,19 +249,33 @@ function isClearlyTransientFinalError(error: string | undefined): boolean {
   );
 }
 
-function canRetryWholeTask(task: ResolvedTask, result: TaskResult): boolean {
+function canRetryWholeTask(
+  task: ResolvedTask,
+  result: TaskResult,
+  hasBashExecution = false,
+): boolean {
   // Whole-task retry can repeat tool side effects. Keep it to stateless fresh
   // tasks, and only when our touched-file accounting says the failed attempt
   // did not write/edit anything. A `model_error` (usage limit, auth, quota) is
   // not transient for the resolved model — retrying with the same model just
   // hits the same wall, so skip it and let the parent resume with a different
-  // model (see the hint in formatFailedTask).
+  // model (see the hint in formatFailedTask). Similarly, once bash executes,
+  // any retry would replay non-idempotent side effects, so suppress it.
+  // We gate on *observed* bash activity plus touchedFiles, not on the tool set
+  // itself: the default tool set includes `bash` for most tasks, and suppressing
+  // retry for every bash-capable task would disable the useful transient-error
+  // retry path even when no side effects occurred. The stricter “any bash tool
+  // → no retry” rule matches the “no filesystem isolation” stance, but is too
+  // restrictive for retry safety — touched-file accounting and observed activity
+  // are the direct side-effect signals.
   return (
     result.failureKind !== "stalled" &&
     result.failureKind !== "model_error" &&
+    result.failureKind !== "deadline_exceeded" &&
     !task.sessionId &&
     !task.resumeFrom &&
     result.touchedFiles.length === 0 &&
+    !hasBashExecution &&
     isClearlyTransientFinalError(result.error)
   );
 }
@@ -215,16 +283,23 @@ function canRetryWholeTask(task: ResolvedTask, result: TaskResult): boolean {
 async function sleepForWholeTaskRetry(
   signal: AbortSignal | undefined,
   delayMs: number,
+  deadlineAt?: number,
 ): Promise<void> {
   if (signal?.aborted) return;
+  if (deadlineAt !== undefined && Date.now() >= deadlineAt) return;
+
+  const retryAt = Date.now() + delayMs;
+  const wakeAt =
+    deadlineAt !== undefined ? Math.min(retryAt, deadlineAt) : retryAt;
+
   await new Promise<void>((resolve) => {
-    let timeout: ReturnType<typeof setTimeout>;
+    let clear: (() => void) | undefined;
     const done = () => {
-      clearTimeout(timeout);
+      clear?.();
       signal?.removeEventListener("abort", done);
       resolve();
     };
-    timeout = setTimeout(done, delayMs);
+    clear = scheduleDeadline(wakeAt, done);
     if (!signal) return;
     signal.addEventListener("abort", done, { once: true });
   });
@@ -244,7 +319,7 @@ async function buildDelegateSession(
   // resource loaders are cached after the first call; provider-configured or
   // allowlisted-extension loaders are deliberately fresh per session because
   // their extension runtime is mutable. The resourceLoader is cwd-scoped (it
-  // scans for AGENTS.md/skills) and the system prompt is per named-agent.
+  // scans for project AGENTS.md/skills) and the system prompt is per named-agent.
   // The custom prompt overrides the default AgentSession system prompt.
   // Pass only the provider needed by this task. This keeps a non-Kilo task
   // from receiving Kilo's provider/auth adapter merely because Kilo is also
@@ -454,7 +529,7 @@ function resolveResumableSessionFile(
  *  Used by both sync (params.async === false) and async (params.async === true) paths.
  *  When task.sessionId is set, the entire acquire/run/close lifecycle runs under
  *  a per-session mutex so concurrent tasks with the same sessionId serialize
- *  cleanly. The lock also covers action='close' and the early-busy/abort paths. */
+ *  cleanly. The lock also covers sessionAction='close' and the early-busy/abort paths. */
 export async function runResolvedTask(
   env: TaskRunEnv,
   task: ResolvedTask,
@@ -496,15 +571,18 @@ async function runResolvedTaskUnlocked(
     p.model = task.model?.id;
 
     // ── Session action handling ───────────────────────────────────────
-    if (task.action === "close") {
+    if (task.sessionAction === "close") {
       if (!task.sessionId) {
         return finishTask(
           env,
           p,
-          failTask(task, "action='close' requires sessionId."),
+          failTask(task, "sessionAction='close' requires sessionId."),
         );
       }
-      const closed = await pool.closePooledAgent(task.sessionId);
+      // The per-session lock for action-based operations is already held by the
+      // outer runResolvedTask() wrapper. Use the internal close helper to avoid a
+      // reentrant deadlock on the same key.
+      const closed = await pool._closePooledAgentWithoutLock(task.sessionId);
       return finishTask(
         env,
         p,
@@ -518,7 +596,7 @@ async function runResolvedTaskUnlocked(
       );
     }
 
-    if (task.action === "list") {
+    if (task.sessionAction === "list") {
       return finishTask(
         env,
         p,
@@ -530,15 +608,53 @@ async function runResolvedTaskUnlocked(
       );
     }
 
+    let hasBashExecution = false;
+    let cumulativeTokens = 0;
+    let cumulativeToolUses = 0;
+    const taskStartedAt = Date.now();
+    const deadlineAt =
+      task.deadlineMs && task.deadlineMs > 0
+        ? taskStartedAt + task.deadlineMs
+        : undefined;
+    let accumulatedUsage = emptyUsage();
+
+    const onAttemptProgress = (u: AgentProgressUpdate): void => {
+      if (
+        !hasBashExecution &&
+        u.activities.some((activity) => activity.name === "bash")
+      ) {
+        hasBashExecution = true;
+      }
+
+      const mapped: AgentProgressUpdate = {
+        ...u,
+        tokens: cumulativeTokens + u.tokens,
+        toolUses: cumulativeToolUses + u.toolUses,
+        durationMs: Date.now() - taskStartedAt,
+      };
+
+      // Keep live totals monotonic across attempts.
+      p.tokens = Math.max(p.tokens, mapped.tokens);
+      p.toolUses = Math.max(p.toolUses, mapped.toolUses);
+      if (mapped.durationMs > p.durationMs) {
+        p.durationMs = mapped.durationMs;
+      }
+      env.onProgress(p, mapped);
+    };
+
     const runAttempt = async (): Promise<TaskResult> => {
+      let attemptToolUsesObserved = 0;
+      const onProgress = (u: AgentProgressUpdate): void => {
+        attemptToolUsesObserved = Math.max(attemptToolUsesObserved, u.toolUses);
+        onAttemptProgress(u);
+      };
+
       // ── Pool / resume / fresh-agent resolution ────────────────────────
       const acquired = await acquireAgentSession(env, task, p);
       if ("error" in acquired) return acquired.error;
 
       // A pool hit is already owned by the pool. Fresh/resumed sessions belong
-      // to this attempt until commit() explicitly transfers ownership. Keeping
-      // this state local makes cleanup a finally invariant rather than a list
-      // of special cases for aborts, stalls, and provider failures.
+      // to this attempt until commit/recordUse/close logic runs.
       let sessionReleased = !acquired.lifecycleOwnsSession;
       try {
         // Re-check abort after acquisition. The pre-acquire check at the top can
@@ -546,22 +662,29 @@ async function runResolvedTaskUnlocked(
         // baseline. runAgentSession re-checks after attaching its listener, but a
         // cancelled ticket should not even start the subagent (no file writes, no
         // pool insert).
-        if (env.signal?.aborted) return failTask(task, "Aborted");
+        if (env.signal?.aborted) {
+          return failTask(task, "Aborted");
+        }
 
         // Snapshot git status before the run so touchedFiles can diff after.
         // AgentSession owns retry/compaction internally — runAgentSession just
-        // drives the prompt and maps events to the progress model.
+        // drives the prompt and maps events to the progress model. Git failures
+        // degrade to an undefined baseline, which tells the runner to skip
+        // git-based attribution entirely; see getGitChangedFiles for the
+        // contract.
         const gitBaseline = await getGitChangedFiles(task.cwd);
-        let r = await runAgentSession(
+        let r = await runAgentSessionForTesting(
           acquired.session,
           task.prompt,
           { cwd: task.cwd },
           env.signal,
-          (u) => env.onProgress(p, u),
+          onProgress,
           gitBaseline,
-          Date.now(),
+          taskStartedAt,
+          deadlineAt,
         );
 
+        cumulativeTokens += r.tokens;
         // The signal can fire after the pre-run check or while the runner is
         // collecting post-prompt evidence. Keep cancellation from looking like
         // success; finally below releases any uncommitted session.
@@ -569,62 +692,86 @@ async function runResolvedTaskUnlocked(
           r = { ...r, error: "Aborted" };
         }
 
-        // A stalled prompt was explicitly aborted and is no longer a safe
-        // continuation. Close a pooled hit; a fresh/resumed session is still
-        // released by the finally below. closePooledAgent removes the pooled
-        // entry before surfacing cleanup errors, so a pool-owned session is
-        // never directly disposed here.
-        if (r.failureKind === "stalled" && task.sessionId) {
-          try {
-            if (await pool.closePooledAgent(task.sessionId)) {
-              sessionReleased = true;
-            }
-          } catch (error) {
-            // Preserve the primary stalled result while logging the cleanup
-            // failure explicitly. A pooled session is already removed by the
-            // pool; a pool miss remains lifecycle-owned and is handled below.
-            console.error(
-              `[delegate] failed to dispose stalled pooled session '${task.sessionId}'`,
-              error,
-            );
-          }
-        }
-
         const sessionFile = resolveResumableSessionFile(
           acquired.sessionFile,
           acquired.sessionManager,
           r.error,
         );
 
-        // Pool bookkeeping. commit() is the sole mutator: it decides
-        // insert-vs-recordUse by map presence (sound because the session lock
-        // serializes same-sessionId tasks). A successful insert transfers
-        // ownership; a failed insert leaves finally responsible for disposal.
-        if (task.sessionId && !r.error) {
-          sessionReleased = pool.commit(task.sessionId, {
-            session: acquired.session,
-            sessionManager: acquired.sessionManager,
-            sessionFile: acquired.sessionFile,
-            frozen: {
-              systemPrompt: task.systemPrompt,
-              model: task.model,
-              thinking: task.thinking,
-              tools: task.tools,
-              cwd: task.cwd,
-            },
-            tokens: r.tokens,
-          });
+        if (task.sessionId) {
+          if (acquired.lifecycleOwnsSession) {
+            // Pool misses (including resumeFrom) transfer ownership only on
+            // successful completion; failures are owned by lifecycle and must
+            // be disposed in this finally path.
+            if (
+              !r.error &&
+              r.failureKind !== "stalled" &&
+              r.failureKind !== "deadline_exceeded"
+            ) {
+              const committed = pool.commit(task.sessionId, {
+                session: acquired.session,
+                sessionManager: acquired.sessionManager,
+                sessionFile: acquired.sessionFile,
+                frozen: {
+                  systemPrompt: task.systemPrompt,
+                  model: task.model,
+                  thinking: task.thinking,
+                  tools: task.tools,
+                  cwd: task.cwd,
+                },
+                tokens: r.tokens,
+              });
+              sessionReleased = sessionReleased || committed;
+            }
+          } else {
+            // A stalled, parent-aborted, or mid-prompt deadline-exceeded pooled
+            // attempt is not safe to keep; the session may have been mutated.
+            // A pre-prompt deadline (runner never called session.prompt()) left
+            // the session in its pre-task state, so return it to the pool intact.
+            if (
+              r.failureKind === "stalled" ||
+              r.error === "Aborted" ||
+              (r.failureKind === "deadline_exceeded" && r.prompted !== false)
+            ) {
+              try {
+                sessionReleased =
+                  (await pool._closePooledAgentWithoutLock(task.sessionId)) ||
+                  sessionReleased;
+              } catch (error) {
+                // Preserve the primary failure result while logging the cleanup
+                // failure explicitly. A pooled session may still be removed by
+                // the pool; a pool-miss remains lifecycle-owned and is handled
+                // by the finally path above.
+                console.error(
+                  `[delegate] failed to dispose aborted, stalled, or deadline-exceeded pooled session '${task.sessionId}'`,
+                  error,
+                );
+              }
+            } else if (r.failureKind !== "deadline_exceeded") {
+              // Pool hits stay owned by the pool, and non-stalled, non-aborted
+              // completions (including failed attempts) must still count usage.
+              pool.recordUse(task.sessionId, r.tokens);
+            }
+            // Pre-prompt deadline (prompted === false): the pooled session was
+            // checked out but never used. Leave it in the pool with no usage
+            // recorded.
+          }
         }
 
+        accumulatedUsage = addUsage(accumulatedUsage, r.usage);
+        cumulativeToolUses += attemptToolUsesObserved;
+
         return {
+          id: task.id,
           agent: task.agentName,
           output: r.output,
           error: r.error,
-          // Classify the failure: the runner sets `stalled` for the inactivity
-          // watchdog; here we add `model_error` for failures attributable to
-          // the resolved model (usage limit, auth, quota) so the parent gets a
-          // "switch model" hint instead of a same-model retry hint, and so
-          // canRetryWholeTask skips the pointless same-model retry.
+          // Classify the failure: the runner sets `stalled` for the
+          // inactivity watchdog; here we add `model_error` for failures
+          // attributable to the resolved model (usage limit, auth, quota) so the
+          // parent gets a "switch model" hint instead of a same-model retry
+          // hint, and so canRetryWholeTask skips the pointless same-model
+          // retry.
           failureKind:
             r.failureKind ??
             (r.error && isModelAttributableError(r.error)
@@ -635,6 +782,7 @@ async function runResolvedTaskUnlocked(
           usage: r.usage,
           sessionFile,
           touchedFiles: r.touchedFiles,
+          attributedFiles: r.attributedFiles ?? [],
         };
       } finally {
         // This runs for ordinary success, normal provider failure, whole-task
@@ -644,29 +792,53 @@ async function runResolvedTaskUnlocked(
       }
     };
 
-    // The session lock is now taken at the top of runResolvedTask (covers the
-    // full acquire/run/close lifecycle), so attempts execute serially per
-    // sessionId without needing an inner lock here.
-    let result = await runAttempt();
-    // Accumulate usage across whole-task retries: the parent pays for every
-    // attempt, including the transient failures that retry. Keep tokens tied
-    // to the accumulated usage as well; otherwise the final result and the
-    // final progress row describe different amounts of work.
-    let accumulatedUsage = result.usage;
-    result = {
-      ...result,
-      tokens: accumulatedUsage.totalTokens,
-      usage: accumulatedUsage,
+    const buildDeadlineExceededResult = (prior?: TaskResult): TaskResult => {
+      const budgetMs = Math.max(0, (deadlineAt ?? 0) - taskStartedAt);
+      return {
+        id: task.id,
+        agent: task.agentName,
+        output: prior?.output ?? "",
+        error: formatDeadlineExceededError(budgetMs),
+        failureKind: "deadline_exceeded",
+        durationMs: prior?.durationMs ?? 0,
+        tokens: prior?.tokens ?? 0,
+        usage: prior?.usage ?? emptyUsage(),
+        sessionFile: prior?.sessionFile,
+        touchedFiles: prior?.touchedFiles ?? [],
+        attributedFiles: prior?.attributedFiles ?? [],
+      };
     };
+
+    let result: TaskResult;
+    try {
+      if (deadlineAt && Date.now() >= deadlineAt) {
+        result = buildDeadlineExceededResult(undefined);
+      } else {
+        result = await runAttempt();
+      }
+    } catch (err) {
+      const failure = failTask(
+        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,
+      });
+    }
+
     const maxRetries = resolvedWholeTaskMaxRetries();
     const baseDelayMs = resolvedWholeTaskBaseDelayMs();
     for (
       let retry = 0;
-      retry < maxRetries && canRetryWholeTask(task, result);
+      retry < maxRetries && canRetryWholeTask(task, result, hasBashExecution);
       retry++
     ) {
       const delayMs = baseDelayMs * 2 ** retry;
-      await sleepForWholeTaskRetry(env.signal, delayMs);
+      p.durationMs = Math.max(p.durationMs, Date.now() - taskStartedAt);
+      await sleepForWholeTaskRetry(env.signal, delayMs, deadlineAt);
       if (env.signal?.aborted) {
         // Preserve any partial output/session path from the last failed attempt
         // while recording that the retry loop was aborted. The task already
@@ -674,27 +846,52 @@ async function runResolvedTaskUnlocked(
         result = {
           ...result,
           error: "Aborted",
+          durationMs: Date.now() - taskStartedAt,
           tokens: accumulatedUsage.totalTokens,
           usage: accumulatedUsage,
+          touchedFiles: result.touchedFiles,
+          sessionFile: result.sessionFile,
         };
         break;
       }
+
+      if (deadlineAt && Date.now() >= deadlineAt) {
+        result = buildDeadlineExceededResult(result);
+        break;
+      }
+
       p.status = "running";
       p.error = undefined;
       p.failureKind = undefined;
       env.onStatusChange?.();
-      result = await runAttempt();
-      accumulatedUsage = addUsage(accumulatedUsage, result.usage);
-      result = {
-        ...result,
-        tokens: accumulatedUsage.totalTokens,
-        usage: accumulatedUsage,
-      };
+      try {
+        result = await runAttempt();
+      } catch (err) {
+        const failure = failTask(
+          task,
+          err instanceof Error ? err.message : String(err),
+        );
+        result = {
+          ...failure,
+          durationMs: Math.max(failure.durationMs, Date.now() - taskStartedAt),
+          tokens: accumulatedUsage.totalTokens,
+          usage: accumulatedUsage,
+        };
+        break;
+      }
     }
-    return finishTask(env, p, result);
+
+    return finishTask(env, p, {
+      ...result,
+      durationMs: Math.max(result.durationMs, Date.now() - taskStartedAt),
+      tokens: Math.max(cumulativeTokens, accumulatedUsage.totalTokens),
+      usage: accumulatedUsage,
+    });
   } catch (err) {
     // Any acquired session is released by runAttempt's finally before an
-    // exception reaches this boundary.
+    // exception reaches this boundary. This outer catch handles unexpected
+    // throws before retry accounting is initialized, so report a minimal
+    // failure without accumulated usage.
     return finishTask(
       env,
       p,