@bastani/atomic 0.8.31-alpha.4 → 0.9.0-alpha.1

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

@@ -0,0 +1,197 @@
+import type { StageOptions } from "../../shared/types.js";
+import type { ConcurrencyLimiter } from "../shared/concurrency.js";
+import type { StageAdapters } from "./stage-runner.js";
+import type { LiveStageRuntime } from "./executor-stage-types.js";
+import { askReadinessViaStageBroker, RESUME_CONTINUATION_PROMPT, shouldInjectResumeContinuation } from "./executor-hil.js";
+import { raceAbort } from "./executor-abort.js";
+import { hasExplicitFastModeCandidate } from "./executor-direct-helpers.js";
+import { applyFailureToStage } from "./executor-lifecycle.js";
+import { isTerminalStage } from "./executor-scheduler.js";
+
+export function createTrackedStageCaller(input: {
+  readonly runtime: LiveStageRuntime;
+  readonly limiter: ConcurrencyLimiter;
+  readonly options: StageOptions | undefined;
+  readonly adapters: StageAdapters;
+  readonly hasContinuation: boolean;
+  readonly hasScopedParents: boolean;
+}): <T>(call: () => Promise<T>, eagerSession?: boolean) => Promise<T> {
+  const { runtime } = input;
+  const readinessGateEnabled = runtime.opts.confirmStageReadiness !== undefined || runtime.opts.usePromptNodesForUi === true;
+  const confirmReadiness = async (): Promise<"advance" | "stay"> => {
+    try {
+      if (runtime.opts.confirmStageReadiness !== undefined) {
+        const ready = await runtime.opts.confirmStageReadiness({
+          runId: runtime.runId,
+          stageId: runtime.stageId,
+          stageName: runtime.name,
+          signal: runtime.signal,
+        });
+        return ready ? "advance" : "stay";
+      }
+      return await askReadinessViaStageBroker(runtime.runId, runtime.stageId, runtime.signal);
+    } catch {
+      return "advance";
+    }
+  };
+
+  const suppressReadinessForCurrentTurn = (): void => {
+    runtime.state.askUserQuestionObservedThisTurn = false;
+    runtime.state.chatAnswerObservedThisTurn = false;
+  };
+
+  const skipResumeContinuationInjection = (): boolean => {
+    if (runtime.state.stageFinalized) return true;
+    if (runtime.state.skippedForParallelFailFast) return true;
+    if (runtime.stageSnapshot.status === "skipped" && runtime.stageSnapshot.skippedReason === "fail-fast") return true;
+    if (isTerminalStage(runtime.stageSnapshot)) return true;
+    if (runtime.stageFailFastScope?.failed === true && runtime.stageFailFastScope.activeStages.has(runtime.stageId)) return true;
+    if (runtime.innerCtx.__structuredOutputFinalized()) return true;
+    return false;
+  };
+
+  const drainResumeContinuations = async <T>(currentResult: T): Promise<T> => {
+    let result = currentResult;
+    while (runtime.state.resumeContinuationPending) {
+      runtime.state.resumeContinuationPending = false;
+      suppressReadinessForCurrentTurn();
+      if (!shouldInjectResumeContinuation({
+        resumeOccurred: true,
+        gateEnabled: readinessGateEnabled,
+        aborted: runtime.signal.aborted,
+      })) {
+        continue;
+      }
+      if (skipResumeContinuationInjection()) continue;
+      result = await raceAbort(runtime.innerCtx.prompt(RESUME_CONTINUATION_PROMPT), runtime.signal) as T;
+    }
+    return result;
+  };
+
+  return async <T>(call: () => Promise<T>, eagerSession = false): Promise<T> => {
+    runtime.exit.throwIfWorkflowExitSelected();
+    await runtime.scheduler.waitForStageRelease(runtime.stageId, runtime.releaseLiveHandle);
+    if (runtime.state.stageFinalized) throw runtime.parallelFailFastError();
+
+    await input.limiter.acquire();
+    try {
+      await runtime.scheduler.waitForStageRelease(runtime.stageId, runtime.releaseLiveHandle);
+      runtime.exit.throwIfWorkflowExitSelected();
+      if (runtime.state.stageFinalized) throw runtime.parallelFailFastError();
+    } catch (err) {
+      input.limiter.release();
+      throw err;
+    }
+
+    if (!input.hasContinuation && runtime.stageSnapshot.startedAt === undefined && !input.hasScopedParents) {
+      const actualParentIds = runtime.scheduler.tracker.currentParents();
+      const sameParents = actualParentIds.length === runtime.stageSnapshot.parentIds.length &&
+        actualParentIds.every((value) => runtime.stageSnapshot.parentIds.includes(value));
+      if (!sameParents) {
+        runtime.scheduler.tracker.replaceParents(runtime.stageId, actualParentIds);
+        runtime.scheduler.setStageParentIds(runtime.stageSnapshot, actualParentIds);
+      }
+    }
+    runtime.stageSnapshot.status = "running";
+    runtime.stageSnapshot.startedAt = Date.now();
+    const hasNoExplicitModelConfig = input.options?.model === undefined && input.options?.fallbackModels === undefined;
+    const promptAdapterHandlesInitialPrompt = input.adapters.prompt !== undefined;
+    if (eagerSession && !promptAdapterHandlesInitialPrompt && (hasNoExplicitModelConfig || await hasExplicitFastModeCandidate({
+      model: input.options?.model,
+      fallbackModels: input.options?.fallbackModels,
+      models: runtime.opts.models,
+    }))) {
+      try {
+        await runtime.innerCtx.__ensureSession();
+        runtime.captureStageSessionMeta();
+      } catch (err) {
+        if (!(err instanceof Error && err.message.includes("prompt adapter not configured"))) throw err;
+      }
+    }
+    runtime.applyModelFallbackMeta(runtime.innerCtx.__modelFallbackMeta());
+    runtime.activeStore.recordStageStart(runtime.runId, runtime.stageSnapshot);
+    runtime.appendStageStartOnce();
+
+    const mcpAllow = input.options?.mcp?.allow ?? null;
+    const mcpDeny = input.options?.mcp?.deny ?? null;
+    const hasMcpScope = mcpAllow !== null || mcpDeny !== null;
+    if (runtime.opts.mcp && hasMcpScope) runtime.opts.mcp.setScope(runtime.stageId, mcpAllow, mcpDeny);
+
+    try {
+      const abortSession = (): void => {
+        void runtime.innerCtx.abort().catch(() => {});
+      };
+      if (runtime.signal.aborted) abortSession();
+      else runtime.signal.addEventListener("abort", abortSession, { once: true });
+      let result: T;
+      try {
+        runtime.state.askUserQuestionObservedThisTurn = false;
+        runtime.state.chatAnswerObservedThisTurn = false;
+        result = await raceAbort(call(), runtime.signal);
+        result = await drainResumeContinuations(result);
+
+        if (!runtime.signal.aborted && readinessGateEnabled) {
+          let resolveNextTurnEnd: (() => void) | null = null;
+          const unsubscribeTurnWatcher = runtime.innerCtx.subscribe((event) => {
+            if ((event as { type?: unknown }).type === "agent_end" && resolveNextTurnEnd) {
+              const resolve = resolveNextTurnEnd;
+              resolveNextTurnEnd = null;
+              resolve();
+            }
+          });
+          try {
+            while (runtime.state.askUserQuestionObservedThisTurn) {
+              const decision = runtime.state.chatAnswerObservedThisTurn ? "stay" : await confirmReadiness();
+              if (decision === "advance") break;
+              if (runtime.signal.aborted) break;
+              runtime.state.askUserQuestionObservedThisTurn = false;
+              runtime.state.chatAnswerObservedThisTurn = false;
+              await raceAbort(new Promise<void>((resolve) => { resolveNextTurnEnd = resolve; }), runtime.signal);
+              if (runtime.signal.aborted) break;
+              result = (runtime.innerCtx.__getLastAssistantText() ?? result) as T;
+              result = await drainResumeContinuations(result);
+            }
+          } finally {
+            resolveNextTurnEnd = null;
+            unsubscribeTurnWatcher();
+          }
+        }
+      } finally {
+        runtime.signal.removeEventListener("abort", abortSession);
+      }
+      runtime.captureStageSessionMeta();
+      runtime.applyModelFallbackMeta(runtime.innerCtx.__modelFallbackMeta());
+      if (runtime.stageFailFastScope?.failed === true && runtime.stageFailFastScope.activeStages.has(runtime.stageId)) {
+        runtime.markSkippedForParallelFailFast();
+        throw runtime.parallelFailFastError();
+      }
+      if (runtime.state.stageFinalized) throw runtime.parallelFailFastError();
+      runtime.stageSnapshot.status = "completed";
+      const assistantText = runtime.innerCtx.__getLastAssistantText();
+      if (assistantText !== undefined) runtime.stageSnapshot.result = assistantText;
+      return result;
+    } catch (err) {
+      const workflowExitAbort = runtime.signal.aborted ? runtime.exit.currentWorkflowExitAbortReason() : undefined;
+      if (workflowExitAbort !== undefined && !runtime.state.skippedForParallelFailFast) {
+        runtime.state.stageClosedByWorkflowExit = true;
+        if (!isTerminalStage(runtime.stageSnapshot)) {
+          runtime.stageSnapshot.status = "skipped";
+          runtime.stageSnapshot.skippedReason = runtime.exit.workflowExitSkippedReason(workflowExitAbort.reason);
+        }
+      } else if (!runtime.signal.aborted && !runtime.state.skippedForParallelFailFast) {
+        applyFailureToStage(runtime.stageSnapshot, runtime.classifyExecutorFailure(err));
+      }
+      throw err;
+    } finally {
+      if (runtime.opts.mcp && hasMcpScope) runtime.opts.mcp.clearScope(runtime.stageId);
+      runtime.captureStageSessionMeta();
+      runtime.finalizeStageSnapshot();
+      if (runtime.state.stageClosedByWorkflowExit || runtime.exit.currentWorkflowExitAbortReason() !== undefined) {
+        await runtime.releaseLiveHandle().catch(() => {});
+      } else {
+        await runtime.dropStageControlForCompletion().catch(() => {});
+      }
+      input.limiter.release();
+    }
+  };
+}

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

@@ -0,0 +1,120 @@
+import type { StageContextWithMeta, StageNoticeInput, LiveStageRuntime } from "./executor-stage-types.js";
+import type { InternalStageContext } from "./stage-runner.js";
+
+function noticeValue(value: unknown): string {
+  if (typeof value === "string") return value;
+  if (value === undefined || value === null) return "";
+  if (typeof value === "object") {
+    const candidate = value as { id?: unknown; name?: unknown; label?: unknown };
+    if (typeof candidate.id === "string") return candidate.id;
+    if (typeof candidate.name === "string") return candidate.name;
+    if (typeof candidate.label === "string") return candidate.label;
+  }
+  return String(value);
+}
+
+function compactionMeta(result: unknown): string | undefined {
+  if (result === undefined || result === null || typeof result !== "object") return undefined;
+  const compaction = result as {
+    stats?: { tokensBefore?: unknown; tokensAfter?: unknown };
+    tokensBefore?: unknown;
+    tokensAfter?: unknown;
+    tokensKept?: unknown;
+  };
+  const beforeRaw = compaction.stats?.tokensBefore ?? compaction.tokensBefore;
+  const keptRaw = compaction.stats?.tokensAfter ?? compaction.tokensKept ?? compaction.tokensAfter;
+  const before = typeof beforeRaw === "number" ? beforeRaw : undefined;
+  const kept = typeof keptRaw === "number" ? keptRaw : undefined;
+  if (before === undefined || kept === undefined) return undefined;
+  return `${(before / 1000).toFixed(1)}k → ${(kept / 1000).toFixed(1)}k`;
+}
+
+export function createStageContext(input: {
+  readonly runtime: LiveStageRuntime;
+  readonly runTrackedStageCall: <T>(call: () => Promise<T>, eagerSession?: boolean) => Promise<T>;
+}): StageContextWithMeta {
+  const { runtime } = input;
+  const recordStageNotice = (notice: StageNoticeInput): void => {
+    runtime.activeStore.recordStageNotice(runtime.runId, runtime.stageId, {
+      id: crypto.randomUUID(),
+      ts: Date.now(),
+      ...notice,
+    });
+  };
+
+  const innerCtx: InternalStageContext = runtime.innerCtx;
+  return {
+    name: innerCtx.name,
+    prompt: (text, promptOptions) => {
+      runtime.throwIfStageMutationBlocked();
+      return input.runTrackedStageCall(() => innerCtx.prompt(text, promptOptions), true);
+    },
+    complete: (text, completeOptions) => {
+      runtime.throwIfStageMutationBlocked();
+      return input.runTrackedStageCall(() => innerCtx.complete(text, completeOptions));
+    },
+    steer: (text) => {
+      runtime.throwIfStageMutationBlocked();
+      return innerCtx.steer(text);
+    },
+    followUp: (text) => {
+      runtime.throwIfStageMutationBlocked();
+      return innerCtx.followUp(text);
+    },
+    subscribe: (listener) => innerCtx.subscribe(listener),
+    get sessionFile() { return innerCtx.sessionFile; },
+    get sessionId() { return innerCtx.sessionId; },
+    setModel: async (model) => {
+      runtime.throwIfStageMutationBlocked();
+      await innerCtx.__ensureSession();
+      runtime.throwIfStageMutationBlocked();
+      recordStageNotice({ kind: "model", from: noticeValue(innerCtx.model), to: noticeValue(model) });
+      await innerCtx.setModel(model);
+    },
+    setThinkingLevel: (level) => {
+      runtime.throwIfStageMutationBlocked();
+      recordStageNotice({ kind: "thinking", from: noticeValue(innerCtx.thinkingLevel), to: noticeValue(level) });
+      innerCtx.setThinkingLevel(level);
+    },
+    cycleModel: async () => {
+      runtime.throwIfStageMutationBlocked();
+      const from = noticeValue(innerCtx.model);
+      const result = await innerCtx.cycleModel();
+      recordStageNotice({ kind: "model", from, to: noticeValue(innerCtx.model) });
+      return result;
+    },
+    cycleThinkingLevel: () => {
+      runtime.throwIfStageMutationBlocked();
+      const from = noticeValue(innerCtx.thinkingLevel);
+      const result = innerCtx.cycleThinkingLevel();
+      recordStageNotice({ kind: "thinking", from, to: noticeValue(innerCtx.thinkingLevel) });
+      return result;
+    },
+    get agent() { return innerCtx.agent; },
+    get model() { return innerCtx.model; },
+    get thinkingLevel() { return innerCtx.thinkingLevel; },
+    get messages() { return innerCtx.messages; },
+    get isStreaming() { return innerCtx.isStreaming; },
+    navigateTree: async (targetId, treeOptions) => {
+      runtime.throwIfStageMutationBlocked();
+      recordStageNotice({ kind: "tree", to: targetId });
+      return innerCtx.navigateTree(targetId, treeOptions);
+    },
+    compact: async () => {
+      runtime.throwIfStageMutationBlocked();
+      const result = await innerCtx.compact();
+      recordStageNotice({ kind: "compaction", to: "compacted", meta: compactionMeta(result) });
+      return result;
+    },
+    abortCompaction: () => {
+      runtime.throwIfStageMutationBlocked();
+      innerCtx.abortCompaction();
+    },
+    abort: async () => {
+      runtime.throwIfStageMutationBlocked();
+      recordStageNotice({ kind: "abort", to: "interrupted" });
+      await innerCtx.abort();
+    },
+    __modelFallbackMeta: () => innerCtx.__modelFallbackMeta(),
+  };
+}

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

@@ -0,0 +1,126 @@
+import type { StageControlHandle, AgentSessionEventListener } from "./stage-control-registry.js";
+import type { LiveStageRuntime } from "./executor-stage-types.js";
+import { isTerminalStage } from "./executor-scheduler.js";
+
+export function createStageControlHandle(runtime: LiveStageRuntime): StageControlHandle {
+  const ensureMessagingSession = async (): Promise<void> => {
+    const meta = runtime.innerCtx.__sessionMeta();
+    if (meta.sessionId !== undefined || meta.sessionFile !== undefined) return;
+    if (runtime.stageSnapshot.sessionFile !== undefined) {
+      await runtime.innerCtx.__ensureSessionFromFile(runtime.stageSnapshot.sessionFile);
+      runtime.captureStageSessionMeta();
+      return;
+    }
+    if (isTerminalStage(runtime.stageSnapshot)) {
+      throw new Error(`atomic-workflows: cannot message stage "${runtime.name}" because no retained session metadata is available.`);
+    }
+  };
+
+  return {
+    runId: runtime.runId,
+    stageId: runtime.stageId,
+    stageName: runtime.name,
+    get status() {
+      return runtime.stageSnapshot.status;
+    },
+    get sessionId() {
+      return runtime.innerCtx.__sessionMeta().sessionId ?? runtime.stageSnapshot.sessionId;
+    },
+    get sessionFile() {
+      return runtime.innerCtx.__sessionMeta().sessionFile ?? runtime.stageSnapshot.sessionFile;
+    },
+    get isStreaming() {
+      return runtime.innerCtx.isStreaming;
+    },
+    get isDisposed() {
+      return runtime.state.liveHandleReleased;
+    },
+    get messages() {
+      return runtime.innerCtx.messages;
+    },
+    get agentSession() {
+      return runtime.innerCtx.__agentSession();
+    },
+    async ensureAttached() {
+      runtime.throwIfStageMutationBlocked();
+      await ensureMessagingSession();
+      await runtime.innerCtx.__ensureSession();
+      runtime.throwIfStageMutationBlocked();
+      runtime.captureStageSessionMeta();
+    },
+    async prompt(text: string) {
+      runtime.throwIfStageMutationBlocked();
+      await ensureMessagingSession();
+      try {
+        await runtime.innerCtx.prompt(text);
+      } finally {
+        runtime.captureStageSessionMeta();
+      }
+      runtime.throwIfStageMutationBlocked();
+    },
+    async steer(text: string) {
+      runtime.throwIfStageMutationBlocked();
+      await ensureMessagingSession();
+      try {
+        await runtime.innerCtx.steer(text);
+      } finally {
+        runtime.captureStageSessionMeta();
+      }
+    },
+    async followUp(text: string) {
+      runtime.throwIfStageMutationBlocked();
+      await ensureMessagingSession();
+      try {
+        await runtime.innerCtx.followUp(text);
+      } finally {
+        runtime.captureStageSessionMeta();
+      }
+    },
+    async pause() {
+      runtime.throwIfStageMutationBlocked();
+      const statusBeforePause = runtime.stageSnapshot.status;
+      const changed = runtime.activeStore.recordStagePaused(runtime.runId, runtime.stageId);
+      if (changed) {
+        runtime.scheduler.ensureReleaseBarrier(runtime.stageId);
+        await runtime.scheduler.cascadePauseFrom(runtime.stageId);
+        const run = runtime.activeStore.runs().find((candidate) => candidate.id === runtime.runId);
+        const stillActive = run?.stages.some(
+          (s) => s.status === "running" && s.id !== runtime.stageId,
+        ) ?? false;
+        if (!stillActive) runtime.activeStore.recordRunPaused(runtime.runId);
+      }
+      if (statusBeforePause === "pending" || statusBeforePause === "running" || runtime.innerCtx.isStreaming) {
+        await runtime.innerCtx.__requestPause();
+      }
+    },
+    async resume(message?: string) {
+      runtime.throwIfStageMutationBlocked();
+      await ensureMessagingSession();
+      const wasPausedBeforeResume = runtime.innerCtx.__isPaused();
+      const hasResumeContinuationMessage = typeof message === "string" && message.trim().length > 0;
+      const previousResumeContinuationPending = runtime.state.resumeContinuationPending;
+      const queuedResumeContinuation = wasPausedBeforeResume && hasResumeContinuationMessage;
+      if (queuedResumeContinuation) runtime.state.resumeContinuationPending = true;
+      try {
+        const changed = runtime.activeStore.recordStageResumed(runtime.runId, runtime.stageId);
+        if (changed) {
+          runtime.scheduler.releaseStageBarrier(runtime.stageId);
+          await runtime.scheduler.cascadeResumeFrom(runtime.stageId);
+          runtime.activeStore.recordRunResumed(runtime.runId);
+        }
+        await runtime.innerCtx.__resume(message);
+      } catch (err) {
+        if (queuedResumeContinuation) runtime.state.resumeContinuationPending = previousResumeContinuationPending;
+        throw err;
+      } finally {
+        runtime.captureStageSessionMeta();
+      }
+    },
+    subscribe(listener: AgentSessionEventListener) {
+      return runtime.innerCtx.subscribe(listener);
+    },
+    async dispose() {
+      await runtime.releaseLiveHandle();
+    },
+  };
+}