oira666_pi-subagent 0.3.4 → 0.3.6

package/README.md

@@ -138,6 +138,8 @@ While a `subagents` tool call is running, mid-stream steering input can be broad
 
 ## Subagent Session Resume
 
+> Requires Pi **0.81.0 or newer**. Crash recovery uses Pi's public full Provider SDK and session-replacement lifecycle.
+
 Subagent subprocesses save sessions in `sessions-subagents`. When a main Pi session is resumed and its latest branch contains an unfinished `subagents` tool call (aborted, errored, or closed by Pi's synthetic unfinished-tool error), the extension can resume that delegation from the saved subagent sessions.
 
 The same detection also runs after navigating the session tree in the TUI (Esc navigation): if you jump back to a point whose branch ends in an unfinished `subagents` call, the extension offers to resume those subagents from their saved sessions.
@@ -146,6 +148,8 @@ The same detection also runs after navigating the session tree in the TUI (Esc n
 - Non-UI modes (`pi -p`, JSON/RPC) resume automatically.
 - Already-finished subagents are reused as completed; unfinished ones continue from their own saved sessions.
 - Nested subagents use the same mechanism recursively.
+- Provider fallback goes through the selected model's effective Pi provider, so custom providers, custom APIs, auth-derived endpoints, headers, and provider-scoped environment are preserved.
+- Pending resume state and delayed callbacks are discarded on `/resume`, `/new`, `/fork`, and `/reload`, preventing stale work from an old runtime from leaking into the replacement session.
 
 | Env Var | Default | Description |
 | --- | --- | --- |

package/index.ts

@@ -10,11 +10,13 @@
 
 import * as fs from "node:fs";
 import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
-import { createFauxCore, fauxAssistantMessage, fauxToolCall } from "@mariozechner/pi-ai";
-// The synthetic resume provider forwards to the real fallback model through a
-// small dispatcher built on stable pi-ai exports. This deliberately avoids the
-// deprecated/temporary `@mariozechner/pi-ai/compat` global `streamSimple`.
-import { streamSimpleForModel as streamModelSimple } from "./resumeStream.js";
+import {
+  createFauxCore,
+  createProvider,
+  fauxAssistantMessage,
+  fauxToolCall,
+  lazyStream,
+} from "@mariozechner/pi-ai";
 import { Type } from "@sinclair/typebox";
 import { type AgentConfig, discoverAgents } from "./agents.js";
 import {
@@ -562,7 +564,6 @@ function isStreamingSteerInput(event: any, ctx: { isIdle: () => boolean }): bool
   return !ctx.isIdle();
 }
 
-const RESUME_STATE_KEY = "__piSubagentResumeState";
 const SUBAGENT_FALLBACK_MODEL_ENV = "PI_SUBAGENT_FALLBACK_MODEL";
 const RESUME_INTERACTIVE_DELAY_MS = 50;
 
@@ -572,21 +573,6 @@ type SyntheticResumeState = {
   trigger: "resumePrompt" | "nextRequest";
 };
 
-function clearSyntheticResumeState(): void {
-  const state = getSyntheticResumeState();
-  state.plans = [];
-  state.phase = "tool";
-  state.trigger = "resumePrompt";
-}
-
-function getSyntheticResumeState(): SyntheticResumeState {
-  const g = globalThis as any;
-  if (!g[RESUME_STATE_KEY]) {
-    g[RESUME_STATE_KEY] = { plans: [], phase: "tool", trigger: "resumePrompt" } satisfies SyntheticResumeState;
-  }
-  return g[RESUME_STATE_KEY] as SyntheticResumeState;
-}
-
 // Model definition for the synthetic subagent-resume provider. Shared between
 // the faux core (which produces the canned assistant turn) and the provider
 // registration below.
@@ -663,23 +649,101 @@ export default function (pi: ExtensionAPI) {
   let lastRestorableModel: any | undefined;
   let latestSessionCtx: any | undefined;
   let pendingInteractiveResumePrompt: string | null = null;
+  let lifecycleGeneration = 0;
+  let sessionActive = false;
+  const scheduledTasks = new Set<ReturnType<typeof setTimeout>>();
+  const resumeState: SyntheticResumeState = {
+    plans: [],
+    phase: "tool",
+    trigger: "resumePrompt",
+  };
+
+  function clearSyntheticResumeState(): void {
+    resumeState.plans = [];
+    resumeState.phase = "tool";
+    resumeState.trigger = "resumePrompt";
+  }
+
+  function scheduleSessionTask(callback: () => void, delayMs: number): void {
+    const expectedGeneration = lifecycleGeneration;
+    const timer = setTimeout(() => {
+      scheduledTasks.delete(timer);
+      if (!sessionActive || expectedGeneration !== lifecycleGeneration) return;
+      callback();
+    }, delayMs);
+    scheduledTasks.add(timer);
+  }
 
-  async function streamWithRealModelFallback(context: any, options: any, fallback: any) {
-    if (!fallback) return null;
-    const auth = resumeModelRegistry
-      ? await resumeModelRegistry.getApiKeyAndHeaders(fallback)
-      : { ok: true, apiKey: undefined, headers: undefined };
-    if (!auth.ok) {
-      throw new Error(auth.error);
+  function mergeProviderHeaders(
+    base: Record<string, string | null> | undefined,
+    override: Record<string, string | null> | undefined,
+  ): Record<string, string | null> | undefined {
+    const merged = new Map<string, [string, string | null]>();
+    for (const headers of [base, override]) {
+      for (const [name, value] of Object.entries(headers ?? {})) {
+        const key = name.toLowerCase();
+        if (value === null) merged.delete(key);
+        else merged.set(key, [name, value]);
+      }
     }
-    return streamModelSimple(fallback, context, {
-      ...options,
-      apiKey: auth.apiKey,
-      headers: auth.headers,
+    return merged.size > 0 ? Object.fromEntries(merged.values()) : undefined;
+  }
+
+  function streamWithRealModelFallback(
+    context: any,
+    options: any,
+    fallback: any,
+    expectedGeneration = lifecycleGeneration,
+  ) {
+    if (!fallback || !resumeModelRegistry) return null;
+
+    // Use pi 0.81's effective Provider instead of dispatching on model.api.
+    // This preserves custom provider streams, provider composition, dynamic
+    // auth base URLs, provider-scoped env, and future/custom API identifiers.
+    return lazyStream(fallback, async () => {
+      if (!sessionActive || expectedGeneration !== lifecycleGeneration) {
+        throw new Error("Subagent resume fallback was cancelled by session replacement.");
+      }
+      const provider = resumeModelRegistry.getProvider?.(fallback.provider);
+      if (!provider || provider.id === RESUME_PROVIDER) {
+        throw new Error(`Subagent resume fallback provider is unavailable: ${fallback.provider}.`);
+      }
+      const [providerResolution, modelResolution] = await Promise.all([
+        resumeModelRegistry.getProviderAuth?.(fallback.provider),
+        resumeModelRegistry.getApiKeyAndHeaders?.(fallback),
+      ]);
+      if (!sessionActive || expectedGeneration !== lifecycleGeneration) {
+        throw new Error("Subagent resume fallback was cancelled by session replacement.");
+      }
+      if (!providerResolution || !modelResolution?.ok) {
+        throw new Error(
+          modelResolution?.error ?? `Provider is not configured: ${fallback.provider}`,
+        );
+      }
+      const providerAuth = providerResolution.auth ?? {};
+      const requestModel = providerAuth.baseUrl
+        ? { ...fallback, baseUrl: providerAuth.baseUrl }
+        : fallback;
+      const requestOptions = {
+        ...options,
+        // Model-aware resolution includes configured/model headers. Never
+        // forward the synthetic provider's no-op credential.
+        apiKey: modelResolution.apiKey,
+        headers: mergeProviderHeaders(modelResolution.headers, options?.headers),
+        env: {
+          ...(providerResolution.env ?? {}),
+          ...(modelResolution.env ?? {}),
+          ...(options?.env ?? {}),
+        },
+      };
+      return provider.streamSimple(requestModel, context, requestOptions);
     });
   }
 
-  async function restoreVisibleModelForResume(): Promise<any | undefined> {
+  async function restoreVisibleModelForResume(
+    expectedGeneration = lifecycleGeneration,
+  ): Promise<any | undefined> {
+    if (!sessionActive || expectedGeneration !== lifecycleGeneration) return undefined;
     const restore = modelToRestoreAfterResume ?? lastRestorableModel;
     if (!restore) return undefined;
     lastRestorableModel = restore;
@@ -687,7 +751,9 @@ export default function (pi: ExtensionAPI) {
       try {
         await pi.setModel(restore);
       } catch (err) {
-        console.error("[pi-subagent] Failed to restore real model during resume:", err);
+        if (sessionActive && expectedGeneration === lifecycleGeneration) {
+          console.error("[pi-subagent] Failed to restore real model during resume:", err);
+        }
       }
     }
     return restore;
@@ -712,12 +778,26 @@ export default function (pi: ExtensionAPI) {
     models: [RESUME_MODEL_DEF],
   });
 
-  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 state = getSyntheticResumeState();
+  const resumeProvider = createProvider({
+    id: RESUME_PROVIDER,
+    name: "Pi Subagent Resume",
+    auth: {
+      apiKey: {
+        name: "Internal synthetic resume provider",
+        async resolve() {
+          return {
+            auth: { apiKey: "pi-subagent-resume-noop-key" },
+            source: "internal synthetic provider",
+          };
+        },
+      },
+    },
+    models: resumeCore.models,
+    api: {
+      stream: resumeCore.stream,
+      streamSimple: (model, context, options) => {
+      const state = resumeState;
+      const expectedGeneration = lifecycleGeneration;
       const discoveredPlans = state.plans.length > 0
         ? state.plans
         : pendingResumePlans.length > 0
@@ -759,7 +839,7 @@ export default function (pi: ExtensionAPI) {
           .result()
           .catch(() => {})
           .finally(() => {
-            void restoreVisibleModelForResume();
+            void restoreVisibleModelForResume(expectedGeneration);
           });
         return stream;
       }
@@ -768,14 +848,28 @@ export default function (pi: ExtensionAPI) {
       // injection turn (e.g. a request raced ahead of the model restore).
       // Forward the request to the real fallback model instead.
       if (phase === "final") {
-        const restore = await restoreVisibleModelForResume();
-        const delegated = await streamWithRealModelFallback(context, options, restore);
-        if (delegated) return delegated;
+        const delegated = streamWithRealModelFallback(
+          context,
+          options,
+          modelToRestoreAfterResume ?? lastRestorableModel,
+          expectedGeneration,
+        );
+        if (delegated) {
+          void restoreVisibleModelForResume(expectedGeneration);
+          return delegated;
+        }
       }
 
-      const restore = await restoreVisibleModelForResume();
-      const fallback = await streamWithRealModelFallback(context, options, restore ?? lastRestorableModel);
-      if (fallback) return fallback;
+      const fallback = streamWithRealModelFallback(
+        context,
+        options,
+        modelToRestoreAfterResume ?? lastRestorableModel,
+        expectedGeneration,
+      );
+      if (fallback) {
+        void restoreVisibleModelForResume(expectedGeneration);
+        return fallback;
+      }
 
       // No real model to fall back to: surface a clear error turn.
       if (!(plans.length > 0 && phase === "tool")) {
@@ -788,9 +882,10 @@ export default function (pi: ExtensionAPI) {
         () => fauxAssistantMessage([], { stopReason: "error", errorMessage: errorText }),
       ]);
       return resumeCore.streamSimple(model, context, options);
+      },
     },
-    models: [{ ...RESUME_MODEL_DEF, api: "openai-responses" }],
   });
+  pi.registerProvider(resumeProvider);
 
   const depthConfig = resolveDelegationDepthConfig(pi);
   const { currentDepth, maxDepth, canDelegate, ancestorAgentStack, preventCycles } =
@@ -1178,6 +1273,7 @@ export default function (pi: ExtensionAPI) {
   }
 
   async function restoreModelAfterResumeFailure(ctx?: { ui?: { notify?: (message: string, type?: "info" | "warning" | "error") => void } }) {
+    if (!sessionActive) return;
     const restore = modelToRestoreAfterResume;
     modelToRestoreAfterResume = undefined;
     pendingResumePlans = [];
@@ -1221,7 +1317,14 @@ export default function (pi: ExtensionAPI) {
 
   // Auto-discover agents on session start
   pi.on("session_start", async (event, ctx) => {
+    lifecycleGeneration += 1;
+    sessionActive = true;
     latestSessionCtx = ctx;
+    resumeModelRegistry = ctx.modelRegistry;
+    clearSyntheticResumeState();
+    pendingResumePlans = [];
+    pendingInteractiveResumePrompt = null;
+    modelToRestoreAfterResume = undefined;
     updateCombinedUsageStatus(ctx);
     try {
       // Always repair sessions left on the synthetic resume model, even in
@@ -1317,6 +1420,24 @@ export default function (pi: ExtensionAPI) {
     }
   });
 
+  // Pi 0.81 replaces and rebinds the entire extension runtime on /resume,
+  // /new, /fork, and /reload. Invalidate every detached callback so it cannot
+  // use stale pi/context objects after the old runtime has been torn down.
+  pi.on("session_shutdown", () => {
+    sessionActive = false;
+    lifecycleGeneration += 1;
+    for (const timer of scheduledTasks) clearTimeout(timer);
+    scheduledTasks.clear();
+    clearSyntheticResumeState();
+    pendingResumePlans = [];
+    pendingInteractiveResumePrompt = null;
+    modelToRestoreAfterResume = undefined;
+    latestSessionCtx = undefined;
+    resumeModelRegistry = undefined;
+    activeSubagentUsageSummaries.clear();
+    activeSubagents.clear();
+  });
+
   /**
    * Detect unfinished subagent calls at the current branch leaf and offer to
    * resume them. Shared between session_start (startup/resume) and
@@ -1373,7 +1494,6 @@ export default function (pi: ExtensionAPI) {
     }
 
     pendingResumePlans = [...plans];
-    const resumeState = getSyntheticResumeState();
     resumeState.plans = [...plans];
     resumeState.phase = "tool";
     // Headless subprocess/RPC subagents cannot answer a visible resume
@@ -1407,7 +1527,7 @@ export default function (pi: ExtensionAPI) {
       // which is the last extension hook before the initial chat render.
       pendingInteractiveResumePrompt = `Resuming ${totalTaskCount} subagents...`;
     } else {
-      setTimeout(() => {
+      scheduleSessionTask(() => {
         try {
           pi.sendUserMessage(`Resuming ${totalTaskCount} subagents...`);
         } catch (err) {
@@ -1452,7 +1572,7 @@ export default function (pi: ExtensionAPI) {
   pi.on("message_end", (_event, ctx) => {
     latestSessionCtx = ctx;
     updateCombinedUsageStatus(ctx);
-    setTimeout(() => updateCombinedUsageStatus(ctx), 0);
+    scheduleSessionTask(() => updateCombinedUsageStatus(ctx), 0);
   });
 
   pi.on("tool_execution_end", (event, ctx) => {
@@ -1460,7 +1580,7 @@ export default function (pi: ExtensionAPI) {
     if (isSubagentToolName(event.toolName)) {
       activeSubagentUsageSummaries.delete(event.toolCallId);
       updateCombinedUsageStatus(ctx);
-      setTimeout(() => updateCombinedUsageStatus(ctx), 0);
+      scheduleSessionTask(() => updateCombinedUsageStatus(ctx), 0);
     }
   });
 
@@ -1468,7 +1588,7 @@ export default function (pi: ExtensionAPI) {
     const prompt = pendingInteractiveResumePrompt;
     if (!prompt) return;
     pendingInteractiveResumePrompt = null;
-    setTimeout(() => {
+    scheduleSessionTask(() => {
       try {
         pi.sendUserMessage(prompt);
       } catch (err) {

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "oira666_pi-subagent",
-  "version": "0.3.4",
+  "version": "0.3.6",
   "description": "Subagent extension for Pi coding agent. Delegate tasks to specialized agents.",
   "type": "module",
   "main": "index.ts",
@@ -9,7 +9,6 @@
     "agents.ts",
     "runner.ts",
     "resume.ts",
-    "resumeStream.ts",
     "names.ts",
     "shared.ts",
     "render.ts",
@@ -53,10 +52,10 @@
     "typescript": "^5.9.3"
   },
   "peerDependencies": {
-    "@mariozechner/pi-agent-core": ">=0.37.0",
-    "@mariozechner/pi-ai": ">=0.37.0",
-    "@mariozechner/pi-coding-agent": ">=0.37.0",
-    "@mariozechner/pi-tui": ">=0.37.0"
+    "@mariozechner/pi-agent-core": ">=0.81.0",
+    "@mariozechner/pi-ai": ">=0.81.0",
+    "@mariozechner/pi-coding-agent": ">=0.81.0",
+    "@mariozechner/pi-tui": ">=0.81.0"
   },
   "peerDependenciesMeta": {
     "@mariozechner/pi-agent-core": {

package/shims.d.ts

@@ -12,6 +12,7 @@ declare module "@mariozechner/pi-ai" {
   export type Message = any;
 
   export function createAssistantMessageEventStream(): any;
+  export function createProvider(options: any): any;
   export function lazyStream(model: any, setup: () => Promise<any>): any;
   export type AssistantMessageEventStream = any;
   export type ProviderStreams = { stream: (...args: any[]) => any; streamSimple: (...args: any[]) => any };
@@ -90,6 +91,7 @@ declare module "@mariozechner/pi-coding-agent" {
   export interface ExtensionAPI {
     registerFlag(name: string, config: any): void;
     getFlag(name: string): unknown;
+    registerProvider(provider: any): void;
     registerProvider(name: string, provider: any): void;
     registerTool(tool: any): void;
     addBeforeAgentStart(hook: (ctx: ExtensionContext) => unknown): void;

package/resumeStream.ts

@@ -1,104 +0,0 @@
-/**
- * Durable model-streaming helper for the synthetic subagent-resume provider.
- *
- * The synthetic resume provider only ever synthesizes an assistant turn that
- * carries the `subagent` tool call(s). But pi may still invoke the provider's
- * `streamSimple` handler in edge/race situations where the real model has not
- * been restored yet (e.g. a request arrives before `pi.setModel(realModel)`
- * has propagated). In that case the handler must forward the request to the
- * real fallback model and return a valid assistant stream.
- *
- * Historically this used the global `streamSimple` dispatcher exported from the
- * `@mariozechner/pi-ai` package root. Pi's provider/model rework moved that
- * dispatcher into the explicitly temporary `@mariozechner/pi-ai/compat`
- * entrypoint ("deleted with the coding-agent ModelManager migration").
- *
- * This module reimplements the same behavior using only stable pi-ai exports:
- *   - the root `lazyStream` helper (returns a stream synchronously while async
- *     setup runs behind it), and
- *   - the per-API `ProviderStreams` factories published under the stable
- *     `@mariozechner/pi-ai/api/*` subpaths.
- *
- * `model.api` selects the concrete API implementation, mirroring exactly what
- * pi core does when it dispatches a stream to the provider that owns a model.
- */
-import { lazyStream } from "@mariozechner/pi-ai";
-import type { AssistantMessageEventStream, ProviderStreams } from "@mariozechner/pi-ai";
-
-type ProviderStreamsFactory = () => ProviderStreams;
-
-/**
- * Lazily import the `ProviderStreams` factory for a given `model.api`.
- *
- * Each entry maps a `KnownApi` id to its stable `/api/*` subpath module and the
- * factory export that module provides. Dynamic `import()` keeps the API modules
- * out of the hot path until a fallback stream is actually needed, and the
- * host's module cache deduplicates repeated loads.
- */
-const API_LOADERS: Record<string, () => Promise<ProviderStreamsFactory>> = {
-  "openai-responses": async () =>
-    (await import("@mariozechner/pi-ai/api/openai-responses.lazy")).openAIResponsesApi,
-  "openai-completions": async () =>
-    (await import("@mariozechner/pi-ai/api/openai-completions.lazy")).openAICompletionsApi,
-  "azure-openai-responses": async () =>
-    (await import("@mariozechner/pi-ai/api/azure-openai-responses.lazy")).azureOpenAIResponsesApi,
-  "openai-codex-responses": async () =>
-    (await import("@mariozechner/pi-ai/api/openai-codex-responses.lazy")).openAICodexResponsesApi,
-  "anthropic-messages": async () =>
-    (await import("@mariozechner/pi-ai/api/anthropic-messages.lazy")).anthropicMessagesApi,
-  "bedrock-converse-stream": async () =>
-    (await import("@mariozechner/pi-ai/api/bedrock-converse-stream.lazy")).bedrockConverseStreamApi,
-  "google-generative-ai": async () =>
-    (await import("@mariozechner/pi-ai/api/google-generative-ai.lazy")).googleGenerativeAIApi,
-  "google-vertex": async () =>
-    (await import("@mariozechner/pi-ai/api/google-vertex.lazy")).googleVertexApi,
-  "mistral-conversations": async () =>
-    (await import("@mariozechner/pi-ai/api/mistral-conversations.lazy")).mistralConversationsApi,
-  "pi-messages": async () =>
-    (await import("@mariozechner/pi-ai/api/pi-messages.lazy")).piMessagesApi,
-};
-
-const providerStreamsCache = new Map<string, ProviderStreams>();
-
-/** API ids for which the resume fallback can forward to a real model. */
-export function getSupportedResumeFallbackApis(): string[] {
-  return Object.keys(API_LOADERS);
-}
-
-/** Load (and cache) the `ProviderStreams` implementation for a `model.api`. */
-export async function resolveProviderStreams(api: string): Promise<ProviderStreams> {
-  const cached = providerStreamsCache.get(api);
-  if (cached) return cached;
-
-  const loader = API_LOADERS[api];
-  if (!loader) {
-    throw new Error(
-      `Subagent resume cannot forward to fallback model: unsupported model API "${api}". ` +
-        `Supported APIs: ${Object.keys(API_LOADERS).join(", ")}.`,
-    );
-  }
-
-  const factory = await loader();
-  const providerStreams = factory();
-  providerStreamsCache.set(api, providerStreams);
-  return providerStreams;
-}
-
-/**
- * Stream a real fallback model through its owning API implementation.
- *
- * Returns synchronously via `lazyStream`; the async API-module load and the
- * underlying provider request run behind the returned stream. `options` is
- * expected to already carry the resolved `apiKey`/`headers` (the extension
- * resolves those through `ctx.modelRegistry.getApiKeyAndHeaders`).
- */
-export function streamSimpleForModel(
-  model: any,
-  context: any,
-  options: any,
-): AssistantMessageEventStream {
-  return lazyStream(model, async () => {
-    const providerStreams = await resolveProviderStreams(model.api);
-    return providerStreams.streamSimple(model, context, options);
-  });
-}