@bermudi/pi-delegate 0.1.18 → 0.1.19

package/extension.ts

@@ -1,11 +1,5 @@
 import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
-import {
-  handleCancel,
-  handlePoll,
-  handleWait,
-  cancelTicketForShutdown,
-  ticketRegistry,
-} from "./tickets.ts";
+import { getDefaultDelegateRuntime, type DelegateRuntime } from "./runtime.ts";
 import { discoverAgents } from "./agents.ts";
 import { getSubagentManualMarkdown } from "./manual.ts";
 import {
@@ -25,7 +19,6 @@ import { hostCompatError } from "./host-compat.ts";
 import { invalidateHostDepsCache } from "./host.ts";
 import { registerProviderExtensionNotifier } from "./provider-extensions.ts";
 import { recordTreeNavigation, resetLeafTracking } from "./leaf.ts";
-import { closeAllPooledAgents } from "./pool.ts";
 import { reconfigureGlobalConcurrency } from "./concurrency.ts";
 import { reloadDelegateConfig, getMaxConcurrent } from "./config.ts";
 import {
@@ -130,8 +123,15 @@ export function _setShutdownDrainTimeoutForTesting(
   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 {
+/** Register the delegate tool and clean up its parent-session resources.
+ *
+ *  Production uses the module default runtime. Tests/embedders may pass a
+ *  fresh runtime as the second argument so all tool execution, status, and
+ *  shutdown operate in an isolated pool/ticket environment. */
+export default function delegateExtension(
+  pi: ExtensionAPI,
+  runtime: DelegateRuntime = getDefaultDelegateRuntime(),
+): 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.
@@ -237,17 +237,17 @@ export default function delegateExtension(pi: ExtensionAPI): void {
 
       // ── Poll action ───────────────────────────────────────────────────
       if (params.ticketAction === "poll") {
-        const result = handlePoll(params, ctx);
+        const result = runtime.tickets.handlePoll(params, ctx);
         succeedCall();
         return result;
       }
 
       // ── Cancel action ─────────────────────────────────────────────────
       if (params.ticketAction === "cancel") {
-        const result = handleCancel(params);
+        const result = runtime.tickets.handleCancel(params);
         // A forced cancel flips the ticket to "cancelling" — keep the
         // footer status in step (deduped; the preview path is a no-op).
-        syncDelegateStatus(ctx);
+        syncDelegateStatus(ctx, runtime);
         succeedCall();
         return result;
       }
@@ -255,7 +255,12 @@ export default function delegateExtension(pi: ExtensionAPI): void {
       // ── Wait action ────────────────────────────────────────────────────
       if (params.ticketAction === "wait") {
         try {
-          const result = await handleWait(params, signal, onUpdate, ctx);
+          const result = await runtime.tickets.handleWait(
+            params,
+            signal,
+            onUpdate,
+            ctx,
+          );
           succeedCall();
           return result;
         } catch (err) {
@@ -283,7 +288,7 @@ export default function delegateExtension(pi: ExtensionAPI): void {
         invalidateHostDepsCache();
         // Keep the footer-status pipeline in step exactly as a normal
         // dispatch would (deduped no-op when nothing is running).
-        syncDelegateStatus(ctx);
+        syncDelegateStatus(ctx, runtime);
         return await dispatchDelegate({
           pi,
           params: { ...params, tasks: [bridgeSessionControlTask(params)] },
@@ -297,6 +302,7 @@ export default function delegateExtension(pi: ExtensionAPI): void {
           signal,
           onUpdate,
           callSpan,
+          runtime,
         });
       }
 
@@ -318,7 +324,7 @@ export default function delegateExtension(pi: ExtensionAPI): void {
       // dispatch narrows it to DelegateToolCtx (which has no `ui`). The
       // status push itself is a deduped no-op here; dispatchAsync re-syncs
       // after registering its ticket.
-      syncDelegateStatus(ctx);
+      syncDelegateStatus(ctx, runtime);
 
       // Keep expensive host deps shared within this dispatch, not indefinitely
       // across dispatches: edits to auth/models/settings/context files must be
@@ -338,6 +344,7 @@ export default function delegateExtension(pi: ExtensionAPI): void {
           signal,
           onUpdate,
           callSpan,
+          runtime,
         });
       } catch (err) {
         failCall();
@@ -360,15 +367,15 @@ export default function delegateExtension(pi: ExtensionAPI): void {
   // 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.
   pi.on("agent_settled", (_event, ctx) => {
-    notifyActiveTicketsOnSettled(ctx);
+    notifyActiveTicketsOnSettled(ctx, runtime);
   });
 
   // Session replacements are cancellable — confirm before killing live work.
   pi.on("session_before_switch", (_event, ctx) =>
-    guardSessionReplacement(ctx, "switch"),
+    guardSessionReplacement(ctx, "switch", runtime),
   );
   pi.on("session_before_fork", (_event, ctx) =>
-    guardSessionReplacement(ctx, "fork"),
+    guardSessionReplacement(ctx, "fork", runtime),
   );
 
   // /tree navigation stays inside the same session: nothing is torn down and
@@ -376,10 +383,12 @@ export default function delegateExtension(pi: ExtensionAPI): void {
   // user moves to. Ask first, and record the new leaf either way so delivery
   // can detect the mismatch (issue #30). `session_tree` also fires for
   // extension-driven ctx.navigateTree, which never reaches the guard.
-  pi.on("session_before_tree", (_event, ctx) => guardTreeNavigation(ctx));
+  pi.on("session_before_tree", (_event, ctx) =>
+    guardTreeNavigation(ctx, runtime),
+  );
   pi.on("session_tree", (event, ctx) => {
     recordTreeNavigation(event.newLeafId);
-    syncDelegateStatus(ctx);
+    syncDelegateStatus(ctx, runtime);
   });
 
   // ── Session shutdown: abort tickets and dispose live pooled sessions ──
@@ -395,7 +404,7 @@ export default function delegateExtension(pi: ExtensionAPI): void {
       // 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();
+      const active = activeTicketSummary(runtime);
       if (active.tickets.length) {
         if (event.reason === "quit") {
           console.error(
@@ -416,15 +425,15 @@ export default function delegateExtension(pi: ExtensionAPI): void {
         }
       }
 
-      for (const ticket of ticketRegistry.values()) {
+      for (const ticket of runtime.tickets.values()) {
         if (ticket.status === "running" || ticket.status === "cancelling") {
-          cancelTicketForShutdown(ticket);
+          runtime.tickets.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);
+      syncDelegateStatus(ctx, runtime);
       // 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
@@ -443,7 +452,7 @@ export default function delegateExtension(pi: ExtensionAPI): void {
       // SQLite still stays open until both cleanup paths finish.
       let poolCleanup: Promise<void>;
       try {
-        poolCleanup = closeAllPooledAgents();
+        poolCleanup = runtime.pool.closeAllPooledAgents();
       } catch (error) {
         console.error("[delegate] pooled-session shutdown start failed", error);
         poolCleanup = Promise.resolve();

package/format.ts

@@ -530,7 +530,10 @@ export function formatCompletedTask(
     if (integration.patchPath)
       parts.push(`full patch: ${integration.patchPath}`);
     if (integration.worktreePath)
-      parts.push(`conflict worktree: ${integration.worktreePath}`);
+      parts.push(`recovery worktree: ${integration.worktreePath}`);
+    if (integration.status === "retained") {
+      parts.push(`not applied: ${integration.reason}`);
+    }
     for (const conflict of integration.conflicts ?? []) {
       parts.push(`conflict: ${conflict.path}: ${conflict.reason}`);
     }

package/host.ts

@@ -85,7 +85,7 @@ export interface HostDepsOptions {
    * host deps are cached per (agentDir + cwd + systemPrompt): the expensive
    * `reload()` (skills, project AGENTS.md discovery) runs once per distinct combo, then
    * is reused across concurrent subagents. Provider-configured or
-   * allowlisted-extension tasks always receive fresh host deps. For ad-hoc
+   * allowlisted-extension tasks always receive fresh host deps. For inline
    * tasks (no named agent) pass undefined to use the discovered prompt.
    */
   systemPrompt?: string;

package/isolated-workspace.ts

@@ -25,6 +25,7 @@ let removeWorktreeHookForTesting:
       destination: string,
     ) => boolean | undefined | Promise<boolean | undefined>)
   | undefined;
+let beforeSourceApplyHookForTesting: (() => void | Promise<void>) | undefined;
 
 /** @internal Keep tests out of the developer's real ~/.pi directory. */
 export function _setIsolatedArtifactRootForTesting(
@@ -44,6 +45,12 @@ export function _setRemoveWorktreeHookForTesting(
   removeWorktreeHookForTesting = hook;
 }
 
+export function _setBeforeSourceApplyHookForTesting(
+  hook: (() => void | Promise<void>) | undefined,
+): void {
+  beforeSourceApplyHookForTesting = hook;
+}
+
 interface CommandResult {
   stdout: string;
   stderr: string;
@@ -698,9 +705,18 @@ interface IsolatedWorker {
   patchPath: string;
 }
 
+export interface IsolatedReconcileOptions {
+  shouldApplySource?: () => boolean;
+  retainedReason?: string;
+  signal?: AbortSignal;
+}
+
 export interface PreparedIsolatedBatch {
   resolved: ResolvedTask[];
-  reconcile(results: TaskResult[]): Promise<TaskResult[]>;
+  reconcile(
+    results: TaskResult[],
+    options?: IsolatedReconcileOptions,
+  ): Promise<TaskResult[]>;
 }
 
 async function restoreAfterFailedApply(
@@ -732,10 +748,45 @@ async function restoreAfterFailedApply(
   }
 }
 
+async function retainAcceptedProposals(
+  group: IsolatedGroup,
+  workers: Map<number, IsolatedWorker>,
+  results: TaskResult[],
+  accepted: Map<number, string[]>,
+  pristineRoot: string,
+  reason: string,
+): Promise<void> {
+  let cleanupIssue:
+    { status: "failed"; reason: string; recoveryPath: string } | undefined;
+  try {
+    await requireWorktreeRemoved(group.sourceRoot, pristineRoot);
+  } catch (error) {
+    cleanupIssue = {
+      status: "failed",
+      reason: error instanceof Error ? error.message : String(error),
+      recoveryPath: pristineRoot,
+    };
+  }
+  for (const [taskIndex, proposedFiles] of accepted) {
+    const worker = workers.get(taskIndex)!;
+    results[taskIndex]!.integration = {
+      status: "retained",
+      reason,
+      proposedFiles,
+      appliedFiles: [],
+      baselineRef: group.baselineRef,
+      proposalRef: worker.proposalRef,
+      patchPath: worker.patchPath,
+      ...(cleanupIssue ? { cleanupIssue } : {}),
+    };
+  }
+}
+
 async function reconcileGroup(
   group: IsolatedGroup,
   workers: Map<number, IsolatedWorker>,
   results: TaskResult[],
+  options: IsolatedReconcileOptions,
 ): Promise<void> {
   let integratedCommit = group.baselineCommit;
   const accepted = new Map<number, string[]>();
@@ -948,6 +999,17 @@ async function reconcileGroup(
     await requireWorktreeRemoved(group.sourceRoot, pristineRoot);
     return;
   }
+  if (options.shouldApplySource && !options.shouldApplySource()) {
+    await retainAcceptedProposals(
+      group,
+      workers,
+      results,
+      accepted,
+      pristineRoot,
+      options.retainedReason ?? "Source application was cancelled.",
+    );
+    return;
+  }
 
   const currentTree = await snapshotTree(
     group.sourceRoot,
@@ -995,6 +1057,18 @@ async function reconcileGroup(
     group.baselineCommit,
     integratedCommit,
   );
+  await beforeSourceApplyHookForTesting?.();
+  if (options.shouldApplySource && !options.shouldApplySource()) {
+    await retainAcceptedProposals(
+      group,
+      workers,
+      results,
+      accepted,
+      pristineRoot,
+      options.retainedReason ?? "Source application was cancelled.",
+    );
+    return;
+  }
   const applyIndex = path.join(group.artifactRoot, "apply.index");
   await fs.promises.rm(applyIndex, { force: true });
   const env = {
@@ -1010,12 +1084,25 @@ async function reconcileGroup(
       cwd: group.sourceRoot,
       env,
     });
+    if (options.shouldApplySource && !options.shouldApplySource()) {
+      await retainAcceptedProposals(
+        group,
+        workers,
+        results,
+        accepted,
+        pristineRoot,
+        options.retainedReason ?? "Source application was cancelled.",
+      );
+      return;
+    }
     await git(["apply", "--binary", "--index", finalPatch], {
       cwd: group.sourceRoot,
       env,
+      signal: options.signal,
     });
   } catch (error) {
     const recoveryRoot = path.join(group.artifactRoot, "failed-apply-files");
+    let rollbackSucceeded = true;
     try {
       await restoreAfterFailedApply(
         group.sourceRoot,
@@ -1024,11 +1111,23 @@ async function reconcileGroup(
         recoveryRoot,
       );
     } catch (rollbackError) {
+      rollbackSucceeded = false;
       console.error(
         "[delegate] isolated apply rollback failed; recovery artifacts retained",
         rollbackError,
       );
     }
+    if (options.signal?.aborted && rollbackSucceeded) {
+      await retainAcceptedProposals(
+        group,
+        workers,
+        results,
+        accepted,
+        pristineRoot,
+        options.retainedReason ?? "Source application was cancelled.",
+      );
+      return;
+    }
     for (const [taskIndex] of accepted) {
       const integration = results[taskIndex]!.integration!;
       const worker = workers.get(taskIndex)!;
@@ -1145,24 +1244,35 @@ async function cleanupCompletedGroupRefs(
   group: IsolatedGroup,
   workers: Map<number, IsolatedWorker>,
   results: readonly TaskResult[],
-): Promise<void> {
+): Promise<boolean> {
   const disposableProposalRefs: string[] = [];
   let retainsRecoveryArtifacts = false;
+  let retainsPrivateRefs = false;
 
   for (const taskIndex of group.taskIndexes) {
-    const status = results[taskIndex]?.integration?.status;
-    if (status === "conflict" || status === "apply_failed") {
+    const integration = results[taskIndex]?.integration;
+    const status = integration?.status;
+    if (!integration) {
+      retainsRecoveryArtifacts = true;
+      retainsPrivateRefs = true;
+    } else if (
+      status === "retained" ||
+      status === "conflict" ||
+      status === "apply_failed"
+    ) {
       retainsRecoveryArtifacts = true;
+      retainsPrivateRefs = true;
     } else if (
       status === "applied_unverified" ||
       status === "no_changes" ||
       status === "discarded"
     ) {
+      if (integration.cleanupIssue) retainsRecoveryArtifacts = true;
       disposableProposalRefs.push(workers.get(taskIndex)!.proposalRef);
     }
   }
 
-  const refs = retainsRecoveryArtifacts
+  const refs = retainsPrivateRefs
     ? disposableProposalRefs
     : [...disposableProposalRefs, group.baselineRef];
   try {
@@ -1172,6 +1282,7 @@ async function cleanupCompletedGroupRefs(
     // turn a successfully applied source change into a reported failure.
     console.error("[delegate] failed to clean completed isolated refs", error);
   }
+  return retainsRecoveryArtifacts;
 }
 
 /** Prepare detached worktrees from one synthetic commit per Git root. The
@@ -1315,11 +1426,14 @@ export async function prepareIsolatedBatch(
 
   return {
     resolved: translated,
-    async reconcile(results: TaskResult[]): Promise<TaskResult[]> {
+    async reconcile(
+      results: TaskResult[],
+      options: IsolatedReconcileOptions = {},
+    ): Promise<TaskResult[]> {
       for (const group of groupsByRoot.values()) {
         try {
           try {
-            await reconcileGroup(group, workers, results);
+            await reconcileGroup(group, workers, results, options);
           } catch (error) {
             console.error(
               "[delegate] isolated group reconciliation failed",
@@ -1332,13 +1446,45 @@ export async function prepareIsolatedBatch(
               error,
             );
           }
-          await cleanupCompletedGroupRefs(group, workers, results);
+          const retainsRecoveryArtifacts = await cleanupCompletedGroupRefs(
+            group,
+            workers,
+            results,
+          );
+          if (!retainsRecoveryArtifacts) {
+            try {
+              await fs.promises.rm(group.artifactRoot, {
+                recursive: true,
+                force: true,
+              });
+            } catch (error) {
+              console.error(
+                `[delegate] failed to remove completed isolated artifacts '${group.artifactRoot}'`,
+                error,
+              );
+            }
+          }
         } finally {
           // Unblocks and serializes any safety-confirmed quarantine cleanups
           // only after all candidate/source/ref work for this group is done.
           group.finishReconciliation();
         }
       }
+      try {
+        await fs.promises.rmdir(batchArtifactRoot);
+      } catch (error) {
+        if (!(
+          error instanceof Error &&
+          "code" in error &&
+          ((error as NodeJS.ErrnoException).code === "ENOENT" ||
+            (error as NodeJS.ErrnoException).code === "ENOTEMPTY")
+        )) {
+          console.error(
+            `[delegate] failed to remove isolated batch directory '${batchArtifactRoot}'`,
+            error,
+          );
+        }
+      }
       return results;
     },
   };

package/lifecycle.ts

@@ -14,8 +14,7 @@ import type {
   TaskRunEnv,
   ToolActivity,
 } from "./types.ts";
-import * as pool from "./pool.ts";
-import { isSessionBusy } from "./tickets.ts";
+import { getDefaultDelegateRuntime, type DelegateRuntime } from "./runtime.ts";
 import {
   createSubagentSessionManager,
   persistSessionHeader,
@@ -50,10 +49,6 @@ let runAgentSessionForTesting: RunAgentSession = runAgentSession;
 type CreateScratchWorkspace = typeof createScratchWorkspace;
 let createScratchWorkspaceForTesting: CreateScratchWorkspace =
   createScratchWorkspace;
-type DetachQuarantinedPooledSession =
-  typeof pool._quarantinePooledAgentWithoutDisposal;
-let detachQuarantinedPooledSessionForTesting: DetachQuarantinedPooledSession =
-  pool._quarantinePooledAgentWithoutDisposal;
 
 export function _setRunAgentSessionForTesting(
   override: RunAgentSession | undefined,
@@ -70,10 +65,11 @@ export function _setCreateScratchWorkspaceForTesting(
 
 /** @internal Simulate a pooled-detachment invariant failure in lifecycle tests. */
 export function _setQuarantinePooledSessionDetachForTesting(
-  override: DetachQuarantinedPooledSession | undefined,
+  override:
+    ((sessionId: string, expectedSession: AgentSession) => boolean) | undefined,
+  runtime: DelegateRuntime = getDefaultDelegateRuntime(),
 ): void {
-  detachQuarantinedPooledSessionForTesting =
-    override ?? pool._quarantinePooledAgentWithoutDisposal;
+  runtime.pool._setQuarantinePooledAgentWithoutDisposalForTesting(override);
 }
 
 /**
@@ -242,14 +238,16 @@ function disposeSession(session: AgentSession, description: string): void {
 /** Detach an abandoned session from every owner immediately, then dispose it
  * only after runner's background termination monitor proves quiescence. */
 function quarantineAcquiredSession(
+  env: TaskRunEnv,
   task: ResolvedTask,
   acquired: AcquiredSession,
   quarantine: SessionQuarantine,
 ): void {
   let mayDisposeAfterSafety = acquired.lifecycleOwnsSession;
   if (!acquired.lifecycleOwnsSession) {
+    const runtime = env.runtime!;
     const detached = task.sessionId
-      ? detachQuarantinedPooledSessionForTesting(
+      ? runtime.pool.quarantinePooledAgentWithoutDisposal(
           task.sessionId,
           acquired.session,
         )
@@ -617,11 +615,12 @@ type AcquireResult = AcquiredSession | { error: TaskResult };
  * checkout is pure (no lastUsed bump); lastUsed is bumped by commit().
  */
 function checkoutPooledSession(
+  env: TaskRunEnv,
   task: ResolvedTask,
   p: TaskProgress,
 ): AcquireResult | undefined {
   if (!task.sessionId) return undefined;
-  const co = pool.checkout(task.sessionId, {
+  const co = env.runtime!.pool.checkout(task.sessionId, {
     cwd: task.cwd,
     thinking: task.thinking,
     tools: task.tools,
@@ -762,7 +761,7 @@ async function acquireAgentSession(
   p: TaskProgress,
 ): Promise<AcquireResult> {
   if (task.sessionId) {
-    const pooled = checkoutPooledSession(task, p);
+    const pooled = checkoutPooledSession(env, task, p);
     if (pooled) return pooled;
   }
   if (task.resumeFrom) return resumeFromSessionFile(env, task, task.resumeFrom);
@@ -813,6 +812,10 @@ export async function runResolvedTask(
   p: TaskProgress,
   taskIndex: number,
 ): Promise<TaskResult> {
+  // Ensure every lifecycle call operates on an explicit runtime. Older callers
+  // (and some test fixtures) do not inject one, so the default runtime is the
+  // backward-compatible fallback.
+  env.runtime ??= getDefaultDelegateRuntime();
   return withResumeTranscriptLock(task.resumeFrom, async (transcript) => {
     let executionTask = task;
     if (task.resumeFrom) {
@@ -879,8 +882,9 @@ export async function runResolvedTask(
       return runResolvedTaskUnlocked(env, executionTask, p, taskIndex);
     };
 
+    const runtime = env.runtime!;
     return executionTask.sessionId
-      ? pool.withSessionLock(executionTask.sessionId, runLocked)
+      ? runtime.pool.withSessionLock(executionTask.sessionId, runLocked)
       : runLocked();
   });
 }
@@ -1198,7 +1202,10 @@ async function applySessionAction(
     // 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);
+    const runtime = env.runtime!;
+    const closed = await runtime.pool.closePooledAgentWithoutLock(
+      task.sessionId,
+    );
     return finishTask(
       env,
       p,
@@ -1213,12 +1220,13 @@ async function applySessionAction(
   }
 
   if (task.sessionAction === "list") {
+    const runtime = env.runtime!;
     return finishTask(
       env,
       p,
       completeSessionAction(
         task,
-        `Active sessions:\n${pool.listPooledAgents().join("\n")}`,
+        `Active sessions:\n${runtime.pool.listPooledAgents().join("\n")}`,
         Date.now() - env.delegateStartedAt,
       ),
     );
@@ -1231,7 +1239,7 @@ function busySessionConflict(
   task: ResolvedTask,
 ): TaskResult | undefined {
   if (!task.sessionId) return undefined;
-  const busyTicketId = isSessionBusy(task.sessionId);
+  const busyTicketId = env.runtime!.tickets.isSessionBusy(task.sessionId);
   if (busyTicketId && busyTicketId !== env.ticketId) {
     return failTask(
       task,
@@ -1303,6 +1311,7 @@ function noteAttemptProgress(
 
 /** Commit, record, or evict a pooled session after one prompt attempt. */
 async function settlePooledAttempt(
+  env: TaskRunEnv,
   task: ResolvedTask,
   acquired: AcquiredSession,
   r: {
@@ -1313,6 +1322,7 @@ async function settlePooledAttempt(
   },
   sessionReleased: boolean,
 ): Promise<boolean> {
+  const runtime = env.runtime!;
   if (!task.sessionId) return sessionReleased;
   if (acquired.lifecycleOwnsSession) {
     // Pool misses (including resumeFrom) transfer ownership only on
@@ -1323,7 +1333,7 @@ async function settlePooledAttempt(
       r.failureKind !== "stalled" &&
       r.failureKind !== "deadline_exceeded"
     ) {
-      const committed = pool.commit(task.sessionId, {
+      const committed = runtime.pool.commit(task.sessionId, {
         session: acquired.session,
         sessionManager: acquired.sessionManager,
         sessionFile: acquired.sessionFile,
@@ -1353,7 +1363,7 @@ async function settlePooledAttempt(
   ) {
     try {
       return (
-        (await pool._closePooledAgentWithoutLock(task.sessionId)) ||
+        (await runtime.pool.closePooledAgentWithoutLock(task.sessionId)) ||
         sessionReleased
       );
     } catch (error) {
@@ -1371,7 +1381,7 @@ async function settlePooledAttempt(
   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);
+    runtime.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
@@ -1558,12 +1568,13 @@ async function runTaskAttempt(
         );
 
     if (quarantine) {
-      quarantineAcquiredSession(task, acquired, quarantine);
+      quarantineAcquiredSession(env, task, acquired, quarantine);
       // Neither lifecycle nor pool owns it now. The deferred safety callback is
       // the sole owner and finally below must not dispose it early.
       sessionReleased = true;
     } else {
       sessionReleased = await settlePooledAttempt(
+        env,
         task,
         acquired,
         r,