oira666_pi-subagent 0.2.4 → 0.2.5

package/index.ts

@@ -9,6 +9,10 @@
  */
 
 import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
+import {
+  createAssistantMessageEventStream,
+  streamSimple as streamModelSimple,
+} from "@mariozechner/pi-ai";
 import { Type } from "@sinclair/typebox";
 import { type AgentConfig, discoverAgents } from "./agents.js";
 import { renderCall, renderResult } from "./render.js";
@@ -17,7 +21,6 @@ import {
   SUBAGENT_RESUME_DISABLE_ENV,
   SUBAGENT_RESUME_PROMPT_ENV,
   buildSubagentSessionDir,
-  ensureDir,
   findLatestResumableSubagentCall,
   getDefaultSubagentSessionRoot,
   isFinishedResult,
@@ -366,11 +369,67 @@ function hasCliInitialPrompt(argv: string[]): boolean {
   return false;
 }
 
+const RESUME_PROVIDER = "pi-subagent-resume";
+const RESUME_MODEL_ID = "synthetic-tool-call";
+const RESUME_STATE_KEY = "__piSubagentResumeState";
+
+type SyntheticResumeState = {
+  plan: ResumableSubagentCall | null;
+  phase: "tool" | "final";
+};
+
+function clearSyntheticResumeState(): void {
+  const state = getSyntheticResumeState();
+  state.plan = null;
+  state.phase = "tool";
+}
+
+function getSyntheticResumeState(): SyntheticResumeState {
+  const g = globalThis as any;
+  if (!g[RESUME_STATE_KEY]) {
+    g[RESUME_STATE_KEY] = { plan: null, phase: "tool" } satisfies SyntheticResumeState;
+  }
+  return g[RESUME_STATE_KEY] as SyntheticResumeState;
+}
+
+function emptyModelUsage() {
+  return {
+    input: 0,
+    output: 0,
+    cacheRead: 0,
+    cacheWrite: 0,
+    totalTokens: 0,
+    cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+  };
+}
+
+function getMessageText(message: any): string {
+  const content = message?.content;
+  if (typeof content === "string") return content;
+  if (!Array.isArray(content)) return "";
+  return content
+    .map((part) => (part?.type === "text" && typeof part.text === "string" ? part.text : ""))
+    .join("");
+}
+
+function isSyntheticResumePrompt(context: any, taskCount: number): boolean {
+  const messages = Array.isArray(context?.messages) ? context.messages : [];
+  const lastUser = [...messages].reverse().find((message) => message?.role === "user");
+  return getMessageText(lastUser).trim() === `Resuming ${taskCount} subagents...`;
+}
+
+function formatModelFlag(model: any): string | undefined {
+  if (!model?.id || !model?.provider) return undefined;
+  return `${model.provider}/${model.id}`;
+}
+
 // ---------------------------------------------------------------------------
 // Extension entry point
 // ---------------------------------------------------------------------------
 
 export default function (pi: ExtensionAPI) {
+  let resumeModelRegistry: any | undefined;
+
   pi.registerFlag("subagent-max-depth", {
     description: "Maximum allowed subagent delegation depth (default: 3).",
     type: "string",
@@ -381,6 +440,97 @@ export default function (pi: ExtensionAPI) {
     type: "boolean",
   });
 
+  pi.registerProvider(RESUME_PROVIDER, {
+    baseUrl: "http://127.0.0.1/pi-subagent-resume",
+    api: "openai-responses",
+    apiKey: "pi-subagent-resume-noop-key",
+    streamSimple: async (model, context, options) => {
+      const stream = createAssistantMessageEventStream();
+      const state = getSyntheticResumeState();
+      const plan = state.plan;
+      const phase = state.phase;
+      if (plan && phase === "tool" && isSyntheticResumePrompt(context, plan.tasks.length)) {
+        state.phase = "final";
+        const toolCall = {
+          type: "toolCall" as const,
+          id: `resume_subagent_${Date.now()}`,
+          name: "subagent",
+          arguments: { tasks: plan.tasks },
+        };
+        const message = {
+          role: "assistant" as const,
+          content: [toolCall],
+          api: model.api,
+          provider: model.provider,
+          model: model.id,
+          usage: emptyModelUsage(),
+          stopReason: "toolUse" as const,
+          timestamp: Date.now(),
+        };
+        queueMicrotask(() => {
+          stream.push({ type: "start", partial: message });
+          stream.push({ type: "toolcall_start", contentIndex: 0, partial: message });
+          stream.push({ type: "toolcall_end", contentIndex: 0, toolCall, partial: message });
+          stream.push({ type: "done", reason: "toolUse", message });
+          stream.end(message);
+        });
+        return stream;
+      }
+
+      if (phase === "final" && modelToRestoreAfterResume) {
+        const restore = modelToRestoreAfterResume;
+        const auth = resumeModelRegistry
+          ? await resumeModelRegistry.getApiKeyAndHeaders(restore)
+          : { ok: true, apiKey: undefined, headers: undefined };
+        if (!auth.ok) {
+          throw new Error(auth.error);
+        }
+        return streamModelSimple(restore, context, {
+          ...options,
+          apiKey: auth.apiKey,
+          headers: auth.headers,
+        });
+      }
+
+      if (!(plan && phase === "tool")) {
+        state.plan = null;
+        state.phase = "tool";
+      }
+      const message = {
+        role: "assistant" as const,
+        content: [],
+        api: model.api,
+        provider: model.provider,
+        model: model.id,
+        usage: emptyModelUsage(),
+        stopReason: "stop" as const,
+        timestamp: Date.now(),
+      };
+      queueMicrotask(() => {
+        stream.push({ type: "start", partial: message });
+        stream.push({ type: "done", reason: "stop", message });
+        stream.end(message);
+      });
+      return stream;
+    },
+    models: [
+      {
+        id: RESUME_MODEL_ID,
+        name: "Pi Subagent Resume",
+        api: "openai-responses",
+        reasoning: false,
+        input: ["text"],
+        cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
+        // Keep this large so Pi's pre-prompt auto-compaction does not call the
+        // synthetic provider before the visible resume prompt is appended. That
+        // would consume the tool-call phase during compaction and the real
+        // resume turn would only see the final text message.
+        contextWindow: 1_000_000,
+        maxTokens: 16,
+      },
+    ],
+  });
+
   const depthConfig = resolveDelegationDepthConfig(pi);
   const { currentDepth, maxDepth, canDelegate, ancestorAgentStack, preventCycles } =
     depthConfig;
@@ -392,8 +542,24 @@ export default function (pi: ExtensionAPI) {
   let currentSessionId = "ephemeral";
   let currentSubagentSessionRoot = "";
   let pendingResumePlan: ResumableSubagentCall | null = null;
+  let modelToRestoreAfterResume: any | undefined;
   const approvedProjectAgentDirsForSession = new Set<string>();
 
+  async function restoreModelAfterResumeFailure(ctx?: { ui?: { notify?: (message: string, type?: "info" | "warning" | "error") => void } }) {
+    const restore = modelToRestoreAfterResume;
+    modelToRestoreAfterResume = undefined;
+    pendingResumePlan = null;
+    clearSyntheticResumeState();
+    if (!restore) return;
+    try {
+      await pi.setModel(restore);
+    } catch (err) {
+      const message = err instanceof Error ? err.message : String(err);
+      ctx?.ui?.notify?.(`Failed to restore model after subagent resume error: ${message}`, "error");
+      console.error("[pi-subagent] Failed to restore model after resume error:", err);
+    }
+  }
+
   // Auto-discover agents on session start
   pi.on("session_start", async (event, ctx) => {
     if (!canDelegate) return;
@@ -430,13 +596,23 @@ export default function (pi: ExtensionAPI) {
       if (!shouldResume) return;
 
       pendingResumePlan = plan;
+      const resumeState = getSyntheticResumeState();
+      resumeState.plan = plan;
+      resumeState.phase = "tool";
+      modelToRestoreAfterResume = ctx.model;
+      resumeModelRegistry = ctx.modelRegistry;
       ensureSubagentToolActive(pi);
+      const resumeModel = ctx.modelRegistry.find(RESUME_PROVIDER, RESUME_MODEL_ID);
+      if (!resumeModel || !(await pi.setModel(resumeModel))) {
+        ctx.ui.notify("Failed to switch to synthetic subagent resume model.", "error");
+        await restoreModelAfterResumeFailure(ctx);
+        return;
+      }
 
       // In print/json subprocesses there is already an initial CLI prompt about
-      // to be sent. Starting another prompt from session_start races with it and
-      // Pi correctly reports "agent is already processing". In that case we only
-      // seed pendingResumePlan; before_agent_start injects the exact subagent
-      // call instruction into the upcoming turn.
+      // to be sent. That prompt will be answered by the synthetic provider with
+      // a real assistant subagent tool call. In interactive mode, submit a short
+      // visible prompt that triggers the same synthetic provider path.
       if (hasCliInitialPrompt(process.argv)) {
         if (ctx.hasUI) ctx.ui.notify(`Resuming ${plan.tasks.length} subagents...`, "info");
       } else {
@@ -444,9 +620,14 @@ export default function (pi: ExtensionAPI) {
       }
     } catch (err) {
       console.error("[pi-subagent] Error in session_start:", err);
+      await restoreModelAfterResumeFailure(ctx);
     }
   });
 
+  pi.on("agent_end", async () => {
+    await restoreModelAfterResumeFailure();
+  });
+
   // Inject available agents into the system prompt
   pi.on("before_agent_start", async (event) => {
     try {
@@ -456,15 +637,9 @@ export default function (pi: ExtensionAPI) {
       const agentList = discoveredAgents
         .map((a) => `- **${a.name}**: ${a.description}`)
         .join("\n");
-      if (pendingResumePlan) ensureSubagentToolActive(pi);
-      const resumeInstruction = pendingResumePlan
-        ? `\n\n## Interrupted Subagent Resume\n\nThe user approved resuming an interrupted subagent delegation. The \`subagent\` tool has been enabled for this turn. You MUST call the \`subagent\` tool exactly once now with these exact arguments and no other tool calls first:\n\n\`\`\`json\n${JSON.stringify({ tasks: pendingResumePlan.tasks }, null, 2)}\n\`\`\`\n`
-        : "";
-
       return {
         systemPrompt:
           event.systemPrompt +
-          resumeInstruction +
           `\n\n## Available Subagents
 
 The following subagents are available via the \`subagent\` tool:
@@ -630,7 +805,9 @@ This guard prevents self-recursion and cyclic handoffs (for example A -> B -> A)
             pendingResumePlan && sameTasks(pendingResumePlan.tasks, tasks)
               ? pendingResumePlan
               : null;
-          if (resumePlan) pendingResumePlan = null;
+          if (resumePlan) {
+            pendingResumePlan = null;
+          }
 
           if (tasks.length === 1) {
             const [task] = tasks;
@@ -646,6 +823,7 @@ This guard prevents self-recursion and cyclic handoffs (for example A -> B -> A)
               resumePlan?.details?.results[0],
               getSessionDirForTask(resumePlan?.previousToolCallId ?? toolCallId, 0),
               !!resumePlan,
+              formatModelFlag(modelToRestoreAfterResume),
             );
           }
 
@@ -659,6 +837,7 @@ This guard prevents self-recursion and cyclic handoffs (for example A -> B -> A)
             resumePlan?.details?.results,
             (index) => getSessionDirForTask(resumePlan?.previousToolCallId ?? toolCallId, index),
             !!resumePlan,
+            formatModelFlag(modelToRestoreAfterResume),
           );
         } catch (err) {
           const msg = err instanceof Error ? err.message : String(err);
@@ -672,7 +851,7 @@ This guard prevents self-recursion and cyclic handoffs (for example A -> B -> A)
 
       },
 
-      renderCall: (args, theme) => renderCall(args, theme),
+      renderCall: (args, theme, context) => renderCall(args, theme, context),
       renderResult: (result, { expanded }, theme) =>
         renderResult(result, expanded, theme),
     });
@@ -680,9 +859,7 @@ This guard prevents self-recursion and cyclic handoffs (for example A -> B -> A)
 
   function getSessionDirForTask(toolCallId: string, index: number): string {
     const root = currentSubagentSessionRoot || pathlessSubagentRootFallback();
-    const dir = buildSubagentSessionDir(root, currentSessionId, toolCallId, index);
-    ensureDir(dir);
-    return dir;
+    return buildSubagentSessionDir(root, currentSessionId, toolCallId, index);
   }
 
   function pathlessSubagentRootFallback(): string {
@@ -705,6 +882,7 @@ This guard prevents self-recursion and cyclic handoffs (for example A -> B -> A)
     previousResult: SingleResult | undefined,
     sessionDir: string,
     resumeExistingSession: boolean,
+    fallbackModel?: string,
   ) {
     if (previousResult && isFinishedResult(previousResult)) {
       return {
@@ -732,8 +910,10 @@ This guard prevents self-recursion and cyclic handoffs (for example A -> B -> A)
       onUpdate,
       makeDetails: makeDetails("single"),
       sessionDir: previousResult?.sessionDir ?? sessionDir,
+      sessionRoot: currentSubagentSessionRoot,
       resumeSession: resumeExistingSession,
       initialResult: previousResult,
+      fallbackModel,
     });
 
     if (isResultError(result)) {
@@ -774,6 +954,7 @@ This guard prevents self-recursion and cyclic handoffs (for example A -> B -> A)
     resumeResults: SingleResult[] | undefined,
     getSessionDir: (index: number) => string,
     resumeExistingSessions: boolean,
+    fallbackModel?: string,
   ) {
     return executeParallelSubprocess(
       tasks,
@@ -789,6 +970,8 @@ This guard prevents self-recursion and cyclic handoffs (for example A -> B -> A)
       resumeResults,
       (index) => getSessionDir(index),
       resumeExistingSessions,
+      currentSubagentSessionRoot,
+      fallbackModel,
     );
   }
 }

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "oira666_pi-subagent",
-  "version": "0.2.4",
+  "version": "0.2.5",
   "description": "Subagent extension for Pi coding agent. Delegate tasks to specialized agents.",
   "type": "module",
   "main": "index.ts",

package/render.ts

@@ -400,14 +400,23 @@ function topLevelSummary(details: SubagentDetails, counts: TreeCounts): string {
 // renderCall — shown while the tool is being invoked
 // ---------------------------------------------------------------------------
 
-export function renderCall(args: Record<string, any>, theme: { fg: ThemeFg; bold: (s: string) => string }): Text {
+export function renderCall(
+	args: Record<string, any>,
+	theme: { fg: ThemeFg; bold: (s: string) => string },
+	context?: { isPartial?: boolean; isError?: boolean },
+): Text {
 	const tasks = Array.isArray(args.tasks) ? args.tasks : [];
 	const count = tasks.length;
+	const icon = context?.isPartial === false
+		? context.isError
+			? theme.fg("error", "❌")
+			: theme.fg("success", "✅")
+		: theme.fg("warning", "⏳");
 	let text = `${theme.fg("toolTitle", theme.bold("subagent "))}${theme.fg("accent", `${count} task${count === 1 ? "" : "s"}`)}`;
 	for (const task of tasks.slice(0, 6)) {
 		const agent = typeof task?.agent === "string" ? task.agent : "...";
 		const preview = typeof task?.task === "string" ? ` ${truncate(task.task, 56)}` : "";
-		text += `\n  ${theme.fg("warning", "⏳")} ${theme.fg("accent", agent)}${theme.fg("dim", preview)}`;
+		text += `\n  ${icon} ${theme.fg("accent", agent)}${theme.fg("dim", preview)}`;
 	}
 	if (tasks.length > 6) text += `\n  ${theme.fg("muted", `... +${tasks.length - 6} more`)}`;
 	return new Text(text, 0, 0);

package/resume.ts

@@ -6,6 +6,7 @@ import { isResultError, isSubagentDetails, type SingleResult, type SubagentDetai
 
 export const SUBAGENT_RESUME_PROMPT_ENV = "PI_SUBAGENT_RESUME_PROMPT";
 export const SUBAGENT_RESUME_DISABLE_ENV = "PI_SUBAGENT_DISABLE_RESUME";
+export const SUBAGENT_SESSION_ROOT_ENV = "PI_SUBAGENT_SESSION_ROOT";
 
 type SessionEntry = ReturnType<ExtensionContext["sessionManager"]["getEntries"]>[number];
 
@@ -25,6 +26,9 @@ export function parseBooleanEnv(raw: unknown): boolean | null {
 }
 
 export function getDefaultSubagentSessionRoot(ctx: ExtensionContext): string {
+  const inheritedRoot = process.env[SUBAGENT_SESSION_ROOT_ENV];
+  if (inheritedRoot) return inheritedRoot;
+
   const mainSessionDir = ctx.sessionManager.getSessionDir?.();
   if (typeof mainSessionDir === "string" && mainSessionDir.length > 0) {
     return path.join(path.dirname(mainSessionDir), "sessions-subagents");

package/runner.ts

@@ -21,6 +21,7 @@ import {
   getFinalOutput,
   getNestedSubagentErrorSummary,
 } from "./types.js";
+import { SUBAGENT_SESSION_ROOT_ENV } from "./resume.js";
 import {
   DEFAULT_MAX_PARALLEL_TASKS,
   DEFAULT_MAX_CONCURRENCY,
@@ -51,6 +52,7 @@ const SUBAGENT_DEPTH_ENV = "PI_SUBAGENT_DEPTH";
 const SUBAGENT_MAX_DEPTH_ENV = "PI_SUBAGENT_MAX_DEPTH";
 const SUBAGENT_STACK_ENV = "PI_SUBAGENT_STACK";
 const SUBAGENT_PREVENT_CYCLES_ENV = "PI_SUBAGENT_PREVENT_CYCLES";
+const SUBAGENT_FALLBACK_MODEL_ENV = "PI_SUBAGENT_FALLBACK_MODEL";
 // PI_OFFLINE intentionally removed: setting it on child processes blocks all API
 // calls and renders subagents unable to do any LLM work. Children inherit the
 // parent's PI_OFFLINE value via process.env spread if needed.
@@ -305,7 +307,7 @@ export function processJsonLine(line: string, result: SingleResult): boolean {
         result.usage.cost += usage.cost?.total || 0;
         result.usage.contextTokens = usage.totalTokens || 0;
       }
-      if (!result.model && msg.model) result.model = msg.model;
+      if (msg.model && msg.model !== "synthetic-tool-call") result.model = msg.model;
       if (msg.stopReason) result.stopReason = msg.stopReason;
       if (msg.errorMessage) result.errorMessage = msg.errorMessage;
     }
@@ -367,6 +369,7 @@ function buildPiArgs(
   task: string,
   sessionDir: string | undefined,
   resumeSession: boolean,
+  fallbackModelOverride?: string,
 ): string[] {
   const args: string[] = [
     "--mode",
@@ -380,7 +383,7 @@ function buildPiArgs(
   args.push("-p");
 
   // Agent config takes priority; fall back to parent CLI value
-  const model = agent.model ?? _inheritedCliArgs.fallbackModel;
+  const model = agent.model ?? fallbackModelOverride ?? process.env[SUBAGENT_FALLBACK_MODEL_ENV] ?? _inheritedCliArgs.fallbackModel;
   if (model) args.push("--model", model);
 
   const thinking = agent.thinking ?? _inheritedCliArgs.fallbackThinking;
@@ -444,10 +447,14 @@ export interface RunAgentOptions {
   makeDetails: (results: SingleResult[]) => SubagentDetails;
   /** Dedicated session directory for this subagent process. */
   sessionDir?: string;
+  /** Top-level root for all subagent session directories in this delegation tree. */
+  sessionRoot?: string;
   /** Continue the most recent session in sessionDir instead of creating a new one. */
   resumeSession?: boolean;
   /** Previously captured state for this same subagent, used to render resumed nested trees. */
   initialResult?: SingleResult;
+  /** Fallback model to use when the agent config does not pin one. */
+  fallbackModel?: string;
 }
 
 /**
@@ -470,8 +477,10 @@ export async function runAgentSubprocess(opts: RunAgentOptions): Promise<SingleR
     onUpdate,
     makeDetails,
     sessionDir,
+    sessionRoot,
     resumeSession = false,
     initialResult,
+    fallbackModel,
   } = opts;
 
   const agent = agents.find((a) => a.name === agentName);
@@ -540,6 +549,7 @@ export async function runAgentSubprocess(opts: RunAgentOptions): Promise<SingleR
       task,
       sessionDir,
       resumeSession,
+      fallbackModel,
     );
     let wasAborted = false;
 
@@ -564,6 +574,8 @@ export async function runAgentSubprocess(opts: RunAgentOptions): Promise<SingleR
           [SUBAGENT_MAX_DEPTH_ENV]: String(propagatedMaxDepth),
           [SUBAGENT_STACK_ENV]: JSON.stringify(propagatedStack),
           [SUBAGENT_PREVENT_CYCLES_ENV]: preventCycles ? "1" : "0",
+          ...(sessionRoot ? { [SUBAGENT_SESSION_ROOT_ENV]: sessionRoot } : {}),
+          ...(fallbackModel ? { [SUBAGENT_FALLBACK_MODEL_ENV]: fallbackModel } : {}),
           // PI_OFFLINE is NOT forced here — see explanation near PI_OFFLINE_ENV.
         },
       });
@@ -761,6 +773,8 @@ export async function executeParallelSubprocess(
   resumeResults?: SingleResult[],
   getSessionDir?: (index: number, task: { agent: string; task: string; cwd?: string }) => string | undefined,
   resumeExistingSessions = false,
+  sessionRoot?: string,
+  fallbackModel?: string,
 ): Promise<{
   content: Array<{ type: "text"; text: string }>;
   details: SubagentDetails;
@@ -855,8 +869,10 @@ export async function executeParallelSubprocess(
         preventCycles,
         signal,
         sessionDir,
+        sessionRoot,
         resumeSession: resumeExistingSessions && !!sessionDir,
         initialResult: previousResult,
+        fallbackModel,
         onUpdate: (partial) => {
           if (partial.details?.results[0]) {
             allResults[index] = partial.details.results[0];