diff --git a/core/src/common.ts b/core/src/common.ts index 026811635..f5fb03658 100644 --- a/core/src/common.ts +++ b/core/src/common.ts @@ -285,10 +285,12 @@ export type { VertexAiSearchToolParams, } from './tools/vertex_ai_search_tool.js'; export {VertexRagRetrievalTool} from './tools/vertex_rag_retrieval_tool.js'; +export {AsyncQueue} from './utils/async_queue.js'; export {getClientLabels, runWithClientLabel} from './utils/client_labels.js'; export {LogLevel, getLogger, setLogLevel, setLogger} from './utils/logger.js'; export type {Logger} from './utils/logger.js'; export {isGemini2OrAbove, isGemini3xFlashLive} from './utils/model_name.js'; +export type {SchemaLike} from './utils/schema.js'; export {zodObjectToSchema} from './utils/simple_zod_to_json.js'; export {Task} from './utils/task.js'; export type {TaskExecutable} from './utils/task.js'; @@ -331,6 +333,65 @@ export { createRestApiTool, } from './tools/openapi_tool/rest_api_tool.js'; +// Workflow (parity port of google/adk-python `google/adk/workflow`). Named +// explicitly (not `export *`) so the top-level surface stays intentional and +// collisions are compile errors; keep this in sync with `./workflow/index.js`. +export { + BaseNode, + BranchPath, + DEFAULT_ROUTE, + Edge, + FunctionNode, + Graph, + JoinNode, + NodeContext, + NodeStatus, + NodeTimeoutError, + ParallelWorker, + RequestInput, + START, + ToolNode, + Workflow, + WorkflowAgent, + WorkflowNode, + commonPrefixOf, + createNodeState, + createSubBranch, + isNodeState, + isRequestInput, + node, + normalizeRetryExceptions, + prepareRetryConfig, +} from './workflow/index.js'; +export type { + BaseNodeConfig, + BuildNodeOptions, + ChainElement, + DynamicEntry, + EdgeItem, + ErrorClass, + FunctionNodeConfig, + FunctionNodeHandler, + FunctionNodeResult, + NodeContextOptions, + NodeLike, + NodeOptions, + NodeResult, + NodeState, + ParallelWorkerConfig, + PreparedRetryConfig, + RequestInputParams, + RetryConfig, + RouteValue, + RoutingMap, + RunNodeOptions, + ScheduleDynamicNode, + ScheduleDynamicNodeOptions, + ToolNodeConfig, + WorkflowAgentConfig, + WorkflowConfig, +} from './workflow/index.js'; + export * from './apps/app.js'; export * from './artifacts/base_artifact_service.js'; export * from './features/feature_registry.js'; diff --git a/core/src/events/event.ts b/core/src/events/event.ts index 9e20794e5..6d39dc3e1 100644 --- a/core/src/events/event.ts +++ b/core/src/events/event.ts @@ -365,10 +365,12 @@ const PRESERVE_KEYS_CAMEL_CASE = [ 'customMetadata', 'content.parts.functionCall.args', 'content.parts.functionResponse.response', - // Workflow: arbitrary node output and checkpointed node state carry - // user-defined keys that must survive round-trips verbatim (a node's original - // input is stashed under `actions.agentState` for HITL resume). + // Workflow: arbitrary node output, emitted route(s), and checkpointed node + // state carry user-defined keys that must survive round-trips verbatim (a + // node's original input is stashed under `actions.agentState` for HITL + // resume, and rehydration reads `output`/`route`/`actions.agentState` back). 'output', + 'route', 'actions.agentState', ]; @@ -389,9 +391,10 @@ const PRESERVE_KEYS_SNAKE_CASE = [ 'custom_metadata', 'content.parts.function_call.args', 'content.parts.function_response.response', - // Workflow: arbitrary node output and checkpointed node state (see the - // camelCase list above). + // Workflow: arbitrary node output, emitted route(s), and checkpointed node + // state (see the camelCase list above). 'output', + 'route', 'actions.agent_state', ]; diff --git a/core/src/utils/schema.ts b/core/src/utils/schema.ts index 7c4cb236d..b7dccd4fa 100644 --- a/core/src/utils/schema.ts +++ b/core/src/utils/schema.ts @@ -26,8 +26,8 @@ import { * A schema accepted by ADK APIs, expressed as a Zod v3 type, a Zod v4 type, or * a genai `Schema`. * - * Use {@link parseWithSchema} to validate a value against one, and - * {@link toJsonSchema} to render one as a plain JSON Schema. + * Use `parseWithSchema` to validate a value against one, and `toJsonSchema` to + * render one as a plain JSON Schema. */ export type SchemaLike = z3.ZodType | z4.ZodType | Schema; diff --git a/core/src/workflow/base_node.ts b/core/src/workflow/base_node.ts index 7423a85e6..119e40014 100644 --- a/core/src/workflow/base_node.ts +++ b/core/src/workflow/base_node.ts @@ -76,7 +76,7 @@ export interface BaseNodeConfig { * {@link Event}s consumed by the engine. */ export abstract class BaseNode { - /** Brand identifying this object as a {@link BaseNode} (see {@link isBaseNode}). */ + /** Brand identifying this object as a {@link BaseNode} (see `isBaseNode`). */ readonly [BASE_NODE_SIGNATURE_SYMBOL] = true; readonly name: string; @@ -163,7 +163,7 @@ export abstract class BaseNode { /** * Validates node input against `inputSchema` (Content passes through). Only * enforced for Zod schemas; a genai `Schema` is left unvalidated (see - * {@link parseWithSchema}). + * `parseWithSchema`). */ protected validateInput(input: TInput): TInput { if (isContent(input)) { @@ -175,7 +175,7 @@ export abstract class BaseNode { /** * Validates node output against `outputSchema` (Content passes through). Only * enforced for Zod schemas; a genai `Schema` is left unvalidated (see - * {@link parseWithSchema}). + * `parseWithSchema`). */ protected validateOutput(output: unknown): unknown { if (isContent(output)) { diff --git a/core/src/workflow/dynamic_node_scheduler.ts b/core/src/workflow/dynamic_node_scheduler.ts new file mode 100644 index 000000000..7cfbe1ec0 --- /dev/null +++ b/core/src/workflow/dynamic_node_scheduler.ts @@ -0,0 +1,151 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {BaseNode} from './base_node.js'; +import {NodeContext, NodeResult} from './node_context.js'; +import {executeChildNode} from './node_runner.js'; +import {createNodeState} from './node_state.js'; +import {NodeStatus} from './node_status.js'; +import { + DynamicNodeRun, + DynamicNodeState, + ScheduleDynamicNode, + ScheduleDynamicNodeOptions, +} from './schedule_dynamic_node.js'; +import { + isFastForwardable, + makeFastForwardResult, + reconstructNodeStatesByPath, +} from './utils/rehydration_utils.js'; + +/** + * Handles `ctx.runNode()` calls for a {@link Workflow} subtree. + * + * Ported (Phase 4 subset) from `google/adk-python` + * `workflow/_dynamic_node_scheduler.py`. Implemented now: fresh execution and + * deduplication of concurrent calls to the same node path. Resumption from + * session events (rehydration + replay interception) is added in Phase 5 at the + * marked hook point. + */ +export class DynamicNodeScheduler implements ScheduleDynamicNode { + /** + * @param state Shared dynamic-node bookkeeping for this workflow subtree. + * @param abortSignal Workflow-scoped cancellation signal, forwarded to each + * dynamic child so a workflow shutting down on error can cancel in-flight + * `ctx.runNode()` children too. + */ + constructor( + private readonly state: DynamicNodeState, + private readonly abortSignal?: AbortSignal, + ) {} + + async schedule( + ctx: NodeContext, + node: BaseNode, + input: unknown, + options: ScheduleDynamicNodeOptions, + ): Promise { + const name = options.nodeName ?? node.name; + const runId = options.runId; + const nodePath = ctx.nodePath + ? `${ctx.nodePath}.${name}@${runId}` + : `${name}@${runId}`; + + const existing = this.state.runs.get(nodePath); + if (existing?.task) { + // Deduplicate concurrent calls: await the in-flight task. + return existing.task; + } + + // Cross-turn resume: rehydrate this dynamic run from prior session events. + if (!this.state.runs.has(nodePath)) { + const prior = reconstructNodeStatesByPath(ctx.session?.events ?? []).get( + nodePath, + ); + if (prior && !node.rerunOnResume && isFastForwardable(prior)) { + // Completed in a prior turn -> return cached output, do not re-execute. + this.state.runs.set(nodePath, { + state: createNodeState({ + status: NodeStatus.COMPLETED, + runId, + parentRunId: ctx.runId, + }), + output: prior.output, + }); + if (options.useAsOutput) { + ctx.output = prior.output; + ctx.route = prior.route; + } + return makeFastForwardResult(ctx, prior); + } + // Otherwise (waiting/unresolved): resume inputs were already merged into + // ctx.resumeInputs by the Workflow; fall through to a fresh run. + } + + return this.runFresh(ctx, node, input, name, runId, nodePath, options); + } + + private async runFresh( + ctx: NodeContext, + node: BaseNode, + input: unknown, + name: string, + runId: string, + nodePath: string, + options: ScheduleDynamicNodeOptions, + ): Promise { + const run: DynamicNodeRun = { + state: createNodeState({ + status: NodeStatus.RUNNING, + input, + runId, + parentRunId: ctx.runId, + }), + }; + this.state.runs.set(nodePath, run); + + run.task = executeChildNode({ + parent: ctx, + node, + input, + abortSignal: this.abortSignal, + options: { + nodeName: name, + runId, + overrideNodePath: nodePath, + useAsOutput: options.useAsOutput, + useSubBranch: options.useSubBranch, + overrideBranch: options.overrideBranch, + overrideIsolationScope: options.overrideIsolationScope, + }, + }); + + const childCtx = await run.task; + this.recordResult(run, childCtx, node); + return childCtx; + } + + private recordResult( + run: DynamicNodeRun, + childCtx: NodeContext, + node: BaseNode, + ): void { + if (childCtx.interruptIds.length > 0) { + run.state.status = NodeStatus.WAITING; + run.state.interrupts = [...childCtx.interruptIds]; + childCtx.interruptIds.forEach((id) => this.state.interruptIds.add(id)); + } else if ( + node.waitForOutput && + childCtx.output === undefined && + childCtx.route === undefined + ) { + run.state.status = NodeStatus.WAITING; + } else { + run.state.status = NodeStatus.COMPLETED; + run.output = childCtx.output; + } + } +} diff --git a/core/src/workflow/graph.ts b/core/src/workflow/graph.ts index a537f4dc9..1d29c1909 100644 --- a/core/src/workflow/graph.ts +++ b/core/src/workflow/graph.ts @@ -13,7 +13,7 @@ import {validateGraph} from './utils/graph_validation.js'; /** * A unique symbol branding {@link Edge} instances. * - * {@link isEdge} matches on this brand rather than `instanceof` so an edge built + * `isEdge` matches on this brand rather than `instanceof` so an edge built * by another copy of adk-js in the same runtime is still recognised (an * `instanceof` check fails across package copies) — mirroring the * `Symbol.for('google.adk.*')` brands used across ADK. @@ -78,7 +78,7 @@ export type EdgeItem = Edge | ChainElement[]; * Mirrors `google/adk-python` `workflow/_graph.py::Edge`. */ export class Edge { - /** Brand identifying this object as an {@link Edge} (see {@link isEdge}). */ + /** Brand identifying this object as an {@link Edge} (see `isEdge`). */ readonly [EDGE_SIGNATURE_SYMBOL] = true; constructor( diff --git a/core/src/workflow/index.ts b/core/src/workflow/index.ts new file mode 100644 index 000000000..ced77ea5b --- /dev/null +++ b/core/src/workflow/index.ts @@ -0,0 +1,75 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * The new ADK workflow module (parity port of `google/adk-python` + * `google/adk/workflow`). Public surface mirrors Python's `__all__`, plus the + * TypeScript-specific `WorkflowAgent` adapter and the types needed to use the + * API from TypeScript. + */ + +// --- Core graph / workflow --- +export {Workflow} from './workflow.js'; +export type {DynamicEntry, WorkflowConfig} from './workflow.js'; +export {WorkflowAgent} from './workflow_agent.js'; +export type {WorkflowAgentConfig} from './workflow_agent.js'; + +// --- Nodes --- +export {BaseNode, START} from './base_node.js'; +export type {BaseNodeConfig} from './base_node.js'; +export {WorkflowNode, node} from './node.js'; +export type {NodeOptions} from './node.js'; +export {FunctionNode} from './nodes/function_node.js'; +export type { + FunctionNodeConfig, + FunctionNodeHandler, + FunctionNodeResult, +} from './nodes/function_node.js'; +export {JoinNode} from './nodes/join_node.js'; +export {ParallelWorker} from './nodes/parallel_worker.js'; +export type {ParallelWorkerConfig} from './nodes/parallel_worker.js'; +export {ToolNode} from './nodes/tool_node.js'; +export type {ToolNodeConfig} from './nodes/tool_node.js'; +export type {BuildNodeOptions} from './utils/workflow_graph_utils.js'; +// LLMAgentWrapper and NodeTool are exported by Part 7 (LLM node). + +// --- Graph model --- +export {DEFAULT_ROUTE, Edge, Graph} from './graph.js'; +export type { + ChainElement, + EdgeItem, + NodeLike, + RouteValue, + RoutingMap, +} from './graph.js'; + +// --- Execution context & state --- +export {BranchPath, commonPrefixOf, createSubBranch} from './branch_path.js'; +export {NodeContext} from './node_context.js'; +export type {NodeContextOptions, NodeResult} from './node_context.js'; +export type {RunNodeOptions} from './node_runner.js'; +export {createNodeState, isNodeState} from './node_state.js'; +export type {NodeState} from './node_state.js'; +export {NodeStatus} from './node_status.js'; +export type { + ScheduleDynamicNode, + ScheduleDynamicNodeOptions, +} from './schedule_dynamic_node.js'; + +// --- HITL --- +export {RequestInput, isRequestInput} from './request_input.js'; +export type {RequestInputParams} from './request_input.js'; + +// --- Retry --- +export {normalizeRetryExceptions, prepareRetryConfig} from './retry_config.js'; +export type { + ErrorClass, + PreparedRetryConfig, + RetryConfig, +} from './retry_config.js'; + +// --- Errors --- +export {NodeTimeoutError} from './errors.js'; diff --git a/core/src/workflow/node.ts b/core/src/workflow/node.ts new file mode 100644 index 000000000..f61f839d3 --- /dev/null +++ b/core/src/workflow/node.ts @@ -0,0 +1,68 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {Event} from '../events/event.js'; +import {experimental} from '../utils/experimental.js'; +import {BaseNode} from './base_node.js'; +import {NodeLike} from './graph.js'; +import {NodeContext} from './node_context.js'; +import {buildNode, BuildNodeOptions} from './utils/workflow_graph_utils.js'; + +/** Options accepted by {@link node}. */ +export type NodeOptions = BuildNodeOptions; + +/** + * Wraps a {@link NodeLike} (function, tool, agent, or existing node) into a + * {@link BaseNode}, optionally overriding its properties. + * + * The TypeScript form is a plain function (there is no `@node` decorator form, + * unlike Python). Examples: + * + * ```ts + * const a = node(myFunction, {name: 'classify'}); + * const b = node(myTool); + * ``` + * + * Ported from `google/adk-python` `workflow/_node.py::node`. + */ +export function node(nodeLike: NodeLike, options: NodeOptions = {}): BaseNode { + return buildNode(nodeLike, options); +} + +/** + * A base class designed for subclassing. Implement {@link runNodeImpl} to + * provide node logic; subclasses inherit the schema/retry/timeout machinery of + * {@link BaseNode}. + * + * Named `WorkflowNode` (not `Node`) to read consistently with `WorkflowAgent` + * and `WorkflowConfig`, and to avoid shadowing the DOM / `@types/node` `Node` + * global in the flat `@google/adk` namespace. The `node()` factory remains the + * ergonomic way to wrap a function/tool/agent. + * + * Mirrors `google/adk-python` `workflow/_node.py::Node`. The `parallel_worker` + * capability is added in Phase 6. + */ +@experimental +export abstract class WorkflowNode< + TInput = unknown, + TOutput = unknown, +> extends BaseNode { + /** + * Implement node execution logic here. May yield `Event`s, raw values, or + * `null` (normalized by {@link BaseNode.run}). + */ + protected abstract runNodeImpl( + ctx: NodeContext, + input: TInput, + ): AsyncGenerator; + + protected async *runImpl( + ctx: NodeContext, + input: TInput, + ): AsyncGenerator { + yield* this.runNodeImpl(ctx, input); + } +} diff --git a/core/src/workflow/node_context.ts b/core/src/workflow/node_context.ts index 72711f015..00670b20e 100644 --- a/core/src/workflow/node_context.ts +++ b/core/src/workflow/node_context.ts @@ -14,6 +14,28 @@ import type {RouteValue} from './graph.js'; import {executeChildNode, RunNodeOptions} from './node_runner.js'; import type {ScheduleDynamicNode} from './schedule_dynamic_node.js'; +/** + * The result of running a node: the fields a caller (and the engine's + * completion handling) reads off a finished run — its `output`, emitted + * `route`, the `branch` it ran on, and any raised interrupt ids. + * + * A node that actually executes returns its full {@link NodeContext} (which + * satisfies this shape). A node that is *fast-forwarded* on resume (its output + * was cached in a prior turn, so its body is not re-run) returns a bare + * `NodeResult` with no live context behaviour — so callers of `ctx.runNode()` + * should treat the result as a `NodeResult` and read only these fields. + */ +export interface NodeResult { + /** The structured output the node produced (if any). */ + output: unknown; + /** The route key(s) the node emitted, if any (array = multi-route). */ + route?: RouteValue | RouteValue[]; + /** The branch the node ran on. */ + branch?: string; + /** Interrupt ids the node is blocked on (empty when it completed). */ + interruptIds: string[]; +} + /** * Options for constructing a {@link NodeContext}. */ @@ -59,12 +81,15 @@ export class NodeContext { interruptIds: string[] = []; /** - * Abort signal for the current node run, set by the engine while a node that - * declares a `timeout` is executing. It fires when the timeout elapses (or the - * invocation itself is aborted). Cooperative node bodies can observe - * `ctx.abortSignal` to cancel their own in-flight work (e.g. pass it to a - * model/tool call); the engine also stops consuming the node's events once it - * fires, so nothing is pushed past the deadline. + * Abort signal for the current node run, set by the engine while the node is + * executing under a deadline or an external cancellation signal — i.e. when + * the node declares a `timeout`, when the invocation itself can be aborted, or + * when it runs inside a Workflow (whose signal fires if a sibling fails). + * Cooperative node bodies can observe `ctx.abortSignal` to wind down their own + * in-flight work (e.g. pass it to a model/tool call); the engine also stops + * consuming the node's events once it fires, so nothing is pushed past + * cancellation. A fired `timeout` surfaces as a `NodeTimeoutError`; an external + * abort stops the node without raising. */ abortSignal?: AbortSignal; @@ -125,8 +150,11 @@ export class NodeContext { /** * Runs a child node programmatically, streaming its events through the same - * channel and resolving to the child's {@link NodeContext} (carrying its - * `output`, `route`, and `interruptIds`). + * channel and resolving to the child's result — a full {@link NodeContext} + * for a node that actually ran, or a bare {@link NodeResult} for one that was + * fast-forwarded from cached output on resume. Either way the caller can read + * `output`, `route`, `branch`, and `interruptIds`; only a `NodeContext` + * offers live behaviour (`emit`, `state`, nested `runNode`). * * When a dynamic-node {@link scheduler} is set (inside a Workflow subtree), * the call routes through it for dedup/resume; otherwise the child runs @@ -136,7 +164,7 @@ export class NodeContext { node: BaseNode, input?: unknown, options?: RunNodeOptions, - ): Promise { + ): Promise { if (this.scheduler) { const nodeName = options?.nodeName ?? node.name; let runId = options?.runId; diff --git a/core/src/workflow/node_runner.ts b/core/src/workflow/node_runner.ts index bbe819c04..d655587b8 100644 --- a/core/src/workflow/node_runner.ts +++ b/core/src/workflow/node_runner.ts @@ -8,7 +8,11 @@ import {InvocationContext} from '../agents/invocation_context.js'; import {Event} from '../events/event.js'; import {BaseNode} from './base_node.js'; import {createSubBranch} from './branch_path.js'; -import {InvocationAbortedError, NodeTimeoutError} from './errors.js'; +import { + InvocationAbortedError, + isInvocationAbortedError, + NodeTimeoutError, +} from './errors.js'; import {NodeContext} from './node_context.js'; import {createNodeState} from './node_state.js'; import {NodeStatus} from './node_status.js'; @@ -48,6 +52,12 @@ export interface ExecuteChildNodeParams { input: unknown; /** Options controlling this run. */ options?: RunNodeOptions; + /** + * Engine-supplied cancellation signal that overrides the parent invocation's + * for this child (used by a Workflow to cancel in-flight siblings when a node + * fails). Defaults to the parent invocation's abort signal. + */ + abortSignal?: AbortSignal; } /** @@ -71,6 +81,7 @@ export async function executeChildNode({ node, input, options = {}, + abortSignal, }: ExecuteChildNodeParams): Promise { const nodeName = options.nodeName ?? node.name; const runId = options.runId ?? nodeName; @@ -91,10 +102,20 @@ export async function executeChildNode({ const isolationScope = options.overrideIsolationScope ?? parent.isolationScope; + // The child observes the engine-supplied abort signal when given (a Workflow + // uses it to cancel siblings on failure), otherwise the parent invocation's. + const effectiveAbortSignal = + abortSignal ?? parent.invocationContext.abortSignal; + const childIc = - branch === parent.invocationContext.branch + branch === parent.invocationContext.branch && + effectiveAbortSignal === parent.invocationContext.abortSignal ? parent.invocationContext - : new InvocationContext({...parent.invocationContext, branch}); + : new InvocationContext({ + ...parent.invocationContext, + branch, + abortSignal: effectiveAbortSignal, + }); const child = new NodeContext({ invocationContext: childIc, @@ -121,6 +142,11 @@ export async function executeChildNode({ await runOnce({node, child, input, nodeName, branch, isolationScope}); succeeded = true; } catch (err) { + // Cancellation is terminal: an aborted invocation (or a sibling failure + // that cancelled this node) is never retried. + if (isInvocationAbortedError(err)) { + throw err; + } // Check retry eligibility with the attempt that just failed, compute its // backoff delay, THEN advance the counter (matches Python semantics). const retryConfig = node.preparedRetryConfig; @@ -132,7 +158,7 @@ export async function executeChildNode({ } const delaySeconds = getRetryDelaySeconds({retryConfig, nodeState}); nodeState.attemptCount += 1; - await delay(delaySeconds * 1000, parent.invocationContext.abortSignal); + await delay(delaySeconds * 1000, effectiveAbortSignal); } } @@ -181,13 +207,19 @@ interface RunOnceParams { * Drives one attempt of `node.run()`, enriching and pushing each event and * tracking the child's output/route. * - * When the node declares a `timeout`, execution is driven step-by-step and - * raced against a deadline: on timeout the engine stops consuming events (so - * nothing is pushed past the deadline — which would otherwise leak into a retry - * or the next node), closes the generator so its `finally` blocks run, and - * aborts `child.abortSignal` so a cooperative node body can cancel its own - * in-flight work. Mirrors the cancellation semantics of Python's + * When the node declares a `timeout` OR an external abort signal is present + * (the invocation's, or the workflow-scoped one used to cancel siblings when + * another node fails), execution is driven step-by-step and raced against those + * conditions: a fired deadline raises {@link NodeTimeoutError}; any other abort + * raises {@link InvocationAbortedError}. Either way the engine stops consuming + * events (so nothing is pushed past cancellation — which would otherwise leak + * into a retry or the next node), closes the generator so its `finally` blocks + * run, and aborts `child.abortSignal` so a cooperative node body can cancel its + * own in-flight work. Mirrors the cancellation semantics of Python's * `asyncio.wait_for`. + * + * When there is neither a deadline nor an abort signal, a plain `for await` + * fast path is used. */ async function runOnce({ node, @@ -223,30 +255,60 @@ async function runOnce({ child.channel.push(event); }; - if (!(typeof node.timeout === 'number' && node.timeout > 0)) { + const parentSignal = child.invocationContext.abortSignal; + const hasTimeout = typeof node.timeout === 'number' && node.timeout > 0; + + // Fast path: no per-node deadline and no external cancellation to observe. + if (!hasTimeout && !parentSignal) { for await (const event of node.run(child, input)) { consume(event); } return; } + // Cooperative cancellation (external abort, no deadline): expose the abort + // signal as `ctx.abortSignal` so a cooperative node can wind down its own work + // (e.g. a Workflow child stopping when a sibling fails), then drain normally. + // We do NOT force-stop: a node that ignores the signal runs to completion + // (best-effort), and a node that fails still surfaces its error — with the + // retry backoff observing the same signal. + if (!hasTimeout) { + child.abortSignal = parentSignal; + try { + for await (const event of node.run(child, input)) { + consume(event); + } + } finally { + child.abortSignal = undefined; + } + return; + } + + // Deadline path: drive the node step-by-step and race each step against the + // timeout (and any external abort). On the deadline (or abort) the engine + // stops consuming events, closes the generator so its `finally` runs, and + // aborts `child.abortSignal` so a cooperative body can cancel its in-flight + // work; the run rejects with NodeTimeoutError. Mirrors Python's + // `asyncio.wait_for`. const timeoutSeconds = node.timeout; const controller = new AbortController(); - const parentSignal = child.invocationContext.abortSignal; const onParentAbort = () => controller.abort(); if (parentSignal?.aborted) { controller.abort(); } else { parentSignal?.addEventListener('abort', onParentAbort, {once: true}); } - const timer = setTimeout(() => controller.abort(), timeoutSeconds * 1000); + const timer = setTimeout( + () => controller.abort(), + (timeoutSeconds ?? 0) * 1000, + ); child.abortSignal = controller.signal; // A single promise that rejects once the deadline (or external abort) fires; // reused across iterations so we don't leak a listener per step. const aborted = new Promise((_, reject) => { const fail = () => - reject(new NodeTimeoutError({nodeName, timeout: timeoutSeconds})); + reject(new NodeTimeoutError({nodeName, timeout: timeoutSeconds ?? 0})); if (controller.signal.aborted) { fail(); } else { diff --git a/core/src/workflow/nodes/parallel_worker.ts b/core/src/workflow/nodes/parallel_worker.ts index 453a41041..02a7d7dbd 100644 --- a/core/src/workflow/nodes/parallel_worker.ts +++ b/core/src/workflow/nodes/parallel_worker.ts @@ -19,7 +19,7 @@ const DEFAULT_MAX_PARALLEL_WORKERS = 8; export interface ParallelWorkerConfig { /** * Maximum number of items processed concurrently. Defaults to - * {@link DEFAULT_MAX_PARALLEL_WORKERS}; pass `Infinity` for unbounded. + * `DEFAULT_MAX_PARALLEL_WORKERS` (8); pass `Infinity` for unbounded. */ maxParallelWorkers?: number; } diff --git a/core/src/workflow/request_input.ts b/core/src/workflow/request_input.ts index baf432ca4..0962a2bcd 100644 --- a/core/src/workflow/request_input.ts +++ b/core/src/workflow/request_input.ts @@ -62,7 +62,7 @@ export class RequestInput { /** * Type guard for {@link RequestInput}. * - * Matches on the {@link REQUEST_INPUT_SIGNATURE_SYMBOL} brand rather than + * Matches on the `REQUEST_INPUT_SIGNATURE_SYMBOL` brand rather than * `instanceof` so it stays correct across package copies (see the brand's doc). */ export function isRequestInput(value: unknown): value is RequestInput { diff --git a/core/src/workflow/schedule_dynamic_node.ts b/core/src/workflow/schedule_dynamic_node.ts index 0e0defd43..571aee3ad 100644 --- a/core/src/workflow/schedule_dynamic_node.ts +++ b/core/src/workflow/schedule_dynamic_node.ts @@ -5,7 +5,7 @@ */ import type {BaseNode} from './base_node.js'; -import type {NodeContext} from './node_context.js'; +import type {NodeContext, NodeResult} from './node_context.js'; import {NodeState} from './node_state.js'; /** @@ -39,7 +39,7 @@ export interface ScheduleDynamicNode { node: BaseNode, input: unknown, options: ScheduleDynamicNodeOptions, - ): Promise; + ): Promise; } /** diff --git a/core/src/workflow/utils/rehydration_utils.ts b/core/src/workflow/utils/rehydration_utils.ts new file mode 100644 index 000000000..3b3bf87c0 --- /dev/null +++ b/core/src/workflow/utils/rehydration_utils.ts @@ -0,0 +1,211 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Reconstructs workflow node state from prior session events, so a resumed + * workflow can fast-forward completed nodes and resolve pending interrupts. + * + * Ported (static-graph subset) from `google/adk-python` + * `workflow/utils/_rehydration_utils.py`. The chronological sequence barrier + * for deterministic parallel/dynamic replay is a Phase 5b continuation. + */ + +import {Event} from '../../events/event.js'; +import {RouteValue} from '../graph.js'; +import type {NodeContext, NodeResult} from '../node_context.js'; + +const RESULT_KEY = 'result'; + +/** Reconstructed state for a single node, keyed by node name. */ +export interface RehydratedNode { + /** The node's cached output from a prior run, if it produced one. */ + output?: unknown; + /** The route(s) the node emitted, if any (array = multi-route). */ + route?: RouteValue | RouteValue[]; + /** The branch the node ran on. */ + branch?: string; + /** The input the node was invoked with (captured when it interrupted). */ + input?: unknown; + /** Interrupt ids the node raised. */ + interruptIds: Set; + /** Resolved interrupt responses, keyed by interrupt id. */ + resolvedResponses: Map; +} + +/** + * Scans session events and reconstructs per-node state (outputs, routes, raised + * interrupts, and resolved interrupt responses). + * + * When `parentPath` is given, reconstruction is scoped to that path's DIRECT + * children (keyed by child name), so nested workflows whose nodes share a name + * do not collide on resume. When omitted, nodes are keyed by their path leaf + * (utility mode, robust to author rewrite). + */ +export function reconstructNodeStates( + events: Event[], + parentPath?: string, +): Map { + if (parentPath) { + return reconstruct(events, (event) => + event.nodeInfo?.path + ? directChildName(event.nodeInfo.path, parentPath) + : undefined, + ); + } + return reconstruct(events, (event) => + event.nodeInfo?.path ? nodeNameFromPath(event.nodeInfo.path) : event.author, + ); +} + +/** + * Like {@link reconstructNodeStates} but keyed by the full node path + * (`wf.node@runId`), so distinct dynamic (`ctx.runNode`) iterations are tracked + * separately for per-run resume/dedup. + */ +export function reconstructNodeStatesByPath( + events: Event[], +): Map { + return reconstruct(events, (event) => event.nodeInfo?.path ?? event.author); +} + +/** Shared scan that groups node events by the key returned by `keyFor`. */ +function reconstruct( + events: Event[], + keyFor: (event: Event) => string | undefined, +): Map { + const nodes = new Map(); + const interruptOwner = new Map(); + + const getNode = (name: string): RehydratedNode => { + let node = nodes.get(name); + if (!node) { + node = {interruptIds: new Set(), resolvedResponses: new Map()}; + nodes.set(name, node); + } + return node; + }; + + for (const event of events) { + // 1. User function responses resolving prior interrupts. + if (event.author === 'user' && event.content?.parts) { + for (const part of event.content.parts) { + const fr = part.functionResponse; + if (fr?.id && interruptOwner.has(fr.id)) { + const owner = interruptOwner.get(fr.id)!; + getNode(owner).resolvedResponses.set( + fr.id, + unwrapResponse(fr.response), + ); + } + } + continue; + } + + // 2. Node events. + const key = keyFor(event); + if (!key) { + continue; + } + const node = getNode(key); + if (event.output !== undefined) { + node.output = event.output; + node.branch = event.branch; + } + if (event.route !== undefined) { + node.route = event.route as RouteValue | RouteValue[]; + } + for (const id of event.longRunningToolIds ?? []) { + node.interruptIds.add(id); + interruptOwner.set(id, key); + } + // Capture the node's original input, stashed on the interrupt event, so a + // resumed waiting node re-runs with it (not the resume message). Guard the + // read since `agentState` is an unknown, arbitrarily-shaped payload. + const agentState = event.actions?.agentState; + if (isRecord(agentState) && 'input' in agentState) { + node.input = agentState.input; + } + } + + return nodes; +} + +/** + * Whether a rehydrated node can be fast-forwarded on resume: it produced an + * output and all of its raised interrupts have been resolved. + */ +export function isFastForwardable(node: RehydratedNode): boolean { + if (node.output === undefined) { + return false; + } + for (const id of node.interruptIds) { + if (!node.resolvedResponses.has(id)) { + return false; + } + } + return true; +} + +/** + * Builds the completion result for a fast-forwarded (cached) node on resume. + * The node's body is not re-run and its events are NOT re-emitted (they already + * exist in the session), so this returns a bare {@link NodeResult} rather than a + * live {@link NodeContext} — the honest type for "cached output, no behaviour". + */ +export function makeFastForwardResult( + parent: NodeContext, + prior: RehydratedNode, +): NodeResult { + return { + output: prior.output, + route: prior.route, + branch: prior.branch ?? parent.branch, + interruptIds: [], + }; +} + +/** Extracts the node name (leaf, without run id) from a dotted node path. */ +export function nodeNameFromPath(path: string): string { + const leaf = path.split(/[./]/).pop() ?? path; + return leaf.split('@')[0]; +} + +/** + * Returns the child node name if `path` is a DIRECT child of `parentPath` + * (e.g. `parent.child` -> `child`, `parent.child@2` -> `child`), or `undefined` + * for a non-descendant or a deeper descendant (e.g. `parent.sub.child`). Used to + * scope rehydration to a single workflow's own nodes. + */ +function directChildName(path: string, parentPath: string): string | undefined { + const prefix = `${parentPath}.`; + if (!path.startsWith(prefix)) { + return undefined; + } + const rest = path.slice(prefix.length); + if (rest.includes('.')) { + return undefined; // a deeper descendant, not a direct child + } + return rest.split('@')[0]; +} + +/** Narrows an unknown value to a plain (non-array) record. */ +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null && !Array.isArray(value); +} + +/** Unwraps a `{result: value}` FunctionResponse envelope to the bare value. */ +export function unwrapResponse(response: unknown): unknown { + if ( + response && + typeof response === 'object' && + !Array.isArray(response) && + Object.keys(response).length === 1 && + RESULT_KEY in response + ) { + return (response as Record)[RESULT_KEY]; + } + return response; +} diff --git a/core/src/workflow/workflow.ts b/core/src/workflow/workflow.ts new file mode 100644 index 000000000..e3152916f --- /dev/null +++ b/core/src/workflow/workflow.ts @@ -0,0 +1,645 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {Event} from '../events/event.js'; +import {experimental} from '../utils/experimental.js'; +import {BaseNode, BaseNodeConfig} from './base_node.js'; +import {commonPrefixOf} from './branch_path.js'; +import {DynamicNodeScheduler} from './dynamic_node_scheduler.js'; +import { + createGraphFromEdgeItems, + EdgeItem, + Graph, + RouteValue, +} from './graph.js'; +import {NodeContext, NodeResult} from './node_context.js'; +import {executeChildNode} from './node_runner.js'; +import {createNodeState, NodeState} from './node_state.js'; +import {NodeStatus} from './node_status.js'; +import {DynamicNodeState} from './schedule_dynamic_node.js'; +import {Trigger} from './trigger.js'; +import { + isFastForwardable, + makeFastForwardResult, + reconstructNodeStates, + RehydratedNode, +} from './utils/rehydration_utils.js'; + +/** + * An imperative workflow entry point. Receives the workflow's node context and + * input, drives execution via `ctx.runNode(...)`, and returns the workflow + * output. Mutually exclusive with `edges`. + */ +export type DynamicEntry = ( + ctx: NodeContext, + input: unknown, +) => unknown | Promise; + +/** + * Configuration for a {@link Workflow}. + * + * A workflow is driven either by a static `edges` graph or by an imperative + * `dynamicEntry` function — exactly one is required, and the two are mutually + * exclusive. The type is a discriminated union so that constraint is enforced at + * compile time (the runtime constructor keeps the equivalent throws for JS + * callers). + */ +export type WorkflowConfig = BaseNodeConfig & { + /** + * Maximum number of graph-scheduled nodes running in parallel. `undefined` + * means unlimited; must be a positive integer otherwise. Does not throttle + * dynamic (`ctx.runNode`) children. + */ + maxConcurrency?: number; +} & ( + | { + /** Edge definitions used to build the workflow graph. */ + edges: EdgeItem[]; + dynamicEntry?: never; + } + | { + /** + * An imperative entry function driving execution via `ctx.runNode(...)`. + * Mutually exclusive with `edges`. + */ + dynamicEntry: DynamicEntry; + edges?: never; + } + ); + +/** + * Mutable, in-memory state for a single {@link Workflow} run. Not persisted; + * discarded when `runImpl` returns. (Replay/checkpoint fields are added in + * Phase 5.) + */ +class LoopState { + readonly nodes = new Map(); + readonly nodeOutputs = new Map(); + readonly nodeBranches = new Map(); + readonly triggerBuffer = new Map(); + readonly pending = new Map>(); + readonly interruptIds = new Set(); + /** Per-node state reconstructed from prior session events (resume). */ + rehydrated: Map = new Map(); + errorShutDown = false; + /** + * Workflow-scoped abort signal handed to each scheduled node so a failure can + * cancel its in-flight siblings (see {@link Workflow.cleanupPending}). + */ + abortSignal?: AbortSignal; +} + +interface CompletedTask { + name: string; + /** + * The finished node's result: a live {@link NodeContext} for a node that ran, + * or a bare {@link NodeResult} for one fast-forwarded from cached output on + * resume. Completion handling reads only the shared fields. + */ + childCtx?: NodeContext | NodeResult; + error?: unknown; +} + +/** + * A graph-based workflow node. `runImpl()` IS the orchestration loop: + * SETUP (seed START triggers) → LOOP (schedule ready nodes, handle + * completions) → FINALIZE (collect the terminal output). + * + * Ported (Phase 2 subset) from `google/adk-python` `workflow/_workflow.py`. + * Replay/checkpointing, dynamic scheduling, and task/chat isolation scopes are + * added in later phases; hook points are marked with TODO(phase-N). + */ +@experimental +export class Workflow extends BaseNode { + readonly graph?: Graph; + readonly dynamicEntry?: DynamicEntry; + readonly maxConcurrency?: number; + + constructor(config: WorkflowConfig) { + super({...config, rerunOnResume: config.rerunOnResume ?? true}); + const hasEdges = !!config.edges && config.edges.length > 0; + if (hasEdges && config.dynamicEntry) { + throw new Error( + `Workflow "${this.name}": "edges" and "dynamicEntry" are mutually exclusive.`, + ); + } + if (!hasEdges && !config.dynamicEntry) { + throw new Error( + `Workflow "${this.name}" requires either "edges" or "dynamicEntry".`, + ); + } + if ( + config.maxConcurrency !== undefined && + (!Number.isInteger(config.maxConcurrency) || config.maxConcurrency < 1) + ) { + throw new Error( + `Workflow "${this.name}": "maxConcurrency" must be a positive integer ` + + `(got ${config.maxConcurrency}).`, + ); + } + this.maxConcurrency = config.maxConcurrency; + this.dynamicEntry = config.dynamicEntry; + if (config.edges && config.edges.length > 0) { + // createGraphFromEdgeItems validates as part of construction. + this.graph = createGraphFromEdgeItems(config.edges); + } + } + + // eslint-disable-next-line require-yield -- child events stream out via ctx.channel/ctx.runNode; this orchestration generator itself yields nothing + protected async *runImpl( + ctx: NodeContext, + nodeInput: unknown, + ): AsyncGenerator { + // Child events are streamed through ctx.channel by ctx.runNode(), so this + // orchestration generator itself yields nothing. + const dynamicState = new DynamicNodeState(); + + // Workflow-scoped cancellation: a controller chained to the invocation's + // abort signal. It is aborted when a node fails (see cleanupPending) so any + // in-flight siblings stop cooperatively instead of running to completion, + // and disposed in the finally so we don't leak the parent-abort listener. + const abort = createWorkflowAbort(ctx.invocationContext.abortSignal); + ctx.scheduler = new DynamicNodeScheduler( + dynamicState, + abort.controller.signal, + ); + + try { + await this.orchestrate(ctx, nodeInput, dynamicState, abort.controller); + } finally { + abort.dispose(); + } + } + + /** + * The orchestration body, wrapped by {@link runImpl} so its workflow-scoped + * abort controller is always disposed. + */ + private async orchestrate( + ctx: NodeContext, + nodeInput: unknown, + dynamicState: DynamicNodeState, + abortController: AbortController, + ): Promise { + // --- REHYDRATE (resume) --- + // Reconstruct node state from prior session events and surface resolved + // interrupt responses so waiting nodes can resume. Scope to this workflow's + // own direct children (by path) so nested workflows with same-named nodes + // don't collide. + const rehydrated = reconstructNodeStates( + ctx.session?.events ?? [], + ctx.nodePath || undefined, + ); + this.applyResumeInputs(ctx, rehydrated); + + if (this.dynamicEntry) { + await this.runDynamicEntry(ctx, nodeInput, dynamicState); + return; + } + + const loop = new LoopState(); + loop.rehydrated = rehydrated; + loop.abortSignal = abortController.signal; + + // --- SETUP --- + this.seedStartTriggers(loop, nodeInput); + + // --- LOOP --- + await this.runLoop(loop, ctx, abortController); + + if (loop.errorShutDown) { + return; + } + + this.collectRemainingInterrupts(loop); + // Fold in interrupts raised by dynamic (ctx.runNode) children. + for (const id of dynamicState.interruptIds) { + loop.interruptIds.add(id); + } + + // --- FINALIZE --- + this.finalize(loop, ctx); + } + + /** + * Runs an imperative `dynamicEntry` workflow. The entry drives execution via + * `ctx.runNode(...)` (routed through the scheduler) and returns the output. + */ + private async runDynamicEntry( + ctx: NodeContext, + nodeInput: unknown, + dynamicState: DynamicNodeState, + ): Promise { + const output = await this.dynamicEntry!(ctx, nodeInput); + if (dynamicState.interruptIds.size > 0) { + ctx.interruptIds = [...dynamicState.interruptIds]; + return; + } + if (output !== undefined) { + ctx.output = output; + } + } + + /** + * Merges resolved interrupt responses from prior session events into + * `ctx.resumeInputs`, so waiting nodes (which read `ctx.resumeInputs[id]`) + * resume with the user's response. Shared by child contexts via propagation. + */ + private applyResumeInputs( + ctx: NodeContext, + rehydrated: Map, + ): void { + for (const node of rehydrated.values()) { + for (const [interruptId, response] of node.resolvedResponses) { + ctx.resumeInputs[interruptId] = response; + } + } + } + + // --- SETUP --- + + private seedStartTriggers(loop: LoopState, nodeInput: unknown): void { + const startEdges = this.graph!.edges.filter( + (e) => e.fromNode.name === '__START__', + ); + const useSubBranch = startEdges.length > 1; + for (const edge of startEdges) { + this.pushTrigger(loop, edge.toNode.name, { + input: nodeInput, + useSubBranch, + }); + } + } + + // --- LOOP --- + + private async runLoop( + loop: LoopState, + ctx: NodeContext, + abortController: AbortController, + ): Promise { + for (;;) { + this.scheduleReadyNodes(loop, ctx); + + if (loop.pending.size === 0) { + break; + } + + const result = await Promise.race(loop.pending.values()); + loop.pending.delete(result.name); + + if (result.error) { + const nodeState = loop.nodes.get(result.name); + if (nodeState) { + nodeState.status = NodeStatus.FAILED; + } + loop.errorShutDown = true; + await this.cleanupPending(loop, abortController); + throw result.error; + } + + await this.handleCompletion(loop, result.name, result.childCtx!); + } + } + + // --- Scheduling --- + + private scheduleReadyNodes(loop: LoopState, ctx: NodeContext): void { + for (const nodeName of [...loop.triggerBuffer.keys()]) { + if (loop.pending.has(nodeName)) { + continue; + } + const state = loop.nodes.get(nodeName); + if (state) { + if (state.status === NodeStatus.RUNNING) { + continue; + } + if ( + state.status === NodeStatus.WAITING && + state.interrupts.length > 0 + ) { + continue; + } + } + if (this.atConcurrencyLimit(loop)) { + break; + } + + const trigger = this.popTrigger(loop, nodeName); + if (!trigger) { + continue; + } + this.prepareNodeStateForStarting(loop, nodeName, trigger); + this.startNodeTask(loop, ctx, nodeName, trigger); + } + } + + private atConcurrencyLimit(loop: LoopState): boolean { + return ( + this.maxConcurrency !== undefined && + loop.pending.size >= this.maxConcurrency + ); + } + + private prepareNodeStateForStarting( + loop: LoopState, + nodeName: string, + trigger: Trigger, + ): void { + const existing = loop.nodes.get(nodeName); + // Fresh NodeState for each run, preserving the run counter. + const state = createNodeState({ + runCounter: existing?.runCounter ?? 0, + }); + state.input = trigger.input; + state.status = NodeStatus.RUNNING; + loop.nodes.set(nodeName, state); + } + + private startNodeTask( + loop: LoopState, + ctx: NodeContext, + nodeName: string, + trigger: Trigger, + ): void { + const node = this.getStaticNode(nodeName); + const nodeState = loop.nodes.get(nodeName)!; + + // Resume: fast-forward a node that already completed in a prior run + // (cached output, all interrupts resolved), unless it must rerun on resume. + const prior = loop.rehydrated.get(nodeName); + if (prior && !node.rerunOnResume && isFastForwardable(prior)) { + loop.pending.set( + nodeName, + Promise.resolve({ + name: nodeName, + childCtx: makeFastForwardResult(ctx, prior), + }), + ); + return; + } + + // Resume with rerun_on_resume=false: a node that interrupted last turn + // (raised interrupts, produced no output) does NOT re-run its body. Instead + // it completes with the resolved resume value(s) as its output, feeding the + // next node. This is Python's two-node request-input pattern, where one node + // yields RequestInput and its successor receives the human's reply as input. + if ( + prior && + !node.rerunOnResume && + prior.output === undefined && + prior.interruptIds.size > 0 + ) { + const values = [...prior.interruptIds].map((id) => ctx.resumeInputs[id]); + if (values.every((v) => v !== undefined)) { + const output = values.length === 1 ? values[0] : values; + const resumeResult: NodeResult = { + output, + route: undefined, + branch: prior.branch ?? ctx.branch, + interruptIds: [], + }; + loop.pending.set( + nodeName, + Promise.resolve({name: nodeName, childCtx: resumeResult}), + ); + return; + } + } + + let runId = nodeState.runId; + if (!runId) { + nodeState.runCounter += 1; + runId = String(nodeState.runCounter); + nodeState.runId = runId; + } + + // On resume, a waiting node (it interrupted last turn) re-runs with its + // ORIGINAL input, not the trigger's (which carries the resume message). + const resuming = + prior !== undefined && + prior.interruptIds.size > 0 && + prior.input !== undefined; + const nodeInput = resuming ? prior.input : trigger.input; + + // Static graph nodes are managed by this loop directly, bypassing the + // dynamic scheduler (which serves user-initiated ctx.runNode() calls). The + // workflow-scoped abort signal lets a sibling's failure cancel this node. + const task: Promise = executeChildNode({ + parent: ctx, + node, + input: nodeInput, + abortSignal: loop.abortSignal, + options: { + runId, + useSubBranch: trigger.useSubBranch, + overrideBranch: trigger.branch, + overrideIsolationScope: trigger.isolationScope, + }, + }).then( + (childCtx) => ({name: nodeName, childCtx}), + (error) => ({name: nodeName, error}), + ); + loop.pending.set(nodeName, task); + } + + // --- Completion handling --- + + private async handleCompletion( + loop: LoopState, + nodeName: string, + childCtx: NodeContext | NodeResult, + ): Promise { + const nodeState = loop.nodes.get(nodeName)!; + const node = this.getStaticNode(nodeName); + + if (childCtx.interruptIds.length > 0) { + nodeState.status = NodeStatus.WAITING; + nodeState.interrupts = [...childCtx.interruptIds]; + childCtx.interruptIds.forEach((id) => loop.interruptIds.add(id)); + return; + } + + if ( + node.waitForOutput && + childCtx.output === undefined && + childCtx.route === undefined + ) { + nodeState.status = NodeStatus.WAITING; + return; + } + + nodeState.status = NodeStatus.COMPLETED; + if (childCtx.output !== undefined) { + loop.nodeOutputs.set(nodeName, childCtx.output); + } + loop.nodeBranches.set(nodeName, childCtx.branch ?? ''); + + this.bufferDownstreamTriggers( + loop, + nodeName, + childCtx.output, + childCtx.route, + childCtx.branch, + ); + } + + private bufferDownstreamTriggers( + loop: LoopState, + nodeName: string, + output: unknown, + route: RouteValue | RouteValue[] | undefined, + branch: string | undefined, + ): void { + const nextNodes = this.graph!.getNextPendingNodes(nodeName, route ?? null); + const useSubBranch = nextNodes.length > 1; + + for (const targetName of nextNodes) { + const targetNode = this.getStaticNode(targetName); + + if (targetNode.requiresAllPredecessors) { + const predecessors = new Set( + this.graph!.edges.filter((e) => e.toNode.name === targetName).map( + (e) => e.fromNode.name, + ), + ); + const allCompleted = [...predecessors].every( + (p) => loop.nodes.get(p)?.status === NodeStatus.COMPLETED, + ); + if (allCompleted) { + const outputs: Record = {}; + for (const p of predecessors) { + outputs[p] = loop.nodeOutputs.get(p); + } + const branches = [...predecessors].map( + (p) => loop.nodeBranches.get(p) ?? '', + ); + const commonBranch = commonPrefixOf(branches); + this.pushTrigger(loop, targetName, { + input: outputs, + useSubBranch: false, + branch: commonBranch || undefined, + }); + } + } else { + this.pushTrigger(loop, targetName, { + input: output, + useSubBranch, + branch, + }); + } + } + } + + private collectRemainingInterrupts(loop: LoopState): void { + for (const nodeState of loop.nodes.values()) { + if ( + nodeState.status === NodeStatus.WAITING && + nodeState.interrupts.length > 0 + ) { + nodeState.interrupts.forEach((id) => loop.interruptIds.add(id)); + } + } + } + + // --- FINALIZE --- + + private finalize(loop: LoopState, ctx: NodeContext): void { + if (loop.interruptIds.size > 0) { + ctx.interruptIds = [...loop.interruptIds]; + return; + } + + const terminalOutputs = [...this.graph!.terminalNodeNames] + .filter((name) => loop.nodeOutputs.has(name)) + .map((name) => loop.nodeOutputs.get(name)); + + if (terminalOutputs.length === 1) { + ctx.output = terminalOutputs[0]; + } else if (terminalOutputs.length > 1) { + throw new Error( + `Workflow ${this.name}: multiple terminal nodes produced output ` + + `(${terminalOutputs.length}). A workflow must have at most one terminal output.`, + ); + } + } + + // --- Utilities --- + + private pushTrigger( + loop: LoopState, + nodeName: string, + trigger: Trigger, + ): void { + const buffer = loop.triggerBuffer.get(nodeName); + if (buffer) { + buffer.push(trigger); + } else { + loop.triggerBuffer.set(nodeName, [trigger]); + } + } + + private popTrigger(loop: LoopState, nodeName: string): Trigger | undefined { + const buffer = loop.triggerBuffer.get(nodeName); + if (!buffer || buffer.length === 0) { + return undefined; + } + const trigger = buffer.shift()!; + if (buffer.length === 0) { + loop.triggerBuffer.delete(nodeName); + } + return trigger; + } + + private getStaticNode(name: string): BaseNode { + const node = this.graph!.nodes.find((n) => n.name === name); + if (!node) { + throw new Error(`Node ${name} not found in graph.`); + } + return node; + } + + private async cleanupPending( + loop: LoopState, + abortController: AbortController, + ): Promise { + // Signal in-flight siblings to stop: cooperative nodes observe + // `ctx.abortSignal`, and the node runner stops consuming their events once + // the signal fires (see node_runner). Then await the outstanding tasks so + // their cleanup runs and events flush; failures are swallowed because the + // workflow is already shutting down on error. + abortController.abort(); + const outstanding = [...loop.pending.values()]; + loop.pending.clear(); + await Promise.allSettled(outstanding); + } +} + +/** + * Creates the workflow-scoped {@link AbortController} that cancels in-flight + * nodes when the workflow shuts down on error. It is chained to the invocation's + * own abort signal (if any) so an invocation-level cancel still propagates to + * nodes; `dispose` detaches that listener to avoid a leak. + */ +function createWorkflowAbort(parentSignal?: AbortSignal): { + controller: AbortController; + dispose: () => void; +} { + const controller = new AbortController(); + if (!parentSignal) { + return {controller, dispose: () => {}}; + } + if (parentSignal.aborted) { + controller.abort(); + return {controller, dispose: () => {}}; + } + const onParentAbort = () => controller.abort(); + parentSignal.addEventListener('abort', onParentAbort, {once: true}); + return { + controller, + dispose: () => parentSignal.removeEventListener('abort', onParentAbort), + }; +} diff --git a/core/src/workflow/workflow_agent.ts b/core/src/workflow/workflow_agent.ts new file mode 100644 index 000000000..a16168eb4 --- /dev/null +++ b/core/src/workflow/workflow_agent.ts @@ -0,0 +1,163 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {Content} from '@google/genai'; +import {BaseAgent} from '../agents/base_agent.js'; +import {InvocationContext} from '../agents/invocation_context.js'; +import {createEvent, Event} from '../events/event.js'; +import {AsyncQueue} from '../utils/async_queue.js'; +import {experimental} from '../utils/experimental.js'; +import {toContent} from './base_node.js'; +import {NodeContext} from './node_context.js'; +import {reconstructNodeStates} from './utils/rehydration_utils.js'; +import {Workflow} from './workflow.js'; + +/** Options for a {@link WorkflowAgent}. */ +export interface WorkflowAgentConfig { + name?: string; + description?: string; +} + +/** + * Adapts a {@link Workflow} (a `BaseNode`) into a `BaseAgent` so it can be run by + * the standard ADK `Runner`. + * + * It sets up the event channel bridge: the workflow's node execution pushes + * events into the channel while this agent's `runAsyncImpl` drains and re-yields + * them to the runtime. The user message (`ctx.userContent`) becomes the + * workflow input. + */ +@experimental +export class WorkflowAgent extends BaseAgent { + readonly workflow: Workflow; + + constructor(workflow: Workflow, config: WorkflowAgentConfig = {}) { + super({ + name: config.name ?? workflow.name, + description: config.description ?? workflow.description, + }); + this.workflow = workflow; + } + + protected async *runAsyncImpl( + ic: InvocationContext, + ): AsyncGenerator { + const channel = new AsyncQueue(); + const root = new NodeContext({ + invocationContext: ic, + channel, + nodePath: '', + runId: this.name, + // Interactive resume: if the workflow is paused on an interrupt and the + // user replies with plain text (not a structured function response), feed + // that text to the pending interrupt(s). Structured function responses are + // still resolved by the workflow's own rehydration. + resumeInputs: resumeInputsFromPlainText(ic), + }); + + const input = extractWorkflowInput(ic.userContent); + + const settle = (async () => { + try { + const wfCtx = await root.runNode(this.workflow, input, { + useAsOutput: true, + }); + // Surface the workflow's final output as an event so consumers (and + // the Runner) can observe it — important for dynamicEntry workflows + // whose return value differs from the last node's event. + if (wfCtx.interruptIds.length === 0 && root.output !== undefined) { + channel.push( + createEvent({ + author: this.name, + invocationId: ic.invocationId, + branch: ic.branch, + content: toContent(root.output), + output: root.output, + }), + ); + } + channel.close(); + } catch (err) { + channel.fail(err); + } + })(); + + try { + for await (const event of channel) { + yield event; + } + await settle; + } finally { + // Ensure a single exit path if the consumer stops early (breaks its + // for-await, or the Runner cancels the invocation): close the channel so + // the workflow's producer stops pushing into a queue nobody drains, and + // await `settle` so its cleanup runs and errors surface. Idempotent on the + // normal path (the channel is already closed and `settle` resolved). + channel.close(); + await settle; + } + } + + // eslint-disable-next-line require-yield -- runLiveImpl must be an AsyncGenerator per BaseAgent, but live mode is unsupported so it only throws + protected async *runLiveImpl(): AsyncGenerator { + throw new Error('WorkflowAgent does not support live mode.'); + } +} + +/** + * When the workflow is paused on exactly one unresolved interrupt and the + * incoming message is plain text (not a structured function response), maps that + * text to the single pending interrupt id so an interactive client (e.g. `adk + * run`) can resume a HITL/auth pause by simply typing a reply. + * + * If more than one interrupt is pending, a plain-text reply is ambiguous — it + * would be broadcast to every pause and at least one node would resume with data + * the user never gave it — so it is ignored here. Addressing a specific pause in + * a multi-interrupt workflow requires structured function responses (resolved by + * the workflow's own rehydration). + */ +function resumeInputsFromPlainText( + ic: InvocationContext, +): Record { + const parts = ic.userContent?.parts ?? []; + const isPlainText = + parts.length > 0 && parts.every((p) => typeof p.text === 'string'); + if (!isPlainText) { + return {}; + } + const text = parts.map((p) => p.text).join(''); + + const pending = new Set(); + for (const node of reconstructNodeStates(ic.session?.events ?? []).values()) { + for (const id of node.interruptIds) { + if (!node.resolvedResponses.has(id)) { + pending.add(id); + } + } + } + + // Only the unambiguous single-pause case is resumable by plain text. + if (pending.size !== 1) { + return {}; + } + const [id] = pending; + return {[id]: text}; +} + +/** + * Derives the workflow input from the user message: plain text when the content + * is text-only, otherwise the raw `Content` (nodes coerce as needed). + */ +function extractWorkflowInput(content?: Content): unknown { + if (!content) { + return undefined; + } + const parts = content.parts ?? []; + if (parts.length > 0 && parts.every((p) => typeof p.text === 'string')) { + return parts.map((p) => p.text).join(''); + } + return content; +} diff --git a/core/test/workflow/auth_gate_test.ts b/core/test/workflow/auth_gate_test.ts new file mode 100644 index 000000000..02693b03e --- /dev/null +++ b/core/test/workflow/auth_gate_test.ts @@ -0,0 +1,149 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {describe, expect, it} from 'vitest'; +import { + AuthCredential, + AuthCredentialTypes, +} from '../../src/auth/auth_credential.js'; +import {AuthScheme} from '../../src/auth/auth_schemes.js'; +import {AuthConfig} from '../../src/auth/auth_tool.js'; +import {Event} from '../../src/events/event.js'; +import {Runner} from '../../src/runner/runner.js'; +import {InMemorySessionService} from '../../src/sessions/in_memory_session_service.js'; +import {NodeContext} from '../../src/workflow/node_context.js'; +import {FunctionNode} from '../../src/workflow/nodes/function_node.js'; +import {hasAuthRequestFunctionCall} from '../../src/workflow/utils/hitl_utils.js'; +import {Workflow} from '../../src/workflow/workflow.js'; +import {WorkflowAgent} from '../../src/workflow/workflow_agent.js'; +import {createIc, driveWorkflow} from './test_helpers.js'; + +const CREDENTIAL_KEY = 'my_api'; + +function apiKeyAuthConfig(): AuthConfig { + return { + authScheme: {type: 'apiKey', in: 'header', name: 'X-API-Key'} as AuthScheme, + rawAuthCredential: {authType: AuthCredentialTypes.API_KEY}, + credentialKey: CREDENTIAL_KEY, + }; +} + +async function collect(gen: AsyncGenerator): Promise { + const out: Event[] = []; + for await (const e of gen) { + out.push(e); + } + return out; +} + +describe('Phase 5b-cont — FunctionNode auth gate', () => { + it('requests credentials, then runs after they are supplied on resume', async () => { + let runs = 0; + let sawApiKey: string | undefined; + + // An auth-gated node must RE-RUN on resume so it can store the supplied + // credential and then run its body (this is why Python's auth samples set + // rerun_on_resume=True). Without it, the default two-node resume semantics + // would complete the node with the raw credential response as its output. + const secured = new FunctionNode( + 'secured', + (ctx: NodeContext) => { + runs++; + const cred = ctx.state.get('temp:' + CREDENTIAL_KEY); + sawApiKey = cred?.apiKey; + return `data(${cred?.apiKey})`; + }, + {authConfig: apiKeyAuthConfig(), rerunOnResume: true}, + ); + + const wf = new Workflow({name: 'auth_wf', edges: [['START', secured]]}); + const agent = new WorkflowAgent(wf); + const sessionService = new InMemorySessionService(); + const session = await sessionService.createSession({ + appName: 'test_app', + userId: 'u1', + }); + const runner = new Runner({appName: 'test_app', agent, sessionService}); + + // Turn 1: no credential -> auth request interrupt, handler NOT run. + const turn1 = await collect( + runner.runAsync({ + userId: 'u1', + sessionId: session.id, + newMessage: {role: 'user', parts: [{text: 'go'}]}, + }), + ); + expect(runs).toBe(0); + expect(turn1.some(hasAuthRequestFunctionCall)).toBe(true); + + // Turn 2: supply the credential (as a filled AuthConfig) and resume. + const credentialResponse: AuthConfig = { + authScheme: { + type: 'apiKey', + in: 'header', + name: 'X-API-Key', + } as AuthScheme, + credentialKey: CREDENTIAL_KEY, + exchangedAuthCredential: { + authType: AuthCredentialTypes.API_KEY, + apiKey: 'secret-123', + }, + }; + const turn2 = await collect( + runner.runAsync({ + userId: 'u1', + sessionId: session.id, + newMessage: { + role: 'user', + parts: [ + { + functionResponse: { + id: CREDENTIAL_KEY, + name: 'adk_request_credential', + response: credentialResponse as unknown as Record< + string, + unknown + >, + }, + }, + ], + }, + }), + ); + + // The node ran once, saw the supplied API key, and produced output. + expect(runs).toBe(1); + expect(sawApiKey).toBe('secret-123'); + expect(turn2.some((e) => e.output === 'data(secret-123)')).toBe(true); + }); + + it('runs immediately when the credential already exists in state', async () => { + let runs = 0; + const secured = new FunctionNode( + 'secured', + () => { + runs++; + return 'ok'; + }, + {authConfig: apiKeyAuthConfig()}, + ); + const wf = new Workflow({name: 'auth_wf2', edges: [['START', secured]]}); + + // Pre-seed the credential directly in the session state. + const {events, output} = await driveWorkflow(wf, 'go', { + ic: createIc({ + ['temp:' + CREDENTIAL_KEY]: { + authType: AuthCredentialTypes.API_KEY, + apiKey: 'pre-existing', + }, + }), + }); + + expect(runs).toBe(1); + expect(output).toBe('ok'); + expect(events.some(hasAuthRequestFunctionCall)).toBe(false); + }); +}); diff --git a/core/test/workflow/dynamic_resume_test.ts b/core/test/workflow/dynamic_resume_test.ts new file mode 100644 index 000000000..f3f0735bf --- /dev/null +++ b/core/test/workflow/dynamic_resume_test.ts @@ -0,0 +1,100 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {describe, expect, it} from 'vitest'; +import {Event} from '../../src/events/event.js'; +import {Runner} from '../../src/runner/runner.js'; +import {InMemorySessionService} from '../../src/sessions/in_memory_session_service.js'; +import {NodeContext} from '../../src/workflow/node_context.js'; +import {FunctionNode} from '../../src/workflow/nodes/function_node.js'; +import {RequestInput} from '../../src/workflow/request_input.js'; +import {hasRequestInputFunctionCall} from '../../src/workflow/utils/hitl_utils.js'; +import {Workflow} from '../../src/workflow/workflow.js'; +import {WorkflowAgent} from '../../src/workflow/workflow_agent.js'; + +async function collect(gen: AsyncGenerator): Promise { + const out: Event[] = []; + for await (const e of gen) { + out.push(e); + } + return out; +} + +describe('Phase 5b-cont — dynamic (ctx.runNode) resume via the Runner', () => { + it('dedups a completed dynamic node and resumes a waiting one', async () => { + let stepRuns = 0; + let askRuns = 0; + + const step = new FunctionNode('step', (_c, input) => { + stepRuns++; + return `step(${input})`; + }); + const ask = new FunctionNode('ask', (ctx: NodeContext) => { + askRuns++; + const answer = ctx.resumeInputs['confirm']; + if (answer === undefined) { + return new RequestInput({interruptId: 'confirm', message: 'confirm?'}); + } + return `confirmed:${answer}`; + }); + + // Imperative workflow: run `step` (completes), then `ask` (interrupts). + const wf = new Workflow({ + name: 'dyn_resume_wf', + dynamicEntry: async (ctx, input) => { + const s = await ctx.runNode(step, input); + const a = await ctx.runNode(ask); + return {step: s.output, ask: a.output}; + }, + }); + + const agent = new WorkflowAgent(wf); + const sessionService = new InMemorySessionService(); + const session = await sessionService.createSession({ + appName: 'test_app', + userId: 'u1', + }); + const runner = new Runner({appName: 'test_app', agent, sessionService}); + + // Turn 1: step runs, ask interrupts. + const turn1 = await collect( + runner.runAsync({ + userId: 'u1', + sessionId: session.id, + newMessage: {role: 'user', parts: [{text: 'x'}]}, + }), + ); + expect(stepRuns).toBe(1); + expect(turn1.some(hasRequestInputFunctionCall)).toBe(true); + + // Turn 2: provide the confirmation and resume. + const turn2 = await collect( + runner.runAsync({ + userId: 'u1', + sessionId: session.id, + newMessage: { + role: 'user', + parts: [ + { + functionResponse: { + id: 'confirm', + name: 'adk_request_input', + response: {result: 'yes'}, + }, + }, + ], + }, + }), + ); + + // `step` was fast-forwarded (cached) -> NOT re-executed. + expect(stepRuns).toBe(1); + // `ask` re-ran with the resolved resume input and completed. + expect(askRuns).toBe(2); + expect(turn2.some((e) => e.output === 'step(x)')).toBe(false); + expect(turn2.some((e) => e.output === 'confirmed:yes')).toBe(true); + }); +}); diff --git a/core/test/workflow/dynamic_workflow_test.ts b/core/test/workflow/dynamic_workflow_test.ts new file mode 100644 index 000000000..587b46a8f --- /dev/null +++ b/core/test/workflow/dynamic_workflow_test.ts @@ -0,0 +1,106 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {describe, expect, it} from 'vitest'; +import {FunctionNode} from '../../src/workflow/nodes/function_node.js'; +import {Workflow} from '../../src/workflow/workflow.js'; +import {driveWorkflow} from './test_helpers.js'; + +describe('Phase 4 — dynamic (imperative) workflows', () => { + it('runs an imperative dynamicEntry driving ctx.runNode()', async () => { + const step = new FunctionNode('step', (_c, input) => `step(${input})`); + const wf = new Workflow({ + name: 'dyn', + dynamicEntry: async (ctx, input) => { + const child = await ctx.runNode(step, input); + return `wrapped[${child.output}]`; + }, + }); + expect((await driveWorkflow(wf, 'x')).output).toBe('wrapped[step(x)]'); + }); + + it('supports a bounded loop (the cycle case that used to hang)', async () => { + // Increment until >= 3; a natural JS loop, terminated by user code. + const inc = new FunctionNode('inc', (_c, n: number) => (n as number) + 1); + const wf = new Workflow({ + name: 'loop', + dynamicEntry: async (ctx, input) => { + let value = input as number; + let iterations = 0; + while (value < 3) { + const child = await ctx.runNode(inc, value); + value = child.output as number; + iterations++; + } + return {value, iterations}; + }, + }); + expect((await driveWorkflow(wf, 0)).output).toEqual({ + value: 3, + iterations: 3, + }); + }); + + it('assigns distinct run ids to repeated dynamic calls (streams each event)', async () => { + const emit = new FunctionNode('emit', (_c, n) => `emit(${n})`); + const wf = new Workflow({ + name: 'repeat', + dynamicEntry: async (ctx) => { + const outs: unknown[] = []; + for (let i = 0; i < 3; i++) { + outs.push((await ctx.runNode(emit, i)).output); + } + return outs; + }, + }); + const {events, output} = await driveWorkflow(wf); + expect(output).toEqual(['emit(0)', 'emit(1)', 'emit(2)']); + // Each iteration streamed its own event. + expect(events.filter((e) => e.author === 'emit')).toHaveLength(3); + }); + + it('deduplicates concurrent ctx.runNode() calls to the same run', async () => { + let executions = 0; + const slow = new FunctionNode('slow', async () => { + executions++; + await new Promise((r) => setTimeout(r, 10)); + return 'done'; + }); + const wf = new Workflow({ + name: 'dedup', + dynamicEntry: async (ctx) => { + // Same explicit runId => same run path => deduped. + const [a, b] = await Promise.all([ + ctx.runNode(slow, undefined, {runId: 'shared'}), + ctx.runNode(slow, undefined, {runId: 'shared'}), + ]); + return {a: a.output, b: b.output, executions}; + }, + }); + expect((await driveWorkflow(wf)).output).toEqual({ + a: 'done', + b: 'done', + executions: 1, + }); + }); + + it('supports the node-as-tool pattern (a node calls a sub-node)', async () => { + const adder = new FunctionNode( + 'adder', + (_c, args: {a: number; b: number}) => args.a + args.b, + ); + const orchestrator = new FunctionNode('orchestrator', async (ctx) => { + const r1 = await ctx.runNode(adder, {a: 2, b: 3}); + const r2 = await ctx.runNode(adder, {a: 10, b: r1.output as number}); + return r2.output; + }); + const wf = new Workflow({ + name: 'node_as_tool', + edges: [['START', orchestrator]], + }); + expect((await driveWorkflow(wf)).output).toBe(15); + }); +}); diff --git a/core/test/workflow/event_model_test.ts b/core/test/workflow/event_model_test.ts index d78d54c71..ce8150c6e 100644 --- a/core/test/workflow/event_model_test.ts +++ b/core/test/workflow/event_model_test.ts @@ -64,6 +64,31 @@ describe('Phase 0 — workflow event-model extensions', () => { expect(back.actions.agentState).toEqual({status: 3}); expect(back.actions.endOfAgent).toBe(true); }); + + it('preserves user-defined keys in output/route/agentState across persistence', () => { + // A persistent session (e.g. VertexAiSessionService) round-trips events + // through snake_case. The workflow's arbitrary payloads — node output, the + // node input stashed under actions.agentState for HITL resume, and the + // emitted route — carry user-defined keys that must survive verbatim. Without + // the preserve-key allowlists, `{userName: 'Ada'}` would come back as + // `{user_name: 'Ada'}` and a resumed node would re-run with mangled keys. + const ev = createEvent({ + author: 'node_a', + output: {userName: 'Ada', nested: {maxRetries: 3}}, + route: 'needsReview', + actions: { + agentState: {input: {firstName: 'Ada', lastName: 'Lovelace'}}, + }, + }); + + const back = transformToCamelCaseEvent(transformToSnakeCaseEvent(ev)); + + expect(back.output).toEqual({userName: 'Ada', nested: {maxRetries: 3}}); + expect(back.route).toBe('needsReview'); + expect(back.actions.agentState).toEqual({ + input: {firstName: 'Ada', lastName: 'Lovelace'}, + }); + }); }); describe('isEvent — signature-symbol brand', () => { diff --git a/core/test/workflow/hitl_flow_test.ts b/core/test/workflow/hitl_flow_test.ts new file mode 100644 index 000000000..2fbd1514b --- /dev/null +++ b/core/test/workflow/hitl_flow_test.ts @@ -0,0 +1,130 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {describe, expect, it} from 'vitest'; +import {Event} from '../../src/events/event.js'; +import {FunctionNode} from '../../src/workflow/nodes/function_node.js'; +import {RequestInput} from '../../src/workflow/request_input.js'; +import { + hasRequestInputFunctionCall, + REQUEST_INPUT_FUNCTION_CALL_NAME, +} from '../../src/workflow/utils/hitl_utils.js'; +import {Workflow} from '../../src/workflow/workflow.js'; +import {driveWorkflow} from './test_helpers.js'; + +/** + * Drives a workflow, optionally supplying resume inputs (keyed by interrupt id). + * Thin wrapper over the shared {@link driveWorkflow} for this file's positional + * resume-input call sites. + */ +function drive( + wf: Workflow, + input?: unknown, + resumeInputs: Record = {}, +): Promise<{output: unknown; interruptIds: string[]; events: Event[]}> { + return driveWorkflow(wf, input, {resumeInputs}); +} + +describe('Phase 5 — HITL (pause / resume)', () => { + it('pauses on RequestInput and surfaces the interrupt id', async () => { + const approval = new FunctionNode('approval', (ctx) => { + const answer = ctx.resumeInputs['approve-1']; + if (answer === undefined) { + return new RequestInput({ + interruptId: 'approve-1', + message: 'Approve?', + }); + } + return `decided:${answer}`; + }); + const wf = new Workflow({name: 'hitl', edges: [['START', approval]]}); + + // Run 1: no resume input → interrupt. + const paused = await drive(wf, undefined); + expect(paused.interruptIds).toEqual(['approve-1']); + expect(paused.output).toBeUndefined(); + // The interrupt surfaced as a request_input function-call event. + expect(paused.events.some(hasRequestInputFunctionCall)).toBe(true); + const fc = paused.events + .flatMap((e) => e.content?.parts ?? []) + .find((p) => p.functionCall?.name === REQUEST_INPUT_FUNCTION_CALL_NAME); + expect(fc?.functionCall?.id).toBe('approve-1'); + }); + + it('resumes and completes when the resume input is provided', async () => { + const approval = new FunctionNode('approval', (ctx) => { + const answer = ctx.resumeInputs['approve-1']; + if (answer === undefined) { + return new RequestInput({ + interruptId: 'approve-1', + message: 'Approve?', + }); + } + return `decided:${answer}`; + }); + const wf = new Workflow({name: 'hitl', edges: [['START', approval]]}); + + // Run 2: provide the resume input → completes. + const resumed = await drive(wf, undefined, {'approve-1': 'yes'}); + expect(resumed.interruptIds).toEqual([]); + expect(resumed.output).toBe('decided:yes'); + }); + + it('propagates an interrupt from a mid-graph node and halts downstream', async () => { + const ran: string[] = []; + const a = new FunctionNode('a', (_c, input) => { + ran.push('a'); + return `a:${input}`; + }); + const gate = new FunctionNode('gate', (ctx, input) => { + ran.push('gate'); + const answer = ctx.resumeInputs['gate-1']; + if (answer === undefined) { + return new RequestInput({interruptId: 'gate-1', message: 'continue?'}); + } + return `${input}|gate:${answer}`; + }); + const c = new FunctionNode('c', (_c, input) => { + ran.push('c'); + return `c:${input}`; + }); + const wf = new Workflow({name: 'chain', edges: [['START', a, gate, c]]}); + + const paused = await drive(wf, 'x'); + expect(paused.interruptIds).toEqual(['gate-1']); + // Downstream node c must NOT have run while gate is waiting. + expect(ran).toEqual(['a', 'gate']); + + const resumed = await drive(wf, 'x', {'gate-1': 'ok'}); + expect(resumed.output).toBe('c:a:x|gate:ok'); + }); + + it('supports HITL in an imperative dynamicEntry workflow', async () => { + const ask = new FunctionNode('ask', (ctx) => { + const answer = ctx.resumeInputs['name']; + if (answer === undefined) { + return new RequestInput({interruptId: 'name', message: 'Your name?'}); + } + return answer; + }); + const wf = new Workflow({ + name: 'dyn_hitl', + dynamicEntry: async (ctx) => { + const child = await ctx.runNode(ask); + if (child.interruptIds.length > 0) { + return undefined; // still waiting + } + return `hello ${child.output}`; + }, + }); + + const paused = await drive(wf); + expect(paused.interruptIds).toEqual(['name']); + + const resumed = await drive(wf, undefined, {name: 'Ada'}); + expect(resumed.output).toBe('hello Ada'); + }); +}); diff --git a/core/test/workflow/node_execution_test.ts b/core/test/workflow/node_execution_test.ts index 629243e2f..e8f0f78d5 100644 --- a/core/test/workflow/node_execution_test.ts +++ b/core/test/workflow/node_execution_test.ts @@ -16,6 +16,7 @@ import { isNodeTimeoutError, } from '../../src/workflow/errors.js'; import {NodeContext} from '../../src/workflow/node_context.js'; +import {executeChildNode} from '../../src/workflow/node_runner.js'; import {createIc, driveNode, FnNode} from './test_helpers.js'; // --- Tests ---------------------------------------------------------------- @@ -96,7 +97,14 @@ describe('Phase 1 — node execution & the push/pull bridge', () => { nodePath: '', runId: 'root', }); - const child = await root.runNode(node, undefined, {useAsOutput: true}); + // executeChildNode returns the concrete child NodeContext (runNode's return + // type widens to NodeContext | NodeResult for the resume fast-forward case). + const child = await executeChildNode({ + parent: root, + node, + input: undefined, + options: {useAsOutput: true}, + }); expect(child.state.get('counter')).toBe(7); expect(child.actions.stateDelta['counter']).toBe(7); }); @@ -165,7 +173,12 @@ describe('Phase 1 — node execution & the push/pull bridge', () => { nodePath: '', runId: 'root', }); - const child = await root.runNode(node, 'x', {useAsOutput: true}); + const child = await executeChildNode({ + parent: root, + node, + input: 'x', + options: {useAsOutput: true}, + }); expect(child.output).toBe('ok'); // The failed first attempt's write must not survive into the committed diff --git a/core/test/workflow/parallel_test.ts b/core/test/workflow/parallel_test.ts new file mode 100644 index 000000000..1b7c4b35d --- /dev/null +++ b/core/test/workflow/parallel_test.ts @@ -0,0 +1,135 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {describe, expect, it} from 'vitest'; +import { + branchPathFromString, + commonPrefixOf, + createSubBranch, +} from '../../src/workflow/branch_path.js'; +import {node} from '../../src/workflow/node.js'; +import {NodeContext} from '../../src/workflow/node_context.js'; +import {FunctionNode} from '../../src/workflow/nodes/function_node.js'; +import {ParallelWorker} from '../../src/workflow/nodes/parallel_worker.js'; +import {Workflow} from '../../src/workflow/workflow.js'; +import {driveWorkflow} from './test_helpers.js'; + +describe('Phase 6 — BranchPath', () => { + it('creates sub-branches with and without run ids', () => { + expect(createSubBranch('parent', {name: 'child', runId: '1'})).toBe( + 'parent.child@1', + ); + expect(createSubBranch(undefined, {name: 'agent'})).toBe('agent'); + }); + + it('computes the common prefix of branches', () => { + expect(commonPrefixOf(['a@1.b@2', 'a@1.c@3'])).toBe('a@1'); + expect(commonPrefixOf(['a@1', 'b@1'])).toBe(''); + expect(commonPrefixOf([])).toBe(''); + }); + + it('detects descendants', () => { + const parent = branchPathFromString('a@1'); + expect(branchPathFromString('a@1.b@2').isDescendantOf(parent)).toBe(true); + expect(branchPathFromString('a@1').isDescendantOf(parent)).toBe(false); + expect(branchPathFromString('x@1.b@2').isDescendantOf(parent)).toBe(false); + }); +}); + +describe('Phase 6 — ParallelWorker', () => { + it('maps a list input across the inner node, preserving order', async () => { + const doubler = new FunctionNode('double', (_c, n: number) => n * 2); + const worker = new ParallelWorker(doubler); + const wf = new Workflow({name: 'pw', edges: [['START', worker]]}); + expect((await driveWorkflow(wf, [1, 2, 3, 4])).output).toEqual([ + 2, 4, 6, 8, + ]); + }); + + it('wraps a single (non-list) input as a one-element list', async () => { + const worker = new ParallelWorker(new FunctionNode('id', (_c, n) => n)); + const wf = new Workflow({name: 'pw1', edges: [['START', worker]]}); + expect((await driveWorkflow(wf, 'solo')).output).toEqual(['solo']); + }); + + it('returns [] for an empty list', async () => { + const worker = new ParallelWorker(new FunctionNode('id', (_c, n) => n)); + const wf = new Workflow({name: 'pw0', edges: [['START', worker]]}); + expect((await driveWorkflow(wf, [])).output).toEqual([]); + }); + + it('respects maxParallelWorkers (bounded concurrency)', async () => { + let active = 0; + let peak = 0; + const slow = new FunctionNode('slow', async (_c, n: number) => { + active++; + peak = Math.max(peak, active); + await new Promise((r) => setTimeout(r, 5)); + active--; + return n; + }); + const worker = new ParallelWorker(slow, {maxParallelWorkers: 2}); + const wf = new Workflow({name: 'bounded', edges: [['START', worker]]}); + const {output: out} = await driveWorkflow(wf, [1, 2, 3, 4, 5, 6]); + expect(out).toEqual([1, 2, 3, 4, 5, 6]); + expect(peak).toBeLessThanOrEqual(2); + }); + + it('cancels remaining work and propagates the first error', async () => { + const flaky = new FunctionNode('flaky', (_c, n: number) => { + if (n === 3) { + throw new Error('boom at 3'); + } + return n; + }); + const worker = new ParallelWorker(flaky, {maxParallelWorkers: 1}); + const wf = new Workflow({name: 'err', edges: [['START', worker]]}); + await expect(driveWorkflow(wf, [1, 2, 3, 4, 5])).rejects.toThrow( + 'boom at 3', + ); + }); + + it('is produced by node(fn, {parallelWorker: true})', async () => { + const n = node((_c: NodeContext, x: number) => x + 1, { + name: 'inc', + parallelWorker: true, + maxParallelWorkers: 3, + }); + expect(n).toBeInstanceOf(ParallelWorker); + const wf = new Workflow({name: 'pwnode', edges: [['START', n]]}); + expect((await driveWorkflow(wf, [10, 20, 30])).output).toEqual([ + 11, 21, 31, + ]); + }); + + it('rejects maxParallelWorkers without parallelWorker', () => { + expect(() => + node((_c: NodeContext, x: unknown) => x, { + name: 'x', + maxParallelWorkers: 2, + }), + ).toThrow(/maxParallelWorkers/); + }); + + it('assigns run ids by item index for deterministic resume', async () => { + // Each child stamps its own run id into the output. With bounded + // concurrency the run id must still equal the item index (not the + // call/completion order), so resume can fast-forward each item correctly. + const worker = new ParallelWorker( + node((ctx: NodeContext, item: string) => `${item}#${ctx.runId}`, { + name: 'w', + }), + {maxParallelWorkers: 2}, + ); + const wf = new Workflow({name: 'pw_ids', edges: [['START', worker]]}); + expect((await driveWorkflow(wf, ['a', 'b', 'c', 'd'])).output).toEqual([ + 'a#0', + 'b#1', + 'c#2', + 'd#3', + ]); + }); +}); diff --git a/core/test/workflow/resume_test.ts b/core/test/workflow/resume_test.ts new file mode 100644 index 000000000..d4212460f --- /dev/null +++ b/core/test/workflow/resume_test.ts @@ -0,0 +1,218 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {describe, expect, it} from 'vitest'; +import { + createEvent, + Event, + transformToCamelCaseEvent, + transformToSnakeCaseEvent, +} from '../../src/events/event.js'; +import {createEventActions} from '../../src/events/event_actions.js'; +import {Runner} from '../../src/runner/runner.js'; +import {InMemorySessionService} from '../../src/sessions/in_memory_session_service.js'; +import {node} from '../../src/workflow/node.js'; +import {NodeContext} from '../../src/workflow/node_context.js'; +import {RequestInput} from '../../src/workflow/request_input.js'; +import {hasRequestInputFunctionCall} from '../../src/workflow/utils/hitl_utils.js'; +import { + isFastForwardable, + reconstructNodeStates, +} from '../../src/workflow/utils/rehydration_utils.js'; +import {Workflow} from '../../src/workflow/workflow.js'; +import {WorkflowAgent} from '../../src/workflow/workflow_agent.js'; + +describe('Phase 5b — rehydration utility', () => { + it('reconstructs completed outputs and unresolved interrupts', () => { + const events: Event[] = [ + createEvent({author: 'a', nodeInfo: {path: 'wf.a'}, output: 'A(x)'}), + createEvent({ + author: 'gate', + nodeInfo: {path: 'wf.gate'}, + content: { + role: 'model', + parts: [{functionCall: {name: 'adk_request_input', id: 'gate-1'}}], + }, + longRunningToolIds: ['gate-1'], + }), + ]; + const states = reconstructNodeStates(events); + + expect(states.get('a')?.output).toBe('A(x)'); + expect(isFastForwardable(states.get('a')!)).toBe(true); + expect([...states.get('gate')!.interruptIds]).toEqual(['gate-1']); + // gate has no output and an unresolved interrupt -> not fast-forwardable. + expect(isFastForwardable(states.get('gate')!)).toBe(false); + }); + + it('resolves an interrupt from a user function response', () => { + const events: Event[] = [ + createEvent({ + author: 'gate', + nodeInfo: {path: 'wf.gate'}, + longRunningToolIds: ['gate-1'], + }), + createEvent({ + author: 'user', + content: { + role: 'user', + parts: [ + { + functionResponse: { + id: 'gate-1', + name: 'adk_request_input', + response: {result: 'approved'}, + }, + }, + ], + }, + }), + ]; + const states = reconstructNodeStates(events); + expect(states.get('gate')?.resolvedResponses.get('gate-1')).toBe( + 'approved', + ); + }); + + it('recovers structured output and interrupt input after a DB serialization round-trip', () => { + const events: Event[] = [ + createEvent({ + author: 'lookup', + nodeInfo: {path: 'wf.lookup'}, + output: {cityName: 'Paris', timeInfo: '10:10 AM'}, + }), + createEvent({ + author: 'gate', + nodeInfo: {path: 'wf.gate'}, + longRunningToolIds: ['gate-1'], + // The engine stashes the waiting node's original input here so it + // re-runs with it on resume (see node_runner runOnce). + actions: createEventActions({ + agentState: {input: {userId: 42, requestedItems: ['a', 'b']}}, + }), + }), + ]; + + // Simulate what a persistent (DB/Vertex) session store does on write+read: + // snake_case on save, camelCase on load. Without the preserve-list fix this + // mangles the arbitrary output/agentState keys. + const persisted = events.map( + (e) => transformToCamelCaseEvent(transformToSnakeCaseEvent(e)) as Event, + ); + + const states = reconstructNodeStates(persisted); + expect(states.get('lookup')?.output).toEqual({ + cityName: 'Paris', + timeInfo: '10:10 AM', + }); + expect(states.get('gate')?.input).toEqual({ + userId: 42, + requestedItems: ['a', 'b'], + }); + }); + + it('scopes reconstruction to direct children so nested same-named nodes do not collide', () => { + const events: Event[] = [ + createEvent({nodeInfo: {path: 'root.process'}, output: 'OUTER'}), + createEvent({nodeInfo: {path: 'root.inner.process'}, output: 'INNER'}), + ]; + const outer = reconstructNodeStates(events, 'root'); + const inner = reconstructNodeStates(events, 'root.inner'); + expect(outer.get('process')?.output).toBe('OUTER'); + expect(inner.get('process')?.output).toBe('INNER'); + // The outer scope must not absorb the nested (grandchild) node. + expect(outer.size).toBe(1); + }); + + it('keys by leaf name when no parent path is given (utility mode)', () => { + const events: Event[] = [ + createEvent({nodeInfo: {path: 'wf.a'}, output: 'A'}), + ]; + expect(reconstructNodeStates(events).get('a')?.output).toBe('A'); + }); +}); + +describe('Phase 5b — HITL resume via the Runner', () => { + it('resumes an interrupted workflow without re-running completed nodes', async () => { + let aRuns = 0; + const a = node( + (_c: NodeContext, input: unknown) => { + aRuns++; + return `A(${input})`; + }, + {name: 'a'}, + ); + // A single-node HITL gate that RE-RUNS on resume to read its answer from + // ctx.resumeInputs. In the faithful (Python) model this is rerun_on_resume= + // true; the default (false) is the two-node pattern where the node does not + // re-run and its output becomes the resume value. + const gate = node( + (ctx: NodeContext, input: unknown) => { + const answer = ctx.resumeInputs['gate-1']; + if (answer === undefined) { + return new RequestInput({interruptId: 'gate-1', message: 'approve?'}); + } + return `${input}|${answer}`; + }, + {name: 'gate', rerunOnResume: true}, + ); + const c = node((_c: NodeContext, input: unknown) => `C(${input})`, { + name: 'c', + }); + const wf = new Workflow({ + name: 'resume_wf', + edges: [['START', a, gate, c]], + }); + + const agent = new WorkflowAgent(wf); + const sessionService = new InMemorySessionService(); + const session = await sessionService.createSession({ + appName: 'test_app', + userId: 'u1', + }); + const runner = new Runner({appName: 'test_app', agent, sessionService}); + + // --- Turn 1: run until the gate interrupts --- + const turn1: Event[] = []; + for await (const event of runner.runAsync({ + userId: 'u1', + sessionId: session.id, + newMessage: {role: 'user', parts: [{text: 'x'}]}, + })) { + turn1.push(event); + } + expect(aRuns).toBe(1); + expect(turn1.some(hasRequestInputFunctionCall)).toBe(true); + // c must not have produced output yet. + expect(turn1.some((e) => e.output === 'C(A(x)|approved)')).toBe(false); + + // --- Turn 2: provide the interrupt response and resume --- + const turn2: Event[] = []; + for await (const event of runner.runAsync({ + userId: 'u1', + sessionId: session.id, + newMessage: { + role: 'user', + parts: [ + { + functionResponse: { + id: 'gate-1', + name: 'adk_request_input', + response: {result: 'approved'}, + }, + }, + ], + }, + })) { + turn2.push(event); + } + + // A was fast-forwarded (cached), NOT re-executed. + expect(aRuns).toBe(1); + // The workflow resumed through the gate and completed at c. + expect(turn2.some((e) => e.output === 'C(A(x)|approved)')).toBe(true); + }); +}); diff --git a/core/test/workflow/routing_test.ts b/core/test/workflow/routing_test.ts new file mode 100644 index 000000000..0fa5db644 --- /dev/null +++ b/core/test/workflow/routing_test.ts @@ -0,0 +1,127 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {describe, expect, it} from 'vitest'; +import {createEvent} from '../../src/events/event.js'; +import {BaseNode} from '../../src/workflow/base_node.js'; +import {DEFAULT_ROUTE, Edge, RouteValue} from '../../src/workflow/graph.js'; +import {JoinNode} from '../../src/workflow/nodes/join_node.js'; +import {Workflow} from '../../src/workflow/workflow.js'; +import {driveNode, FnNode} from './test_helpers.js'; + +const emit = (name: string, route: RouteValue): BaseNode => + new FnNode(name, () => createEvent({route, output: route})); + +const echo = (name: string): BaseNode => + new FnNode(name, (_c, i) => `${name}(${i})`); + +describe('workflow routing values', () => { + it('matches numeric route keys', async () => { + const router = emit('router', 2); + const a = echo('a'); + const b = echo('b'); + const wf = new Workflow({ + name: 'numeric_route', + edges: [ + ['START', router], + [router, {1: a, 2: b}], + ], + }); + expect((await driveNode(wf, 'x')).output).toBe('b(2)'); + }); + + it('matches a list route on an explicit Edge (any listed value)', async () => { + const router = emit('router', 3); + const target = echo('target'); + const other = echo('other'); + const wf = new Workflow({ + name: 'list_route', + edges: [ + ['START', router], + new Edge(router, target, [2, 3]), + new Edge(router, other, 9), + ], + }); + expect((await driveNode(wf, 'x')).output).toBe('target(3)'); + }); + + it('fans out from a single route to multiple nodes, then joins', async () => { + const router = emit('router', 'go'); + const a = echo('a'); + const b = echo('b'); + const join = new JoinNode({name: 'join'}); + const wf = new Workflow({ + name: 'route_fan_out', + edges: [ + ['START', router], + [router, {go: [a, b]}], + [[a, b], join], + ], + }); + expect((await driveNode(wf, 'x')).output).toEqual({ + a: 'a(go)', + b: 'b(go)', + }); + }); + + it('uses DEFAULT_ROUTE only when no specific route matches', async () => { + // Specific route matches -> takes the specific branch. + const router = emit('router', 'known'); + const wfMatch = new Workflow({ + name: 'default_route_match', + edges: [ + ['START', router], + [router, {known: echo('a'), [DEFAULT_ROUTE]: echo('fb')}], + ], + }); + expect((await driveNode(wfMatch, 'x')).output).toBe('a(known)'); + + // No specific route matches -> falls back to DEFAULT_ROUTE. + const router2 = emit('router2', 'unknown'); + const wfFallback = new Workflow({ + name: 'default_route_fallback', + edges: [ + ['START', router2], + [router2, {known: echo('a2'), [DEFAULT_ROUTE]: echo('fb2')}], + ], + }); + expect((await driveNode(wfFallback, 'x')).output).toBe('fb2(unknown)'); + }); + + it('matches a boolean route key in a routing map', async () => { + const router = emit('router', true); + const wf = new Workflow({ + name: 'bool_route', + edges: [ + ['START', router], + [router, {true: echo('yes'), false: echo('no')}], + ], + }); + expect((await driveNode(wf, 'x')).output).toBe('yes(true)'); + }); + + it('fires multiple branches when a node emits an array of routes', async () => { + const router = new FnNode('router', () => + createEvent({route: ['a', 'b'], output: 'msg'}), + ); + const a = echo('a'); + const b = echo('b'); + const c = echo('c'); // present in the map but not emitted -> must not run + const join = new JoinNode({name: 'join'}); + const wf = new Workflow({ + name: 'multi_route', + edges: [ + ['START', router], + [router, {a, b, c}], + [[a, b], join], + ], + }); + expect((await driveNode(wf, 'x')).output).toEqual({ + a: 'a(msg)', + b: 'b(msg)', + }); + }); +}); diff --git a/core/test/workflow/runner_integration_test.ts b/core/test/workflow/runner_integration_test.ts new file mode 100644 index 000000000..fa8742fb3 --- /dev/null +++ b/core/test/workflow/runner_integration_test.ts @@ -0,0 +1,115 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {describe, expect, it} from 'vitest'; +import {createEvent, Event} from '../../src/events/event.js'; +import {Runner} from '../../src/runner/runner.js'; +import {InMemorySessionService} from '../../src/sessions/in_memory_session_service.js'; +import {DEFAULT_ROUTE} from '../../src/workflow/graph.js'; +import {node} from '../../src/workflow/node.js'; +import {NodeContext} from '../../src/workflow/node_context.js'; +import {Workflow} from '../../src/workflow/workflow.js'; +import {WorkflowAgent} from '../../src/workflow/workflow_agent.js'; + +async function runViaRunner( + workflow: Workflow, + text: string, +): Promise { + const agent = new WorkflowAgent(workflow); + const sessionService = new InMemorySessionService(); + const session = await sessionService.createSession({ + appName: 'test_app', + userId: 'u1', + }); + const runner = new Runner({appName: 'test_app', agent, sessionService}); + + const events: Event[] = []; + for await (const event of runner.runAsync({ + userId: 'u1', + sessionId: session.id, + newMessage: {role: 'user', parts: [{text}]}, + })) { + events.push(event); + } + return events; +} + +describe('Phase 8 — WorkflowAgent via the real Runner', () => { + it('runs a single-node workflow end-to-end', async () => { + const wf = new Workflow({ + name: 'greet_wf', + edges: [ + [ + 'START', + node((_c: NodeContext, input: string) => `hello ${input}`, { + name: 'greet', + }), + ], + ], + }); + const events = await runViaRunner(wf, 'world'); + expect(events.some((e) => e.output === 'hello world')).toBe(true); + }); + + it('runs a linear sequence end-to-end (input threads through)', async () => { + const a = node((_c: NodeContext, i: string) => `${i}->A`, {name: 'a'}); + const b = node((_c: NodeContext, i: string) => `${i}->B`, {name: 'b'}); + const wf = new Workflow({name: 'seq_wf', edges: [['START', a, b]]}); + + const events = await runViaRunner(wf, 'INIT'); + const outputs = events + .filter((e) => e.output !== undefined) + .map((e) => e.output); + expect(outputs).toContain('INIT->A'); + expect(outputs).toContain('INIT->A->B'); + }); + + it('runs a routed workflow end-to-end', async () => { + const route = node( + (_c: NodeContext, input: string) => + createEvent({route: input.includes('?') ? 'q' : 's', output: input}), + {name: 'route'}, + ); + const q = node((_c: NodeContext, i: string) => `Q:${i}`, {name: 'q'}); + const s = node((_c: NodeContext, i: string) => `S:${i}`, {name: 's'}); + const wf = new Workflow({ + name: 'route_wf', + edges: [ + ['START', route], + [route, {q, s}], + ], + }); + + const events = await runViaRunner(wf, 'hi?'); + expect(events.some((e) => e.output === 'Q:hi?')).toBe(true); + }); + + it('falls back to DEFAULT_ROUTE end-to-end', async () => { + const check = node( + (_c: NodeContext, input: string) => + createEvent( + input === 'skip' ? {output: input} : {route: 'go', output: input}, + ), + {name: 'check'}, + ); + const go = node((_c: NodeContext, i: string) => `GO:${i}`, { + name: 'go_node', + }); + const fallback = node((_c: NodeContext, i: string) => `DEFAULT:${i}`, { + name: 'fallback', + }); + const wf = new Workflow({ + name: 'default_wf', + edges: [ + ['START', check], + [check, {go, [DEFAULT_ROUTE]: fallback}], + ], + }); + + const events = await runViaRunner(wf, 'skip'); + expect(events.some((e) => e.output === 'DEFAULT:skip')).toBe(true); + }); +}); diff --git a/core/test/workflow/test_helpers.ts b/core/test/workflow/test_helpers.ts index 41a758d59..622cba6b5 100644 --- a/core/test/workflow/test_helpers.ts +++ b/core/test/workflow/test_helpers.ts @@ -71,6 +71,47 @@ export async function driveNode( return {events, output: root.output, ctx: root}; } +/** Options for {@link driveWorkflow}. */ +export interface DriveWorkflowOptions { + /** InvocationContext to run under (defaults to a fresh {@link createIc}). */ + ic?: InvocationContext; + /** Resume inputs keyed by interrupt id (for HITL/auth resume). */ + resumeInputs?: Record; +} + +/** + * Drives a workflow (or any node) to completion and returns its streamed events, + * final output, and the interrupt ids it is paused on — the shared harness for + * the workflow-level tests (replaces the per-file `createIc`/`driveWorkflow` + * copies that reached for `as unknown as Session/BaseAgent`). + */ +export async function driveWorkflow( + wf: BaseNode, + input?: unknown, + options: DriveWorkflowOptions = {}, +): Promise<{events: Event[]; output: unknown; interruptIds: string[]}> { + const channel = new AsyncQueue(); + const root = new NodeContext({ + invocationContext: options.ic ?? createIc(), + channel, + nodePath: '', + runId: 'root', + resumeInputs: options.resumeInputs, + }); + const events: Event[] = []; + const resultPromise = root.runNode(wf, input, {useAsOutput: true}); + const settle = resultPromise.then( + () => channel.close(), + (err) => channel.fail(err), + ); + for await (const ev of channel) { + events.push(ev); + } + await settle; + const result = await resultPromise; + return {events, output: root.output, interruptIds: result.interruptIds}; +} + /** A node whose behavior is a plain function returning a value or Event. */ export class FnNode extends BaseNode { constructor( diff --git a/core/test/workflow/workflow_advanced_test.ts b/core/test/workflow/workflow_advanced_test.ts new file mode 100644 index 000000000..1ceb1f7ce --- /dev/null +++ b/core/test/workflow/workflow_advanced_test.ts @@ -0,0 +1,240 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {describe, expect, it} from 'vitest'; +import {createEvent, Event} from '../../src/events/event.js'; +import {BaseNode} from '../../src/workflow/base_node.js'; +import {FunctionNode} from '../../src/workflow/nodes/function_node.js'; +import {JoinNode} from '../../src/workflow/nodes/join_node.js'; +import {Workflow} from '../../src/workflow/workflow.js'; +import {createIc, driveNode, FnNode} from './test_helpers.js'; + +describe('workflow — maxConcurrency', () => { + it('bounds the number of concurrently running graph nodes', async () => { + let active = 0; + let peak = 0; + const slow = (name: string): BaseNode => + new FunctionNode(name, async () => { + active++; + peak = Math.max(peak, active); + await new Promise((r) => setTimeout(r, 10)); + active--; + return name; + }); + const join = new JoinNode({name: 'join'}); + const wf = new Workflow({ + name: 'bounded', + maxConcurrency: 2, + edges: [['START', [slow('a'), slow('b'), slow('c'), slow('d')], join]], + }); + + await driveNode(wf, 'x'); + expect(peak).toBeLessThanOrEqual(2); + expect(peak).toBeGreaterThan(0); + }); + + it('rejects maxConcurrency below 1 (0 is not "unlimited")', () => { + const n = new FunctionNode('n', () => 'x'); + expect( + () => + new Workflow({name: 'bad0', edges: [['START', n]], maxConcurrency: 0}), + ).toThrow(/positive integer/); + }); + + it('rejects a non-integer maxConcurrency', () => { + const n = new FunctionNode('n', () => 'x'); + expect( + () => + new Workflow({ + name: 'bad_frac', + edges: [['START', n]], + maxConcurrency: 1.5, + }), + ).toThrow(/positive integer/); + }); +}); + +describe('workflow — sibling cancellation on failure', () => { + it('cancels an in-flight cooperative sibling when another node fails', async () => { + let siblingCancelled = false; + // A cooperative node that waits, but bails out early if the workflow-scoped + // abort signal fires (which happens when its sibling throws). + const patient = new FunctionNode('patient', async (ctx) => { + await new Promise((resolve) => { + const timer = setTimeout(resolve, 2000); + ctx.abortSignal?.addEventListener( + 'abort', + () => { + siblingCancelled = true; + clearTimeout(timer); + resolve(); + }, + {once: true}, + ); + }); + return 'patient-done'; + }); + const boom = new FunctionNode('boom', async () => { + // Let `patient` start and attach its abort listener first, then fail. + await new Promise((r) => setTimeout(r, 10)); + throw new Error('boom'); + }); + const wf = new Workflow({ + name: 'cancel_siblings', + edges: [['START', [patient, boom]]], + }); + + await expect(driveNode(wf, 'x')).rejects.toThrow('boom'); + expect(siblingCancelled).toBe(true); + }); +}); + +describe('workflow — error propagation', () => { + it('fails the workflow when a node throws (no retry)', async () => { + const boom = new FunctionNode('boom', () => { + throw new Error('kaboom'); + }); + const wf = new Workflow({name: 'err', edges: [['START', boom]]}); + await expect(driveNode(wf, 'x')).rejects.toThrow('kaboom'); + }); + + it('does not run downstream nodes after an upstream failure', async () => { + let downstreamRan = false; + const boom = new FunctionNode('boom', () => { + throw new Error('stop'); + }); + const after = new FunctionNode('after', () => { + downstreamRan = true; + return 'after'; + }); + const wf = new Workflow({ + name: 'err_chain', + edges: [['START', boom, after]], + }); + await expect(driveNode(wf, 'x')).rejects.toThrow('stop'); + expect(downstreamRan).toBe(false); + }); +}); + +describe('workflow — join with three predecessors', () => { + it('waits for all predecessors before the join runs', async () => { + const a = new FnNode('a', (_c, i) => `A(${i})`); + const b = new FnNode('b', (_c, i) => `B(${i})`); + const c = new FnNode('c', (_c, i) => `C(${i})`); + const join = new JoinNode({name: 'join'}); + const wf = new Workflow({ + name: 'triple_join', + edges: [['START', [a, b, c], join]], + }); + expect((await driveNode(wf, 'x')).output).toEqual({ + a: 'A(x)', + b: 'B(x)', + c: 'C(x)', + }); + }); +}); + +describe('workflow — retry with exception allow-list', () => { + it('retries only listed error types', async () => { + let attempts = 0; + const node = new FunctionNode( + 'typed', + () => { + attempts++; + if (attempts < 2) { + throw new TypeError('transient'); + } + return 'ok'; + }, + { + retryConfig: { + maxAttempts: 4, + initialDelay: 0.001, + jitter: 0, + exceptions: [TypeError], + }, + }, + ); + const wf = new Workflow({name: 'typed_retry', edges: [['START', node]]}); + expect((await driveNode(wf, 'x')).output).toBe('ok'); + expect(attempts).toBe(2); + }); + + it('does not retry an unlisted error type', async () => { + let attempts = 0; + const node = new FunctionNode( + 'typed2', + () => { + attempts++; + throw new RangeError('nope'); + }, + { + retryConfig: { + maxAttempts: 4, + initialDelay: 0.001, + jitter: 0, + exceptions: [TypeError], + }, + }, + ); + const wf = new Workflow({name: 'typed_retry2', edges: [['START', node]]}); + await expect(driveNode(wf, 'x')).rejects.toThrow('nope'); + expect(attempts).toBe(1); + }); +}); + +describe('workflow — rerunOnResume', () => { + it('re-runs a rerunOnResume node on resume instead of fast-forwarding', async () => { + let runs = 0; + // A node that already "completed" in a prior turn (output event in session) + // but is marked rerunOnResume, so it must run again. + const node = new FnNode( + 'always', + () => { + runs++; + return 'fresh'; + }, + {rerunOnResume: true}, + ); + const wf = new Workflow({name: 'rerun', edges: [['START', node]]}); + + // Seed a session as if `node` completed in a prior turn. + const priorEvent: Event = createEvent({ + author: 'always', + nodeInfo: {path: 'rerun.always'}, + output: 'stale', + }); + const ic = createIc(); + ic.session.events.push(priorEvent); + + const {output} = await driveNode(wf, 'x', ic); + // Because rerunOnResume is true, it re-executed rather than using 'stale'. + expect(runs).toBe(1); + expect(output).toBe('fresh'); + }); + + it('fast-forwards a completed node that is NOT rerunOnResume', async () => { + let runs = 0; + const node = new FnNode('once', () => { + runs++; + return 'fresh'; + }); + const wf = new Workflow({name: 'ff', edges: [['START', node]]}); + + const priorEvent: Event = createEvent({ + author: 'once', + nodeInfo: {path: 'ff.once'}, + output: 'cached', + }); + const ic = createIc(); + ic.session.events.push(priorEvent); + + const {output} = await driveNode(wf, 'x', ic); + // Not rerunOnResume + has cached output -> fast-forwarded, not re-run. + expect(runs).toBe(0); + expect(output).toBe('cached'); + }); +}); diff --git a/core/test/workflow/workflow_agent_test.ts b/core/test/workflow/workflow_agent_test.ts new file mode 100644 index 000000000..571e6a4c7 --- /dev/null +++ b/core/test/workflow/workflow_agent_test.ts @@ -0,0 +1,81 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {describe, expect, it} from 'vitest'; +import {InvocationContext} from '../../src/agents/invocation_context.js'; +import {Event} from '../../src/events/event.js'; +import {PluginManager} from '../../src/plugins/plugin_manager.js'; +import {createSession} from '../../src/sessions/session.js'; +import {NodeContext} from '../../src/workflow/node_context.js'; +import {RequestInput} from '../../src/workflow/request_input.js'; +import {createRequestInputEvent} from '../../src/workflow/utils/hitl_utils.js'; +import {Workflow} from '../../src/workflow/workflow.js'; +import {WorkflowAgent} from '../../src/workflow/workflow_agent.js'; + +/** A session event standing in for a node that raised an unresolved interrupt. */ +function pendingInterruptEvent(id: string): Event { + const event = createRequestInputEvent( + new RequestInput({interruptId: id, message: '?'}), + ); + event.author = id; + return event; +} + +/** + * Runs the agent against a session pre-seeded with `pendingIds` unresolved + * interrupts and a plain-text user reply, returning the `resumeInputs` the + * workflow was driven with (captured via a dynamicEntry). + */ +async function resumeInputsFor( + pendingIds: string[], + replyText: string, +): Promise> { + let captured: Record = {}; + const wf = new Workflow({ + name: 'capture', + dynamicEntry: async (ctx: NodeContext) => { + captured = {...ctx.resumeInputs}; + return 'done'; + }, + }); + const agent = new WorkflowAgent(wf); + + const session = createSession({ + id: 's1', + appName: 'app', + userId: 'u', + lastUpdateTime: Date.now(), + }); + for (const id of pendingIds) { + session.events.push(pendingInterruptEvent(id)); + } + + const ic = new InvocationContext({ + invocationId: 'inv-1', + session, + agent, + userContent: {role: 'user', parts: [{text: replyText}]}, + pluginManager: new PluginManager(), + }); + + for await (const _ of agent.runAsync(ic)) { + void _; + } + return captured; +} + +describe('WorkflowAgent — plain-text resume', () => { + it('feeds a plain-text reply to the single pending interrupt', async () => { + expect(await resumeInputsFor(['only'], 'yes')).toEqual({only: 'yes'}); + }); + + it('ignores a plain-text reply when multiple interrupts are pending', async () => { + // Broadcasting one answer to every pause would resume a node with data the + // user never gave it; the ambiguous case is dropped (structured function + // responses are required to address a specific interrupt). + expect(await resumeInputsFor(['first', 'second'], 'yes')).toEqual({}); + }); +}); diff --git a/core/test/workflow/workflow_test.ts b/core/test/workflow/workflow_test.ts new file mode 100644 index 000000000..6cab53d52 --- /dev/null +++ b/core/test/workflow/workflow_test.ts @@ -0,0 +1,130 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {describe, expect, it} from 'vitest'; +import {createEvent} from '../../src/events/event.js'; +import {BaseNode} from '../../src/workflow/base_node.js'; +import {DEFAULT_ROUTE} from '../../src/workflow/graph.js'; +import {NodeContext} from '../../src/workflow/node_context.js'; +import {Workflow} from '../../src/workflow/workflow.js'; +import {driveWorkflow, FnNode} from './test_helpers.js'; + +// Fan-in barrier: waits for all predecessors, then emits the aggregated inputs. +class JoinNode extends BaseNode { + override get requiresAllPredecessors(): boolean { + return true; + } + protected async *runImpl(_ctx: NodeContext, input: unknown) { + yield input; + } +} + +describe('Phase 2 — Workflow orchestration', () => { + it('runs a linear sequence and threads input downstream (baseline bug fix)', async () => { + const a = new FnNode('step_a', (_c, input) => `${input}->A`); + const b = new FnNode('step_b', (_c, input) => `${input}->B`); + const c = new FnNode('step_c', (_c, input) => `${input}->C`); + const wf = new Workflow({name: 'seq', edges: [['START', a, b, c]]}); + + const {output, events} = await driveWorkflow(wf, 'INIT'); + + // The initial input reaches the first node (previously 'undefined->A'). + expect(output).toBe('INIT->A->B->C'); + const outputs = events + .filter((e) => e.output !== undefined) + .map((e) => e.output); + expect(outputs).toEqual(['INIT->A', 'INIT->A->B', 'INIT->A->B->C']); + }); + + it('routes conditionally via a routing map', async () => { + const router = new FnNode('router', (_c, input) => + createEvent({ + route: (input as string).endsWith('?') ? 'question' : 'statement', + output: input, + }), + ); + const q = new FnNode('answer', (_c, input) => `Q:${input}`); + const s = new FnNode('comment', (_c, input) => `S:${input}`); + const wf = new Workflow({ + name: 'router_wf', + edges: [ + ['START', router], + [router, {question: q, statement: s}], + ], + }); + + expect((await driveWorkflow(wf, 'what?')).output).toBe('Q:what?'); + expect((await driveWorkflow(wf, 'hello')).output).toBe('S:hello'); + }); + + it('falls back to DEFAULT_ROUTE when no specific route matches', async () => { + const check = new FnNode('check', (_c, input) => + createEvent( + input === 'jane' + ? {output: input} // no route -> DEFAULT + : {route: 'retry', output: input}, + ), + ); + const retry = new FnNode('retry_node', (_c, input) => `RETRY:${input}`); + const gen = new FnNode('generate', (_c, input) => `GEN:${input}`); + const wf = new Workflow({ + name: 'default_route_wf', + edges: [ + ['START', check], + [check, {retry, [DEFAULT_ROUTE]: gen}], + ], + }); + + expect((await driveWorkflow(wf, 'john')).output).toBe('RETRY:john'); + expect((await driveWorkflow(wf, 'jane')).output).toBe('GEN:jane'); + }); + + it('fans out to parallel branches and joins them at a barrier', async () => { + const a = new FnNode('A', (_c, input) => `A(${input})`); + const b = new FnNode('B', (_c, input) => `B(${input})`); + const join = new JoinNode({name: 'join'}); + const wf = new Workflow({ + name: 'fan_wf', + edges: [['START', [a, b], join]], + }); + + const {output} = await driveWorkflow(wf, 'x'); + expect(output).toEqual({A: 'A(x)', B: 'B(x)'}); + }); + + it('rejects an unconditional cycle at construction', () => { + const a = new FnNode('cyc_a', (_c, i) => i); + const b = new FnNode('cyc_b', (_c, i) => i); + expect( + () => + new Workflow({ + name: 'cycle_wf', + edges: [ + ['START', a], + [a, b], + [b, a], + ], + }), + ).toThrow(/cycle/i); + }); + + it('rejects an unreachable node', () => { + const a = new FnNode('reach_a', (_c, i) => i); + const orphan = new FnNode('orphan', (_c, i) => i); + const b = new FnNode('reach_b', (_c, i) => i); + expect( + () => + new Workflow({ + name: 'unreachable_wf', + // orphan -> b is never reachable from START. + edges: [ + ['START', a], + [orphan, b], + ], + }), + ).toThrow(/unreachable/i); + }); +});