@bastani/atomic 0.8.28 → 0.8.29-alpha.3

package/dist/builtin/workflows/builtin/shared-prompts.ts

@@ -9,3 +9,10 @@ export const WORKER_PREFLIGHT_CONTRACT = [
   "If setup requirements cannot be determined confidently, delegate a focused discovery task before implementation instead of guessing.",
   "If setup remains blocked after evidence-based discovery and setup attempts, report the blocker with commands tried and the exact evidence needed to continue.",
 ].join("\n");
+
+export const E2E_VERIFICATION_GUIDANCE = [
+  "Verify correctness end-to-end whenever practical for user-visible behavior; do not rely only on code inspection, unit tests, or stage summaries when an executable user scenario can prove the outcome.",
+  "For web or frontend flows — including frontend changes whose correctness depends on backend/API behavior — use the browser skill, or delegate to a subagent with `skill: \"browser\"`, to drive the application like a user and capture screenshot, DOM, or network evidence when that proves the objective.",
+  "For TUI or terminal-app flows, use the tmux skill, or delegate to a subagent with `skill: \"tmux\"`, to launch the app in an isolated tmux session, send keys, capture pane output, and simulate the scenario end to end.",
+  "If end-to-end verification is not practical in this checkout, record what was attempted, the smallest missing prerequisite, and the narrower validation that was run instead; do not claim end-to-end proof when it was not performed.",
+].join("\n");

package/dist/builtin/workflows/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@bastani/workflows",
-  "version": "0.8.28",
+  "version": "0.8.29-alpha.3",
   "private": true,
   "description": "Atomic extension for multi-stage workflow authoring and execution.",
   "contributors": [
@@ -83,7 +83,7 @@
   },
   "peerDependencies": {
     "@bastani/atomic": "*",
-    "@earendil-works/pi-tui": "^0.78.1"
+    "@earendil-works/pi-tui": "^0.79.3"
   },
   "peerDependenciesMeta": {
     "@bastani/atomic": {

package/dist/builtin/workflows/skills/research-codebase/SKILL.md

@@ -65,10 +65,24 @@ The user's research question/request is: **$ARGUMENTS**
         - The agent fetches live web content using the **browser** skill's `browse` CLI (or `npx browse` / `curl`). Instruct it to apply the token-efficient fetch order: (1) try `curl https://<site>/llms.txt` for an AI-friendly index (see [llmstxt.org](https://llmstxt.org/llms.txt)), (2) try `curl <url> -H "Accept: text/markdown"` to get pre-converted Markdown (supported on Cloudflare-hosted docs via [Markdown for Agents](https://developers.cloudflare.com/fundamentals/reference/markdown-for-agents/)), (3) fall back to HTML parsing via `browse`
         - Instruct the agent to return LINKS with their findings and INCLUDE those links in the research document
         - The agent should persist reusable source documents under `research/web/<YYYY-MM-DD>-<kebab-case-topic>.md` (with frontmatter noting `source_url`, `fetched_at`, and `fetch_method`) so future research can reuse them without re-fetching
-        - Output directory for the synthesized research artifact: `research/docs/`
+        - Output directory for the synthesized web research artifacts: `research/web/`:
+
+          When you fetch a document that is worth keeping for future sessions (reference docs, API schemas, SDK guides, release notes, troubleshooting writeups, architecture articles), `write` it to `research/web/<YYYY-MM-DD>-<kebab-case-topic>.md` with frontmatter capturing:
+          
+          ```markdown
+          ---
+          source_url: <original URL>
+          fetched_at: <YYYY-MM-DD>
+          fetch_method: read | llms.txt | markdown-accept-header | browser | browse
+          topic: <short description>
+          ---
+          ```
+
+        - Followed by the extracted content (trimmed of nav chrome, ads, and irrelevant boilerplate). This lets future work reuse the lookup without re-fetching. Before fetching anything, quickly `find research/web/` for an existing, recent copy.
+
         - Examples:
-            - If researching `Redis` locks usage, the agent might find relevant usage and create a document `research/docs/2024-01-15-redis-locks-usage.md` with internal links to Redis docs and code references (and cache the fetched Redis docs under `research/web/`)
-            - If researching `OAuth` flows, the agent might find relevant external articles and create a document `research/docs/2024-01-16-oauth-flows.md` with links to those articles
+            - If researching `Redis` locks usage, the agent might find relevant usage and create a document `research/web/2024-01-15-redis-locks-usage.md` with internal links to Redis docs and code references (and cache the fetched Redis docs under `research/web/`)
+            - If researching `OAuth` flows, the agent might find relevant external articles and create a document `research/web/2024-01-16-oauth-flows.md` with links to those articles
 
     The key is to use these agents intelligently:
     - Start with locator agents to find what exists

package/dist/builtin/workflows/src/extension/wiring.ts

@@ -22,7 +22,7 @@
  */
 
 import { basename } from "node:path";
-import type { ChatMessageRenderOptions, CreateAgentSessionOptions } from "@bastani/atomic";
+import type { ChatMessageRenderOptions, CreateAgentSessionOptions, PackageSource } from "@bastani/atomic";
 import type { StageAdapters, StageSessionCreateResult, StageSessionRuntime } from "../runs/foreground/stage-runner.js";
 import type { StageExecutionMeta, StageOptions } from "../shared/types.js";
 import { stageUiBroker, type StageUiBroker } from "../shared/stage-ui-broker.js";
@@ -109,7 +109,7 @@ export interface PiCodingAgentSdk {
     cwd: string;
     agentDir: string;
     settingsManager?: PiSdkSettingsManager;
-    builtinPackagePaths?: string[];
+    builtinPackagePaths?: PackageSource[];
   }) => PiSdkResourceLoader;
   createAgentSession(options?: AtomicCreateAgentSessionOptions): Promise<{ session: StageSessionRuntime }>;
 }
@@ -156,7 +156,7 @@ export async function prepareAtomicStageSessionOptions(
     settingsManager,
     builtinPackagePaths: stageBuiltinPackagePaths(sdk.getBuiltinPackagePaths?.() ?? []),
   });
-  await resourceLoader.reload();
+  await reloadWorkflowStageResources(resourceLoader);
 
   return {
     ...atomicOptions,
@@ -167,13 +167,60 @@ export async function prepareAtomicStageSessionOptions(
   };
 }
 
-function stageBuiltinPackagePaths(paths: readonly string[]): string[] {
+function stageBuiltinPackagePaths(paths: readonly string[]): PackageSource[] {
   // Workflow stages are child AgentSessions owned by the workflow extension.
   // Loading the workflows extension again inside that child session replays its
   // `session_start` lifecycle and clears/kills the parent workflow store. Keep
-  // the other builtin packages (subagents, mcp, web-access, intercom), but do
-  // not recursively install workflows into workflow stage sessions.
-  return paths.filter((path) => basename(path) !== "workflows");
+  // the workflows package itself so its bundled skills/prompts/resources remain
+  // available, but disable only its extension entry for stage sessions.
+  return paths.map((path) =>
+    basename(path) === "workflows" ? { source: path, extensions: [] } : path,
+  );
+}
+
+const SUBAGENT_CHILD_EXTENSION_ENV_KEYS = [
+  "ATOMIC_SUBAGENT_CHILD",
+  "ATOMIC_SUBAGENT_FANOUT_CHILD",
+  "PI_SUBAGENT_CHILD",
+  "PI_SUBAGENT_FANOUT_CHILD",
+] as const;
+
+let workflowStageResourceReloadQueue: Promise<void> = Promise.resolve();
+
+async function reloadWorkflowStageResources(resourceLoader: PiSdkResourceLoader): Promise<void> {
+  const queuedReload = workflowStageResourceReloadQueue.then(() =>
+    reloadWorkflowStageResourcesWithEnvIsolation(resourceLoader),
+  );
+  workflowStageResourceReloadQueue = queuedReload.catch(() => undefined);
+  return queuedReload;
+}
+
+async function reloadWorkflowStageResourcesWithEnvIsolation(resourceLoader: PiSdkResourceLoader): Promise<void> {
+  // Workflow stage sessions are already governed by an orchestration context
+  // that disables recursive workflow tools and caps nested subagent depth. When
+  // a workflow itself runs inside a subagent child process, inherited subagent
+  // child env flags would otherwise make the bundled subagents extension skip
+  // registering its `subagent` tool before the stage session exists. Isolate
+  // extension discovery from those parent-process flags so an explicit
+  // `tools: ["subagent"]` allowlist works the same in workflow stages everywhere.
+  // The isolation mutates process-global env, so serialize the full
+  // save/delete/reload/restore section. Without this queue, overlapping workflow
+  // stage session creation can snapshot an already-cleared env and restore that
+  // stale snapshot after another reload restores the real parent values.
+  const previousValues = new Map<string, string | undefined>();
+  for (const key of SUBAGENT_CHILD_EXTENSION_ENV_KEYS) {
+    previousValues.set(key, process.env[key]);
+    delete process.env[key];
+  }
+  try {
+    await resourceLoader.reload();
+  } finally {
+    for (const key of SUBAGENT_CHILD_EXTENSION_ENV_KEYS) {
+      const previousValue = previousValues.get(key);
+      if (previousValue === undefined) delete process.env[key];
+      else process.env[key] = previousValue;
+    }
+  }
 }
 
 async function createPiSdkAgentSession(
@@ -253,7 +300,7 @@ async function createTestAgentSession(_options?: CreateAgentSessionOptions): Pro
 function stripWorkflowOnlyOptions(options: (StageOptions | CreateAgentSessionOptions) | undefined): CreateAgentSessionOptions | undefined {
   if (!options) return options;
   const maybeWorkflowOptions = options as StageOptions;
-  const { mcp: _mcp, fallbackModels: _fallbackModels, ...sessionOptions } = maybeWorkflowOptions;
+  const { schema: _schema, mcp: _mcp, fallbackModels: _fallbackModels, ...sessionOptions } = maybeWorkflowOptions;
   return sessionOptions as CreateAgentSessionOptions;
 }
 
@@ -265,7 +312,7 @@ function makeWorkflowStageOrchestrationContext(meta: StageExecutionMeta): NonNul
     workflowStageName: meta.stageName,
     constraints: {
       disableWorkflowTool: true,
-      maxSubagentDepth: 1,
+      maxSubagentDepth: 2,
     },
   };
 }
@@ -499,6 +546,14 @@ export interface PiOverlayHandle {
  * (`overlay-adapter.ts`); inline pickers leave it unset and dismiss
  * via the factory `done()` callback.
  */
+export interface PiHostCustomUiState {
+  blockingInlineCustomUiDepth: number;
+  blockingInlineCustomUiActive: boolean;
+  blockingInlineCustomUiFocusDeferred?: boolean;
+}
+
+export type PiHostCustomUiStateListener = (state: PiHostCustomUiState) => void;
+
 export interface PiCustomOverlayOptions {
   /**
    * `true` mounts a floating popup; `false` mounts a focused
@@ -506,6 +561,8 @@ export interface PiCustomOverlayOptions {
    * place of the editor until the factory's `done()` callback fires.
    */
   overlay: boolean;
+  /** Keep host inline custom UI pending in the background while this overlay is visible. */
+  deferInlineCustomUiFocus?: boolean;
   /**
    * Geometry / anchoring intended for pi-tui's `resolveOverlayLayout`.
    * NOT forwarded by current pi interactive `custom()` — see
@@ -636,6 +693,12 @@ export interface PiUISurface {
   setTitle?: (title: string) => void;
   /** Show a custom component or overlay. */
   custom?: PiCustomOverlayFunction;
+  /** Get host-owned inline custom UI focus state, if exposed by the host. */
+  getHostCustomUiState?: () => PiHostCustomUiState;
+  /** Observe host-owned inline custom UI focus state changes, if exposed by the host. */
+  onHostCustomUiStateChange?: (listener: PiHostCustomUiStateListener) => () => void;
+  /** Move focus to a mounted host-owned inline custom UI, if one is pending. */
+  focusHostInlineCustomUi?: () => boolean;
   pasteToEditor?: (text: string) => void;
   setEditorText?: (text: string) => void;
   getEditorText?: () => string;

package/dist/builtin/workflows/src/extension/workflow-schema.ts

@@ -37,6 +37,39 @@ const McpOptionsSchema = Type.Object({
   deny: Type.Optional(Type.Array(Type.String())),
 });
 
+const JsonSchemaObjectTypeValue = {
+  anyOf: [
+    { const: "object" },
+    { type: "array", minItems: 1, maxItems: 1, items: { const: "object" } },
+  ],
+};
+
+const JsonSchemaExplicitObjectDescriptor = {
+  type: "object",
+  required: ["type"],
+  properties: { type: JsonSchemaObjectTypeValue },
+  additionalProperties: true,
+};
+
+const JsonSchemaObject = Type.Unsafe<Record<string, unknown>>({
+  description: "Top-level object JSON Schema used as structured_output tool arguments for this workflow item.",
+  anyOf: [
+    JsonSchemaExplicitObjectDescriptor,
+    {
+      type: "object",
+      required: ["allOf"],
+      properties: {
+        allOf: {
+          type: "array",
+          minItems: 1,
+          items: JsonSchemaExplicitObjectDescriptor,
+        },
+      },
+      additionalProperties: true,
+    },
+  ],
+});
+
 const BashCommandRuleSchema = Type.Union([
   Type.String(),
   Type.Object({ prefix: Type.String() }, { additionalProperties: false }),
@@ -55,6 +88,7 @@ const BashCommandPolicySchema = Type.Object({
 }, { additionalProperties: false });
 
 const StageSessionOptionProperties = {
+  schema: Type.Optional(JsonSchemaObject),
   cwd: Type.Optional(Type.String()),
   agentDir: Type.Optional(Type.String()),
   authStorage: Type.Optional(SdkSessionOptionSchema("authStorage")),

package/dist/builtin/workflows/src/runs/foreground/executor.ts

@@ -1267,6 +1267,15 @@ function truncateByBytes(text: string, maxBytes: number): { text: string; trunca
   return { text: text.slice(0, low), truncated: true };
 }
 
+function structuredTaskOutputText(value: unknown): string {
+  if (typeof value === "string") return value;
+  try {
+    return JSON.stringify(value, null, 2);
+  } catch (error) {
+    throw new Error(`atomic-workflows: structured task output is not JSON-serializable: ${error instanceof Error ? error.message : String(error)}`);
+  }
+}
+
 function truncateTaskOutput(text: string, maxOutput: WorkflowMaxOutput | undefined): string {
   const limits = normalizeMaxOutput(maxOutput);
   const byLines = truncateByLines(text, limits.lines);
@@ -4827,11 +4836,12 @@ export async function run<TInputs extends WorkflowInputValues>(
           taskStageOptions(resolvedTaskOptions),
           stageFailFastScope,
         );
-        const rawText = await stage.prompt(
+        const rawOutput = await stage.prompt(
           applyTaskContext(`${taskReadInstruction(resolvedTaskOptions)}${taskPrompt(resolvedTaskOptions)}`, taskPrevious(resolvedTaskOptions)),
           taskPromptOptions(resolvedTaskOptions),
         );
-        const text = truncateTaskOutput(rawText, resolvedTaskOptions.maxOutput);
+        const structured = typeof rawOutput === "string" ? undefined : rawOutput;
+        const text = truncateTaskOutput(structuredTaskOutputText(rawOutput), resolvedTaskOptions.maxOutput);
         const sessionId = (() => {
           try {
             return stage.sessionId;
@@ -4844,6 +4854,7 @@ export async function run<TInputs extends WorkflowInputValues>(
           name,
           stageName: name,
           text,
+          ...(structured !== undefined ? { structured: structured as WorkflowSerializableValue } : {}),
           ...(sessionId !== undefined ? { sessionId } : {}),
           ...(stage.sessionFile !== undefined ? { sessionFile: stage.sessionFile } : {}),
           ...(stageMeta.model !== undefined ? { model: stageMeta.model } : {}),

package/dist/builtin/workflows/src/runs/foreground/stage-runner.ts

@@ -10,11 +10,14 @@
 import { mkdir, writeFile } from "node:fs/promises";
 import { dirname, isAbsolute, resolve } from "node:path";
 import {
+  createStructuredOutputCapture,
+  createStructuredOutputTool,
   shouldApplyCodexFastModeForScope,
   SessionManager,
   type AgentSession,
   type CreateAgentSessionOptions,
   type PromptOptions,
+  type StructuredOutputCapture,
 } from "@bastani/atomic";
 import type {
   CompleteStageOpts,
@@ -28,6 +31,7 @@ import type {
   WorkflowExecutionMode,
   WorkflowModelCatalogPort,
 } from "../../shared/types.js";
+import type { Static, TSchema } from "typebox";
 import {
   buildModelCandidatesFromCatalog,
   errorMessage,
@@ -167,6 +171,7 @@ export interface InternalStageContext extends StageContext {
 function stripWorkflowOnlyOptions(options: StageOptions | undefined): CreateAgentSessionOptions {
   if (!options) return {};
   const {
+    schema: _schema,
     mcp: _mcp,
     fallbackModels: _fallbackModels,
     fallbackThinkingLevels: _fallbackThinkingLevels,
@@ -530,6 +535,43 @@ function splitPromptOptions(options: StagePromptOptions | undefined): {
   };
 }
 
+const STRUCTURED_OUTPUT_TOOL_NAME = "structured_output";
+
+function structuredOutputPrompt(text: string): string {
+  return `${text}\n\nFinal output contract:\n- Your final action MUST be a structured_output tool call.\n- Pass the schema fields directly as tool arguments; do not wrap them in { value: ... } unless the schema explicitly defines a top-level value field.\n- Do not emit a prose final answer instead of structured_output.\n- If you need to inspect files or run commands first, do so, then call structured_output exactly once.`;
+}
+
+function stringifyStructuredOutputValue(value: unknown): string {
+  try {
+    return JSON.stringify(value, null, 2);
+  } catch (error) {
+    throw new Error(`atomic-workflows: structured_output returned a non-serializable value: ${error instanceof Error ? error.message : String(error)}`);
+  }
+}
+
+function stageOptionsWithStructuredOutput(
+  options: StageOptions | undefined,
+  capture: StructuredOutputCapture<unknown> | undefined,
+): StageOptions | undefined {
+  if (!options?.schema || !capture) return options;
+  const tools = options.tools === undefined
+    ? undefined
+    : Array.from(new Set([...options.tools, STRUCTURED_OUTPUT_TOOL_NAME]));
+  const excludedTools = options.excludedTools?.filter((toolName) => toolName !== STRUCTURED_OUTPUT_TOOL_NAME);
+  return {
+    ...options,
+    ...(tools !== undefined ? { tools } : {}),
+    ...(excludedTools !== undefined ? { excludedTools } : {}),
+    customTools: [
+      ...(options.customTools ?? []),
+      createStructuredOutputTool({
+        schema: options.schema as TSchema,
+        capture: capture as StructuredOutputCapture<Static<TSchema>>,
+      }),
+    ],
+  };
+}
+
 function validatePromptOutputOptions(outputOptions: StageOutputOptions): void {
   if (outputOptions.outputMode === "file-only" && (typeof outputOptions.output !== "string" || outputOptions.output.length === 0)) {
     throw new Error(
@@ -564,7 +606,9 @@ async function finalizePromptOutput(
 
 export function createStageContext(opts: StageRunnerOpts): InternalStageContext {
   const { stageId, stageName, adapters, runId, signal, stageOptions, executionMode } = opts;
-  const meta: StageExecutionMeta = { runId, stageId, stageName, signal, stageOptions, executionMode };
+  const structuredOutputCapture = stageOptions?.schema ? createStructuredOutputCapture<unknown>() : undefined;
+  const effectiveStageOptions = stageOptionsWithStructuredOutput(stageOptions, structuredOutputCapture);
+  const meta: StageExecutionMeta = { runId, stageId, stageName, signal, stageOptions: effectiveStageOptions, executionMode };
   let session: StageSessionRuntime | undefined;
   let sessionPromise: Promise<StageSessionRuntime> | undefined;
   let lastAssistantText: string | undefined;
@@ -633,7 +677,7 @@ export function createStageContext(opts: StageRunnerOpts): InternalStageContext
   }
 
   const hasExplicitModelFallbackConfig =
-    stageOptions?.model !== undefined || (stageOptions?.fallbackModels?.length ?? 0) > 0;
+    effectiveStageOptions?.model !== undefined || (effectiveStageOptions?.fallbackModels?.length ?? 0) > 0;
   let candidatesPromise: Promise<WorkflowResolvedModelCandidate[]> | undefined;
   let activeCandidateIndex: number | undefined;
   let selectedModel: string | undefined;
@@ -653,9 +697,9 @@ export function createStageContext(opts: StageRunnerOpts): InternalStageContext
   function modelCandidates(): Promise<WorkflowResolvedModelCandidate[]> {
     if (!candidatesPromise) {
       candidatesPromise = buildModelCandidatesFromCatalog({
-        primaryModel: stageOptions?.model,
-        fallbackModels: stageOptions?.fallbackModels,
-        fallbackThinkingLevels: stageOptions?.fallbackThinkingLevels,
+        primaryModel: effectiveStageOptions?.model,
+        fallbackModels: effectiveStageOptions?.fallbackModels,
+        fallbackThinkingLevels: effectiveStageOptions?.fallbackThinkingLevels,
         catalog: modelCatalog,
       });
     }
@@ -663,9 +707,9 @@ export function createStageContext(opts: StageRunnerOpts): InternalStageContext
   }
 
   function stageOptionsForCandidate(candidate: WorkflowResolvedModelCandidate | undefined): StageOptions | undefined {
-    if (candidate === undefined) return stageOptions;
+    if (candidate === undefined) return effectiveStageOptions;
     return {
-      ...(stageOptions ?? {}),
+      ...(effectiveStageOptions ?? {}),
       model: candidate.value,
       ...(candidate.reasoningLevel !== undefined ? { thinkingLevel: candidate.reasoningLevel } : {}),
       fallbackModels: undefined,
@@ -677,7 +721,7 @@ export function createStageContext(opts: StageRunnerOpts): InternalStageContext
 
   function isWorkflowFastModeEnabled(): boolean | undefined {
     const model = session?.model;
-    const settingsManager = sessionSettingsManager ?? stageOptions?.settingsManager;
+    const settingsManager = sessionSettingsManager ?? effectiveStageOptions?.settingsManager;
     if (model === undefined || settingsManager === undefined) return undefined;
     return shouldApplyCodexFastModeForScope(model, settingsManager.getCodexFastModeSettings(), "workflow");
   }
@@ -705,7 +749,7 @@ export function createStageContext(opts: StageRunnerOpts): InternalStageContext
   }
 
   function effectiveCandidateReasoning(candidate: WorkflowResolvedModelCandidate): StageOptions["thinkingLevel"] | undefined {
-    return candidate.reasoningLevel ?? stageOptions?.thinkingLevel;
+    return candidate.reasoningLevel ?? effectiveStageOptions?.thinkingLevel;
   }
 
   function modelAttemptReasoning(candidate: WorkflowResolvedModelCandidate): Pick<WorkflowModelAttempt, "reasoningLevel"> {
@@ -715,7 +759,7 @@ export function createStageContext(opts: StageRunnerOpts): InternalStageContext
 
   function applyCandidateThinking(candidate: WorkflowResolvedModelCandidate | undefined): void {
     pendingThinkingLevel = candidate === undefined
-      ? stageOptions?.thinkingLevel
+      ? effectiveStageOptions?.thinkingLevel
       : effectiveCandidateReasoning(candidate);
   }
 
@@ -843,6 +887,13 @@ export function createStageContext(opts: StageRunnerOpts): InternalStageContext
     }
 
     let index = activeCandidateIndex ?? 0;
+    const capturedStructuredOutputForAttempt = (): boolean =>
+      structuredOutputCapture?.called === true && signal?.aborted !== true;
+    const recordSuccessfulAttempt = (candidate: WorkflowResolvedModelCandidate): void => {
+      modelAttempts.push({ model: candidate.id, success: true, ...modelAttemptReasoning(candidate) });
+      pendingFallbackWarnings.length = 0;
+    };
+
     while (index < candidates.length) {
       const candidate = candidates[index]!;
       const activeSession = session && activeCandidateIndex === index
@@ -855,13 +906,20 @@ export function createStageContext(opts: StageRunnerOpts): InternalStageContext
         const { terminalScanStartIndex } = await promptWithPauseResume(activeSession, text, sdkOptions);
         const terminalFailure = latestTerminalAssistantFailureSince(activeSession.messages, terminalScanStartIndex);
         if (terminalFailure !== undefined) {
+          if (capturedStructuredOutputForAttempt()) {
+            recordSuccessfulAttempt(candidate);
+            return;
+          }
           throw new WorkflowPromptModelFailure(terminalFailure);
         }
-        modelAttempts.push({ model: candidate.id, success: true, ...modelAttemptReasoning(candidate) });
-        pendingFallbackWarnings.length = 0;
+        recordSuccessfulAttempt(candidate);
         return;
       } catch (err) {
         const message = errorMessage(err);
+        if (capturedStructuredOutputForAttempt() && isRetryableModelFailure(err)) {
+          recordSuccessfulAttempt(candidate);
+          return;
+        }
         modelAttempts.push({ model: candidate.id, success: false, ...modelAttemptReasoning(candidate), error: message });
         if (signal?.aborted || !isRetryableModelFailure(err) || index === candidates.length - 1) {
           modelWarnings.push(...pendingFallbackWarnings);
@@ -887,15 +945,29 @@ export function createStageContext(opts: StageRunnerOpts): InternalStageContext
 
     async prompt(text, options) {
       const { sdkOptions, outputOptions } = splitPromptOptions(options);
-      const runtimeCwd = typeof stageOptions?.cwd === "string" ? stageOptions.cwd : process.cwd();
+      const runtimeCwd = typeof effectiveStageOptions?.cwd === "string" ? effectiveStageOptions.cwd : process.cwd();
       validatePromptOutputOptions(outputOptions);
+      if (structuredOutputCapture?.called) {
+        throw new Error("atomic-workflows: stage schema supports one prompt() call per stage context because structured_output may be called exactly once. Create a new ctx.stage(...) for each additional schema-backed prompt.");
+      }
       if (adapters.prompt) {
+        if (structuredOutputCapture) {
+          throw new Error("atomic-workflows: stage schema requires an AgentSessionAdapter so the structured_output tool can be registered.");
+        }
         const rawText = await adapters.prompt.prompt(text, meta);
         lastAssistantText = await finalizePromptOutput(rawText, outputOptions, runtimeCwd);
         adapterMessages = assistantMessage(lastAssistantText);
         return lastAssistantText;
       }
-      await promptWithFallback(text, sdkOptions);
+      await promptWithFallback(structuredOutputCapture ? structuredOutputPrompt(text) : text, sdkOptions);
+      if (structuredOutputCapture) {
+        if (!structuredOutputCapture.called) {
+          throw new Error("atomic-workflows: stage configured with schema must finish by calling structured_output.");
+        }
+        const rawStructuredText = stringifyStructuredOutputValue(structuredOutputCapture.value);
+        lastAssistantText = await finalizePromptOutput(rawStructuredText, outputOptions, runtimeCwd);
+        return structuredOutputCapture.value as never;
+      }
       const rawText = lastAssistantTextFromSession(session, lastAssistantText, terminatingToolCallIds) ?? "";
       lastAssistantText = await finalizePromptOutput(rawText, outputOptions, runtimeCwd);
       return lastAssistantText;

package/dist/builtin/workflows/src/shared/authoring-contract.d.ts

@@ -48,6 +48,7 @@ export interface WorkflowModelFallbackFields {
     readonly fallbackThinkingLevels?: readonly string[];
 }
 export type WorkflowModelValue = string | object;
+export type WorkflowStageResult<TSchemaDef extends TSchema | undefined = undefined> = [TSchemaDef] extends [TSchema] ? Static<TSchemaDef> : string;
 export interface WorkflowModelUsage extends WorkflowSerializableObject {
     readonly input?: number;
     readonly output?: number;
@@ -111,7 +112,9 @@ export interface WorkflowFastModeSettings extends WorkflowSerializableObject {
 export interface WorkflowFastModeSettingsManager {
     getCodexFastModeSettings(): WorkflowFastModeSettings;
 }
-export interface StageOptions extends WorkflowModelFallbackFields {
+export interface StageOptions<TSchemaDef extends TSchema | undefined = TSchema | undefined> extends WorkflowModelFallbackFields {
+    /** Optional structured final-answer schema. When set, the stage receives a schema-specific `structured_output` tool and must finish by calling it. */
+    readonly schema?: TSchemaDef;
     readonly model?: WorkflowModelValue;
     readonly mcp?: StageMcpOptions;
     readonly tools?: readonly string[];
@@ -231,9 +234,9 @@ export interface StageAdapters {
     readonly prompt?: PromptAdapter;
     readonly complete?: CompleteAdapter;
 }
-export interface StageContext {
+export interface StageContext<TSchemaDef extends TSchema | undefined = undefined> {
     readonly name: string;
-    prompt(text: string, options?: StagePromptOptions): Promise<string>;
+    prompt(text: string, options?: StagePromptOptions): Promise<WorkflowStageResult<TSchemaDef>>;
     complete(text: string, options?: CompleteStageOpts): Promise<string>;
     steer(text: string): Promise<void>;
     followUp(text: string): Promise<void>;
@@ -279,6 +282,8 @@ export interface WorkflowTaskContext extends WorkflowSerializableObject {
 export type WorkflowTaskContextInput = string | WorkflowTaskContext | WorkflowTaskResult;
 export interface WorkflowTaskResult extends WorkflowTaskContext {
     readonly stageName: string;
+    /** Parsed structured value when the task/stage was configured with `schema`. */
+    readonly structured?: WorkflowSerializableValue;
     readonly sessionId?: string;
     readonly sessionFile?: string;
     readonly artifacts?: readonly WorkflowArtifact[];
@@ -400,6 +405,9 @@ export interface WorkflowRunContext<TInputs extends WorkflowInputValues = Workfl
     readonly inputs: Readonly<TInputs>;
     readonly cwd?: string;
     exit(options?: WorkflowExitOptions<TOutputs>): never;
+    stage<TSchemaDef extends TSchema>(name: string, options: StageOptions<TSchemaDef> & {
+        readonly schema: TSchemaDef;
+    }): StageContext<TSchemaDef>;
     stage(name: string, options?: StageOptions): StageContext;
     task(name: string, options: WorkflowTaskOptions): Promise<WorkflowTaskResult>;
     chain(steps: readonly WorkflowTaskStep[], options?: WorkflowChainOptions): Promise<WorkflowTaskResult[]>;

package/dist/builtin/workflows/src/shared/types.ts

@@ -153,9 +153,11 @@ export interface StageMcpOptions extends AuthoringContract.StageMcpOptions {
  * All pi SDK createAgentSession options are forwarded to the stage session;
  * workflow-owned options such as `mcp` and `gitWorktreeDir` are stripped before SDK session creation.
  */
-export interface StageOptions
+export interface StageOptions<TSchemaDef extends TSchema | undefined = TSchema | undefined>
   extends Omit<CreateAgentSessionOptions, "model" | keyof AuthoringContract.StageOptions>,
-    Omit<Mutable<AuthoringContract.StageOptions>, "sessionManager" | "settingsManager"> {
+    Omit<Mutable<AuthoringContract.StageOptions<TSchemaDef>>, "sessionManager" | "settingsManager"> {
+  /** Optional structured final-answer schema. When set, the stage receives a schema-specific `structured_output` tool and must finish by calling it. */
+  schema?: TSchemaDef;
   /** Model id or pi SDK model object used as the primary stage model. */
   model?: WorkflowModelValue;
   /** Per-stage MCP server gating. No-op when no WorkflowMcpPort is configured. */
@@ -231,6 +233,7 @@ export interface WorkflowPersistencePort {
 export type WorkflowTaskContext = AuthoringContract.WorkflowTaskContext;
 export type WorkflowTaskContextInput = AuthoringContract.WorkflowTaskContextInput;
 export type WorkflowTaskResult = AuthoringContract.WorkflowTaskResult;
+export type WorkflowStageResult<TSchemaDef extends TSchema | undefined = undefined> = AuthoringContract.WorkflowStageResult<TSchemaDef>;
 
 /**
  * Higher-level task API: create a tracked stage, optionally inject prior task
@@ -276,12 +279,12 @@ export interface WorkflowDirectOptions extends StageOptions, Omit<Mutable<Author
  * This exposes the supported subset of pi's SDK AgentSession. The workflow
  * executor owns disposal and wraps prompt() with stage lifecycle tracking.
  */
-export interface StageContext {
+export interface StageContext<TSchemaDef extends TSchema | undefined = undefined> {
   /** Human-readable name for this stage (used in TUI + persistence). */
   readonly name: string;
 
   /** Send a prompt and wait for completion. */
-  prompt(text: string, options?: StagePromptOptions): Promise<string>;
+  prompt(text: string, options?: StagePromptOptions): Promise<WorkflowStageResult<TSchemaDef>>;
   complete(text: string, options?: CompleteStageOpts): Promise<string>;
 
   /** Queue messages during streaming. */
@@ -344,6 +347,7 @@ export interface WorkflowRunContext<
    * @param name   Human-readable stage name (used in TUI + persistence).
    * @param options Optional per-stage configuration (mcp allow/deny, etc.).
    */
+  stage<TSchemaDef extends TSchema>(name: string, options: StageOptions<TSchemaDef> & { schema: TSchemaDef }): StageContext<TSchemaDef>;
   stage(name: string, options?: StageOptions): StageContext;
   /**
    * Safe high-level task primitive. Equivalent to creating a named stage and