@mystilleef/pi-subagent 0.10.2 → 0.11.0

package/src/child/process.ts

@@ -5,38 +5,24 @@
  */
 
 import { type ChildProcess, spawn } from "node:child_process";
-import * as fs from "node:fs";
 import readline from "node:readline";
-import {
-  clampThinkingLevel,
-  getModel,
-  getSupportedThinkingLevels,
-  type Message,
-  type ModelThinkingLevel,
-} from "@earendil-works/pi-ai";
+import type { Message } from "@earendil-works/pi-ai";
 import type { AgentConfig, ThinkingLevel } from "../agent/agents.js";
 import { getFinalOutput } from "../output/ui.js";
-import { makeToolPreview, renderToolActivity } from "../progress/progress.js";
-import { SENSITIVE_PATTERN } from "../progress/progress-format.js";
-import { isToolCallPart } from "../progress/progress-state.js";
+import { serializeSamplingParams } from "../shared/sampling.js";
 import {
   type OnUpdateCallback,
   type SingleResult,
-  type StreamingProgress,
   type SubagentDetails,
   TOOL_RESULT_FAILED_MESSAGE,
-  type ToolActivity,
 } from "../shared/types.js";
 import {
   detectMessageError,
-  findLastAssistantTextMessage,
   getPiInvocation,
   getSubagentDepth,
   getSubagentRuntimeLimits,
   resolveAgentSkillArgs,
   subagentDepthEnv,
-  truncateOutput,
-  writePromptToTempFile,
 } from "../shared/utils.js";
 import {
   type ChildEventParseResult,
@@ -44,7 +30,34 @@ import {
   parseChildEventLine,
   TOOL_EXECUTION_UPDATE_EVENT,
 } from "./child-events.js";
+import { getLatestOutcomeFromMessages } from "./complete-outcome.js";
+import {
+  buildModelDisplay,
+  type ChildModelSettings,
+  resolveEffectiveChildModelSettings,
+  resolveThinkingLevel,
+} from "./model-resolution.js";
+import {
+  appendWithByteLimit,
+  resolveCompleteExtensionPath,
+  resolveSamplingExtensionPath,
+} from "./process-utils.js";
 import { appendSubagentResultContract } from "./prompt-contract.js";
+import {
+  beginPromptSetup,
+  cleanupPromptSetupResult,
+  cleanupTempPrompt,
+} from "./prompt-setup.js";
+import {
+  addMessageToResult,
+  createErrorResult,
+  errorForDepthLimit,
+  errorForUnknownAgent,
+  initRuntimeResult,
+  type RuntimeResult,
+  rebuildResultFromMessages,
+} from "./result-builder.js";
+import { type EmitUpdateFn, makeEmitUpdate } from "./streaming-progress.js";
 import {
   acquireChildSleepInhibitor,
   getProcessTreeSpawnOptions,
@@ -54,35 +67,13 @@ import {
   terminateChildProcess,
 } from "./termination.js";
 
-export function resolveThinkingLevel(
-  requested: ThinkingLevel,
-  provider: string,
-  modelId: string,
-): { level: ThinkingLevel; warning?: string } {
-  const model = getModel(provider as never, modelId as never);
-  if (!model) return { level: requested };
-  const mkWarning = (effective: ThinkingLevel) =>
-    `Thinking level "${requested}" not supported by model "${provider}/${modelId}"; using "${effective}" instead`;
-  if (model.reasoning === false) {
-    return { level: "off", warning: mkWarning("off") };
-  }
-  if (!model.thinkingLevelMap) return { level: requested };
-  const supported = getSupportedThinkingLevels(model);
-  if (supported.length === 0) return { level: requested };
-  const clamped = clampThinkingLevel(
-    model,
-    requested as ModelThinkingLevel,
-  ) as ThinkingLevel;
-  if (clamped === requested) return { level: requested };
-  return { level: clamped, warning: mkWarning(clamped) };
-}
+const COMPLETE_EXTENSION_PATH = resolveCompleteExtensionPath();
+const SAMPLING_EXTENSION_PATH = resolveSamplingExtensionPath();
+
+export { resolveThinkingLevel } from "./model-resolution.js";
+export { makeEmitUpdate } from "./streaming-progress.js";
 
 type RuntimeLimits = ReturnType<typeof getSubagentRuntimeLimits>;
-type RuntimeResult = SingleResult & { messages: Message[] };
-type ChildModelSettings = {
-  provider?: string | undefined;
-  id?: string | undefined;
-};
 type SleepInhibitorAcquirer = (pid: number) => Promise<SleepInhibitorHandle>;
 
 type RunSingleAgentOptions = {
@@ -90,18 +81,9 @@ type RunSingleAgentOptions = {
   getOrchestratorPid?: () => unknown;
 };
 
-export class SubagentAbortError extends Error {
-  readonly result: SingleResult;
-  constructor(result: SingleResult) {
-    super("Subagent was aborted");
-    this.name = "SubagentAbortError";
-    this.result = result;
-  }
-}
-
-type TempPrompt = { dir: string; filePath: string };
-
-type PromptSetupResult = { tmpPrompt: TempPrompt | null } | { error: unknown };
+export type RunSingleAgentResult =
+  | { kind: "completed"; result: SingleResult }
+  | { kind: "aborted"; result: SingleResult };
 
 interface SubagentState {
   result: RuntimeResult;
@@ -112,53 +94,6 @@ interface SubagentState {
   terminationPromise?: Promise<unknown>;
 }
 
-function appendWithByteLimit(
-  current: string,
-  data: string | Buffer,
-  max: number,
-): string {
-  const currentBytes = Buffer.from(current, "utf-8");
-  if (currentBytes.length >= max) return current;
-  const incomingBytes = Buffer.isBuffer(data)
-    ? data
-    : Buffer.from(data, "utf-8");
-  const combined = Buffer.concat([currentBytes, incomingBytes]);
-  if (combined.length <= max) return combined.toString("utf-8");
-  return truncateValidUtf8(combined, max);
-}
-
-function truncateValidUtf8(buffer: Buffer, max: number): string {
-  let end = Math.min(max, buffer.length);
-  while (end > 0) {
-    const candidate = buffer.subarray(0, end).toString("utf-8");
-    if (!candidate.endsWith("�")) return candidate;
-    end -= 1;
-  }
-  return "";
-}
-
-/**
- * Rationale: Subagent usage reporting needs context window awareness to provide
- * meaningful "context full" indicators to the parent.
- */
-function resolveContextWindowTokens(msg: Message): number | undefined {
-  const m = msg as unknown as Record<string, unknown>;
-  if (typeof m["provider"] !== "string" || typeof m["model"] !== "string")
-    return;
-  try {
-    const contextWindow = getModel(
-      m["provider"] as never,
-      m["model"] as never,
-    )?.contextWindow;
-    return Number.isFinite(contextWindow) && contextWindow > 0
-      ? contextWindow
-      : undefined;
-  } catch {
-    /* model lookup failures return undefined to skip context window tracking */
-    return;
-  }
-}
-
 function getAbortReason(signal: AbortSignal): string {
   const { reason } = signal;
   if (reason instanceof Error && reason.message) return reason.message;
@@ -216,15 +151,13 @@ function startSleepInhibitorRelease(
   return acquisitionPromise.then(releaseSleepInhibitor, () => {});
 }
 
-function hasCompletedAgentOutput(result: RuntimeResult): boolean {
+export function hasCompletedAgentOutput(
+  result: RuntimeResult,
+  outcome?: string,
+): boolean {
   if (result.finalOutput.trim()) return true;
-  return result.messages.some(
-    (msg) =>
-      msg.role === "assistant" &&
-      msg.content.some(
-        (part) => part.type === "text" && Boolean(part.text?.trim()),
-      ),
-  );
+  if (outcome?.trim()) return true;
+  return getFinalOutput(result.messages).trim().length > 0;
 }
 
 /**
@@ -232,15 +165,16 @@ function hasCompletedAgentOutput(result: RuntimeResult): boolean {
  * we force-kill them after a grace period and treat it as success (0) if they
  * actually produced output.
  */
-function getAgentEndTimeoutExitCode(
+export function getAgentEndTimeoutExitCode(
   result: RuntimeResult,
   spawnError: Error | undefined,
+  outcome: string | undefined,
 ): number | undefined {
   if (result.termination?.cancelReason !== "agent_end_timeout") return;
   if (spawnError) return;
   if (result.stopReason === "error" || result.stopReason === "aborted") return;
   if (result.errorMessage?.trim()) return;
-  return hasCompletedAgentOutput(result) ? 0 : 1;
+  return hasCompletedAgentOutput(result, outcome) ? 0 : 1;
 }
 
 /**
@@ -303,287 +237,6 @@ async function waitForSubagentProcess(
   });
 }
 
-function buildModelDisplay(
-  effectiveModel: ChildModelSettings,
-  thinking: ThinkingLevel,
-): string | undefined {
-  const parts: string[] = [];
-  if (effectiveModel.provider) {
-    parts.push(effectiveModel.provider);
-  }
-  if (effectiveModel.id) {
-    parts.push(effectiveModel.id);
-  }
-  if (thinking) {
-    parts.push(thinking);
-  }
-  return parts.length > 0 ? parts.join(" ･ ") : undefined;
-}
-
-const EMPTY_USAGE = {
-  input: 0,
-  output: 0,
-  cacheRead: 0,
-  cacheWrite: 0,
-  cost: 0,
-  contextTokens: 0,
-  turns: 0,
-};
-
-function initRuntimeResult(
-  agentName: string,
-  source: "user" | "project" | "unknown",
-  task: string,
-  modelDisplay: string | undefined,
-): RuntimeResult {
-  return {
-    agent: agentName,
-    agentSource: source,
-    task,
-    exitCode: 0,
-    finalOutput: "",
-    messages: [],
-    stderr: "",
-    usage: { ...EMPTY_USAGE },
-    model: modelDisplay,
-  };
-}
-
-function accumulateUsage(result: RuntimeResult, msg: Message): void {
-  if (msg.role !== "assistant") return;
-  result.usage.turns++;
-  const { usage } = msg;
-  if (!usage) return;
-  result.usage.input += usage.input || 0;
-  result.usage.output += usage.output || 0;
-  result.usage.cacheRead += usage.cacheRead || 0;
-  result.usage.cacheWrite += usage.cacheWrite || 0;
-  result.usage.cost += usage.cost?.total || 0;
-  result.usage.contextTokens = usage.totalTokens || 0;
-  const ctxWindowTokens = resolveContextWindowTokens(msg);
-  if (ctxWindowTokens !== undefined)
-    result.usage.contextWindowTokens = ctxWindowTokens;
-}
-
-function addMessageToResult(result: RuntimeResult, msg: Message): void {
-  result.messages.push(msg);
-  result.finalOutput = truncateOutput(getFinalOutput(result.messages));
-  if (msg.role === "toolResult" && msg.isError) {
-    result.errorMessage ||= TOOL_RESULT_FAILED_MESSAGE;
-  } else if (result.errorMessage === TOOL_RESULT_FAILED_MESSAGE) {
-    delete result.errorMessage;
-  }
-  if (msg.role === "assistant") {
-    accumulateUsage(result, msg);
-    if (!result.model && msg.model) result.model = msg.model;
-    if (msg.stopReason) result.stopReason = msg.stopReason;
-    if (msg.errorMessage) result.errorMessage = msg.errorMessage;
-  }
-}
-
-function createErrorResult(
-  agent: string,
-  source: "user" | "project" | "unknown",
-  task: string,
-  error: string,
-  model?: string,
-): SingleResult {
-  return {
-    agent,
-    agentSource: source,
-    task,
-    exitCode: 1,
-    finalOutput: "",
-    stderr: error,
-    usage: { ...EMPTY_USAGE },
-    model,
-  };
-}
-
-function errorForUnknownAgent(
-  agentName: string,
-  agents: AgentConfig[],
-  task: string,
-): SingleResult {
-  const available = agents.map((a) => `"${a.name}"`).join(", ") || "none";
-  return createErrorResult(
-    agentName,
-    "unknown",
-    task,
-    `Unknown agent: "${agentName}". Available agents: ${available}.`,
-  );
-}
-
-function errorForDepthLimit(
-  agentName: string,
-  source: "user" | "project" | "unknown",
-  task: string,
-  depth: number,
-  maxDepth: number,
-  model?: string,
-): SingleResult {
-  return createErrorResult(
-    agentName,
-    source,
-    task,
-    `Subagent nesting limit reached (depth ${depth}/${maxDepth}).`,
-    model,
-  );
-}
-
-async function cleanupTempPrompt(tmpPrompt: TempPrompt): Promise<void> {
-  try {
-    await fs.promises.unlink(tmpPrompt.filePath);
-    await fs.promises.rmdir(tmpPrompt.dir);
-  } catch {
-    /* temp file cleanup failures are non-fatal; OS will clean up eventually */
-  }
-}
-
-function beginPromptSetup(agent: AgentConfig): Promise<PromptSetupResult> {
-  if (!agent.systemPrompt.trim()) return Promise.resolve({ tmpPrompt: null });
-  return writePromptToTempFile(agent.name, agent.systemPrompt).then(
-    (tmpPrompt) => ({ tmpPrompt }),
-    (error: unknown) => ({ error }),
-  );
-}
-
-async function cleanupPromptSetupResult(
-  setup: PromptSetupResult,
-): Promise<void> {
-  if ("tmpPrompt" in setup && setup.tmpPrompt) {
-    await cleanupTempPrompt(setup.tmpPrompt);
-  }
-}
-
-function findRecentMessagesAnchor(messages: Message[]): number {
-  return findLastAssistantTextMessage(messages);
-}
-
-function deriveStreamingProgress(messages: Message[]): StreamingProgress {
-  const toolCalls: { id: string; preview: string }[] = [];
-  let lastToolPreview: string | undefined;
-  let activeToolActivity: ToolActivity | undefined;
-  for (const msg of messages) {
-    if (msg.role !== "assistant" || !Array.isArray(msg.content)) continue;
-    for (const part of msg.content) {
-      if (!isToolCallPart(part)) continue;
-      const preview = sanitizeProgressPreview(
-        makeToolPreview(part.name, part.arguments),
-        part.name,
-      );
-      toolCalls.push({ id: part.id, preview });
-      lastToolPreview = preview;
-      activeToolActivity = { toolName: part.name, inputSummary: preview };
-    }
-  }
-  const activityText = renderToolActivity(activeToolActivity);
-  return {
-    activeToolActivity,
-    activityText,
-    toolCalls,
-    lastToolPreview,
-  };
-}
-
-function sanitizeProgressPreview(preview: string, toolName: string): string {
-  return SENSITIVE_PATTERN.test(preview) ? toolName : preview;
-}
-
-/**
- * Merge incoming tool activity with existing progress activity.
- * When tool names match, prefer richer inputSummary from incoming.
- * Otherwise, replace entirely with incoming activity.
- */
-function mergeToolActivity(
-  existing: ToolActivity | undefined,
-  incoming: ToolActivity,
-): ToolActivity {
-  if (existing && existing.toolName === incoming.toolName) {
-    const incomingSummary = incoming.inputSummary;
-    const preferIncoming =
-      incomingSummary && incomingSummary !== incoming.toolName;
-    return {
-      ...existing,
-      inputSummary: preferIncoming ? incomingSummary : existing.inputSummary,
-      instanceName: incoming.instanceName ?? existing.instanceName,
-      child: incoming.child ?? existing.child,
-    };
-  }
-  return incoming;
-}
-
-/**
- * Apply tool activity and result completion updates to progress state.
- * Handles merging of child events with parent activity tree.
- */
-function applyActivityUpdates(
-  progress: StreamingProgress,
-  options: { toolActivity?: ToolActivity; toolResultCompleted?: boolean },
-  previousActivity?: ToolActivity,
-): void {
-  // Preserve stored activity tree for tool-result completion signals
-  // so the parent retains nested context until newer activity arrives
-  if (options.toolResultCompleted && previousActivity) {
-    progress.activeToolActivity = previousActivity;
-    const renderedText = renderToolActivity(previousActivity);
-    if (renderedText !== undefined) progress.activityText = renderedText;
-    else delete progress.activityText;
-  }
-  // Handle parsed tool activity from child events
-  // Merge with parent activity if this is a nested update
-  if (options.toolActivity) {
-    progress.activeToolActivity = mergeToolActivity(
-      progress.activeToolActivity,
-      options.toolActivity,
-    );
-    const renderedActivity = renderToolActivity(progress.activeToolActivity);
-    if (renderedActivity !== undefined)
-      progress.activityText = renderedActivity;
-    else delete progress.activityText;
-  }
-  if (options.toolResultCompleted) {
-    progress.toolResultCompleted = true;
-  }
-}
-
-export function makeEmitUpdate(
-  result: RuntimeResult,
-  onUpdate: OnUpdateCallback | undefined,
-  makeDetails: (
-    results: RuntimeResult[],
-    options?: { includeMessages?: boolean; recentMessages?: Message[] },
-  ) => SubagentDetails,
-): (options?: {
-  toolActivity?: ToolActivity;
-  toolResultCompleted?: boolean;
-}) => void {
-  return (options) => {
-    const msgs = result.messages;
-    const anchorIdx = findRecentMessagesAnchor(msgs);
-    const recentMessages =
-      anchorIdx >= 0 ? msgs.slice(anchorIdx) : msgs.slice(-5);
-    const progress = deriveStreamingProgress(msgs);
-    if (options) {
-      applyActivityUpdates(
-        progress,
-        options,
-        result.progress?.activeToolActivity,
-      );
-    }
-    result.progress = progress;
-    onUpdate?.({
-      content: [
-        { type: "text", text: progress.activityText ?? "(running...)" },
-      ],
-      details: makeDetails([result], {
-        includeMessages: true,
-        recentMessages,
-      }),
-    });
-  };
-}
-
 function makeRequestTerminator(
   proc: ChildProcess,
   terminateOptions: {
@@ -613,10 +266,7 @@ function clearGraceTimer(state: SubagentState): void {
 function handleMessageEvent(
   event: ChildKnownEvent,
   state: SubagentState,
-  emitUpdate: (options?: {
-    toolActivity?: ToolActivity;
-    toolResultCompleted?: boolean;
-  }) => void,
+  emitUpdate: EmitUpdateFn,
 ): void {
   if (event.type !== "message_end" && event.type !== "tool_result_end") return;
   if (event.message) {
@@ -628,10 +278,7 @@ function handleMessageEvent(
 
 function handleToolExecutionUpdateEvent(
   event: ChildKnownEvent,
-  emitUpdate: (options?: {
-    toolActivity?: ToolActivity;
-    toolResultCompleted?: boolean;
-  }) => void,
+  emitUpdate: EmitUpdateFn,
 ): void {
   if (event.type !== TOOL_EXECUTION_UPDATE_EVENT) return;
   emitUpdate({ toolActivity: event.toolActivity });
@@ -640,17 +287,12 @@ function handleToolExecutionUpdateEvent(
 function handleAgentEndEvent(
   event: ChildKnownEvent,
   state: SubagentState,
-  emitUpdate: (options?: {
-    toolActivity?: ToolActivity;
-    toolResultCompleted?: boolean;
-  }) => void,
+  emitUpdate: EmitUpdateFn,
   requestTermination: (reason: string) => Promise<unknown>,
 ): void {
   if (event.type !== "agent_end") return;
-  if (state.result.messages.length === 0 && Array.isArray(event.messages)) {
-    for (const msg of event.messages as Message[]) {
-      addMessageToResult(state.result, msg);
-    }
+  if (Array.isArray(event.messages) && event.messages.length > 0) {
+    rebuildResultFromMessages(state.result, event.messages as Message[]);
     emitUpdate();
   }
   if (state.agentEndGraceTimer || state.terminationPromise) return;
@@ -677,10 +319,7 @@ function formatUnknownEventDiagnostic(
 function processEventLine(
   line: string,
   state: SubagentState,
-  emitUpdate: (options?: {
-    toolActivity?: ToolActivity;
-    toolResultCompleted?: boolean;
-  }) => void,
+  emitUpdate: EmitUpdateFn,
   requestTermination: (reason: string) => Promise<unknown>,
   debugEventDiagnostics: boolean,
 ): void {
@@ -719,16 +358,11 @@ function setupAbortHandler(
   return onAbort;
 }
 
-function resolveEffectiveChildModelSettings(
-  agent: AgentConfig,
-  parentModel: ChildModelSettings | undefined,
-): ChildModelSettings {
-  return {
-    provider: agent.provider ?? parentModel?.provider,
-    id:
-      agent.model ??
-      (agent.provider === undefined ? parentModel?.id : undefined),
-  };
+function buildSamplingEnv(agent: AgentConfig): string | undefined {
+  return serializeSamplingParams({
+    temperature: agent.temperature,
+    topP: agent.topP,
+  });
 }
 
 function buildPiArgs(
@@ -744,11 +378,16 @@ function buildPiArgs(
     args.push("--provider", effectiveModel.provider);
   if (effectiveModel.id) args.push("--model", effectiveModel.id);
   args.push("--thinking", thinking);
-  if (agent.tools?.length) args.push("--tools", agent.tools.join(","));
+  if (agent.tools) {
+    const tools = new Set(agent.tools);
+    tools.add("complete");
+    args.push("--tools", [...tools].join(","));
+  }
   if (agent.skills) args.push("--no-skills", ...resolvedSkills.args);
   if (tmpPrompt) {
     args.push("--append-system-prompt", tmpPrompt.filePath);
   }
+  args.push("--extension", COMPLETE_EXTENSION_PATH);
   const taskPrompt = task
     ? `Task: ${task}`
     : "Run according to your system prompt. If no explicit task was provided, use the default context described there.";
@@ -756,13 +395,19 @@ function buildPiArgs(
   return args;
 }
 
+function buildChildEnv(samplingEnv: string | undefined): NodeJS.ProcessEnv {
+  const { PI_SAMPLING_PARAMS: _, ...parentEnv } = process.env;
+  return {
+    ...parentEnv,
+    ...subagentDepthEnv(),
+    ...(samplingEnv ? { PI_SAMPLING_PARAMS: samplingEnv } : {}),
+  };
+}
+
 function setupChildProcess(
   proc: ChildProcess,
   state: SubagentState,
-  emitUpdate: (options?: {
-    toolActivity?: ToolActivity;
-    toolResultCompleted?: boolean;
-  }) => void,
+  emitUpdate: EmitUpdateFn,
   requestTermination: (reason: string) => Promise<unknown>,
   debugEventDiagnostics: boolean,
 ): void {
@@ -801,17 +446,16 @@ function setupChildProcess(
 async function finalizeResult(
   state: SubagentState,
   startedAt: number,
-): Promise<SingleResult> {
+): Promise<RunSingleAgentResult> {
   state.result.durationMs = Date.now() - startedAt;
   clearGraceTimer(state);
   if (state.terminationPromise) await state.terminationPromise;
   if (state.spawnError) state.result.exitCode = 1;
-  if (detectMessageError(state.result.messages)) {
-    state.result.errorMessage ||= TOOL_RESULT_FAILED_MESSAGE;
-  }
+  const outcome = getLatestOutcomeFromMessages(state.result.messages);
   const agentEndTimeoutExitCode = getAgentEndTimeoutExitCode(
     state.result,
     state.spawnError,
+    outcome,
   );
   if (agentEndTimeoutExitCode !== undefined) {
     state.result.exitCode = agentEndTimeoutExitCode;
@@ -819,11 +463,25 @@ async function finalizeResult(
   if (state.result.termination?.cancelReason === "agent_end_timeout") {
     state.result.stderr = state.result.stderr.replace(/^Terminated\r?\n/gm, "");
   }
+  // Skip unresolved-error detection when complete succeeded: intermediate tool
+  // errors (e.g. a linter exiting non-zero) are expected and already addressed.
+  if (!outcome && detectMessageError(state.result.messages)) {
+    state.result.errorMessage ||= TOOL_RESULT_FAILED_MESSAGE;
+  }
   if (state.wasAborted) {
     state.result.stderr = "";
-    throw new SubagentAbortError(state.result);
+    return { kind: "aborted", result: state.result };
+  }
+  const isSemanticallySuccessful =
+    !state.spawnError &&
+    state.result.exitCode === 0 &&
+    (Boolean(outcome) ||
+      (!state.result.errorMessage?.trim() &&
+        !detectMessageError(state.result.messages)));
+  if (isSemanticallySuccessful && outcome) {
+    state.result.outcome = outcome;
   }
-  return state.result;
+  return { kind: "completed", result: state.result };
 }
 
 /**
@@ -854,19 +512,26 @@ export async function runSingleAgent(
   parentThinking: ThinkingLevel,
   debugEventDiagnostics = false,
   options: RunSingleAgentOptions = {},
-): Promise<SingleResult> {
+): Promise<RunSingleAgentResult> {
   const agent = agents.find((a) => a.name === agentName);
-  if (!agent) return errorForUnknownAgent(agentName, agents, task);
+  if (!agent)
+    return {
+      kind: "completed",
+      result: errorForUnknownAgent(agentName, agents, task),
+    };
   const runtimeLimits = getSubagentRuntimeLimits();
   const depth = getSubagentDepth();
   if (depth >= runtimeLimits.maxDepth) {
-    return errorForDepthLimit(
-      agentName,
-      agent.source,
-      task,
-      depth,
-      runtimeLimits.maxDepth,
-    );
+    return {
+      kind: "completed",
+      result: errorForDepthLimit(
+        agentName,
+        agent.source,
+        task,
+        depth,
+        runtimeLimits.maxDepth,
+      ),
+    };
   }
   const requestedThinking = agent.thinking ?? parentThinking;
   const effectiveModel = resolveEffectiveChildModelSettings(agent, parentModel);
@@ -888,13 +553,16 @@ export async function runSingleAgent(
   if ("error" in resolvedSkills) {
     const promptSetup = await promptSetupPromise;
     await cleanupPromptSetupResult(promptSetup);
-    return createErrorResult(
-      agentName,
-      agent.source,
-      task,
-      resolvedSkills.error,
-      modelDisplay,
-    );
+    return {
+      kind: "completed",
+      result: createErrorResult(
+        agentName,
+        agent.source,
+        task,
+        resolvedSkills.error,
+        modelDisplay,
+      ),
+    };
   }
   const promptSetup = await promptSetupPromise;
   if ("error" in promptSetup) throw promptSetup.error;
@@ -906,6 +574,7 @@ export async function runSingleAgent(
   };
   if (thinkingWarning) state.result.thinkingWarning = thinkingWarning;
   const tmpPrompt = promptSetup.tmpPrompt;
+  const samplingEnv = buildSamplingEnv(agent);
   try {
     const args = buildPiArgs(
       agent,
@@ -915,17 +584,21 @@ export async function runSingleAgent(
       resolvedSkills,
       tmpPrompt,
     );
+    if (samplingEnv) {
+      args.push("--extension", SAMPLING_EXTENSION_PATH);
+    }
     const invocation = getPiInvocation(args);
     const terminateOptions = {
       tree: true,
       platform: process.platform,
       processTreeDetached: process.platform !== "win32",
     };
+    const childEnv = buildChildEnv(samplingEnv);
     const proc = spawn(invocation.command, invocation.args, {
       cwd: defaultCwd,
       shell: invocation.command === "pi" && process.platform === "win32",
       stdio: ["ignore", "pipe", "pipe"],
-      env: { ...process.env, ...subagentDepthEnv() },
+      env: childEnv,
       ...getProcessTreeSpawnOptions(terminateOptions.tree),
     });
     const processDone = waitForSubagentProcess(proc);