add persisted pi workflow extension

extensions/_workflows/runner.ts

@@ -12,10 +12,10 @@
 
 import {
   createAgentSession,
-  DefaultResourceLoader,
+  type DefaultResourceLoader,
   defineTool,
   SessionManager,
-  SettingsManager,
+  type SettingsManager,
   type AgentSession,
   type AgentSessionEvent,
   type AgentSessionEventListener,
@@ -29,6 +29,7 @@ import {
   childToolPolicy,
   createChildResources,
   shutdownAndDisposeChildSession,
+  waitForChildOperation,
 } from "../shared/child-session.ts";
 import { createToolCallTimeoutGuard } from "../shared/tool-call-timeout.ts";
 import { emptyUsage, type AgentUsage, type TranscriptEntry } from "./model.ts";
@@ -68,7 +69,10 @@ export interface AgentOutcome {
   error?: string;
   aborted: boolean;
   usage: AgentUsage;
+  provider?: string;
   model?: string;
+  thinkingLevel?: ThinkingLevel;
+  childSessionId?: string;
   contextWindow?: number;
   transcript: TranscriptEntry[];
 }
@@ -91,7 +95,15 @@ export interface RunAgentOptions {
   settingsManager: SettingsManager;
   modelRegistry: ExtensionContext["modelRegistry"];
   signal?: AbortSignal;
+  tools?: string[];
+  customTools?: ToolDefinition<any, any, any>[];
   onProgress?: (progress: AgentProgress) => void;
+  onSessionCreated?: (metadata: {
+    childSessionId: string;
+    provider?: string;
+    model?: string;
+    thinkingLevel: ThinkingLevel;
+  }) => void;
   /** Test-only override for the per-tool execution timeout. */
   toolCallTimeoutMs?: number;
   /** Test-only override for the first assistant response-event timeout. */
@@ -430,18 +442,24 @@ export async function runAgent(
   options: RunAgentOptions,
 ): Promise<AgentOutcome> {
   let structured: unknown;
-  let customTools: ToolDefinition[] | undefined;
+  let customTools: ToolDefinition<any, any, any>[];
   let session: AgentSession | undefined;
   let unsubscribeToolTimeout: (() => void) | undefined;
   try {
-    customTools =
-      options.schema !== undefined
-        ? [
+    customTools = [
+      ...(options.customTools ?? []),
+      ...(options.schema === undefined
+        ? []
+        : [
             makeStructuredOutputTool(options.schema, (value) => {
               structured = value;
             }),
-          ]
-        : undefined;
+          ]),
+    ];
+    const customToolNames = customTools.map((tool) => tool.name);
+    if (new Set(customToolNames).size !== customToolNames.length) {
+      throw new Error("Workflow child custom tool names must be unique");
+    }
     ({ session } = await createAgentSession({
       cwd: options.cwd,
       ...(options.model ? { model: options.model } : {}),
@@ -451,7 +469,8 @@ export async function runAgent(
       resourceLoader: options.loader,
       settingsManager: options.settingsManager,
       sessionManager: SessionManager.inMemory(options.cwd),
-      ...(customTools ? { customTools } : {}),
+      ...(customTools.length > 0 ? { customTools } : {}),
+      ...(options.tools ? { tools: options.tools } : {}),
       ...childToolPolicy(),
     }));
     await bindChildSessionExtensions(session);
@@ -468,15 +487,30 @@ export async function runAgent(
       error: `Failed to create agent session: ${errorText(error)}`,
       aborted: false,
       usage: emptyUsage(),
+      provider: options.model?.provider,
       model: options.model?.id,
+      thinkingLevel: options.thinkingLevel,
       contextWindow: options.model?.contextWindow,
       transcript: [],
     };
   }
 
   const childSession = session;
+  const childSessionId = childSession.sessionManager.getSessionId();
+  const thinkingLevel = childSession.thinkingLevel;
   let usage = emptyUsage();
+  let provider = childSession.model?.provider ?? options.model?.provider;
   let modelId = childSession.model?.id ?? options.model?.id;
+  try {
+    options.onSessionCreated?.({
+      childSessionId,
+      provider,
+      model: modelId,
+      thinkingLevel,
+    });
+  } catch {
+    // Metadata observers must not disrupt or leak a successfully created child.
+  }
   let contextWindow = childSession.model?.contextWindow;
   let stopReason: string | undefined;
   let errorMessage: string | undefined;
@@ -487,6 +521,7 @@ export async function runAgent(
     usage = computeUsage(messages);
 
     const sessionModel = childSession.model;
+    provider = sessionModel?.provider ?? provider;
     modelId = sessionModel?.id ?? modelId;
     contextWindow = sessionModel?.contextWindow ?? contextWindow;
     const context = childSession.getContextUsage();
@@ -515,6 +550,7 @@ export async function runAgent(
         !sessionModel ||
         (msg.provider === sessionModel.provider &&
           msg.model === sessionModel.id);
+      provider = msg.provider;
       const reportedId = msg.responseModel ?? msg.model;
       const reportedModel = responseMatchesSession
         ? options.modelRegistry.find(msg.provider, reportedId)
@@ -555,9 +591,14 @@ export async function runAgent(
 
   let aborted = false;
   let abortPromise: Promise<void> | undefined;
+  let resolveCancellation = () => {};
+  const cancellation = new Promise<void>((resolve) => {
+    resolveCancellation = resolve;
+  });
   const onAbort = () => {
     aborted = true;
-    abortPromise ??= childSession.abort().catch(() => {});
+    resolveCancellation();
+    abortPromise ??= waitForChildOperation(childSession.abort());
   };
   if (options.signal) {
     if (options.signal.aborted) onAbort();
@@ -573,9 +614,12 @@ export async function runAgent(
         model: modelId,
       });
       markFirstResponse = watchdog.markResponse;
-      await watchdog.waitFor(
-        childSession.prompt(buildWorkflowAgentPrompt(options.prompt)),
-      );
+      await Promise.race([
+        watchdog.waitFor(
+          childSession.prompt(buildWorkflowAgentPrompt(options.prompt)),
+        ),
+        cancellation,
+      ]);
     }
   } catch (error) {
     errorMessage = errorMessage ?? errorText(error);
@@ -602,7 +646,10 @@ export async function runAgent(
       error: "Agent was aborted",
       aborted: true,
       usage,
+      provider,
       model: modelId,
+      thinkingLevel,
+      childSessionId,
       contextWindow,
       transcript,
     };
@@ -617,7 +664,10 @@ export async function runAgent(
       error: errorMessage ?? "Agent failed",
       aborted: false,
       usage,
+      provider,
       model: modelId,
+      thinkingLevel,
+      childSessionId,
       contextWindow,
       transcript,
     };
@@ -631,7 +681,10 @@ export async function runAgent(
         "Agent finished without calling structured_output; no structured result matching the schema was produced.",
       aborted: false,
       usage,
+      provider,
       model: modelId,
+      thinkingLevel,
+      childSessionId,
       contextWindow,
       transcript,
     };
@@ -643,7 +696,10 @@ export async function runAgent(
     structured,
     aborted: false,
     usage,
+    provider,
     model: modelId,
+    thinkingLevel,
+    childSessionId,
     contextWindow,
     transcript,
   };