wave-agent-sdk 0.17.1 → 0.17.2

package/src/workflow/workflowApis.ts

@@ -0,0 +1,412 @@
+import type { SubagentManager } from "../managers/subagentManager.js";
+import * as fs from "fs";
+import * as path from "path";
+import { ConcurrencyLimiter } from "./concurrencyLimiter.js";
+import { BudgetTracker } from "./budgetTracker.js";
+import { ProgressReporter } from "./progressReporter.js";
+import { Journal } from "./journal.js";
+import type { BudgetInfo, AgentMeta, WorkflowProgressEvent } from "./types.js";
+import {
+  createStructuredOutputPrompt,
+  createStructuredOutputTool,
+  extractStructuredResult,
+} from "./structuredOutput.js";
+import { AGENT_TOOL_NAME, WORKFLOW_TOOL_NAME } from "../constants/tools.js";
+import { logger } from "../utils/globalLogger.js";
+
+/**
+ * Attempt to parse args if it was passed as a JSON string.
+ * LLMs may serialize array/object args as a string (e.g. "[3,8,15,42,7]")
+ * which would cause `for...of` to iterate character-by-character.
+ */
+function parseArgsIfNeeded(args: unknown): unknown {
+  if (typeof args === "string") {
+    try {
+      const parsed = JSON.parse(args);
+      // Only accept parsed arrays/objects, not primitives like "42" → 42
+      if (typeof parsed === "object" && parsed !== null) {
+        return parsed;
+      }
+    } catch {
+      // Not valid JSON — use as-is (string)
+    }
+  }
+  return args;
+}
+
+const MAX_TOTAL_AGENTS = 1000;
+const MAX_ITEMS_PER_CALL = 4096;
+
+interface WorkflowApiContext {
+  subagentManager: SubagentManager;
+  concurrencyLimiter: ConcurrencyLimiter;
+  budgetTracker: BudgetTracker;
+  progressReporter: ProgressReporter;
+  journal: Journal;
+  abortSignal: AbortSignal;
+  args: unknown;
+  onLog?: (message: string) => void;
+  /** When resuming, offset the agent counter by this many existing entries */
+  initialAgentCount?: number;
+  /** Session directory for computing transcript paths */
+  sessionDir: string;
+  /** Per-run directory for writing agent metadata sidecars */
+  runDir: string;
+  /** Per-agent abort controllers for kill/skip support */
+  agentControllers: Map<number, AbortController>;
+  /** Optional callback for progress events */
+  onProgress?: (event: WorkflowProgressEvent) => void;
+}
+
+interface AgentOpts {
+  label?: string;
+  phase?: string;
+  schema?: object;
+  model?: string;
+  isolation?: string;
+  agentType?: string;
+}
+
+export interface WorkflowApis {
+  agent: (prompt: string, opts?: AgentOpts) => Promise<unknown>;
+  parallel: (thunks: Array<() => Promise<unknown>>) => Promise<unknown[]>;
+  pipeline: (
+    items: unknown[],
+    ...stages: Array<
+      (prev: unknown, item: unknown, index: number) => Promise<unknown>
+    >
+  ) => Promise<unknown[]>;
+  phase: (title: string) => void;
+  log: (message: string) => void;
+  args: unknown;
+  budget: BudgetInfo;
+}
+
+export function createWorkflowApis(ctx: WorkflowApiContext): WorkflowApis {
+  let agentCounter = ctx.initialAgentCount ?? 0;
+
+  const agent = async (prompt: string, opts?: AgentOpts): Promise<unknown> => {
+    const index = agentCounter++;
+
+    // Check agent limit
+    if (index >= MAX_TOTAL_AGENTS) {
+      throw new Error(
+        `Workflow exceeded maximum agent count of ${MAX_TOTAL_AGENTS}`,
+      );
+    }
+
+    // Check abort
+    if (ctx.abortSignal.aborted) {
+      return null;
+    }
+
+    // Check budget
+    if (ctx.budgetTracker.isExceeded()) {
+      throw new Error("Workflow token budget exceeded");
+    }
+
+    // Check journal for cached result (resume)
+    const cached = ctx.journal.getCachedResult(index);
+    if (cached !== undefined) {
+      logger.debug(`[Workflow] agent(${index}): using cached result`);
+      return cached;
+    }
+
+    // Acquire concurrency slot
+    await ctx.concurrencyLimiter.acquire();
+
+    // Create per-agent abort controller linked to the run's signal
+    const agentController = new AbortController();
+    ctx.agentControllers.set(index, agentController);
+    // If the run is already aborted, abort this agent immediately
+    if (ctx.abortSignal.aborted) {
+      agentController.abort();
+    }
+    // Propagate run abort to agent abort
+    const onRunAbort = () => agentController.abort();
+    ctx.abortSignal.addEventListener("abort", onRunAbort);
+
+    try {
+      // Resolve subagent type
+      const subagentType = opts?.agentType || "general-purpose";
+      let configuration = await ctx.subagentManager.findSubagent(subagentType);
+
+      if (!configuration) {
+        logger.warn(
+          `[Workflow] agent(${index}): subagent type "${subagentType}" not found, falling back to general-purpose`,
+        );
+        configuration =
+          await ctx.subagentManager.findSubagent("general-purpose");
+        if (!configuration) {
+          throw new Error(`No subagent type available for agent call`);
+        }
+      }
+
+      // Build the effective prompt
+      let effectivePrompt = prompt;
+      if (opts?.schema) {
+        effectivePrompt += createStructuredOutputPrompt(opts.schema);
+      }
+
+      // Set phase if specified
+      if (opts?.phase) {
+        ctx.progressReporter.setPhase(opts.phase);
+      }
+      ctx.progressReporter.agentStarted();
+
+      // Create subagent instance
+      const instance = await ctx.subagentManager.createInstance(configuration, {
+        description: opts?.label || `workflow-agent-${index}`,
+        prompt: effectivePrompt,
+        subagent_type: subagentType,
+        model: opts?.model,
+      });
+
+      // Capture subagent linkage info
+      const subagentId = instance.subagentId;
+      const transcriptPath = instance.messageManager.getTranscriptPath();
+
+      // If schema provided, register StructuredOutput tool on the subagent's tool manager
+      // and force the model to call it via tool_choice
+      if (opts?.schema) {
+        const structuredTool = createStructuredOutputTool(opts.schema);
+        instance.toolManager.register(structuredTool);
+        instance.aiManager.toolChoiceOverride = {
+          type: "function",
+          function: { name: "StructuredOutput" },
+        };
+      }
+
+      // Deny Agent and Workflow tools in workflow subagents
+      // (prevent infinite recursion)
+      instance.permissionManager.addTemporaryRules([
+        `${AGENT_TOOL_NAME}:deny`,
+        `${WORKFLOW_TOOL_NAME}:deny`,
+      ]);
+
+      // Execute agent
+      const result = await ctx.subagentManager.executeAgent(
+        instance,
+        effectivePrompt,
+        agentController.signal,
+        false,
+      );
+
+      // Track token usage — sum from assistant messages' usage field
+      // (getUsages() is empty because onUsageAdded targets the parent agent)
+      const messages = instance.messageManager.getMessages();
+      const tokens = messages.reduce((sum, msg) => {
+        if (msg.role !== "assistant" || !msg.usage) return sum;
+        const u = msg.usage;
+        return (
+          sum +
+          (u.total_tokens || 0) +
+          (u.cache_read_input_tokens || 0) +
+          (u.cache_creation_input_tokens || 0)
+        );
+      }, 0);
+      ctx.budgetTracker.addUsage(tokens);
+      ctx.progressReporter.agentCompleted(tokens);
+
+      // Extract structured result if schema was provided
+      let finalResult: unknown;
+      if (opts?.schema) {
+        const messages = instance.messageManager.getMessages();
+        // Build the format extractStructuredResult expects:
+        // Look for StructuredOutput tool calls in ToolBlocks
+        const structuredToolBlocks = messages
+          .filter((m) => m.role === "assistant")
+          .flatMap((m) => m.blocks)
+          .filter(
+            (b): b is import("../types/messaging.js").ToolBlock =>
+              b.type === "tool" &&
+              b.name === "StructuredOutput" &&
+              b.stage === "end",
+          );
+
+        if (structuredToolBlocks.length > 0) {
+          // Parse the StructuredOutput tool call parameters
+          const lastBlock =
+            structuredToolBlocks[structuredToolBlocks.length - 1];
+          try {
+            const parsed = JSON.parse(lastBlock.parameters || "{}");
+            finalResult = parsed;
+          } catch {
+            // Parameters parse failed, try result
+            try {
+              const parsed = JSON.parse(lastBlock.result || "null");
+              finalResult = parsed;
+            } catch {
+              finalResult = result;
+            }
+          }
+        } else {
+          // Fallback: try extractStructuredResult with message content
+          finalResult = extractStructuredResult(
+            messages.map((m) => {
+              const textBlock = m.blocks.find(
+                (b): b is import("../types/messaging.js").TextBlock =>
+                  b.type === "text",
+              );
+              return {
+                role: m.role,
+                content: textBlock?.content,
+                tool_calls: (
+                  m as unknown as {
+                    tool_calls?: Array<{
+                      function: { name: string; arguments: string };
+                    }>;
+                  }
+                ).tool_calls,
+              };
+            }),
+            opts.schema,
+          );
+          if (finalResult === null) {
+            // Schema enforcement failed — return the raw text
+            finalResult = result;
+          }
+        }
+      } else {
+        finalResult = result;
+      }
+
+      // Append to journal
+      ctx.journal.append({
+        agentIndex: index,
+        prompt,
+        opts: { ...opts } as Record<string, unknown>,
+        result: finalResult,
+        tokens,
+        subagentId,
+        transcriptPath,
+      });
+
+      // Write agent metadata sidecar
+      try {
+        const agentMeta: AgentMeta = {
+          agentType: subagentType,
+          subagentId,
+          transcriptPath,
+          label: opts?.label,
+          phase: opts?.phase,
+        };
+        const metaPath = path.join(ctx.runDir, "agents", `${index}.meta.json`);
+        await fs.promises.writeFile(
+          metaPath,
+          JSON.stringify(agentMeta, null, 2),
+          "utf-8",
+        );
+      } catch (metaErr) {
+        logger.warn(
+          `[Workflow] Failed to write agent meta for ${index}: ${metaErr instanceof Error ? metaErr.message : String(metaErr)}`,
+        );
+      }
+
+      // Cleanup
+      ctx.subagentManager.cleanupInstance(instance.subagentId);
+
+      return finalResult;
+    } catch (error) {
+      // Agent errors are logged but don't crash the workflow
+      // Return null so the caller can filter with .filter(Boolean)
+      const errorMsg = error instanceof Error ? error.message : String(error);
+      logger.warn(`[Workflow] agent(${index}) failed: ${errorMsg}`);
+
+      // Write agent_failed entry to journal for skip/retry tracking
+      ctx.journal.append({
+        type: "agent_failed",
+        agentIndex: index,
+        error: errorMsg,
+      });
+
+      // Emit progress event for agent failure
+      ctx.progressReporter.agentFailed(index);
+
+      return null;
+    } finally {
+      ctx.agentControllers.delete(index);
+      ctx.abortSignal.removeEventListener("abort", onRunAbort);
+      ctx.concurrencyLimiter.release();
+    }
+  };
+
+  const parallel = async (
+    thunks: Array<() => Promise<unknown>>,
+  ): Promise<unknown[]> => {
+    if (thunks.length > MAX_ITEMS_PER_CALL) {
+      throw new Error(
+        `parallel() accepts at most ${MAX_ITEMS_PER_CALL} items, got ${thunks.length}`,
+      );
+    }
+
+    // Note: thunks typically call agent() which acquires its own slot,
+    // so parallel does NOT acquire a slot per thunk to avoid deadlock.
+    const results = await Promise.allSettled(
+      thunks.map(async (thunk) => {
+        try {
+          return await thunk();
+        } catch {
+          return null;
+        }
+      }),
+    );
+
+    // Convert rejected promises to null
+    return results.map((r) => (r.status === "fulfilled" ? r.value : null));
+  };
+
+  const pipeline = async (
+    items: unknown[],
+    ...stages: Array<
+      (prev: unknown, item: unknown, index: number) => Promise<unknown>
+    >
+  ): Promise<unknown[]> => {
+    if (items.length > MAX_ITEMS_PER_CALL) {
+      throw new Error(
+        `pipeline() accepts at most ${MAX_ITEMS_PER_CALL} items, got ${items.length}`,
+      );
+    }
+
+    // Run each item through all stages, items are independent.
+    // Note: agent() inside stages acquires its own concurrency slot,
+    // so pipeline does NOT acquire a slot per item to avoid deadlock.
+    const results = await Promise.allSettled(
+      items.map(async (item, index) => {
+        try {
+          let result: unknown = undefined;
+          for (const stage of stages) {
+            result = await stage(result, item, index);
+          }
+          return result;
+        } catch (error) {
+          logger.warn(
+            `[Workflow] pipeline item ${index} failed: ${error instanceof Error ? error.message : String(error)}`,
+          );
+          return null;
+        }
+      }),
+    );
+
+    return results.map((r) => (r.status === "fulfilled" ? r.value : null));
+  };
+
+  const phase = (title: string): void => {
+    ctx.progressReporter.setPhase(title);
+  };
+
+  const log = (message: string): void => {
+    ctx.journal.appendLog(message);
+    ctx.onLog?.(message);
+  };
+
+  return {
+    agent,
+    parallel,
+    pipeline,
+    phase,
+    log,
+    args: parseArgsIfNeeded(ctx.args),
+    budget: ctx.budgetTracker.toBudgetInfo(),
+  };
+}