pi-workflow-engine 0.11.0 → 0.13.0

package/.pi/extensions/pi-workflow-engine/src/agent-runner.ts

@@ -1,10 +1,12 @@
 import { assertWorkflowBudgetAvailable } from "./budget.ts";
-import { throwIfAborted } from "./cancellation.ts";
+import { WorkflowAgentTimeoutError } from "./agent-limits.ts";
+import { abortReason, linkAbortSignal, throwIfAborted } from "./cancellation.ts";
 import { executeAgentAttempt } from "./agent-attempt.ts";
 import {
   createAgentReplayPlan,
   isReplayEnabled,
   settleAgentAttempt,
+  type AgentAttemptResult,
   type AgentReplayPlan,
 } from "./agent-replay.ts";
 import {
@@ -18,13 +20,16 @@ import type {
 } from "./agent-runner-types.ts";
 import type { AgentResumeBaseContext } from "./resume-context.ts";
 import { unknownErrorMessage } from "./unknown-error.ts";
+import {
+  agentRetryDelayMs,
+  WorkflowProviderError,
+} from "./agent-retry.ts";
+import { resolveAgentModelProfile } from "./model-profiles.ts";
 
 export type {
   AgentExecutionOptions,
   AgentProgress,
-  AgentRunnerEvent,
   AgentRunnerSession,
-  AgentRunnerToolInfo,
   CreateAgentSession,
   RunContext,
 } from "./agent-runner-types.ts";
@@ -35,7 +40,7 @@ export type { ResolvedAgentModel, ResolvedAgentModelRequest } from "./agent-sess
  * Run one subagent to completion in an isolated in-memory session.
  *
  * Session construction, replay validation, and workspace cleanup live in focused
- * modules; this boundary owns only admission, progress, and the one cache-to-live retry.
+ * modules; this boundary owns admission, progress, and bounded cache/provider retries.
  */
 export async function runAgent(
   rc: RunContext,
@@ -53,35 +58,63 @@ export async function runAgent(
 
   return await rc.perf.time("agent.total_ms", async () => {
     throwIfAborted(rc.signal);
-    const model = resolveModel(rc, opts, label);
-    const replay = createAgentReplayPlan(prompt, opts);
+    const routing = resolveAgentRouting(rc, opts, label);
+    const effectiveOpts = routing.thinkingLevel === opts.thinkingLevel
+      ? opts
+      : { ...opts, thinkingLevel: routing.thinkingLevel };
+    const replay = createAgentReplayPlan(prompt, effectiveOpts);
     if (!isReplayEnabled(replay)) assertWorkflowBudgetAvailable(rc.budget);
 
     const rowId = rc.progress.agentQueued(opts.phase, label);
+    const liveScope = createAgentLiveScope(rc, label);
     try {
       return await rc.semaphore.run(
         async () => {
-          throwIfAborted(rc.signal);
-          if (!isReplayEnabled(replay)) assertWorkflowBudgetAvailable(rc.budget);
+          const agentRc: RunContext = { ...rc, signal: liveScope.signal };
+          throwIfAborted(agentRc.signal);
+          if (!isReplayEnabled(replay)) assertWorkflowBudgetAvailable(agentRc.budget);
           rc.progress.agentStart(opts.phase, label, rowId);
           if (replay.kind === "disabled") {
             rc.progress.log(`${label}: resume disabled for this call (${replay.reason})`);
           }
 
           let attemptPlan: AgentReplayPlan = replay;
+          let providerRetries = 0;
           while (true) {
-            const outcome = await executeAgentAttempt({
-              rc,
-              prompt,
-              opts,
-              resumeBaseContext,
-              model,
-              replay: attemptPlan,
-              label,
-              rowId,
-              tags,
-            });
-            const settlement = await settleAgentAttempt({ rc, label, tags, replay: attemptPlan, outcome });
+            let outcome: AgentAttemptResult;
+            try {
+              outcome = await executeAgentAttempt({
+                rc: agentRc,
+                prompt,
+                opts: effectiveOpts,
+                resumeBaseContext,
+                model: routing.model,
+                replay: attemptPlan,
+                label,
+                rowId,
+                tags,
+                admitLiveAgent: liveScope.admit,
+              });
+            } catch (error) {
+              if (!(error instanceof WorkflowProviderError) || !error.retryable) throw error;
+              if (providerRetries >= agentRc.agentRetries) {
+                if (agentRc.agentRetries > 0) {
+                  rc.progress.log(`${label}: transient provider failure after ${providerRetries} retries; giving up`);
+                }
+                throw error;
+              }
+
+              assertWorkflowBudgetAvailable(agentRc.budget);
+              providerRetries++;
+              const delayMs = agentRetryDelayMs(providerRetries);
+              rc.progress.log(`${label}: transient provider failure; retry ${providerRetries}/${agentRc.agentRetries} in ${delayMs}ms`);
+              rc.perf.counter("agent.provider_retry", 1, tags);
+              rc.perf.observe("agent.provider_retry_delay_ms", delayMs, tags);
+              attemptPlan = { kind: "off" };
+              await agentRc.retryScheduler.sleep(delayMs, agentRc.signal);
+              continue;
+            }
+            const settlement = await settleAgentAttempt({ rc: agentRc, label, tags, replay: attemptPlan, outcome });
             if (settlement.kind === "retry-live") {
               attemptPlan = { kind: "off" };
               continue;
@@ -96,20 +129,56 @@ export async function runAgent(
         },
       );
     } catch (error) {
-      rc.progress.agentFailed(label, error, rowId);
-      rc.progress.log(`${label} failed: ${unknownErrorMessage(error)}`);
-      throw error;
+      let failure = error;
+      if (liveScope.signal.aborted) failure = abortReason(liveScope.signal);
+      if (rc.signal?.aborted) failure = abortReason(rc.signal);
+      rc.progress.agentFailed(label, failure, rowId);
+      rc.progress.log(`${label} failed: ${unknownErrorMessage(failure)}`);
+      throw failure;
+    } finally {
+      liveScope.dispose();
     }
   }, tags);
 }
 
-function resolveModel(
+function createAgentLiveScope(rc: RunContext, label: string) {
+  const controller = new AbortController();
+  const unlink = linkAbortSignal(rc.signal, controller);
+  let timer: ReturnType<typeof setTimeout> | undefined;
+  let timerStarted = false;
+
+  return {
+    signal: controller.signal,
+    admit() {
+      rc.agentLimiter.admit(controller.signal);
+      if (timerStarted) return;
+      timerStarted = true;
+      timer = setTimeout(
+        () => controller.abort(new WorkflowAgentTimeoutError(label, rc.agentTimeoutMs)),
+        rc.agentTimeoutMs,
+      );
+    },
+    dispose() {
+      if (timer) clearTimeout(timer);
+      unlink();
+    },
+  };
+}
+
+function resolveAgentRouting(
   rc: RunContext,
   opts: AgentExecutionOptions,
   label: string,
-): ResolvedAgentModel["model"] {
+): { readonly model: ResolvedAgentModel["model"]; readonly thinkingLevel: AgentExecutionOptions["thinkingLevel"] } {
   try {
-    return resolveAgentModel(opts.model, rc.modelRegistry, rc.hostModel).model;
+    return resolveAgentModelProfile(
+      {
+        request: opts,
+        profiles: rc.modelProfiles,
+        resolveExplicitModel: (modelRef) => resolveAgentModel(modelRef, rc.modelRegistry, rc.hostModel).model,
+        hostModel: rc.hostModel,
+      },
+    );
   } catch (error) {
     rc.progress.agentFailed(label, error);
     throw error;

package/.pi/extensions/pi-workflow-engine/src/agent-session-providers.ts

@@ -0,0 +1,50 @@
+import type { Api, Model } from "@earendil-works/pi-ai";
+import type { ModelRegistry, ModelRuntime } from "@earendil-works/pi-coding-agent";
+
+type HostProviderRegistry = Pick<
+  ModelRegistry,
+  | "getApiKeyForProvider"
+  | "getProviderAuthStatus"
+  | "getRegisteredProviderConfig"
+  | "getRegisteredProviderIds"
+  | "isUsingOAuth"
+>;
+
+/** Mirror the host's live provider and selected-model auth state into a child runtime. */
+export async function synchronizeWorkflowModelRuntime(input: {
+  readonly host: HostProviderRegistry;
+  readonly child: ModelRuntime;
+  readonly selectedModel: Model<Api> | undefined;
+  readonly removeChildOnlyProviders: boolean;
+}): Promise<void> {
+  const { host, child, selectedModel, removeChildOnlyProviders } = input;
+  const hostProviderIds = new Set(host.getRegisteredProviderIds());
+
+  if (removeChildOnlyProviders) {
+    for (const providerId of child.getRegisteredProviderIds()) {
+      if (!hostProviderIds.has(providerId)) child.unregisterProvider(providerId);
+    }
+  }
+
+  for (const providerId of hostProviderIds) {
+    const config = host.getRegisteredProviderConfig(providerId);
+    if (!config) continue;
+    child.unregisterProvider(providerId);
+    child.registerProvider(providerId, config);
+  }
+
+  const selectedProvider = selectedModel?.provider;
+  if (!selectedModel || !selectedProvider) return;
+
+  const hostAuth = host.getProviderAuthStatus(selectedProvider);
+  const childAuth = child.getProviderAuthStatus(selectedProvider);
+  if (hostAuth.configured && !childAuth.configured && host.isUsingOAuth(selectedModel)) {
+    throw new Error(
+      `Workflow subagents cannot inherit OAuth credentials for "${selectedProvider}" from a host-only credential store; configure OAuth in Pi's shared agent directory.`,
+    );
+  }
+  if (hostAuth.source === "runtime" || (hostAuth.configured && !childAuth.configured)) {
+    const apiKey = await host.getApiKeyForProvider(selectedProvider);
+    if (apiKey !== undefined) await child.setRuntimeApiKey(selectedProvider, apiKey);
+  }
+}