From 5a80d068fc9e31de46c7ef72b1a905465a725736 Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Wed, 15 Jul 2026 13:55:19 -0700 Subject: [PATCH 01/41] feat(workflow): implement Python ADK workflow parity with strict TypeScript types --- core/src/agents/invocation_context.ts | 16 + core/src/events/event.ts | 17 ++ core/src/index.ts | 1 + core/src/workflow/base_node.ts | 74 +++++ core/src/workflow/dynamic_node_scheduler.ts | 150 ++++++++++ core/src/workflow/index.ts | 47 +++ core/src/workflow/node_runner.ts | 282 ++++++++++++++++++ core/src/workflow/node_state.ts | 92 ++++++ core/src/workflow/nodes/function_node.ts | 134 +++++++++ core/src/workflow/nodes/join_node.ts | 129 ++++++++ core/src/workflow/nodes/llm_agent_wrapper.ts | 90 ++++++ core/src/workflow/nodes/tool_node.ts | 66 ++++ core/src/workflow/parallel_worker.ts | 132 ++++++++ core/src/workflow/retry_config.ts | 67 +++++ core/src/workflow/run_node.ts | 96 ++++++ core/src/workflow/trigger.ts | 101 +++++++ core/src/workflow/utils/graph_parser.ts | 176 +++++++++++ core/src/workflow/utils/graph_validation.ts | 92 ++++++ core/src/workflow/utils/hitl_utils.ts | 89 ++++++ core/src/workflow/utils/rehydration_utils.ts | 114 +++++++ core/src/workflow/utils/replay_manager.ts | 90 ++++++ core/src/workflow/utils/retry_utils.ts | 114 +++++++ core/src/workflow/workflow.ts | 141 +++++++++ core/test/workflow/dynamic_workflow_test.ts | 107 +++++++ core/test/workflow/graph_parser_test.ts | 105 +++++++ .../workflow/hitl_and_rehydration_test.ts | 189 ++++++++++++ .../workflow/join_node_and_parallel_test.ts | 130 ++++++++ core/test/workflow/node_runner_test.ts | 177 +++++++++++ core/test/workflow/workflow_agent_test.ts | 141 +++++++++ 29 files changed, 3159 insertions(+) create mode 100644 core/src/workflow/base_node.ts create mode 100644 core/src/workflow/dynamic_node_scheduler.ts create mode 100644 core/src/workflow/index.ts create mode 100644 core/src/workflow/node_runner.ts create mode 100644 core/src/workflow/node_state.ts create mode 100644 core/src/workflow/nodes/function_node.ts create mode 100644 core/src/workflow/nodes/join_node.ts create mode 100644 core/src/workflow/nodes/llm_agent_wrapper.ts create mode 100644 core/src/workflow/nodes/tool_node.ts create mode 100644 core/src/workflow/parallel_worker.ts create mode 100644 core/src/workflow/retry_config.ts create mode 100644 core/src/workflow/run_node.ts create mode 100644 core/src/workflow/trigger.ts create mode 100644 core/src/workflow/utils/graph_parser.ts create mode 100644 core/src/workflow/utils/graph_validation.ts create mode 100644 core/src/workflow/utils/hitl_utils.ts create mode 100644 core/src/workflow/utils/rehydration_utils.ts create mode 100644 core/src/workflow/utils/replay_manager.ts create mode 100644 core/src/workflow/utils/retry_utils.ts create mode 100644 core/src/workflow/workflow.ts create mode 100644 core/test/workflow/dynamic_workflow_test.ts create mode 100644 core/test/workflow/graph_parser_test.ts create mode 100644 core/test/workflow/hitl_and_rehydration_test.ts create mode 100644 core/test/workflow/join_node_and_parallel_test.ts create mode 100644 core/test/workflow/node_runner_test.ts create mode 100644 core/test/workflow/workflow_agent_test.ts diff --git a/core/src/agents/invocation_context.ts b/core/src/agents/invocation_context.ts index 8a6c4e19d..6b3312481 100644 --- a/core/src/agents/invocation_context.ts +++ b/core/src/agents/invocation_context.ts @@ -38,6 +38,8 @@ export interface InvocationContextParams { activeStreamingTools?: Record; pluginManager: PluginManager; abortSignal?: AbortSignal; + agentStates?: Record; + endOfAgents?: Record; } /** @@ -185,6 +187,17 @@ export class InvocationContext { readonly abortSignal?: AbortSignal; + /** + * Checkpointed states for workflow nodes under this invocation. + */ + agentStates: Record; + + /** + * Tracks whether specific agents or workflows have reached the end of their execution. + */ + + endOfAgents: Record; + /** * @param params The parameters for creating an invocation context. */ @@ -203,7 +216,10 @@ export class InvocationContext { this.activeStreamingTools = params.activeStreamingTools; this.pluginManager = params.pluginManager; this.abortSignal = params.abortSignal; + this.agentStates = params.agentStates ?? {}; + this.endOfAgents = params.endOfAgents ?? {}; // Inherit the parent invocation's cost manager when one is available. + // Child contexts created for sub-agents, agent transfers and loop // iterations (via createInvocationContext / createBranchCtxForSubAgent) // carry the parent context's fields over, so reusing its cost manager diff --git a/core/src/events/event.ts b/core/src/events/event.ts index 4b039732d..f5fcae31c 100644 --- a/core/src/events/event.ts +++ b/core/src/events/event.ts @@ -208,6 +208,23 @@ export function pruneThoughts(event: Event): Event { const ASCII_LETTERS_AND_NUMBERS = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789'; +/** + * Type guard to check if an object is an instance of Event. + * + * @param obj The object to check. + * @returns True if the object matches the Event structure. + */ +export function isEvent(obj: unknown): obj is Event { + return ( + typeof obj === 'object' && + obj !== null && + 'invocationId' in obj && + typeof (obj as Event).invocationId === 'string' && + 'actions' in obj && + typeof (obj as Event).actions === 'object' + ); +} + /** * Generates a new unique ID for the event. */ diff --git a/core/src/index.ts b/core/src/index.ts index b2a05a2c9..23829507a 100644 --- a/core/src/index.ts +++ b/core/src/index.ts @@ -55,3 +55,4 @@ export * from './telemetry/setup.js'; export * from './tools/mcp/mcp_session_manager.js'; export * from './tools/mcp/mcp_tool.js'; export * from './tools/mcp/mcp_toolset.js'; +export * from './workflow/index.js'; diff --git a/core/src/workflow/base_node.ts b/core/src/workflow/base_node.ts new file mode 100644 index 000000000..e762eb9f6 --- /dev/null +++ b/core/src/workflow/base_node.ts @@ -0,0 +1,74 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {InvocationContext} from '../agents/invocation_context.js'; +import {Event} from '../events/event.js'; +import {RetryConfig, normalizeRetryConfig} from './retry_config.js'; + +/** + * Options for configuring a BaseNode. + */ +export interface BaseNodeOptions { + /** + * If true, the node will re-execute when a workflow is resumed even if + * historical completed state exists in `InvocationContext.agentStates`. + * Default is false. + */ + rerunOnResume?: boolean; + + /** + * Optional retry configuration for handling transient errors during execution. + */ + retryConfig?: RetryConfig; +} + +/** + * Abstract base class for all nodes in an ADK Workflow. + * A node represents a discrete unit of execution within a static graph or dynamic chain. + */ +export abstract class BaseNode { + /** + * The canonical name of the node. Must be unique within a workflow graph. + */ + readonly name: string; + + /** + * Whether this node should re-execute when resuming a paused or rehydrated workflow. + */ + readonly rerunOnResume: boolean; + + /** + * The normalized retry configuration for this node, if any. + */ + readonly retryConfig?: Required; + + /** + * Optional cached output payload stored on the instance during generator execution. + */ + lastOutputPayload?: unknown; + + constructor(name: string, options?: BaseNodeOptions) { + if (!name || typeof name !== 'string' || name.trim().length === 0) { + throw new Error('Node name must be a non-empty string.'); + } + this.name = name.trim(); + this.rerunOnResume = options?.rerunOnResume ?? false; + this.retryConfig = normalizeRetryConfig(options?.retryConfig); + } + + /** + * Core execution contract for a node. + * + * @param ctx The invocation context of the current workflow execution. + * @param input Optional input payload passed from upstream nodes or dynamic scheduler. + * @yields Events generated during node execution (including partial output or route events). + * @returns The final output payload of this node. + */ + abstract run( + ctx: InvocationContext, + input?: TInput, + ): AsyncGenerator; +} diff --git a/core/src/workflow/dynamic_node_scheduler.ts b/core/src/workflow/dynamic_node_scheduler.ts new file mode 100644 index 000000000..277430539 --- /dev/null +++ b/core/src/workflow/dynamic_node_scheduler.ts @@ -0,0 +1,150 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {InvocationContext} from '../agents/invocation_context.js'; +import {Event, isEvent} from '../events/event.js'; +import {BaseNode} from './base_node.js'; +import { + consumeGenerator, + generateExecutionId, + getOrInitAgentStates, +} from './node_runner.js'; +import {NodeState, NodeStatus, isNodeState} from './node_state.js'; +import {FunctionNode, FunctionNodeHandler} from './nodes/function_node.js'; + +/** + * Type for the dynamic workflow entry point. + */ +export type DynamicEntry = + | BaseNode + | FunctionNodeHandler; + +/** + * Options for the DynamicNodeScheduler. + */ +export interface DynamicNodeSchedulerOptions { + /** + * Key inside `InvocationContext.agentStates` where the final output of the + * dynamic entry point should be saved upon completion. + */ + outputKey?: string; +} + +/** + * Coordinates and executes a dynamic workflow where control flow (`async/await`, loops, conditionals) + * is driven programmatically by Python/TS code calling `ctx.runNode(...)`. + * Manages deterministic ID counters (`exec_node__`) and checkpoint skip-on-resume. + */ +export class DynamicNodeScheduler { + readonly entryNode: BaseNode; + readonly options: DynamicNodeSchedulerOptions; + + /** + * @param entry A BaseNode instance or a function handler to serve as the root of the dynamic workflow. + * @param options Optional configuration (outputKey). + */ + constructor(entry: DynamicEntry, options?: DynamicNodeSchedulerOptions) { + if (typeof entry === 'function') { + this.entryNode = new FunctionNode('dynamic_entry_node', entry); + } else if (isBaseNode(entry)) { + this.entryNode = entry; + } else { + throw new Error( + 'DynamicNodeScheduler requires a valid BaseNode instance or function handler.', + ); + } + this.options = options || {}; + } + + /** + * Runs the dynamic workflow entry node, intercepting events and handling checkpointing. + */ + async *runAsync( + ctx: InvocationContext, + initialInput?: unknown, + ): AsyncGenerator { + const agentStates = getOrInitAgentStates(ctx); + const execId = generateExecutionId(ctx, this.entryNode.name); + + const existingState = agentStates[execId] as NodeState | undefined; + if ( + existingState && + isNodeState(existingState) && + existingState.status === NodeStatus.COMPLETED && + !this.entryNode.rerunOnResume + ) { + if (this.options.outputKey) { + agentStates[this.options.outputKey] = existingState.outputPayload; + } + return; + } + + const stateRecord: NodeState = { + executionId: execId, + nodeName: this.entryNode.name, + status: NodeStatus.RUNNING, + inputPayload: initialInput, + timestamp: Date.now(), + }; + agentStates[execId] = stateRecord; + + try { + const generator = this.entryNode.run(ctx, initialInput); + const yieldedEvents: Event[] = []; + + const {output, isPausedHitl} = await consumeGenerator( + generator, + async (ev) => { + if (isEvent(ev)) { + yieldedEvents.push(ev); + } + }, + ); + + for (const ev of yieldedEvents) { + yield ev; + } + + if (isPausedHitl || ctx.endInvocation || ctx.abortSignal?.aborted) { + stateRecord.status = NodeStatus.PAUSED_HITL; + stateRecord.timestamp = Date.now(); + ctx.endInvocation = true; + return; + } + + const finalResult = + output !== undefined + ? output + : (this.entryNode.lastOutputPayload ?? + stateRecord.lastOutputPayload ?? + initialInput); + stateRecord.status = NodeStatus.COMPLETED; + stateRecord.outputPayload = finalResult; + stateRecord.timestamp = Date.now(); + + if (this.options.outputKey) { + agentStates[this.options.outputKey] = finalResult; + } + } catch (error: unknown) { + stateRecord.status = NodeStatus.FAILED; + stateRecord.errorMessage = + error instanceof Error ? error.message : String(error); + stateRecord.timestamp = Date.now(); + throw error; + } + } +} + +function isBaseNode(obj: unknown): obj is BaseNode { + return ( + typeof obj === 'object' && + obj !== null && + 'name' in obj && + typeof (obj as BaseNode).name === 'string' && + 'run' in obj && + typeof (obj as BaseNode).run === 'function' + ); +} diff --git a/core/src/workflow/index.ts b/core/src/workflow/index.ts new file mode 100644 index 000000000..6668ad581 --- /dev/null +++ b/core/src/workflow/index.ts @@ -0,0 +1,47 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +export {BaseNode, type BaseNodeOptions} from './base_node.js'; +export { + DynamicNodeScheduler, + type DynamicEntry, + type DynamicNodeSchedulerOptions, +} from './dynamic_node_scheduler.js'; +export { + NodeRunner, + generateExecutionId, + getOrInitAgentStates, + type NodeRunnerOptions, +} from './node_runner.js'; +export {NodeStatus, isNodeState, type NodeState} from './node_state.js'; +export {FunctionNode, type FunctionNodeHandler} from './nodes/function_node.js'; +export {JoinNode, type JoinNodeOptions} from './nodes/join_node.js'; +export {LLMAgentWrapper} from './nodes/llm_agent_wrapper.js'; +export {ToolNode} from './nodes/tool_node.js'; +export {runInParallel, type ParallelRunOptions} from './parallel_worker.js'; +export {normalizeRetryConfig, type RetryConfig} from './retry_config.js'; +export {runNode, type RunNodeOptions} from './run_node.js'; +export {Trigger, type TriggerPredicate} from './trigger.js'; +export { + ParsedGraph, + parseGraph, + type AdjacencyEdge, + type EdgeElement, + type GraphEdge, +} from './utils/graph_parser.js'; +export {validateGraph} from './utils/graph_validation.js'; +export { + createRequestInputEvent, + injectHitlResumptionInput, + type RequestInputOptions, +} from './utils/hitl_utils.js'; +export { + persistAgentStatesToSession, + rehydrateAgentStates, +} from './utils/rehydration_utils.js'; +export {ReplayManager} from './utils/replay_manager.js'; +export {runWithRetry} from './utils/retry_utils.js'; +export {Workflow, isWorkflow, type WorkflowConfig} from './workflow.js'; diff --git a/core/src/workflow/node_runner.ts b/core/src/workflow/node_runner.ts new file mode 100644 index 000000000..ae3e02a6a --- /dev/null +++ b/core/src/workflow/node_runner.ts @@ -0,0 +1,282 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {InvocationContext} from '../agents/invocation_context.js'; +import {Event} from '../events/event.js'; +import {BaseNode} from './base_node.js'; +import {NodeState, NodeStatus, isNodeState} from './node_state.js'; +import {GraphEdge, ParsedGraph, parseGraph} from './utils/graph_parser.js'; +import {validateGraph} from './utils/graph_validation.js'; +import {runWithRetry} from './utils/retry_utils.js'; + +/** + * Options for configuring the NodeRunner. + */ +export interface NodeRunnerOptions { + /** + * Whether to allow cycles in the graph during validation. + * Default is false. + */ + allowCycles?: boolean; + + /** + * Key inside `InvocationContext.agentStates` where the final leaf node outputs + * should also be written or aggregated, if requested by the workflow. + */ + outputKey?: string; +} + +interface QueueItem { + readonly node: BaseNode; + readonly inputPayload?: unknown; + readonly sourceNodeName: string; +} + +/** + * Consumes an AsyncGenerator to completion, capturing all yielded Events via onEvent + * and extracting the final return value when `done: true`. + * Also detects if any yielded Event signals a Human-in-the-Loop (`RequestInput`) pause condition. + */ +export async function consumeGenerator( + generator: AsyncGenerator, + onEvent?: (event: Event) => void | Promise, +): Promise<{ + output: TOutput | undefined; + isPausedHitl: boolean; + lastEvent?: Event; +}> { + let isPausedHitl = false; + let lastEvent: Event | undefined; + + while (true) { + const {value, done} = await generator.next(); + if (done) { + let output = value as TOutput | undefined; + if ( + output === undefined && + lastEvent?.actions && + typeof lastEvent.actions === 'object' && + 'output' in (lastEvent.actions as Record) + ) { + output = (lastEvent.actions as Record) + .output as TOutput; + } + return {output, isPausedHitl, lastEvent}; + } + + const event = value as Event; + lastEvent = event; + if (onEvent) { + await onEvent(event); + } + if (isHitlPauseEvent(event)) { + isPausedHitl = true; + return {output: undefined, isPausedHitl, lastEvent: event}; + } + } +} + +/** + * Executes a static graph workflow (`edges`) using topological queue-based scheduling, + * evaluating edge triggers upon node completion, checkpointing state in `InvocationContext.agentStates`, + * and handling Human-in-the-Loop (`RequestInput` / `PAUSED_HITL`) interruptions cleanly. + */ +export class NodeRunner { + readonly graph: ParsedGraph; + readonly options: NodeRunnerOptions; + + /** + * @param edgesOrGraph Array of GraphEdge sequences or a pre-parsed ParsedGraph. + * @param options Optional configuration for the runner. + */ + constructor( + edgesOrGraph: GraphEdge[] | ParsedGraph, + options?: NodeRunnerOptions, + ) { + if (edgesOrGraph instanceof ParsedGraph) { + this.graph = edgesOrGraph; + } else { + this.graph = parseGraph(edgesOrGraph); + } + this.options = options || {}; + validateGraph(this.graph, {allowCycles: this.options.allowCycles}); + } + + /** + * Executes the workflow graph from "START" (or from paused/rehydrated checkpoints). + * @param ctx The invocation context for the workflow run. + * @param initialInput Optional initial input payload passed to START nodes. + * @yields All events generated during node execution. + */ + async *runAsync( + ctx: InvocationContext, + initialInput?: unknown, + ): AsyncGenerator { + const agentStates = getOrInitAgentStates(ctx); + const queue: QueueItem[] = []; + + // 1. Initialize queue with edges originating from "START" + const startEdges = this.graph.adjacencyList.get('START') || []; + for (const edge of startEdges) { + queue.push({ + node: edge.target, + inputPayload: initialInput, + sourceNodeName: 'START', + }); + } + + // 2. Queue processing loop + while (queue.length > 0) { + if (ctx.endInvocation || ctx.abortSignal?.aborted) { + break; + } + + const item = queue.shift()!; + const execId = generateExecutionId(ctx, item.node.name); + + const existingState = agentStates[execId] as NodeState | undefined; + let nodeOutput: unknown = undefined; + + if ( + existingState && + isNodeState(existingState) && + existingState.status === NodeStatus.COMPLETED && + !item.node.rerunOnResume + ) { + nodeOutput = existingState.outputPayload; + } else { + const effectiveInput = + existingState?.inputPayload !== undefined + ? existingState.inputPayload + : item.inputPayload; + const stateRecord: NodeState = { + executionId: execId, + nodeName: item.node.name, + status: NodeStatus.RUNNING, + inputPayload: effectiveInput, + timestamp: Date.now(), + }; + agentStates[execId] = stateRecord; + + try { + const generator = runWithRetry( + () => item.node.run(ctx, effectiveInput), + item.node.retryConfig, + ctx.abortSignal, + ); + + const yieldedEvents: Event[] = []; + const {output, isPausedHitl} = await consumeGenerator( + generator, + async (event) => { + yieldedEvents.push(event); + }, + ); + + for (const ev of yieldedEvents) { + yield ev; + } + + if (isPausedHitl || ctx.endInvocation || ctx.abortSignal?.aborted) { + stateRecord.status = NodeStatus.PAUSED_HITL; + stateRecord.timestamp = Date.now(); + ctx.endInvocation = true; + break; + } + + nodeOutput = + output !== undefined + ? output + : (item.node.lastOutputPayload ?? + stateRecord.lastOutputPayload ?? + item.inputPayload); + stateRecord.status = NodeStatus.COMPLETED; + stateRecord.outputPayload = nodeOutput; + stateRecord.timestamp = Date.now(); + } catch (error: unknown) { + stateRecord.status = NodeStatus.FAILED; + stateRecord.errorMessage = + error instanceof Error ? error.message : String(error); + stateRecord.timestamp = Date.now(); + throw error; + } + } + + // 3. Evaluate outgoing edges and enqueue successors whose triggers are satisfied + const outgoingEdges = this.graph.adjacencyList.get(item.node.name) || []; + for (const edge of outgoingEdges) { + let triggerSatisfied = true; + if (edge.trigger) { + triggerSatisfied = await edge.trigger.evaluate(ctx, nodeOutput); + } + + if (triggerSatisfied) { + queue.push({ + node: edge.target, + inputPayload: nodeOutput, + sourceNodeName: item.node.name, + }); + } + } + } + + if (this.options.outputKey) { + const finalStates: Record = {}; + for (const state of Object.values(agentStates)) { + if (isNodeState(state) && state.status === NodeStatus.COMPLETED) { + finalStates[state.nodeName] = state.outputPayload; + } + } + agentStates[this.options.outputKey] = finalStates; + } + } +} + +/** + * Gets or initializes the `agentStates` record on the invocation context. + */ +export function getOrInitAgentStates( + ctx: InvocationContext, +): Record { + const unknownCtx = ctx as unknown as Record; + if (!unknownCtx.agentStates || typeof unknownCtx.agentStates !== 'object') { + unknownCtx.agentStates = {}; + } + return unknownCtx.agentStates as Record; +} + +/** + * Generates a deterministic execution ID for a node based on the context branch and node name. + */ +export function generateExecutionId( + ctx: InvocationContext, + nodeName: string, +): string { + const branchPrefix = ctx.branch ? `${ctx.branch}.` : ''; + return `exec_node_${branchPrefix}${nodeName}`; +} + +/** + * Checks whether an event signals a Human-in-the-Loop pause (`RequestInput`). + */ +export function isHitlPauseEvent(event: Event): boolean { + if (!event) return false; + if (event.actions && typeof event.actions === 'object') { + if ( + 'requestInput' in event.actions && + Boolean((event.actions as Record).requestInput) + ) { + return true; + } + } + if ( + 'requestInput' in event && + Boolean((event as Record).requestInput) + ) { + return true; + } + return false; +} diff --git a/core/src/workflow/node_state.ts b/core/src/workflow/node_state.ts new file mode 100644 index 000000000..4edcd4dfb --- /dev/null +++ b/core/src/workflow/node_state.ts @@ -0,0 +1,92 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {Event} from '../events/event.js'; + +/** + * Represents the execution status of a node in a workflow graph or dynamic chain. + */ +export enum NodeStatus { + PENDING = 'PENDING', + RUNNING = 'RUNNING', + COMPLETED = 'COMPLETED', + PAUSED_HITL = 'PAUSED_HITL', + FAILED = 'FAILED', +} + +/** + * Checkpointed state for a specific node execution. + * Stored inside `InvocationContext.agentStates[executionId]`. + */ +export interface NodeState { + /** + * The deterministic execution ID assigned to this node execution. + */ + executionId: string; + + /** + * The canonical name of the node. + */ + nodeName: string; + + /** + * The current status of the node execution. + */ + status: NodeStatus; + + /** + * The input payload passed into the node during execution. + */ + inputPayload?: TInput; + + /** + * The final output payload yielded or returned by the node upon completion. + */ + outputPayload?: TOutput; + + /** + * Error message if the node execution failed (`status === FAILED`). + */ + errorMessage?: string; + + /** + * Timestamp in milliseconds when this state record was last updated. + */ + timestamp: number; + + /** + * Events emitted by the node during live execution, cached for replaying on resumption. + */ + cachedEvents?: Event[]; + + /** + * Indicates if this node previously paused for Human-in-the-Loop (`RequestInput`). + */ + wasPausedHitl?: boolean; + + /** + * Stores intermediate or final payload before completion status transition. + */ + lastOutputPayload?: unknown; +} + +/** + * Type guard to check if an object is a valid NodeState instance. + * @param obj The object to check. + * @returns True if the object matches the NodeState structure. + */ +export function isNodeState(obj: unknown): obj is NodeState { + return ( + typeof obj === 'object' && + obj !== null && + 'executionId' in obj && + typeof (obj as NodeState).executionId === 'string' && + 'nodeName' in obj && + typeof (obj as NodeState).nodeName === 'string' && + 'status' in obj && + Object.values(NodeStatus).includes((obj as NodeState).status) + ); +} diff --git a/core/src/workflow/nodes/function_node.ts b/core/src/workflow/nodes/function_node.ts new file mode 100644 index 000000000..87d4b6a18 --- /dev/null +++ b/core/src/workflow/nodes/function_node.ts @@ -0,0 +1,134 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {Content} from '@google/genai'; +import {InvocationContext} from '../../agents/invocation_context.js'; +import {createEvent, Event, isEvent} from '../../events/event.js'; +import {BaseNode, BaseNodeOptions} from '../base_node.js'; + +/** + * Type for the function wrapped by a FunctionNode. + * Can return a direct value, a Promise of a value, an Event, or an AsyncGenerator of Events. + */ +export type FunctionNodeHandler = ( + ctx: InvocationContext, + input?: TInput, +) => + | AsyncGenerator + | Promise + | TOutput + | Event; + +/** + * A concrete node that wraps a deterministic JavaScript/TypeScript function. + * Automatically handles generator streams, Event returns, or boxes plain return values into Event outputs. + */ +export class FunctionNode extends BaseNode< + TInput, + TOutput +> { + private readonly handler: FunctionNodeHandler; + + /** + * @param name Unique name for this function node. + * @param handler The execution logic function. + * @param options Optional BaseNode configuration (rerunOnResume, retryConfig). + */ + constructor( + name: string, + handler: FunctionNodeHandler, + options?: BaseNodeOptions, + ) { + super(name, options); + if (typeof handler !== 'function') { + throw new Error( + `FunctionNode "${name}" requires a valid function handler.`, + ); + } + this.handler = handler; + } + + /** + * Executes the wrapped handler function inside the workflow context. + */ + async *run( + ctx: InvocationContext, + input?: TInput, + ): AsyncGenerator { + const resultOrGen = this.handler(ctx, input); + + if (isAsyncGenerator(resultOrGen)) { + const finalVal = yield* resultOrGen; + this.lastOutputPayload = finalVal; + return finalVal; + } + + const res = await Promise.resolve(resultOrGen); + + if (isEvent(res)) { + yield res; + const extracted = + res.content ?? + (typeof res.actions === 'object' && + res.actions !== null && + 'output' in (res.actions as Record) + ? (res.actions as Record).output + : res); + this.lastOutputPayload = extracted; + return extracted as TOutput; + } + + if (res !== undefined && res !== null) { + const boxedEvent = createEvent({ + invocationId: ctx.invocationId, + author: this.name, + branch: ctx.branch, + content: toContent(res), + actions: {output: res}, + }); + yield boxedEvent; + } + + this.lastOutputPayload = res; + return res as TOutput; + } +} + +function isAsyncGenerator( + obj: unknown, +): obj is AsyncGenerator { + return ( + typeof obj === 'object' && + obj !== null && + Symbol.asyncIterator in obj && + 'next' in obj && + typeof (obj as Record).next === 'function' + ); +} + +function toContent(val: unknown): Content | undefined { + if (!val) return undefined; + if (typeof val === 'object' && 'role' in val && 'parts' in val) { + return val as Content; + } + if (typeof val === 'string') { + return { + role: 'model', + parts: [{text: val}], + }; + } + try { + return { + role: 'model', + parts: [{text: JSON.stringify(val)}], + }; + } catch { + return { + role: 'model', + parts: [{text: String(val)}], + }; + } +} diff --git a/core/src/workflow/nodes/join_node.ts b/core/src/workflow/nodes/join_node.ts new file mode 100644 index 000000000..595572aad --- /dev/null +++ b/core/src/workflow/nodes/join_node.ts @@ -0,0 +1,129 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {InvocationContext} from '../../agents/invocation_context.js'; +import {createEvent, Event} from '../../events/event.js'; +import {BaseNode, BaseNodeOptions} from '../base_node.js'; +import {getOrInitAgentStates} from '../node_runner.js'; +import {isNodeState, NodeStatus} from '../node_state.js'; + +/** + * Options for configuring a JoinNode. + */ +export interface JoinNodeOptions extends BaseNodeOptions { + /** + * The number of distinct upstream predecessor nodes that must complete + * before this join node unblocks and emits a combined output. + * Must be >= 1. + */ + upstreamCount: number; + + /** + * Optional array of explicit predecessor node names to wait on. + * If provided, `upstreamCount` must match `predecessors.length`. + */ + predecessors?: string[]; +} + +/** + * A synchronization barrier node used in fan-out/fan-in parallel workflows. + * Waits until `upstreamCount` predecessor branches have reached completion (`COMPLETED`), + * then aggregates their outputs into a dictionary and yields a single combined event. + */ +export class JoinNode< + TInput = unknown, + TOutput = Record, +> extends BaseNode { + readonly upstreamCount: number; + readonly predecessors?: string[]; + + constructor(name: string, options: JoinNodeOptions) { + if ( + !options || + typeof options.upstreamCount !== 'number' || + options.upstreamCount < 1 + ) { + throw new Error( + `JoinNode "${name}" requires a valid upstreamCount >= 1.`, + ); + } + if ( + options.predecessors && + options.predecessors.length !== options.upstreamCount + ) { + throw new Error( + `JoinNode "${name}" upstreamCount (${options.upstreamCount}) does not match predecessors.length (${options.predecessors.length}).`, + ); + } + super(name, options); + this.upstreamCount = options.upstreamCount; + this.predecessors = options.predecessors; + } + + /** + * Evaluates the completion status of upstream predecessor nodes in `InvocationContext.agentStates`. + * Only unblocks and yields combined output when all required predecessors have completed. + */ + async *run( + ctx: InvocationContext, + _input?: TInput, + ): AsyncGenerator { + const agentStates = getOrInitAgentStates(ctx); + + const completedPredecessors: Record = {}; + let count = 0; + + if (this.predecessors && this.predecessors.length > 0) { + for (const predName of this.predecessors) { + for (const state of Object.values(agentStates)) { + if ( + isNodeState(state) && + state.nodeName === predName && + state.status === NodeStatus.COMPLETED + ) { + completedPredecessors[predName] = state.outputPayload; + count++; + break; + } + } + } + } else { + for (const state of Object.values(agentStates)) { + if ( + isNodeState(state) && + state.nodeName !== this.name && + state.status === NodeStatus.COMPLETED && + !(state.nodeName in completedPredecessors) + ) { + completedPredecessors[state.nodeName] = state.outputPayload; + count++; + } + } + } + + if (count < this.upstreamCount) { + return completedPredecessors as unknown as TOutput; + } + + const joinEvent = createEvent({ + invocationId: ctx.invocationId, + author: this.name, + branch: ctx.branch, + actions: { + joinCompleted: { + node: this.name, + upstreamCount: this.upstreamCount, + predecessors: Object.keys(completedPredecessors), + outputs: completedPredecessors, + }, + }, + }); + + yield joinEvent; + this.lastOutputPayload = completedPredecessors; + return completedPredecessors as unknown as TOutput; + } +} diff --git a/core/src/workflow/nodes/llm_agent_wrapper.ts b/core/src/workflow/nodes/llm_agent_wrapper.ts new file mode 100644 index 000000000..919d2968c --- /dev/null +++ b/core/src/workflow/nodes/llm_agent_wrapper.ts @@ -0,0 +1,90 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {BaseAgent} from '../../agents/base_agent.js'; +import { + InvocationContext, + InvocationContextParams, +} from '../../agents/invocation_context.js'; +import {Event} from '../../events/event.js'; +import {BaseNode, BaseNodeOptions} from '../base_node.js'; + +/** + * A concrete node that wraps any ADK BaseAgent (e.g., LlmAgent, SequentialAgent) + * so it can participate as a node inside a workflow graph. + * Enforces single-turn task execution mode and relays generated events. + */ +export class LLMAgentWrapper< + TInput = unknown, + TOutput = unknown, +> extends BaseNode { + readonly agent: BaseAgent; + + /** + * @param agent The BaseAgent instance to wrap. + * @param options Optional BaseNode configuration (name override, rerunOnResume, retryConfig). + */ + constructor(agent: BaseAgent, options?: BaseNodeOptions & {name?: string}) { + if (!agent || typeof agent.runAsync !== 'function') { + throw new Error('LLMAgentWrapper requires a valid BaseAgent instance.'); + } + super(options?.name || agent.name || 'llm_agent_wrapper', options); + this.agent = agent; + } + + /** + * Invokes the wrapped agent via runAsync and relays all produced events. + */ + async *run( + ctx: InvocationContext, + input?: TInput, + ): AsyncGenerator { + let lastOutput: unknown = undefined; + + const childCtxParams: InvocationContextParams & Record = { + ...ctx, + agent: this.agent, + }; + + if (input !== undefined && input !== null) { + if (typeof input === 'string') { + childCtxParams.userContent = { + role: 'user', + parts: [{text: input}], + }; + } else if ( + typeof input === 'object' && + 'role' in input && + 'parts' in input + ) { + childCtxParams.userContent = input as {role: string; parts: unknown[]}; + } + } + + const childCtx = new InvocationContext(childCtxParams); + + for await (const event of this.agent.runAsync(childCtx)) { + yield event; + if (event.content?.parts?.length) { + const texts = event.content.parts.map((p) => p.text).filter(Boolean); + if (texts.length > 0) { + lastOutput = texts.join('\n'); + } + } + if ( + event.actions && + typeof event.actions === 'object' && + 'output' in (event.actions as Record) + ) { + lastOutput = (event.actions as Record).output; + } + } + + const finalVal = lastOutput ?? input; + this.lastOutputPayload = finalVal; + return finalVal as TOutput; + } +} diff --git a/core/src/workflow/nodes/tool_node.ts b/core/src/workflow/nodes/tool_node.ts new file mode 100644 index 000000000..8169eef80 --- /dev/null +++ b/core/src/workflow/nodes/tool_node.ts @@ -0,0 +1,66 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {InvocationContext} from '../../agents/invocation_context.js'; +import {createEvent, Event} from '../../events/event.js'; +import {BaseTool} from '../../tools/base_tool.js'; +import {BaseNode, BaseNodeOptions} from '../base_node.js'; + +/** + * A concrete node that wraps an ADK BaseTool (or FunctionTool) so it can be executed + * directly as a step in a workflow graph without requiring an LLM wrapper. + */ +export class ToolNode< + TInput = Record, + TOutput = unknown, +> extends BaseNode { + readonly tool: BaseTool; + + /** + * @param tool The BaseTool instance to execute when this node runs. + * @param options Optional BaseNode configuration (name override, rerunOnResume, retryConfig). + */ + constructor(tool: BaseTool, options?: BaseNodeOptions & {name?: string}) { + if (!tool || typeof tool.runAsync !== 'function') { + throw new Error( + 'ToolNode requires a valid BaseTool instance with runAsync().', + ); + } + super(options?.name || tool.name || 'tool_node', options); + this.tool = tool; + } + + /** + * Executes the wrapped tool using parameters from the input payload. + */ + async *run( + ctx: InvocationContext, + input?: TInput, + ): AsyncGenerator { + const params = typeof input === 'object' && input !== null ? input : {}; + const result = await this.tool.runAsync( + {invocationContext: ctx}, + params as Record, + ); + + const event = createEvent({ + invocationId: ctx.invocationId, + author: this.name, + branch: ctx.branch, + actions: { + toolExecution: { + name: this.tool.name, + input: params, + output: result, + }, + }, + }); + + yield event; + this.lastOutputPayload = result; + return result as TOutput; + } +} diff --git a/core/src/workflow/parallel_worker.ts b/core/src/workflow/parallel_worker.ts new file mode 100644 index 000000000..3d3a83ead --- /dev/null +++ b/core/src/workflow/parallel_worker.ts @@ -0,0 +1,132 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {InvocationContext} from '../agents/invocation_context.js'; +import {BaseNode} from './base_node.js'; +import {consumeGenerator, getOrInitAgentStates} from './node_runner.js'; +import {NodeState, NodeStatus, isNodeState} from './node_state.js'; +import {FunctionNode, FunctionNodeHandler} from './nodes/function_node.js'; + +/** + * Options for configuring parallel branch execution (`runInParallel`). + */ +export interface ParallelRunOptions { + /** + * Optional prefix for naming the child branches in `InvocationContext.branch`. + * Defaults to the target node name. + */ + branchPrefix?: string; + + /** + * If true, throws an error immediately if any parallel worker branch fails. + * If false, catches branch errors and returns undefined/error markers for failed items. + * Default is true. + */ + stopOnError?: boolean; +} + +/** + * Executes a target node (or function) concurrently across an array of input items using isolated + * child `InvocationContext` branches. Prevents concurrent async tasks from corrupting shared + * session event histories or racing on shared node state checkpoints. + * + * @param ctx The parent invocation context. + * @param nodeOrFunc The BaseNode or function handler to execute for each item. + * @param items Array of input items to process in parallel. + * @param options Optional settings (branchPrefix, stopOnError). + * @returns Promise resolving to an array of output payloads corresponding 1-to-1 with the items array. + */ +export async function runInParallel( + ctx: InvocationContext, + nodeOrFunc: BaseNode | FunctionNodeHandler, + items: TInput[], + options?: ParallelRunOptions, +): Promise { + if (!Array.isArray(items) || items.length === 0) { + return []; + } + + const node = + typeof nodeOrFunc === 'function' + ? new FunctionNode('parallel_worker_node', nodeOrFunc) + : nodeOrFunc; + + const prefix = options?.branchPrefix || node.name; + const stopOnError = options?.stopOnError ?? true; + const parentStates = getOrInitAgentStates(ctx); + + const tasks = items.map(async (item, index) => { + const branchName = ctx.branch + ? `${ctx.branch}.${prefix}_${index}` + : `${prefix}_${index}`; + + const childCtx = new InvocationContext({ + ...ctx, + branch: branchName, + }); + + const childStates: Record = {...parentStates}; + (childCtx as unknown as Record).agentStates = childStates; + + const execId = `exec_node_${branchName}.${node.name}`; + + const existingState = childStates[execId] as NodeState | undefined; + if ( + existingState && + isNodeState(existingState) && + existingState.status === NodeStatus.COMPLETED && + !node.rerunOnResume + ) { + return existingState.outputPayload as TOutput; + } + + const stateRecord: NodeState = { + executionId: execId, + nodeName: `${node.name}_${index}`, + status: NodeStatus.RUNNING, + inputPayload: item, + timestamp: Date.now(), + }; + childStates[execId] = stateRecord; + + try { + const generator = node.run(childCtx, item); + const {output, isPausedHitl} = await consumeGenerator(generator); + + if (isPausedHitl || ctx.endInvocation || ctx.abortSignal?.aborted) { + stateRecord.status = NodeStatus.PAUSED_HITL; + stateRecord.timestamp = Date.now(); + ctx.endInvocation = true; + throw new Error( + `Parallel worker branch "${branchName}" requested HITL pause.`, + ); + } + + const finalVal = + output !== undefined + ? output + : (node.lastOutputPayload ?? stateRecord.lastOutputPayload ?? item); + stateRecord.status = NodeStatus.COMPLETED; + stateRecord.outputPayload = finalVal; + stateRecord.timestamp = Date.now(); + + parentStates[execId] = stateRecord; + return finalVal as TOutput; + } catch (err: unknown) { + stateRecord.status = NodeStatus.FAILED; + stateRecord.errorMessage = + err instanceof Error ? err.message : String(err); + stateRecord.timestamp = Date.now(); + parentStates[execId] = stateRecord; + if (stopOnError) { + throw err; + } + return undefined as unknown as TOutput; + } + }); + + return await Promise.all(tasks); +} diff --git a/core/src/workflow/retry_config.ts b/core/src/workflow/retry_config.ts new file mode 100644 index 000000000..0b6f678e0 --- /dev/null +++ b/core/src/workflow/retry_config.ts @@ -0,0 +1,67 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Configuration options for node execution retries upon transient failures. + */ +export interface RetryConfig { + /** + * Maximum number of execution attempts (including the initial attempt). + * Must be >= 1. + */ + maxAttempts: number; + + /** + * Initial delay in milliseconds before the first retry. + * Default is 1000ms (1 second). + */ + initialDelayMs?: number; + + /** + * Maximum delay in milliseconds between retries. + * Default is 30000ms (30 seconds). + */ + maxDelayMs?: number; + + /** + * Multiplier applied to the delay after each retry attempt (exponential backoff). + * Default is 2.0. + */ + backoffFactor?: number; + + /** + * Optional array of Error constructors or error message patterns that should trigger a retry. + * If not specified, all errors are considered retryable up to `maxAttempts`. + */ + retryableErrors?: Array Error | string | RegExp>; +} + +/** + * Validates and normalizes a RetryConfig into canonical defaults. + * @param config Optional user-provided RetryConfig. + * @returns Normalized RetryConfig or undefined if not provided. + */ +export function normalizeRetryConfig( + config?: RetryConfig, +): Required | undefined { + if (!config) { + return undefined; + } + + if (config.maxAttempts < 1) { + throw new Error( + `RetryConfig.maxAttempts must be at least 1, received: ${config.maxAttempts}`, + ); + } + + return { + maxAttempts: config.maxAttempts, + initialDelayMs: config.initialDelayMs ?? 1000, + maxDelayMs: config.maxDelayMs ?? 30000, + backoffFactor: config.backoffFactor ?? 2.0, + retryableErrors: config.retryableErrors ?? [], + }; +} diff --git a/core/src/workflow/run_node.ts b/core/src/workflow/run_node.ts new file mode 100644 index 000000000..2d8e7a046 --- /dev/null +++ b/core/src/workflow/run_node.ts @@ -0,0 +1,96 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {InvocationContext} from '../agents/invocation_context.js'; +import {BaseNode} from './base_node.js'; +import {consumeGenerator, getOrInitAgentStates} from './node_runner.js'; +import {NodeState, NodeStatus, isNodeState} from './node_state.js'; +import {FunctionNode, FunctionNodeHandler} from './nodes/function_node.js'; + +/** + * Options for `runNode`. + */ +export interface RunNodeOptions { + /** + * Custom execution ID override for this node execution. + * If not provided, a deterministic execution ID based on branch and node name is used. + */ + customExecutionId?: string; +} + +/** + * Programmatically runs a target node (or function) inside the invocation context, + * checking and persisting state checkpoints (`InvocationContext.agentStates[execId]`). + * Essential for dynamic workflows where execution flow is controlled by TS code. + * + * @param ctx The current invocation context. + * @param nodeOrFunc A BaseNode instance or function handler to execute. + * @param input Optional input payload. + * @param options Optional settings (customExecutionId). + * @returns Promise resolving to the final output payload of the node. + */ +export async function runNode( + ctx: InvocationContext, + nodeOrFunc: BaseNode | FunctionNodeHandler, + input?: TInput, + options?: RunNodeOptions, +): Promise { + const node = + typeof nodeOrFunc === 'function' + ? new FunctionNode('run_node_dynamic', nodeOrFunc) + : nodeOrFunc; + + const agentStates = getOrInitAgentStates(ctx); + const execId = + options?.customExecutionId ?? + `exec_node_${ctx.branch ? ctx.branch + '.' : ''}${node.name}`; + + const existingState = agentStates[execId] as NodeState | undefined; + if ( + existingState && + isNodeState(existingState) && + existingState.status === NodeStatus.COMPLETED && + !node.rerunOnResume + ) { + return existingState.outputPayload as TOutput; + } + + const stateRecord: NodeState = { + executionId: execId, + nodeName: node.name, + status: NodeStatus.RUNNING, + inputPayload: input, + timestamp: Date.now(), + }; + agentStates[execId] = stateRecord; + + try { + const generator = node.run(ctx, input); + const {output, isPausedHitl} = await consumeGenerator(generator); + + if (isPausedHitl || ctx.endInvocation || ctx.abortSignal?.aborted) { + stateRecord.status = NodeStatus.PAUSED_HITL; + stateRecord.timestamp = Date.now(); + ctx.endInvocation = true; + throw new Error(`Node "${node.name}" requested HITL pause.`); + } + + const finalVal = + output !== undefined + ? output + : (node.lastOutputPayload ?? stateRecord.lastOutputPayload ?? input); + stateRecord.status = NodeStatus.COMPLETED; + stateRecord.outputPayload = finalVal; + stateRecord.timestamp = Date.now(); + return finalVal as TOutput; + } catch (error: unknown) { + stateRecord.status = NodeStatus.FAILED; + stateRecord.errorMessage = + error instanceof Error ? error.message : String(error); + stateRecord.timestamp = Date.now(); + throw error; + } +} diff --git a/core/src/workflow/trigger.ts b/core/src/workflow/trigger.ts new file mode 100644 index 000000000..9a713c26d --- /dev/null +++ b/core/src/workflow/trigger.ts @@ -0,0 +1,101 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {InvocationContext} from '../agents/invocation_context.js'; +import {Event} from '../events/event.js'; + +/** + * A predicate function evaluated against the event or output yielded by an upstream node. + * @param ctx The current invocation context. + * @param eventOrOutput The event or output payload produced by the source node. + * @returns True if the transition condition is satisfied, false otherwise. + */ +export type TriggerPredicate = ( + ctx: InvocationContext, + eventOrOutput: Event | unknown, +) => boolean | Promise; + +/** + * Represents a conditional trigger attached to a workflow graph edge. + * Determines whether a specific transition route should be taken upon upstream node completion. + */ +export class Trigger { + private readonly routeKey?: string; + private readonly predicate?: TriggerPredicate; + + private constructor(options: { + routeKey?: string; + predicate?: TriggerPredicate; + }) { + this.routeKey = options.routeKey; + this.predicate = options.predicate; + } + + /** + * Creates a route matching trigger. + * The trigger evaluates to true if the emitted `Event.actions.route` (or event payload route) matches `routeKey`. + * @param routeKey The exact string route key to match. + */ + static fromRoute(routeKey: string): Trigger { + if (!routeKey || typeof routeKey !== 'string') { + throw new Error( + 'Trigger.fromRoute requires a non-empty string routeKey.', + ); + } + return new Trigger({routeKey}); + } + + /** + * Creates a predicate-based trigger. + * The trigger evaluates to true if the provided predicate function returns true. + * @param predicate A boolean function evaluated against the context and node output/event. + */ + static fromPredicate(predicate: TriggerPredicate): Trigger { + if (typeof predicate !== 'function') { + throw new Error('Trigger.fromPredicate requires a function predicate.'); + } + return new Trigger({predicate}); + } + + /** + * Evaluates whether this trigger is satisfied by the given event or output payload. + * @param ctx The current invocation context. + * @param eventOrOutput The event or output payload produced by the source node. + * @returns Promise resolving to true if the transition should occur, false otherwise. + */ + async evaluate( + ctx: InvocationContext, + eventOrOutput: Event | unknown, + ): Promise { + if (this.routeKey) { + if (eventOrOutput && typeof eventOrOutput === 'object') { + const obj = eventOrOutput as Record; + if ( + 'actions' in obj && + obj.actions && + typeof obj.actions === 'object' && + 'route' in (obj.actions as Record) && + (obj.actions as Record).route === this.routeKey + ) { + return true; + } + if ('route' in obj && obj.route === this.routeKey) { + return true; + } + } + if (eventOrOutput === this.routeKey) { + return true; + } + return false; + } + + if (this.predicate) { + return await this.predicate(ctx, eventOrOutput); + } + + return true; + } +} diff --git a/core/src/workflow/utils/graph_parser.ts b/core/src/workflow/utils/graph_parser.ts new file mode 100644 index 000000000..566303972 --- /dev/null +++ b/core/src/workflow/utils/graph_parser.ts @@ -0,0 +1,176 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {BaseNode} from '../base_node.js'; +import {Trigger} from '../trigger.js'; + +/** + * A single element inside a GraphEdge array. + */ +export type EdgeElement = + | string // e.g., "START" + | BaseNode // concrete node instance + | Record // route map: { ROUTE_A: nodeA, ROUTE_B: nodeB } + | [Trigger, BaseNode]; // conditional tuple: [trigger, targetNode] + +/** + * A workflow graph edge definition. + * Examples: + * ["START", nodeA, nodeB, nodeC] + * [routerNode, { ROUTE_X: nodeC, ROUTE_Y: nodeD }] + * [nodeA, [Trigger.fromPredicate(...), nodeB]] + */ +export type GraphEdge = EdgeElement[]; + +/** + * Internal representation of a directed edge between two nodes. + */ +export interface AdjacencyEdge { + readonly source: string; // source node name or "START" + readonly target: BaseNode; + readonly trigger?: Trigger; +} + +/** + * A parsed and structured workflow graph ready for execution or validation. + */ +export class ParsedGraph { + readonly nodes = new Map(); + readonly adjacencyList = new Map(); + readonly inboundCounts = new Map(); + + constructor() { + this.adjacencyList.set('START', []); + this.inboundCounts.set('START', 0); + } + + /** + * Registers a node in the graph and initializes its adjacency and inbound counters if new. + * @param node The node to add. + */ + addNode(node: BaseNode): void { + if (!this.nodes.has(node.name)) { + this.nodes.set(node.name, node); + this.adjacencyList.set(node.name, []); + this.inboundCounts.set(node.name, 0); + } + } + + /** + * Adds a directed edge from `source` to `target` with an optional `trigger`. + * @param source Source node name (or "START"). + * @param target Target node instance. + * @param trigger Optional trigger condition. + */ + addEdge(source: string, target: BaseNode, trigger?: Trigger): void { + this.addNode(target); + if (source !== 'START' && !this.nodes.has(source)) { + throw new Error( + `Source node "${source}" must be added to the graph or referenced before defining an edge from it.`, + ); + } + + const edges = this.adjacencyList.get(source) || []; + edges.push({source, target, trigger}); + this.adjacencyList.set(source, edges); + + const currentInbound = this.inboundCounts.get(target.name) || 0; + this.inboundCounts.set(target.name, currentInbound + 1); + } +} + +/** + * Parses an array of user-defined GraphEdge structures into an internal ParsedGraph. + * @param edges Array of GraphEdge sequences or branch definitions. + * @returns A structured ParsedGraph. + */ +export function parseGraph(edges: GraphEdge[]): ParsedGraph { + if (!Array.isArray(edges) || edges.length === 0) { + throw new Error( + 'parseGraph requires a non-empty array of GraphEdge definitions.', + ); + } + + const graph = new ParsedGraph(); + + for (const edgeSeq of edges) { + if (!Array.isArray(edgeSeq) || edgeSeq.length < 2) { + throw new Error( + 'Each GraphEdge definition must be an array with at least 2 elements (e.g., ["START", nodeA]).', + ); + } + + for (let i = 0; i < edgeSeq.length - 1; i++) { + const current = edgeSeq[i]; + const next = edgeSeq[i + 1]; + + // Resolve source name + let sourceName: string; + if (typeof current === 'string' && current === 'START') { + sourceName = 'START'; + } else if (isBaseNode(current)) { + graph.addNode(current); + sourceName = current.name; + } else { + throw new Error( + `Invalid source element at index ${i} in edge sequence. Must be "START" or a BaseNode instance.`, + ); + } + + // Resolve target(s) + if (isBaseNode(next)) { + graph.addEdge(sourceName, next); + } else if (isRouteMap(next)) { + for (const [routeKey, targetNode] of Object.entries(next)) { + if (!isBaseNode(targetNode)) { + throw new Error( + `Target for route "${routeKey}" from source "${sourceName}" must be a BaseNode instance.`, + ); + } + graph.addEdge(sourceName, targetNode, Trigger.fromRoute(routeKey)); + } + } else if (isTriggerTuple(next)) { + const [trigger, targetNode] = next; + graph.addEdge(sourceName, targetNode, trigger); + } else { + throw new Error( + `Invalid target element at index ${i + 1} from source "${sourceName}". Must be a BaseNode, route dictionary, or [Trigger, BaseNode] tuple.`, + ); + } + } + } + + return graph; +} + +function isBaseNode(obj: unknown): obj is BaseNode { + return ( + typeof obj === 'object' && + obj !== null && + 'name' in obj && + typeof (obj as BaseNode).name === 'string' && + 'run' in obj && + typeof (obj as BaseNode).run === 'function' + ); +} + +function isRouteMap(obj: unknown): obj is Record { + return ( + typeof obj === 'object' && + obj !== null && + !isBaseNode(obj) && + !Array.isArray(obj) + ); +} + +function isTriggerTuple(obj: unknown): obj is [Trigger, BaseNode] { + return ( + Array.isArray(obj) && + obj.length === 2 && + obj[0] instanceof Trigger && + isBaseNode(obj[1]) + ); +} diff --git a/core/src/workflow/utils/graph_validation.ts b/core/src/workflow/utils/graph_validation.ts new file mode 100644 index 000000000..e8f75ef17 --- /dev/null +++ b/core/src/workflow/utils/graph_validation.ts @@ -0,0 +1,92 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {ParsedGraph} from './graph_parser.js'; + +/** + * Performs structural validation on a ParsedGraph before execution begins. + * Verifies reachability, checks for unintended cycles in DAG mode, and validates JoinNode upstream counts. + * + * @param graph The ParsedGraph to validate. + * @param options Optional validation settings (e.g. `allowCycles`). + * @throws Error if the graph structure is invalid or malformed. + */ +export function validateGraph( + graph: ParsedGraph, + options?: {allowCycles?: boolean}, +): void { + // 1. Reachability check from "START" + const visited = new Set(['START']); + const queue: string[] = ['START']; + + while (queue.length > 0) { + const current = queue.shift()!; + const edges = graph.adjacencyList.get(current) || []; + for (const edge of edges) { + if (!visited.has(edge.target.name)) { + visited.add(edge.target.name); + queue.push(edge.target.name); + } + } + } + + for (const [nodeName] of graph.nodes) { + if (!visited.has(nodeName)) { + throw new Error( + `Graph validation failed: Node "${nodeName}" is unreachable from "START". Check your edge definitions.`, + ); + } + } + + // 2. Cycle detection (DFS via recursion stack) if !allowCycles + if (!options?.allowCycles) { + const recursionStack = new Set(); + const dfsVisited = new Set(); + + const checkCycles = (nodeName: string): void => { + dfsVisited.add(nodeName); + recursionStack.add(nodeName); + + const edges = graph.adjacencyList.get(nodeName) || []; + for (const edge of edges) { + const targetName = edge.target.name; + if (!dfsVisited.has(targetName)) { + checkCycles(targetName); + } else if (recursionStack.has(targetName)) { + throw new Error( + `Graph validation failed: Cycle detected involving node "${targetName}". If your workflow intentionally contains loops, enable cycle support or use dynamic routing.`, + ); + } + } + + recursionStack.delete(nodeName); + }; + + checkCycles('START'); + } + + // 3. JoinNode upstream predecessor validation + for (const [nodeName, node] of graph.nodes) { + const nodeObj = node as unknown as Record; + if ( + 'upstreamCount' in nodeObj && + typeof nodeObj.upstreamCount === 'number' + ) { + const upstreamCount = nodeObj.upstreamCount as number; + const actualInbound = graph.inboundCounts.get(nodeName) || 0; + if (upstreamCount < 1) { + throw new Error( + `JoinNode "${nodeName}" has invalid upstreamCount: ${upstreamCount}. Must be >= 1.`, + ); + } + if (upstreamCount > actualInbound) { + throw new Error( + `JoinNode "${nodeName}" expects ${upstreamCount} upstream predecessors, but only has ${actualInbound} inbound edges defined in the graph.`, + ); + } + } + } +} diff --git a/core/src/workflow/utils/hitl_utils.ts b/core/src/workflow/utils/hitl_utils.ts new file mode 100644 index 000000000..a19515eec --- /dev/null +++ b/core/src/workflow/utils/hitl_utils.ts @@ -0,0 +1,89 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {InvocationContext} from '../../agents/invocation_context.js'; +import {createEvent, Event} from '../../events/event.js'; +import {isNodeState, NodeStatus} from '../node_state.js'; + +import {getOrInitAgentStates} from '../node_runner.js'; + +/** + * Options when creating a HITL input request. + */ +export interface RequestInputOptions { + /** + * Optional custom prompt or question to display to the user. + */ + prompt?: string; + + /** + * Optional structured schema or options describing what input is required. + */ + schema?: Record; +} + +/** + * Creates an Event that signals a Human-in-the-Loop (`RequestInput`) pause condition to the workflow engine. + * + * @param ctx The current invocation context. + * @param nodeName Name of the node requesting input. + * @param options Optional prompt and schema describing the required input. + */ +export function createRequestInputEvent( + ctx: InvocationContext, + nodeName: string, + options?: RequestInputOptions, +): Event { + return createEvent({ + invocationId: ctx.invocationId, + author: nodeName, + branch: ctx.branch, + content: options?.prompt + ? {role: 'model', parts: [{text: options.prompt}]} + : undefined, + actions: { + requestInput: { + nodeName, + prompt: options?.prompt, + schema: options?.schema, + }, + }, + }); +} + +/** + * Locates any node inside `InvocationContext.agentStates` whose status is `PAUSED_HITL`, + * and injects the resumption input payload so that subsequent workflow execution can proceed from that node. + * + * @param ctx The invocation context being resumed. + * @param resumptionInput The user's input payload provided upon resumption. + * @returns The name and execution ID of the resumed node, or undefined if no paused node was found. + */ +export function injectHitlResumptionInput( + ctx: InvocationContext, + resumptionInput: unknown, +): {nodeName: string; executionId: string} | undefined { + const agentStates = getOrInitAgentStates(ctx); + + for (const [execId, state] of Object.entries(agentStates)) { + if ( + isNodeState(state) && + state.status === NodeStatus.COMPLETED && + state.wasPausedHitl + ) { + continue; + } + if (isNodeState(state) && state.status === NodeStatus.PAUSED_HITL) { + state.status = NodeStatus.RUNNING; + state.inputPayload = resumptionInput; + state.wasPausedHitl = true; + state.timestamp = Date.now(); + return {nodeName: state.nodeName, executionId: execId}; + } + } + + return undefined; +} diff --git a/core/src/workflow/utils/rehydration_utils.ts b/core/src/workflow/utils/rehydration_utils.ts new file mode 100644 index 000000000..cf0abe6d8 --- /dev/null +++ b/core/src/workflow/utils/rehydration_utils.ts @@ -0,0 +1,114 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {InvocationContext} from '../../agents/invocation_context.js'; +import {Session} from '../../sessions/session.js'; +import {getOrInitAgentStates} from '../node_runner.js'; +import {NodeStatus, isNodeState} from '../node_state.js'; + +/** + * Scans historical session events and state metadata to reconstruct `InvocationContext.agentStates` + * and `InvocationContext.endOfAgents` when resuming a durable session from storage (such as database or GCS). + * + * @param session The session loaded from storage containing historical events and state. + * @param ctx The invocation context to populate with rehydrated checkpoints. + */ +export function rehydrateAgentStates( + session: Session, + ctx: InvocationContext, +): void { + const agentStates = getOrInitAgentStates(ctx); + + if (session.state && typeof session.state === 'object') { + const sessionState = session.state as Record; + if ( + 'agentStates' in sessionState && + sessionState.agentStates && + typeof sessionState.agentStates === 'object' + ) { + for (const [key, val] of Object.entries( + sessionState.agentStates as Record, + )) { + if (!agentStates[key] && isNodeState(val)) { + agentStates[key] = val; + } + } + } + if ( + 'endOfAgents' in sessionState && + sessionState.endOfAgents && + typeof sessionState.endOfAgents === 'object' + ) { + for (const [key, val] of Object.entries( + sessionState.endOfAgents as Record, + )) { + if (typeof val === 'boolean') { + ctx.endOfAgents[key] = val; + } + } + } + } + + if (Array.isArray(session.events)) { + for (const event of session.events) { + if (!event || typeof event !== 'object') continue; + + const eventRecord = event as Record; + const actions = eventRecord.actions as + | Record + | undefined; + if (actions && typeof actions === 'object') { + if ( + 'nodeExecution' in actions && + actions.nodeExecution && + typeof actions.nodeExecution === 'object' + ) { + const {executionId, nodeName, status, outputPayload} = + actions.nodeExecution as Record; + if ( + executionId && + typeof executionId === 'string' && + !agentStates[executionId] + ) { + agentStates[executionId] = { + executionId, + nodeName: + typeof nodeName === 'string' ? nodeName : 'unknown_node', + status: + status === 'PAUSED_HITL' + ? NodeStatus.PAUSED_HITL + : NodeStatus.COMPLETED, + outputPayload, + timestamp: + typeof eventRecord.timestamp === 'number' + ? eventRecord.timestamp + : Date.now(), + }; + } + } + } + } + } +} + +/** + * Persists current `InvocationContext.agentStates` and `InvocationContext.endOfAgents` snapshots + * onto the session's state dictionary so they can be securely serialized by session services. + * + * @param ctx The invocation context whose states should be saved. + * @param session The session object to update. + */ +export function persistAgentStatesToSession( + ctx: InvocationContext, + session: Session, +): void { + if (!session.state || typeof session.state !== 'object') { + session.state = {}; + } + const sessionState = session.state as Record; + sessionState.agentStates = {...getOrInitAgentStates(ctx)}; + sessionState.endOfAgents = {...ctx.endOfAgents}; +} diff --git a/core/src/workflow/utils/replay_manager.ts b/core/src/workflow/utils/replay_manager.ts new file mode 100644 index 000000000..d182af851 --- /dev/null +++ b/core/src/workflow/utils/replay_manager.ts @@ -0,0 +1,90 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {InvocationContext} from '../../agents/invocation_context.js'; +import {Event} from '../../events/event.js'; +import {BaseNode} from '../base_node.js'; +import {generateExecutionId, getOrInitAgentStates} from '../node_runner.js'; +import {NodeState, NodeStatus, isNodeState} from '../node_state.js'; + +/** + * Manages event replay when a workflow node is bypassed due to a completed historical checkpoint (`rerunOnResume == false`). + * Yields historical events associated with that node execution so client UIs and event subscribers can reconstruct the full trajectory. + */ +export class ReplayManager { + /** + * Checks if the given node has a historical COMPLETED checkpoint in `InvocationContext.agentStates` + * and yields its historical events (or a synthetic replay event) if `!node.rerunOnResume`. + * + * @param ctx The current invocation context. + * @param node The node being checked for replay. + * @yields Historical or synthetic replay events if checkpoint exists. + * @returns True if the node was successfully replayed (and execution should be skipped), false otherwise. + */ + static async *replayIfCompleted( + ctx: InvocationContext, + node: BaseNode, + ): AsyncGenerator { + if (node.rerunOnResume) { + return false; + } + + const agentStates = getOrInitAgentStates(ctx); + const execId = generateExecutionId(ctx, node.name); + const existingState = agentStates[execId] as NodeState | undefined; + + if ( + existingState && + isNodeState(existingState) && + existingState.status === NodeStatus.COMPLETED + ) { + if (Array.isArray(existingState.cachedEvents)) { + for (const event of existingState.cachedEvents) { + yield event; + } + } else if (existingState.outputPayload !== undefined) { + yield { + invocationId: ctx.invocationId, + author: node.name, + branch: ctx.branch, + actions: { + nodeExecutionReplay: { + executionId: execId, + nodeName: node.name, + outputPayload: existingState.outputPayload, + timestamp: existingState.timestamp, + }, + }, + } as Event; + } + return true; + } + + return false; + } + + /** + * Caches emitted events onto a node's state record during live execution so they can be replayed on subsequent resumptions. + * + * @param ctx The invocation context. + * @param execId The execution ID of the running node. + * @param event The emitted event to cache. + */ + static cacheEventForReplay( + ctx: InvocationContext, + execId: string, + event: Event, + ): void { + const agentStates = getOrInitAgentStates(ctx); + const state = agentStates[execId] as NodeState | undefined; + if (state && isNodeState(state)) { + if (!Array.isArray(state.cachedEvents)) { + state.cachedEvents = []; + } + state.cachedEvents.push(event); + } + } +} diff --git a/core/src/workflow/utils/retry_utils.ts b/core/src/workflow/utils/retry_utils.ts new file mode 100644 index 000000000..104cebd48 --- /dev/null +++ b/core/src/workflow/utils/retry_utils.ts @@ -0,0 +1,114 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {Event} from '../../events/event.js'; +import {RetryConfig, normalizeRetryConfig} from '../retry_config.js'; + +/** + * Sleeps for a specified number of milliseconds unless the abort signal fires. + * @param ms Delay in milliseconds. + * @param abortSignal Optional AbortSignal to cancel sleeping early. + */ +async function sleepWithSignal( + ms: number, + abortSignal?: AbortSignal, +): Promise { + if (abortSignal?.aborted) { + throw new Error('Aborted before retry delay.'); + } + return new Promise((resolve, reject) => { + const timer = setTimeout(() => { + abortSignal?.removeEventListener('abort', onAbort); + resolve(); + }, ms); + const onAbort = () => { + clearTimeout(timer); + reject(new Error('Aborted during retry delay.')); + }; + abortSignal?.addEventListener('abort', onAbort); + }); +} + +/** + * Checks whether an error is retryable according to the RetryConfig. + * @param error The error thrown during execution. + * @param retryableErrors Array of Error constructors, strings, or RegExps. + */ +function isErrorRetryable( + error: unknown, + retryableErrors: Required['retryableErrors'], +): boolean { + if (!retryableErrors || retryableErrors.length === 0) { + return true; + } + + const errObj = + typeof error === 'object' && error !== null + ? (error as Record) + : undefined; + const errMsg = + errObj && typeof errObj.message === 'string' ? errObj.message : undefined; + + return retryableErrors.some((matcher) => { + if (typeof matcher === 'string') { + return errMsg && errMsg.includes(matcher); + } + if (matcher instanceof RegExp) { + return errMsg && matcher.test(errMsg); + } + if (typeof matcher === 'function' && error instanceof matcher) { + return true; + } + return false; + }); +} + +/** + * Wraps an async generator with retry logic according to the provided RetryConfig. + * If the generator throws a retryable error mid-stream or during start, it will back off and retry from the beginning. + * + * @param generatorFactory A factory function that creates a fresh AsyncGenerator for each attempt. + * @param retryConfig Optional RetryConfig or undefined (if undefined, runs once without retrying). + * @param abortSignal Optional AbortSignal to halt retries upon cancellation. + */ +export async function* runWithRetry( + generatorFactory: () => AsyncGenerator, + retryConfig?: RetryConfig, + abortSignal?: AbortSignal, +): AsyncGenerator { + const config = normalizeRetryConfig(retryConfig); + if (!config || config.maxAttempts <= 1) { + return yield* generatorFactory(); + } + + let attempt = 1; + while (true) { + if (abortSignal?.aborted) { + throw new Error('Execution aborted before attempt.'); + } + + const generator = generatorFactory(); + try { + const result = yield* generator; + return result; + } catch (error: unknown) { + if ( + attempt >= config.maxAttempts || + !isErrorRetryable(error, config.retryableErrors) || + abortSignal?.aborted + ) { + throw error; + } + + const delayMs = Math.min( + config.initialDelayMs * Math.pow(config.backoffFactor, attempt - 1), + config.maxDelayMs, + ); + await sleepWithSignal(delayMs, abortSignal); + attempt++; + } + } +} diff --git a/core/src/workflow/workflow.ts b/core/src/workflow/workflow.ts new file mode 100644 index 000000000..fe21f8779 --- /dev/null +++ b/core/src/workflow/workflow.ts @@ -0,0 +1,141 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {BaseAgent, BaseAgentConfig} from '../agents/base_agent.js'; +import {InvocationContext} from '../agents/invocation_context.js'; +import {Event} from '../events/event.js'; +import {BaseNode} from './base_node.js'; +import { + DynamicEntryFunction, + DynamicNodeScheduler, +} from './dynamic_node_scheduler.js'; +import {NodeRunner} from './node_runner.js'; +import {GraphEdge} from './utils/graph_parser.js'; + +/** + * A unique symbol to identify ADK Workflow agent instances. + */ +const WORKFLOW_SIGNATURE_SYMBOL = Symbol.for('google.adk.workflow'); + +/** + * Type guard to check if an object is an instance of Workflow. + * @param obj The object to check. + * @returns True if the object is an instance of Workflow, false otherwise. + */ +export function isWorkflow(obj: unknown): obj is Workflow { + return ( + typeof obj === 'object' && + obj !== null && + WORKFLOW_SIGNATURE_SYMBOL in obj && + (obj as Record)[WORKFLOW_SIGNATURE_SYMBOL] === true + ); +} + +/** + * Configuration options for creating a Workflow agent. + * Workflows must define exactly one of `edges` (for static DAG execution) or `dynamicEntry` (for programmatic execution). + */ +export interface WorkflowConfig extends BaseAgentConfig { + /** + * Static graph edge definitions (e.g., `["START", nodeA, nodeB]` or `[routerNode, { ROUTE_A: nodeC }]`). + * Mutually exclusive with `dynamicEntry`. + */ + edges?: GraphEdge[]; + + /** + * Programmatic entry node or async function handler (`async (ctx, input) => ...`) that coordinates + * child nodes using `ctx.runNode(...)`. Mutually exclusive with `edges`. + */ + dynamicEntry?: BaseNode | DynamicEntryFunction; + + /** + * Optional key inside `InvocationContext.agentStates` where the final output of the workflow + * should be stored upon successful completion. + */ + outputKey?: string; + + /** + * If true, the workflow will force re-execution on resumption even if historical outputs exist. + * Default is false. + */ + rerunOnResume?: boolean; + + /** + * If true, allows directed cycles inside static `edges` graph validation. + * Default is false. + */ + allowCycles?: boolean; +} + +/** + * The top-level Workflow agent in ADK-JS. + * Inherits from `BaseAgent` and orchestrates multi-step node execution using either a static graph DAG (`NodeRunner`) + * or dynamic programmatic scheduling (`DynamicNodeScheduler`). + */ +export class Workflow extends BaseAgent { + readonly [WORKFLOW_SIGNATURE_SYMBOL] = true; + + readonly edges?: GraphEdge[]; + readonly dynamicEntry?: BaseNode | DynamicEntryFunction; + readonly outputKey?: string; + readonly rerunOnResume: boolean; + readonly allowCycles: boolean; + + constructor(config: WorkflowConfig) { + super(config); + if (config.edges && config.dynamicEntry) { + throw new Error( + `Workflow "${this.name}" cannot have both "edges" and "dynamicEntry" defined. They are mutually exclusive.`, + ); + } + if (!config.edges && !config.dynamicEntry) { + throw new Error( + `Workflow "${this.name}" must define either "edges" (for static graphs) or "dynamicEntry" (for dynamic code execution).`, + ); + } + + this.edges = config.edges; + this.dynamicEntry = config.dynamicEntry; + this.outputKey = config.outputKey; + this.rerunOnResume = config.rerunOnResume ?? false; + this.allowCycles = config.allowCycles ?? false; + } + + /** + * Executes the workflow via text-based or programmatic invocation. + */ + protected async *runAsyncImpl( + context: InvocationContext, + ): AsyncGenerator { + if (context.endOfAgents[this.name]) { + return; + } + + if (this.edges) { + const runner = new NodeRunner(this.edges, { + outputKey: this.outputKey, + allowCycles: this.allowCycles, + }); + yield* runner.runAsync(context, context.userContent); + } else if (this.dynamicEntry) { + const scheduler = new DynamicNodeScheduler(this.dynamicEntry, { + outputKey: this.outputKey, + }); + yield* scheduler.runAsync(context, context.userContent); + } + + context.endOfAgents[this.name] = true; + } + + /** + * Executes the workflow via audio/video live streaming invocation. + */ + protected async *runLiveImpl( + context: InvocationContext, + ): AsyncGenerator { + yield* this.runAsyncImpl(context); + } +} diff --git a/core/test/workflow/dynamic_workflow_test.ts b/core/test/workflow/dynamic_workflow_test.ts new file mode 100644 index 000000000..43b219b18 --- /dev/null +++ b/core/test/workflow/dynamic_workflow_test.ts @@ -0,0 +1,107 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {describe, expect, it, vi} from 'vitest'; +import {BaseAgent} from '../../src/agents/base_agent.js'; +import {InvocationContext} from '../../src/agents/invocation_context.js'; +import {PluginManager} from '../../src/plugins/plugin_manager.js'; +import {Session} from '../../src/sessions/session.js'; +import { + DynamicNodeScheduler, + FunctionNode, + NodeStatus, + runNode, +} from '../../src/workflow/index.js'; + +describe('Workflow DynamicNodeScheduler & runNode', () => { + function createTestContext( + params?: Partial, + ): InvocationContext { + const session: Session = { + id: 'session-dyn', + appName: 'test-app', + userId: 'test-user', + events: [], + state: {}, + }; + + return new InvocationContext({ + invocationId: 'inv-dyn', + session, + agent: { + name: 'mock_agent', + runAsync: async function* () {}, + } as unknown as BaseAgent, + pluginManager: new PluginManager(), + ...params, + }); + } + + it('should run dynamic workflows and track individual node checkpoints via runNode', async () => { + const ctx = createTestContext(); + const nodeA = new FunctionNode('calc_a', (_ctx, num: number) => num * 2); + const nodeB = new FunctionNode('calc_b', (_ctx, num: number) => num + 10); + + const dynamicEntry = async (context: InvocationContext, input?: number) => { + const resA = await runNode(context, nodeA, input || 5); + if (resA > 5) { + const resB = await runNode(context, nodeB, resA); + return resB; + } + return resA; + }; + + const scheduler = new DynamicNodeScheduler(dynamicEntry, { + outputKey: 'dynResult', + }); + for await (const _ of scheduler.runAsync(ctx, 4)) { + /* consume events */ + } + + expect(ctx.agentStates['exec_node_calc_a'].status).toBe( + NodeStatus.COMPLETED, + ); + expect(ctx.agentStates['exec_node_calc_a'].outputPayload).toBe(8); + expect(ctx.agentStates['exec_node_calc_b'].status).toBe( + NodeStatus.COMPLETED, + ); + expect(ctx.agentStates['exec_node_calc_b'].outputPayload).toBe(18); + expect(ctx.agentStates['dynResult']).toBe(18); + }); + + it('should skip completed nodes inside dynamic execution on resume', async () => { + const ctx = createTestContext(); + const spyA = vi.fn((_ctx, num: number) => num * 100); + const spyB = vi.fn((_ctx, num: number) => num + 50); + const nodeA = new FunctionNode('node_dyn_a', spyA, {rerunOnResume: false}); + const nodeB = new FunctionNode('node_dyn_b', spyB, {rerunOnResume: true}); + + ctx.agentStates['exec_node_node_dyn_a'] = { + executionId: 'exec_node_node_dyn_a', + nodeName: 'node_dyn_a', + status: NodeStatus.COMPLETED, + outputPayload: 999, + timestamp: Date.now(), + }; + + const dynamicEntry = async (context: InvocationContext, input?: number) => { + const resA = await runNode(context, nodeA, input || 1); + const resB = await runNode(context, nodeB, resA); + return resB; + }; + + const scheduler = new DynamicNodeScheduler(dynamicEntry); + for await (const _ of scheduler.runAsync(ctx, 2)) { + /* consume events */ + } + + // spyA skipped, reused 999 + expect(spyA).not.toHaveBeenCalled(); + // spyB executed with 999 + expect(spyB).toHaveBeenCalledTimes(1); + expect(spyB).toHaveBeenCalledWith(ctx, 999); + }); +}); diff --git a/core/test/workflow/graph_parser_test.ts b/core/test/workflow/graph_parser_test.ts new file mode 100644 index 000000000..e49557045 --- /dev/null +++ b/core/test/workflow/graph_parser_test.ts @@ -0,0 +1,105 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {describe, expect, it} from 'vitest'; +import { + FunctionNode, + JoinNode, + parseGraph, + Trigger, + validateGraph, +} from '../../src/workflow/index.js'; + +describe('Workflow Graph Parser & Validation', () => { + const nodeA = new FunctionNode('node_a', () => 'A'); + const nodeB = new FunctionNode('node_b', () => 'B'); + const nodeC = new FunctionNode('node_c', () => 'C'); + const router = new FunctionNode('router', () => 'router_res'); + + it('should parse sequential edges accurately', () => { + const graph = parseGraph([['START', nodeA, nodeB, nodeC]]); + expect(graph.nodes.size).toBe(3); + expect(graph.nodes.get('node_a')).toBe(nodeA); + expect(graph.nodes.get('node_b')).toBe(nodeB); + expect(graph.nodes.get('node_c')).toBe(nodeC); + + const startEdges = graph.adjacencyList.get('START') || []; + expect(startEdges.length).toBe(1); + expect(startEdges[0].target).toBe(nodeA); + + const aEdges = graph.adjacencyList.get('node_a') || []; + expect(aEdges.length).toBe(1); + expect(aEdges[0].target).toBe(nodeB); + + const bEdges = graph.adjacencyList.get('node_b') || []; + expect(bEdges.length).toBe(1); + expect(bEdges[0].target).toBe(nodeC); + + expect(graph.inboundCounts.get('node_a')).toBe(1); + expect(graph.inboundCounts.get('node_b')).toBe(1); + expect(graph.inboundCounts.get('node_c')).toBe(1); + }); + + it('should parse route maps and trigger tuples', () => { + const customTrigger = Trigger.fromPredicate(() => true); + const graph = parseGraph([ + ['START', router], + [router, {ROUTE_X: nodeA, ROUTE_Y: nodeB}], + [nodeA, [customTrigger, nodeC]], + ]); + + expect(graph.nodes.size).toBe(4); + const routerEdges = graph.adjacencyList.get('router') || []; + expect(routerEdges.length).toBe(2); + expect(routerEdges[0].target.name).toBe('node_a'); + expect(routerEdges[1].target.name).toBe('node_b'); + + const aEdges = graph.adjacencyList.get('node_a') || []; + expect(aEdges.length).toBe(1); + expect(aEdges[0].trigger).toBe(customTrigger); + }); + + it('should throw during validation when a node is unreachable from START', () => { + const graph = parseGraph([ + ['START', nodeA], + [nodeB, nodeC], // nodeB and nodeC have no path from START + ]); + + expect(() => validateGraph(graph)).toThrowError( + /unreachable from "START"/i, + ); + }); + + it('should detect cycles and throw unless allowCycles is true', () => { + const graph = parseGraph([ + ['START', nodeA, nodeB], + [nodeB, nodeA], // cycle nodeB -> nodeA + ]); + + expect(() => validateGraph(graph, {allowCycles: false})).toThrowError( + /Cycle detected/i, + ); + + expect(() => validateGraph(graph, {allowCycles: true})).not.toThrow(); + }); + + it('should validate JoinNode upstreamCount integrity', () => { + const joinNode = new JoinNode('join_node', {upstreamCount: 2}); + const validGraph = parseGraph([ + ['START', nodeA, joinNode], + ['START', nodeB, joinNode], + ]); + + expect(() => validateGraph(validGraph)).not.toThrow(); + + const invalidJoin = new JoinNode('bad_join', {upstreamCount: 5}); + const invalidGraph = parseGraph([['START', nodeA, invalidJoin]]); + + expect(() => validateGraph(invalidGraph)).toThrowError( + /expects 5 upstream predecessors, but only has 1 inbound edges/i, + ); + }); +}); diff --git a/core/test/workflow/hitl_and_rehydration_test.ts b/core/test/workflow/hitl_and_rehydration_test.ts new file mode 100644 index 000000000..1801b05f7 --- /dev/null +++ b/core/test/workflow/hitl_and_rehydration_test.ts @@ -0,0 +1,189 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {describe, expect, it} from 'vitest'; +import {BaseAgent} from '../../src/agents/base_agent.js'; +import {InvocationContext} from '../../src/agents/invocation_context.js'; +import {PluginManager} from '../../src/plugins/plugin_manager.js'; +import {Session} from '../../src/sessions/session.js'; +import { + createRequestInputEvent, + FunctionNode, + injectHitlResumptionInput, + NodeRunner, + NodeStatus, + persistAgentStatesToSession, + rehydrateAgentStates, + ReplayManager, +} from '../../src/workflow/index.js'; + +describe('Workflow HITL & State Rehydration', () => { + function createTestContext(session?: Session): InvocationContext { + const s: Session = session || { + id: 'session-hitl', + appName: 'test-app', + userId: 'test-user', + events: [], + state: {}, + }; + + return new InvocationContext({ + invocationId: 'inv-hitl', + session: s, + agent: { + name: 'mock_agent', + runAsync: async function* () {}, + } as unknown as BaseAgent, + pluginManager: new PluginManager(), + }); + } + + it('should pause workflow execution when a node yields a RequestInput event', async () => { + const ctx = createTestContext(); + const hitlNode = new FunctionNode('approval_step', (context) => { + return createRequestInputEvent(context, 'approval_step', { + prompt: 'Please approve this transaction (yes/no)', + }); + }); + const downstreamNode = new FunctionNode('post_approval', () => 'EXECUTED'); + + const runner = new NodeRunner([['START', hitlNode, downstreamNode]]); + const events: unknown[] = []; + for await (const ev of runner.runAsync(ctx)) { + events.push(ev); + } + + expect(events.length).toBe(1); + expect( + ( + events[0] as unknown as Record< + string, + Record + > + ).actions.requestInput.nodeName, + ).toBe('approval_step'); + expect(ctx.agentStates['exec_node_approval_step'].status).toBe( + NodeStatus.PAUSED_HITL, + ); + + expect(ctx.endInvocation).toBe(true); + // downstreamNode must NOT have run while paused + expect(ctx.agentStates['exec_node_post_approval']).toBeUndefined(); + }); + + it('should inject resumption input and complete the paused node on subsequent turn', async () => { + const ctx = createTestContext(); + const hitlNode = new FunctionNode( + 'approval_step', + (_context, input?: string) => { + if (!input) { + return createRequestInputEvent(_context, 'approval_step', { + prompt: 'Approve?', + }); + } + return `APPROVED_WITH_${input}`; + }, + ); + const downstreamNode = new FunctionNode( + 'post_approval', + (_context, input: string) => `DONE_${input}`, + ); + + const runner = new NodeRunner([['START', hitlNode, downstreamNode]]); + + // Turn 1: Pauses + for await (const _ of runner.runAsync(ctx)) { + /* consume events */ + } + expect(ctx.agentStates['exec_node_approval_step'].status).toBe( + NodeStatus.PAUSED_HITL, + ); + + // Turn 2: User provides input 'YES'. Inject it and run again. + ctx.endInvocation = false; + const resumedInfo = injectHitlResumptionInput(ctx, 'YES'); + expect(resumedInfo?.nodeName).toBe('approval_step'); + expect(ctx.agentStates['exec_node_approval_step'].status).toBe( + NodeStatus.RUNNING, + ); + expect(ctx.agentStates['exec_node_approval_step'].inputPayload).toBe('YES'); + + for await (const _ of runner.runAsync(ctx)) { + /* consume events */ + } + + expect(ctx.agentStates['exec_node_approval_step'].status).toBe( + NodeStatus.COMPLETED, + ); + expect(ctx.agentStates['exec_node_approval_step'].outputPayload).toBe( + 'APPROVED_WITH_YES', + ); + expect(ctx.agentStates['exec_node_post_approval'].status).toBe( + NodeStatus.COMPLETED, + ); + expect(ctx.agentStates['exec_node_post_approval'].outputPayload).toBe( + 'DONE_APPROVED_WITH_YES', + ); + }); + + it('should persist and rehydrate checkpoints from session state accurately', () => { + const session: Session = { + id: 's-rehydrate', + appName: 'app', + userId: 'user', + events: [], + state: {}, + }; + const ctx1 = createTestContext(session); + ctx1.agentStates['exec_node_saved'] = { + executionId: 'exec_node_saved', + nodeName: 'saved', + status: NodeStatus.COMPLETED, + outputPayload: {foo: 'bar'}, + timestamp: 12345, + }; + ctx1.endOfAgents['my_wf'] = true; + + persistAgentStatesToSession(ctx1, session); + + // Now simulate a fresh context loading from the same session + const ctx2 = createTestContext(session); + rehydrateAgentStates(session, ctx2); + + expect(ctx2.agentStates['exec_node_saved'].status).toBe( + NodeStatus.COMPLETED, + ); + expect(ctx2.agentStates['exec_node_saved'].outputPayload).toEqual({ + foo: 'bar', + }); + expect(ctx2.endOfAgents['my_wf']).toBe(true); + }); + + it('should yield historical replay events when ReplayManager inspects completed checkpoints', async () => { + const ctx = createTestContext(); + const node = new FunctionNode('past_node', () => 'historical_res', { + rerunOnResume: false, + }); + ctx.agentStates['exec_node_past_node'] = { + executionId: 'exec_node_past_node', + nodeName: 'past_node', + status: NodeStatus.COMPLETED, + outputPayload: 'historical_res', + timestamp: 99999, + }; + + const gen = ReplayManager.replayIfCompleted(ctx, node); + const ev = await gen.next(); + expect(ev.done).toBe(false); + expect(ev.value.actions.nodeExecutionReplay.outputPayload).toBe( + 'historical_res', + ); + + const res = await gen.next(); + expect(res.done).toBe(true); + expect(res.value).toBe(true); // Replayed successfully, skip real run + }); +}); diff --git a/core/test/workflow/join_node_and_parallel_test.ts b/core/test/workflow/join_node_and_parallel_test.ts new file mode 100644 index 000000000..a78db2ccd --- /dev/null +++ b/core/test/workflow/join_node_and_parallel_test.ts @@ -0,0 +1,130 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {describe, expect, it} from 'vitest'; +import {BaseAgent} from '../../src/agents/base_agent.js'; +import {InvocationContext} from '../../src/agents/invocation_context.js'; +import {PluginManager} from '../../src/plugins/plugin_manager.js'; +import {Session} from '../../src/sessions/session.js'; +import { + FunctionNode, + JoinNode, + NodeStatus, + runInParallel, +} from '../../src/workflow/index.js'; + +describe('Workflow ParallelWorker & JoinNode', () => { + function createTestContext( + params?: Partial, + ): InvocationContext { + const session: Session = { + id: 'session-par', + appName: 'test-app', + userId: 'test-user', + events: [], + state: {}, + }; + + return new InvocationContext({ + invocationId: 'inv-par', + session, + agent: { + name: 'mock_agent', + runAsync: async function* () {}, + } as unknown as BaseAgent, + pluginManager: new PluginManager(), + ...params, + }); + } + + it('should run items in parallel and merge child checkpoints back into parent context', async () => { + const ctx = createTestContext(); + const workerNode = new FunctionNode( + 'worker', + async (_ctx, item: string) => { + return `processed_${item}`; + }, + ); + + const results = await runInParallel(ctx, workerNode, [ + 'alpha', + 'beta', + 'gamma', + ]); + + expect(results).toEqual([ + 'processed_alpha', + 'processed_beta', + 'processed_gamma', + ]); + expect(ctx.agentStates['exec_node_worker_0.worker'].status).toBe( + NodeStatus.COMPLETED, + ); + expect(ctx.agentStates['exec_node_worker_0.worker'].outputPayload).toBe( + 'processed_alpha', + ); + expect(ctx.agentStates['exec_node_worker_1.worker'].status).toBe( + NodeStatus.COMPLETED, + ); + expect(ctx.agentStates['exec_node_worker_1.worker'].outputPayload).toBe( + 'processed_beta', + ); + expect(ctx.agentStates['exec_node_worker_2.worker'].status).toBe( + NodeStatus.COMPLETED, + ); + expect(ctx.agentStates['exec_node_worker_2.worker'].outputPayload).toBe( + 'processed_gamma', + ); + }); + + it('should synchronize at a JoinNode when all upstream predecessors reach COMPLETED', async () => { + const ctx = createTestContext(); + const joinNode = new JoinNode('join_sync', { + upstreamCount: 2, + predecessors: ['branch_1', 'branch_2'], + }); + + // 1. First branch finishes, second has not started + ctx.agentStates['exec_node_branch_1'] = { + executionId: 'exec_node_branch_1', + nodeName: 'branch_1', + status: NodeStatus.COMPLETED, + outputPayload: {data: 100}, + timestamp: Date.now(), + }; + + const gen1 = joinNode.run(ctx); + const res1 = await gen1.next(); + // Since only 1 of 2 predecessors completed, joinNode returns partial state and does NOT yield a joinCompleted event + expect(res1.done).toBe(true); + expect(res1.value).toEqual({branch_1: {data: 100}}); + + // 2. Second branch now finishes + ctx.agentStates['exec_node_branch_2'] = { + executionId: 'exec_node_branch_2', + nodeName: 'branch_2', + status: NodeStatus.COMPLETED, + outputPayload: {data: 200}, + timestamp: Date.now(), + }; + + const gen2 = joinNode.run(ctx); + const ev = await gen2.next(); // Should yield joinCompleted event + expect(ev.done).toBe(false); + expect(ev.value.actions.joinCompleted).toBeDefined(); + expect(ev.value.actions.joinCompleted.outputs).toEqual({ + branch_1: {data: 100}, + branch_2: {data: 200}, + }); + + const finalRes = await gen2.next(); + expect(finalRes.done).toBe(true); + expect(finalRes.value).toEqual({ + branch_1: {data: 100}, + branch_2: {data: 200}, + }); + }); +}); diff --git a/core/test/workflow/node_runner_test.ts b/core/test/workflow/node_runner_test.ts new file mode 100644 index 000000000..868a5f833 --- /dev/null +++ b/core/test/workflow/node_runner_test.ts @@ -0,0 +1,177 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {describe, expect, it, vi} from 'vitest'; +import {BaseAgent} from '../../src/agents/base_agent.js'; +import {InvocationContext} from '../../src/agents/invocation_context.js'; +import {PluginManager} from '../../src/plugins/plugin_manager.js'; +import {Session} from '../../src/sessions/session.js'; +import { + FunctionNode, + NodeRunner, + NodeStatus, +} from '../../src/workflow/index.js'; + +describe('Workflow NodeRunner & Checkpointing', () => { + function createTestContext( + params?: Partial, + ): InvocationContext { + const session: Session = { + id: 'session-123', + appName: 'test-app', + userId: 'test-user', + events: [], + state: {}, + }; + + return new InvocationContext({ + invocationId: 'inv-456', + session, + agent: { + name: 'mock_agent', + runAsync: async function* () {}, + } as unknown as BaseAgent, + pluginManager: new PluginManager(), + ...params, + }); + } + + it('should run sequential nodes and pass input payloads downstream', async () => { + const ctx = createTestContext(); + const nodeA = new FunctionNode( + 'step_1', + (_ctx, input: string) => `${input}_A`, + ); + const nodeB = new FunctionNode( + 'step_2', + (_ctx, input: string) => `${input}_B`, + ); + + const runner = new NodeRunner([['START', nodeA, nodeB]], { + outputKey: 'finalOutput', + }); + const events: unknown[] = []; + for await (const event of runner.runAsync(ctx, 'INITIAL')) { + events.push(event); + } + + expect(ctx.agentStates['exec_node_step_1'].status).toBe( + NodeStatus.COMPLETED, + ); + expect(ctx.agentStates['exec_node_step_1'].outputPayload).toBe('INITIAL_A'); + expect(ctx.agentStates['exec_node_step_2'].status).toBe( + NodeStatus.COMPLETED, + ); + expect(ctx.agentStates['exec_node_step_2'].outputPayload).toBe( + 'INITIAL_A_B', + ); + expect(ctx.agentStates['finalOutput']).toEqual({ + step_1: 'INITIAL_A', + step_2: 'INITIAL_A_B', + }); + }); + + it('should route conditionally when a router node emits a route action or string', async () => { + const ctx = createTestContext(); + const router = new FunctionNode('router', () => 'ROUTE_Y'); + const nodeX = new FunctionNode( + 'node_x', + vi.fn(() => 'X'), + ); + const nodeY = new FunctionNode( + 'node_y', + vi.fn(() => 'Y'), + ); + + const runner = new NodeRunner([ + ['START', router], + [router, {ROUTE_X: nodeX, ROUTE_Y: nodeY}], + ]); + + for await (const _ of runner.runAsync(ctx)) { + /* consume events */ + } + + expect(ctx.agentStates['exec_node_router'].status).toBe( + NodeStatus.COMPLETED, + ); + expect(ctx.agentStates['exec_node_node_y'].status).toBe( + NodeStatus.COMPLETED, + ); + expect(ctx.agentStates['exec_node_node_x']).toBeUndefined(); // ROUTE_X never enqueued + }); + + it('should skip completed nodes on resume unless rerunOnResume is true', async () => { + const ctx = createTestContext(); + const spyA = vi.fn(() => 'fresh_A'); + const spyB = vi.fn(() => 'fresh_B'); + const nodeA = new FunctionNode('node_a', spyA, {rerunOnResume: false}); + const nodeB = new FunctionNode('node_b', spyB, {rerunOnResume: true}); + + // Pre-populate agentStates as if nodeA and nodeB completed in a previous run + ctx.agentStates['exec_node_node_a'] = { + executionId: 'exec_node_node_a', + nodeName: 'node_a', + status: NodeStatus.COMPLETED, + outputPayload: 'cached_A', + timestamp: Date.now() - 10000, + }; + ctx.agentStates['exec_node_node_b'] = { + executionId: 'exec_node_node_b', + nodeName: 'node_b', + status: NodeStatus.COMPLETED, + outputPayload: 'cached_B', + timestamp: Date.now() - 10000, + }; + + const runner = new NodeRunner([['START', nodeA, nodeB]]); + for await (const _ of runner.runAsync(ctx)) { + /* consume events */ + } + + // nodeA should have been skipped (spyA not called), and cached_A passed to nodeB + expect(spyA).not.toHaveBeenCalled(); + // nodeB has rerunOnResume: true, so spyB MUST be called with cached_A + expect(spyB).toHaveBeenCalledTimes(1); + expect(spyB).toHaveBeenCalledWith(ctx, 'cached_A'); + }); + + it('should retry node execution upon transient errors according to retryConfig', async () => { + const ctx = createTestContext(); + let attempts = 0; + const flakyNode = new FunctionNode( + 'flaky_node', + () => { + attempts++; + if (attempts < 3) { + throw new Error('Transient timeout error'); + } + return 'SUCCESS_AFTER_RETRY'; + }, + { + retryConfig: { + maxAttempts: 3, + initialDelayMs: 10, + maxDelayMs: 50, + backoffFactor: 1.5, + }, + }, + ); + + const runner = new NodeRunner([['START', flakyNode]]); + for await (const _ of runner.runAsync(ctx)) { + /* consume events */ + } + + expect(attempts).toBe(3); + expect(ctx.agentStates['exec_node_flaky_node'].status).toBe( + NodeStatus.COMPLETED, + ); + expect(ctx.agentStates['exec_node_flaky_node'].outputPayload).toBe( + 'SUCCESS_AFTER_RETRY', + ); + }); +}); diff --git a/core/test/workflow/workflow_agent_test.ts b/core/test/workflow/workflow_agent_test.ts new file mode 100644 index 000000000..a3a786c2d --- /dev/null +++ b/core/test/workflow/workflow_agent_test.ts @@ -0,0 +1,141 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {describe, expect, it, vi} from 'vitest'; +import {BaseAgent} from '../../src/agents/base_agent.js'; +import {InvocationContext} from '../../src/agents/invocation_context.js'; + +import {PluginManager} from '../../src/plugins/plugin_manager.js'; +import {Session} from '../../src/sessions/session.js'; +import { + FunctionNode, + isWorkflow, + NodeStatus, + Workflow, +} from '../../src/workflow/index.js'; + +describe('Workflow Agent Orchestrator (`Workflow`)', () => { + function createTestContext( + params?: Partial, + ): InvocationContext { + const session: Session = { + id: 'session-wf', + appName: 'test-app', + userId: 'test-user', + events: [], + state: {}, + }; + + return new InvocationContext({ + invocationId: 'inv-wf', + session, + agent: { + name: 'mock_parent', + runAsync: async function* () {}, + } as unknown as BaseAgent, + pluginManager: new PluginManager(), + ...params, + }); + } + + it('should identify Workflow instances accurately via isWorkflow type guard', () => { + const wf = new Workflow({ + name: 'test_wf', + edges: [['START', new FunctionNode('a', () => 'a')]], + }); + + expect(isWorkflow(wf)).toBe(true); + expect(isWorkflow({name: 'not_wf'})).toBe(false); + }); + + it('should throw an error if both or neither of edges and dynamicEntry are defined', () => { + expect(() => new Workflow({name: 'empty_wf'})).toThrowError( + /must define either "edges"/i, + ); + + expect( + () => + new Workflow({ + name: 'conflicted_wf', + edges: [['START', new FunctionNode('a', () => 'a')]], + dynamicEntry: async () => 'b', + }), + ).toThrowError(/cannot have both "edges" and "dynamicEntry" defined/i); + }); + + it('should run a static graph Workflow from runAsync and mark endOfAgents upon completion', async () => { + const ctx = createTestContext(); + const nodeA = new FunctionNode('step_first', () => 'First'); + const nodeB = new FunctionNode( + 'step_second', + (_ctx, input: string) => `${input}_Second`, + ); + + const wf = new Workflow({ + name: 'static_wf', + edges: [['START', nodeA, nodeB]], + outputKey: 'wfResult', + }); + + for await (const _ of wf.runAsync(ctx)) { + /* consume events */ + } + + expect(ctx.agentStates['exec_node_step_first'].status).toBe( + NodeStatus.COMPLETED, + ); + expect(ctx.agentStates['exec_node_step_second'].status).toBe( + NodeStatus.COMPLETED, + ); + expect(ctx.agentStates['wfResult']).toEqual({ + step_first: 'First', + step_second: 'First_Second', + }); + expect(ctx.endOfAgents['static_wf']).toBe(true); + }); + + it('should run a dynamic Workflow from runAsync and mark endOfAgents upon completion', async () => { + const ctx = createTestContext(); + const spy = vi.fn(async (_ctx, input: number) => input * 5); + const dynNode = new FunctionNode('dyn_mul', spy); + + const wf = new Workflow({ + name: 'dynamic_wf', + dynamicEntry: dynNode, + outputKey: 'dynOut', + }); + + // Provide initial input via userContent text + ctx.userContent = {role: 'user', parts: [{text: '10'}]}; + + for await (const _ of wf.runAsync(ctx)) { + /* consume events */ + } + + expect(spy).toHaveBeenCalledTimes(1); + expect(spy.mock.calls[0][1]).toEqual({role: 'user', parts: [{text: '10'}]}); + expect(ctx.agentStates['dynOut']).toBeDefined(); + expect(ctx.endOfAgents['dynamic_wf']).toBe(true); + }); + + it('should skip execution if context.endOfAgents already marks the workflow as true', async () => { + const ctx = createTestContext(); + const spy = vi.fn(() => 'should_not_run'); + const wf = new Workflow({ + name: 'already_done_wf', + edges: [['START', new FunctionNode('step_never', spy)]], + }); + + ctx.endOfAgents['already_done_wf'] = true; + + for await (const _ of wf.runAsync(ctx)) { + /* consume events */ + } + + expect(spy).not.toHaveBeenCalled(); + expect(ctx.agentStates['exec_node_step_never']).toBeUndefined(); + }); +}); From bb615bf725bdd4ac7a928b1f0f71d2a1dfe5f257 Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Wed, 15 Jul 2026 14:17:04 -0700 Subject: [PATCH 02/41] test(workflows): add comprehensive workflow integration tests and parity enhancements (#490) --- core/src/workflow/index.ts | 2 +- core/src/workflow/node_runner.ts | 30 ++- core/src/workflow/nodes/join_node.ts | 22 +- core/src/workflow/trigger.ts | 26 ++- core/src/workflow/utils/graph_parser.ts | 117 +++++++--- .../workflows/dynamic_nodes_workflow_test.ts | 105 +++++++++ .../workflows/fan_out_fan_in_workflow_test.ts | 77 +++++++ .../workflows/loop_workflow_test.ts | 121 ++++++++++ .../workflows/nested_workflow_test.ts | 88 ++++++++ .../workflows/node_as_tool_workflow_test.ts | 116 ++++++++++ .../workflows/route_workflow_test.ts | 207 ++++++++++++++++++ .../workflows/sequence_workflow_test.ts | 157 +++++++++++++ 12 files changed, 1016 insertions(+), 52 deletions(-) create mode 100644 tests/integration/workflows/dynamic_nodes_workflow_test.ts create mode 100644 tests/integration/workflows/fan_out_fan_in_workflow_test.ts create mode 100644 tests/integration/workflows/loop_workflow_test.ts create mode 100644 tests/integration/workflows/nested_workflow_test.ts create mode 100644 tests/integration/workflows/node_as_tool_workflow_test.ts create mode 100644 tests/integration/workflows/route_workflow_test.ts create mode 100644 tests/integration/workflows/sequence_workflow_test.ts diff --git a/core/src/workflow/index.ts b/core/src/workflow/index.ts index 6668ad581..9121a86ba 100644 --- a/core/src/workflow/index.ts +++ b/core/src/workflow/index.ts @@ -24,7 +24,7 @@ export {ToolNode} from './nodes/tool_node.js'; export {runInParallel, type ParallelRunOptions} from './parallel_worker.js'; export {normalizeRetryConfig, type RetryConfig} from './retry_config.js'; export {runNode, type RunNodeOptions} from './run_node.js'; -export {Trigger, type TriggerPredicate} from './trigger.js'; +export {DEFAULT_ROUTE, Trigger, type TriggerPredicate} from './trigger.js'; export { ParsedGraph, parseGraph, diff --git a/core/src/workflow/node_runner.ts b/core/src/workflow/node_runner.ts index ae3e02a6a..1877d633a 100644 --- a/core/src/workflow/node_runner.ts +++ b/core/src/workflow/node_runner.ts @@ -207,13 +207,29 @@ export class NodeRunner { // 3. Evaluate outgoing edges and enqueue successors whose triggers are satisfied const outgoingEdges = this.graph.adjacencyList.get(item.node.name) || []; + let hasRoutingEdges = false; + let matchedSpecificRoute = false; + let defaultRouteEdge: (typeof outgoingEdges)[0] | undefined; + for (const edge of outgoingEdges) { - let triggerSatisfied = true; - if (edge.trigger) { - triggerSatisfied = await edge.trigger.evaluate(ctx, nodeOutput); + if (!edge.trigger) { + queue.push({ + node: edge.target, + inputPayload: nodeOutput, + sourceNodeName: item.node.name, + }); + continue; + } + + hasRoutingEdges = true; + if (edge.trigger.isDefaultRoute()) { + defaultRouteEdge = edge; + continue; } + const triggerSatisfied = await edge.trigger.evaluate(ctx, nodeOutput); if (triggerSatisfied) { + matchedSpecificRoute = true; queue.push({ node: edge.target, inputPayload: nodeOutput, @@ -221,6 +237,14 @@ export class NodeRunner { }); } } + + if (hasRoutingEdges && !matchedSpecificRoute && defaultRouteEdge) { + queue.push({ + node: defaultRouteEdge.target, + inputPayload: nodeOutput, + sourceNodeName: item.node.name, + }); + } } if (this.options.outputKey) { diff --git a/core/src/workflow/nodes/join_node.ts b/core/src/workflow/nodes/join_node.ts index 595572aad..ddcfd772f 100644 --- a/core/src/workflow/nodes/join_node.ts +++ b/core/src/workflow/nodes/join_node.ts @@ -40,27 +40,25 @@ export class JoinNode< readonly upstreamCount: number; readonly predecessors?: string[]; - constructor(name: string, options: JoinNodeOptions) { - if ( - !options || - typeof options.upstreamCount !== 'number' || - options.upstreamCount < 1 - ) { + constructor(name: string, options?: Partial) { + const count = options?.upstreamCount ?? 0; + if (count < 0) { throw new Error( - `JoinNode "${name}" requires a valid upstreamCount >= 1.`, + `JoinNode "${name}" requires a valid upstreamCount >= 0.`, ); } if ( - options.predecessors && - options.predecessors.length !== options.upstreamCount + options?.predecessors && + options.predecessors.length !== count && + count > 0 ) { throw new Error( - `JoinNode "${name}" upstreamCount (${options.upstreamCount}) does not match predecessors.length (${options.predecessors.length}).`, + `JoinNode "${name}" upstreamCount (${count}) does not match predecessors.length (${options.predecessors.length}).`, ); } super(name, options); - this.upstreamCount = options.upstreamCount; - this.predecessors = options.predecessors; + this.upstreamCount = count; + this.predecessors = options?.predecessors; } /** diff --git a/core/src/workflow/trigger.ts b/core/src/workflow/trigger.ts index 9a713c26d..f0da3eb9a 100644 --- a/core/src/workflow/trigger.ts +++ b/core/src/workflow/trigger.ts @@ -7,6 +7,11 @@ import {InvocationContext} from '../agents/invocation_context.js'; import {Event} from '../events/event.js'; +/** + * The default route key used as fallback when no conditional routes match. + */ +export const DEFAULT_ROUTE = '__DEFAULT__'; + /** * A predicate function evaluated against the event or output yielded by an upstream node. * @param ctx The current invocation context. @@ -34,18 +39,31 @@ export class Trigger { this.predicate = options.predicate; } + /** + * Whether this trigger represents the fallback default route (`DEFAULT_ROUTE`). + */ + isDefaultRoute(): boolean { + return this.routeKey === DEFAULT_ROUTE; + } + /** * Creates a route matching trigger. * The trigger evaluates to true if the emitted `Event.actions.route` (or event payload route) matches `routeKey`. * @param routeKey The exact string route key to match. */ - static fromRoute(routeKey: string): Trigger { - if (!routeKey || typeof routeKey !== 'string') { + static fromRoute(routeKey: string | symbol): Trigger { + if (!routeKey && typeof routeKey !== 'symbol') { throw new Error( - 'Trigger.fromRoute requires a non-empty string routeKey.', + 'Trigger.fromRoute requires a non-empty string or symbol routeKey.', ); } - return new Trigger({routeKey}); + const keyStr = + typeof routeKey === 'symbol' + ? routeKey === Symbol.for('DEFAULT_ROUTE') + ? DEFAULT_ROUTE + : routeKey.description || String(routeKey) + : routeKey; + return new Trigger({routeKey: keyStr}); } /** diff --git a/core/src/workflow/utils/graph_parser.ts b/core/src/workflow/utils/graph_parser.ts index 566303972..6835f0ba1 100644 --- a/core/src/workflow/utils/graph_parser.ts +++ b/core/src/workflow/utils/graph_parser.ts @@ -13,13 +13,16 @@ import {Trigger} from '../trigger.js'; export type EdgeElement = | string // e.g., "START" | BaseNode // concrete node instance - | Record // route map: { ROUTE_A: nodeA, ROUTE_B: nodeB } - | [Trigger, BaseNode]; // conditional tuple: [trigger, targetNode] + | Record // route map: { ROUTE_A: nodeA, ROUTE_B: nodeB } + | [Trigger, BaseNode] // conditional tuple: [trigger, targetNode] + | BaseNode[] // fan-out node array: [nodeA, nodeB] + | readonly BaseNode[]; // fan-out node array /** * A workflow graph edge definition. * Examples: * ["START", nodeA, nodeB, nodeC] + * ["START", [nodeA, nodeB], joinNode] * [routerNode, { ROUTE_X: nodeC, ROUTE_Y: nodeD }] * [nodeA, [Trigger.fromPredicate(...), nodeB]] */ @@ -107,38 +110,60 @@ export function parseGraph(edges: GraphEdge[]): ParsedGraph { const current = edgeSeq[i]; const next = edgeSeq[i + 1]; - // Resolve source name - let sourceName: string; - if (typeof current === 'string' && current === 'START') { - sourceName = 'START'; - } else if (isBaseNode(current)) { - graph.addNode(current); - sourceName = current.name; - } else { - throw new Error( - `Invalid source element at index ${i} in edge sequence. Must be "START" or a BaseNode instance.`, - ); - } - - // Resolve target(s) - if (isBaseNode(next)) { - graph.addEdge(sourceName, next); - } else if (isRouteMap(next)) { - for (const [routeKey, targetNode] of Object.entries(next)) { - if (!isBaseNode(targetNode)) { - throw new Error( - `Target for route "${routeKey}" from source "${sourceName}" must be a BaseNode instance.`, - ); + const sources = flattenSource(current, graph, i); + for (const sourceName of sources) { + if (isBaseNode(next)) { + graph.addEdge(sourceName, next); + } else if (Array.isArray(next) && !isTriggerTuple(next)) { + for (const targetNode of next) { + if (!isBaseNode(targetNode)) { + throw new Error( + `All elements in target fan-out array from source "${sourceName}" must be BaseNode instances.`, + ); + } + graph.addEdge(sourceName, targetNode); + } + } else if (isRouteMap(next)) { + const entries: [string | symbol, BaseNode][] = [ + ...Object.entries(next), + ...Object.getOwnPropertySymbols(next).map( + (sym) => + [sym, (next as Record)[sym]] as [ + string | symbol, + BaseNode, + ], + ), + ]; + for (const [routeKey, targetNode] of entries) { + if (!isBaseNode(targetNode)) { + throw new Error( + `Target for route "${String(routeKey)}" from source "${sourceName}" must be a BaseNode instance.`, + ); + } + graph.addEdge(sourceName, targetNode, Trigger.fromRoute(routeKey)); } - graph.addEdge(sourceName, targetNode, Trigger.fromRoute(routeKey)); + } else if (isTriggerTuple(next)) { + const [trigger, targetNode] = next; + graph.addEdge(sourceName, targetNode, trigger); + } else { + throw new Error( + `Invalid target element at index ${i + 1} from source "${sourceName}". Must be a BaseNode, BaseNode array, route dictionary, or [Trigger, BaseNode] tuple.`, + ); } - } else if (isTriggerTuple(next)) { - const [trigger, targetNode] = next; - graph.addEdge(sourceName, targetNode, trigger); - } else { - throw new Error( - `Invalid target element at index ${i + 1} from source "${sourceName}". Must be a BaseNode, route dictionary, or [Trigger, BaseNode] tuple.`, - ); + } + } + } + + for (const [nodeName, node] of graph.nodes.entries()) { + if ( + node && + node.constructor && + node.constructor.name === 'JoinNode' && + (node as unknown as {upstreamCount: number}).upstreamCount === 0 + ) { + const count = graph.inboundCounts.get(nodeName) || 0; + if (count >= 1) { + (node as unknown as {upstreamCount: number}).upstreamCount = count; } } } @@ -146,6 +171,34 @@ export function parseGraph(edges: GraphEdge[]): ParsedGraph { return graph; } +function flattenSource( + element: EdgeElement, + graph: ParsedGraph, + index: number, +): string[] { + if (typeof element === 'string' && element === 'START') { + return ['START']; + } + if (isBaseNode(element)) { + graph.addNode(element); + return [element.name]; + } + if (Array.isArray(element) && !isTriggerTuple(element)) { + return element.map((node, idx) => { + if (!isBaseNode(node)) { + throw new Error( + `Invalid source element inside array at index ${index}[${idx}]. Must be a BaseNode instance.`, + ); + } + graph.addNode(node); + return node.name; + }); + } + throw new Error( + `Invalid source element at index ${index} in edge sequence. Must be "START", a BaseNode, or an array of BaseNode instances.`, + ); +} + function isBaseNode(obj: unknown): obj is BaseNode { return ( typeof obj === 'object' && diff --git a/tests/integration/workflows/dynamic_nodes_workflow_test.ts b/tests/integration/workflows/dynamic_nodes_workflow_test.ts new file mode 100644 index 000000000..99b4b775e --- /dev/null +++ b/tests/integration/workflows/dynamic_nodes_workflow_test.ts @@ -0,0 +1,105 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { + createEvent, + Event, + FunctionNode, + InMemoryRunner, + InvocationContext, + Workflow, +} from '@google/adk'; +import {describe, expect, it} from 'vitest'; + +describe('Workflow Samples: Dynamic Nodes & Dynamic Fan-Out', () => { + it('should run a dynamic entry workflow scheduling downstream nodes via ctx.runNode (dynamic_nodes sample parity)', async () => { + const generatorNode = new FunctionNode( + 'generator', + (_ctx, topic: string) => `Catchy headline about ${topic}`, + ); + + const evaluatorNode = new FunctionNode( + 'evaluator', + (_ctx, headline: string) => `Evaluated: [${headline}] - Grade A`, + ); + + const dynamicEntryNode = new FunctionNode( + 'orchestrate', + async (ctx: InvocationContext, topic: string) => { + // Execute generator using ctx.runNode + const genOutput = await ctx.runNode(generatorNode, topic); + // Pass output to evaluator using ctx.runNode + const evalOutput = await ctx.runNode(evaluatorNode, genOutput); + return createEvent({message: `Final report: ${evalOutput}`}); + }, + ); + + const rootAgent = new Workflow({ + name: 'dynamic_nodes_workflow', + edges: [['START', dynamicEntryNode]], + }); + + const runner = new InMemoryRunner({agent: rootAgent}); + const events: Event[] = []; + + for await (const event of runner.runEphemeral({ + userId: 'test_user', + newMessage: {role: 'user', parts: [{text: 'AI Innovations'}]}, + })) { + events.push(event); + } + + const messages = events + .flatMap((e) => e.content?.parts?.map((p) => p.text) ?? []) + .join(''); + expect(messages).toContain( + 'Final report: Evaluated: [Catchy headline about AI Innovations] - Grade A', + ); + }); + + it('should perform dynamic fan-out and fan-in across items (dynamic_fan_out_fan_in sample parity)', async () => { + const processTopicNode = new FunctionNode( + 'process_topic', + (_ctx, topic: string) => `Processed: ${topic.trim().toUpperCase()}`, + ); + + const dynamicOrchestrator = new FunctionNode( + 'orchestrate_fan_out', + async (ctx: InvocationContext, input: string) => { + const topics = input.split(',').map((t) => t.trim()); + // Dynamic fan-out executing multiple nodes in parallel via Promise.all with ctx.runNode + const results = await Promise.all( + topics.map((topic) => ctx.runNode(processTopicNode, topic)), + ); + return createEvent({ + message: `Aggregated Topics: ${results.join(' | ')}`, + }); + }, + ); + + const rootAgent = new Workflow({ + name: 'dynamic_fanout_workflow', + edges: [['START', dynamicOrchestrator]], + }); + + const runner = new InMemoryRunner({agent: rootAgent}); + const events: Event[] = []; + + for await (const event of runner.runEphemeral({ + userId: 'test_user', + newMessage: {role: 'user', parts: [{text: 'apple, banana, cherry'}]}, + })) { + events.push(event); + } + + const messages = events + .flatMap((e) => e.content?.parts?.map((p) => p.text) ?? []) + .join(''); + expect(messages).toContain( + 'Aggregated Topics: Processed: APPLE | Processed: BANANA | Processed: CHERRY', + ); + }); +}); diff --git a/tests/integration/workflows/fan_out_fan_in_workflow_test.ts b/tests/integration/workflows/fan_out_fan_in_workflow_test.ts new file mode 100644 index 000000000..37b8ac0c7 --- /dev/null +++ b/tests/integration/workflows/fan_out_fan_in_workflow_test.ts @@ -0,0 +1,77 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { + createEvent, + Event, + FunctionNode, + InMemoryRunner, + JoinNode, + Workflow, +} from '@google/adk'; +import {describe, expect, it} from 'vitest'; + +describe('Workflow Samples: Fan Out Fan In with JoinNode', () => { + it('should run parallel nodes and aggregate outputs at JoinNode (fan_out_fan_in sample parity)', async () => { + const makeUppercase = new FunctionNode( + 'make_uppercase', + (_ctx, input: string) => input.toUpperCase(), + ); + + const countCharacters = new FunctionNode( + 'count_characters', + (_ctx, input: string) => input.length, + ); + + const reverseString = new FunctionNode( + 'reverse_string', + (_ctx, input: string) => input.split('').reverse().join(''), + ); + + const joinNode = new JoinNode('join_for_results'); + + const aggregateNode = new FunctionNode( + 'aggregate', + (_ctx, input: Record) => { + return createEvent({ + message: + `Uppercase: ${input['make_uppercase']}\n\n` + + `Character Count: ${input['count_characters']}\n\n` + + `Reversed: ${input['reverse_string']}\n\n`, + }); + }, + ); + + const rootAgent = new Workflow({ + name: 'fan_out_fan_in_workflow', + edges: [ + [ + 'START', + [makeUppercase, countCharacters, reverseString], + joinNode, + aggregateNode, + ], + ], + }); + + const runner = new InMemoryRunner({agent: rootAgent}); + const events: Event[] = []; + + for await (const event of runner.runEphemeral({ + userId: 'test_user', + newMessage: {role: 'user', parts: [{text: 'adk workflow'}]}, + })) { + events.push(event); + } + + const messages = events + .flatMap((e) => e.content?.parts?.map((p) => p.text) ?? []) + .join(''); + expect(messages).toContain('Uppercase: ADK WORKFLOW'); + expect(messages).toContain('Character Count: 12'); + expect(messages).toContain('Reversed: wolfkrow kda'); + }); +}); diff --git a/tests/integration/workflows/loop_workflow_test.ts b/tests/integration/workflows/loop_workflow_test.ts new file mode 100644 index 000000000..90a6131fe --- /dev/null +++ b/tests/integration/workflows/loop_workflow_test.ts @@ -0,0 +1,121 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { + createEvent, + Event, + FunctionNode, + InMemoryRunner, + Workflow, +} from '@google/adk'; +import {describe, expect, it} from 'vitest'; + +describe('Workflow Samples: Loop & Loop Self', () => { + it('should support multi-node looping with conditional exit (loop sample parity)', async () => { + const processInput = new FunctionNode( + 'process_input', + (_ctx, input: string) => + createEvent({state: {topic: input, attempts: 0}}), + ); + + const generateHeadline = new FunctionNode('generate_headline', (ctx) => { + const attempts = ((ctx.session.state['attempts'] as number) || 0) + 1; + ctx.session.state['attempts'] = attempts; + const topic = ctx.session.state['topic'] as string; + const headline = + attempts === 1 + ? `General news about ${topic}` + : `Tech Breakthrough in ${topic}`; + return createEvent({state: {currentHeadline: headline}}); + }); + + const evaluateHeadline = new FunctionNode('evaluate_headline', (ctx) => { + const headline = ctx.session.state['currentHeadline'] as string; + const isTech = headline.includes('Tech'); + return {grade: isTech ? 'tech-related' : 'unrelated', headline}; + }); + + const routeHeadline = new FunctionNode( + 'route_headline', + (_ctx, input: {grade: string}) => createEvent({route: input.grade}), + ); + + const rootAgent = new Workflow({ + name: 'loop_workflow', + edges: [ + [ + 'START', + processInput, + generateHeadline, + evaluateHeadline, + routeHeadline, + ], + [routeHeadline, {unrelated: generateHeadline}], + ], + outputKey: 'loopResult', + allowCycles: true, + }); + + const runner = new InMemoryRunner({agent: rootAgent}); + const events: Event[] = []; + + for await (const event of runner.runEphemeral({ + userId: 'test_user', + newMessage: {role: 'user', parts: [{text: 'Software Engineering'}]}, + })) { + events.push(event); + } + + // It should have cycled: attempts should be 2 when it exits via "tech-related" (no route handler -> end) + expect(events.length).toBeGreaterThanOrEqual(1); + const finalEvent = events[events.length - 1]; + expect(finalEvent).toBeDefined(); + }); + + it('should support a node looping back to itself (loop_self sample parity)', async () => { + let guessCount = 0; + const guessNode = new FunctionNode('guess_node', () => { + guessCount++; + if (guessCount < 3) { + return createEvent({ + message: `Guess ${guessCount}: wrong`, + route: 'guessed_wrong', + }); + } + return createEvent({ + message: `Guess ${guessCount}: correct!`, + route: 'guessed_right', + }); + }); + + const rootAgent = new Workflow({ + name: 'loop_self_workflow', + edges: [ + ['START', guessNode], + [guessNode, {guessed_wrong: guessNode}], + ], + allowCycles: true, + }); + + const runner = new InMemoryRunner({agent: rootAgent}); + const events: Event[] = []; + + for await (const event of runner.runEphemeral({ + userId: 'test_user', + newMessage: {role: 'user', parts: [{text: 'Guess a number'}]}, + })) { + events.push(event); + } + + expect(guessCount).toBe(3); + const messages = events + .flatMap((e) => e.content?.parts?.map((p) => p.text) ?? []) + .filter(Boolean); + expect(messages).toContain('Guess 1: wrong'); + expect(messages).toContain('Guess 2: wrong'); + expect(messages).toContain('Guess 3: correct!'); + }); +}); diff --git a/tests/integration/workflows/nested_workflow_test.ts b/tests/integration/workflows/nested_workflow_test.ts new file mode 100644 index 000000000..ff54e20e5 --- /dev/null +++ b/tests/integration/workflows/nested_workflow_test.ts @@ -0,0 +1,88 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { + createEvent, + Event, + FunctionNode, + InMemoryRunner, + JoinNode, + Workflow, +} from '@google/adk'; +import {describe, expect, it} from 'vitest'; + +describe('Workflow Samples: Nested Workflow Composition', () => { + it('should compose a sub-Workflow inside a parent Workflow (nested_workflow sample parity)', async () => { + const findNameNode = new FunctionNode( + 'find_name', + (_ctx, input: string) => `Person_${input}`, + ); + + const generateBioNode = new FunctionNode( + 'generate_bio', + (_ctx, name: string) => `Bio for ${name}: Famous historical author.`, + ); + + // Sub-workflow wraps two sequential steps + const findFamousPersonWorkflow = new Workflow({ + name: 'find_famous_person_workflow', + edges: [['START', findNameNode, generateBioNode]], + outputKey: 'famousPersonResult', + }); + + const findHistoricalEventNode = new FunctionNode( + 'find_historical_event', + (_ctx, input: string) => + `Historical event in year ${input}: First publication of landmark novel.`, + ); + + const joinNode = new JoinNode('join_for_aggregation'); + + const formatOutputNode = new FunctionNode( + 'format_output', + (_ctx, input: Record) => { + // When a Workflow node completes inside a parent graph, its final output payload is passed to the join node + return createEvent({ + message: + `Person Bio: ${JSON.stringify(input['find_famous_person_workflow'])}\n` + + `Event: ${input['find_historical_event']}`, + }); + }, + ); + + const rootAgent = new Workflow({ + name: 'nested_root_workflow', + edges: [ + [ + 'START', + [findFamousPersonWorkflow, findHistoricalEventNode], + joinNode, + formatOutputNode, + ], + ], + }); + + const runner = new InMemoryRunner({agent: rootAgent}); + const events: Event[] = []; + + for await (const event of runner.runEphemeral({ + userId: 'test_user', + newMessage: {role: 'user', parts: [{text: '1984'}]}, + })) { + events.push(event); + } + + const messages = events + .flatMap((e) => e.content?.parts?.map((p) => p.text) ?? []) + .join(''); + expect(messages).toContain( + 'Bio for Person_1984: Famous historical author.', + ); + expect(messages).toContain( + 'Historical event in year 1984: First publication of landmark novel.', + ); + }); +}); diff --git a/tests/integration/workflows/node_as_tool_workflow_test.ts b/tests/integration/workflows/node_as_tool_workflow_test.ts new file mode 100644 index 000000000..2d7e4b716 --- /dev/null +++ b/tests/integration/workflows/node_as_tool_workflow_test.ts @@ -0,0 +1,116 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { + AgentTool, + createEvent, + FunctionNode, + LlmAgent, + Workflow, +} from '@google/adk'; +import {describe, it} from 'vitest'; +import {runTestCase} from '../test_case_utils.js'; + +describe('Workflow Samples: Node / Workflow as Tool', () => { + it('should allow an LlmAgent to call a Workflow wrapped via AgentTool (node_as_tool sample parity)', async () => { + const lookupDbNode = new FunctionNode( + 'lookup_db', + (_ctx, customerId: string) => { + if (customerId === 'C123') { + return JSON.stringify({name: 'Jane Doe', tier: 'Gold', balance: 450}); + } + return JSON.stringify({error: 'Customer not found'}); + }, + ); + + const customerLookupWorkflow = new Workflow({ + name: 'customer_lookup_workflow', + edges: [['START', lookupDbNode]], + }); + + const workflowTool = new AgentTool({ + agent: customerLookupWorkflow, + }); + + const rootAgent = new LlmAgent({ + name: 'customer_service_agent', + instruction: + 'Use the lookup_customer tool when asked for customer details.', + tools: [workflowTool], + }); + + await runTestCase({ + agent: rootAgent, + turns: [ + { + userPrompt: 'Can you check details for customer C123?', + expectedEvents: [ + createEvent({ + author: 'customer_service_agent', + content: { + role: 'model', + parts: [ + { + functionCall: { + id: 'call_lookup_1', + name: 'lookup_customer', + args: {input: 'C123'}, + }, + }, + ], + }, + }), + createEvent({ + author: 'customer_service_agent', + content: { + role: 'model', + parts: [ + { + text: 'Customer C123 is Jane Doe with Gold tier and a balance of $450.', + }, + ], + }, + }), + ], + }, + ], + modelResponses: [ + { + candidates: [ + { + content: { + role: 'model', + parts: [ + { + functionCall: { + id: 'call_lookup_1', + name: 'lookup_customer', + args: {input: 'C123'}, + }, + }, + ], + }, + }, + ], + }, + { + candidates: [ + { + content: { + role: 'model', + parts: [ + { + text: 'Customer C123 is Jane Doe with Gold tier and a balance of $450.', + }, + ], + }, + }, + ], + }, + ], + }); + }); +}); diff --git a/tests/integration/workflows/route_workflow_test.ts b/tests/integration/workflows/route_workflow_test.ts new file mode 100644 index 000000000..6a92f20c9 --- /dev/null +++ b/tests/integration/workflows/route_workflow_test.ts @@ -0,0 +1,207 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { + createEvent, + DEFAULT_ROUTE, + Event, + FunctionNode, + InMemoryRunner, + LlmAgent, + Workflow, +} from '@google/adk'; +import {describe, expect, it} from 'vitest'; +import {runTestCase} from '../test_case_utils.js'; + +describe('Workflow Samples: Route & Agent in Workflow with DEFAULT_ROUTE', () => { + it('should route conditionally based on event route (route sample parity)', async () => { + const classifyInputNode = new FunctionNode( + 'classify_input', + (_ctx, input: string) => { + const isQuestion = input.endsWith('?'); + return createEvent({ + route: isQuestion ? 'question' : 'statement', + }); + }, + ); + + const answerQuestionAgent = new LlmAgent({ + name: 'answer_question', + instruction: 'Answer the user question concisely.', + }); + + const commentOnStatementAgent = new LlmAgent({ + name: 'comment_on_statement', + instruction: 'Provide a brief comment on the statement.', + }); + + const rootAgent = new Workflow({ + name: 'router_workflow', + edges: [ + ['START', classifyInputNode], + [ + classifyInputNode, + { + question: answerQuestionAgent, + statement: commentOnStatementAgent, + }, + ], + ], + }); + + // Test question route + await runTestCase({ + agent: rootAgent, + turns: [ + { + userPrompt: 'What is ADK?', + expectedEvents: [ + createEvent({ + author: 'answer_question', + content: { + role: 'model', + parts: [{text: 'ADK is the Agent Development Kit.'}], + }, + }), + ], + }, + ], + modelResponses: [ + { + candidates: [ + { + content: { + role: 'model', + parts: [{text: 'ADK is the Agent Development Kit.'}], + }, + }, + ], + }, + ], + }); + + // Test statement route + await runTestCase({ + agent: rootAgent, + turns: [ + { + userPrompt: 'ADK supports workflows and routing.', + expectedEvents: [ + createEvent({ + author: 'comment_on_statement', + content: { + role: 'model', + parts: [{text: 'That is correct and very powerful!'}], + }, + }), + ], + }, + ], + modelResponses: [ + { + candidates: [ + { + content: { + role: 'model', + parts: [{text: 'That is correct and very powerful!'}], + }, + }, + ], + }, + ], + }); + }); + + it('should support DEFAULT_ROUTE fallback when specific route does not match (agent_in_workflow sample parity)', async () => { + const checkIdentityNode = new FunctionNode( + 'check_identity', + (_ctx, name: string) => { + if (name.toLowerCase() !== 'jane doe') { + return createEvent({ + message: `Could not find matching records for ${name}. Let's try again.`, + route: 'retry', + }); + } + return createEvent({ + message: `Hello ${name}! Let me look up your orders.`, + }); + }, + ); + + const retryHandlerNode = new FunctionNode('retry_handler', () => + createEvent({message: 'Retrying identity check...'}), + ); + + const generateInstructionAgent = new LlmAgent({ + name: 'generate_instruction', + instruction: 'Generate preparation instruction for Jane Doe.', + }); + + const rootAgent = new Workflow({ + name: 'agent_in_workflow', + edges: [ + ['START', checkIdentityNode], + [ + checkIdentityNode, + { + retry: retryHandlerNode, + [DEFAULT_ROUTE]: generateInstructionAgent, + }, + ], + ], + }); + + // Test when route="retry" matches specific route in routing table + const runnerRetry = new InMemoryRunner({agent: rootAgent}); + const retryEvents: Event[] = []; + for await (const event of runnerRetry.runEphemeral({ + userId: 'user1', + newMessage: {role: 'user', parts: [{text: 'John Smith'}]}, + })) { + retryEvents.push(event); + } + expect( + retryEvents.some( + (e) => e.content?.parts?.[0].text === 'Retrying identity check...', + ), + ).toBe(true); + + // Test when no route is yielded, taking DEFAULT_ROUTE fallback + await runTestCase({ + agent: rootAgent, + turns: [ + { + userPrompt: 'Jane Doe', + expectedEvents: [ + createEvent({ + author: 'generate_instruction', + content: { + role: 'model', + parts: [ + {text: 'Please fast for 12 hours before your lipid panel.'}, + ], + }, + }), + ], + }, + ], + modelResponses: [ + { + candidates: [ + { + content: { + role: 'model', + parts: [ + {text: 'Please fast for 12 hours before your lipid panel.'}, + ], + }, + }, + ], + }, + ], + }); + }); +}); diff --git a/tests/integration/workflows/sequence_workflow_test.ts b/tests/integration/workflows/sequence_workflow_test.ts new file mode 100644 index 000000000..0d06b84e1 --- /dev/null +++ b/tests/integration/workflows/sequence_workflow_test.ts @@ -0,0 +1,157 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { + createEvent, + Event, + FunctionNode, + InMemoryRunner, + LlmAgent, + Workflow, +} from '@google/adk'; +import {describe, expect, it} from 'vitest'; +import {runTestCase} from '../test_case_utils.js'; + +describe('Workflow Samples: Sequence, Message & State', () => { + it('should run sequential workflow with LLM agents (sequence sample parity)', async () => { + const generateFruitAgent = new LlmAgent({ + name: 'generate_fruit_agent', + instruction: + 'Return the name of a random fruit. Return only the name, nothing else.', + }); + + const generateBenefitAgent = new LlmAgent({ + name: 'generate_benefit_agent', + instruction: 'Tell me a health benefit about the specified fruit.', + }); + + const rootAgent = new Workflow({ + name: 'root_agent', + edges: [['START', generateFruitAgent, generateBenefitAgent]], + }); + + await runTestCase({ + agent: rootAgent, + turns: [ + { + userPrompt: 'Tell me about a fruit.', + expectedEvents: [ + createEvent({ + author: 'generate_fruit_agent', + content: {role: 'model', parts: [{text: 'Apple'}]}, + }), + createEvent({ + author: 'generate_benefit_agent', + content: { + role: 'model', + parts: [{text: 'Apples are rich in fiber and vitamin C.'}], + }, + }), + ], + }, + ], + modelResponses: [ + {candidates: [{content: {role: 'model', parts: [{text: 'Apple'}]}}]}, + { + candidates: [ + { + content: { + role: 'model', + parts: [{text: 'Apples are rich in fiber and vitamin C.'}], + }, + }, + ], + }, + ], + }); + }); + + it('should emit event message from a FunctionNode (message sample parity)', async () => { + const messageNode = new FunctionNode('emit_message', () => + createEvent({ + content: { + role: 'model', + parts: [{text: 'Hello from workflow function node!'}], + }, + }), + ); + + const rootAgent = new Workflow({ + name: 'message_workflow', + edges: [['START', messageNode]], + }); + + const runner = new InMemoryRunner({agent: rootAgent}); + const events: Event[] = []; + + for await (const event of runner.runEphemeral({ + userId: 'test_user', + newMessage: {role: 'user', parts: [{text: 'Start'}]}, + })) { + events.push(event); + } + + expect(events.length).toBeGreaterThanOrEqual(1); + const msgEvents = events.filter((e) => + e.content?.parts?.some( + (p) => p.text === 'Hello from workflow function node!', + ), + ); + expect(msgEvents.length).toBeGreaterThanOrEqual(1); + }); + + it('should read and update session state across FunctionNodes (state sample parity)', async () => { + const initNode = new FunctionNode('init_state', (ctx, input: string) => { + ctx.session.state['topic'] = input; + ctx.session.state['count'] = 1; + return createEvent({ + actions: {stateDelta: {topic: input, count: 1}}, + content: { + role: 'model', + parts: [{text: `Initialized ${input}`}], + }, + }); + }); + + const updateNode = new FunctionNode('update_state', (ctx) => { + const currentCount = (ctx.session.state['count'] as number) || 0; + const topic = (ctx.session.state['topic'] as string) || ''; + const newCount = currentCount + 1; + ctx.session.state['count'] = newCount; + ctx.session.state['lastProcessed'] = `${topic}_processed`; + return createEvent({ + actions: { + stateDelta: {count: newCount, lastProcessed: `${topic}_processed`}, + }, + content: { + role: 'model', + parts: [{text: `Processed ${topic} with count ${newCount}`}], + }, + }); + }); + + const rootAgent = new Workflow({ + name: 'state_workflow', + edges: [['START', initNode, updateNode]], + }); + + const runner = new InMemoryRunner({agent: rootAgent}); + const events: Event[] = []; + + for await (const event of runner.runEphemeral({ + userId: 'test_user', + newMessage: {role: 'user', parts: [{text: 'AI Workflows'}]}, + })) { + events.push(event); + } + + expect(events.length).toBeGreaterThanOrEqual(2); + const lastEvent = events[events.length - 1]; + expect(lastEvent.content?.parts?.[0].text).toContain( + 'Processed AI Workflows with count 2', + ); + }); +}); From e8b82888f61109f04033d89fa61906d319c28534 Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Wed, 15 Jul 2026 14:20:09 -0700 Subject: [PATCH 03/41] fix(workflow): resolve strict TypeScript build and EventAction errors (#490) --- core/src/events/event.ts | 11 +++++++++-- core/src/events/event_actions.ts | 8 ++++++++ core/src/workflow/dynamic_node_scheduler.ts | 7 ++++++- core/src/workflow/index.ts | 1 + core/src/workflow/node_runner.ts | 4 ++-- core/src/workflow/nodes/function_node.ts | 4 ++-- core/src/workflow/nodes/llm_agent_wrapper.ts | 8 +++++--- core/src/workflow/nodes/tool_node.ts | 9 +++++---- core/src/workflow/utils/rehydration_utils.ts | 2 +- core/src/workflow/utils/replay_manager.ts | 8 ++++---- 10 files changed, 43 insertions(+), 19 deletions(-) diff --git a/core/src/events/event.ts b/core/src/events/event.ts index f5fcae31c..cddc49021 100644 --- a/core/src/events/event.ts +++ b/core/src/events/event.ts @@ -64,19 +64,26 @@ export interface Event extends LlmResponse { timestamp: number; } +/** + * Parameters for creating an event with partial fields. + */ +export interface CreateEventParams extends Omit, 'actions'> { + actions?: Partial; +} + /** * Creates an event from a partial event. * * @param params The partial event to create the event from. * @returns The event. */ -export function createEvent(params: Partial = {}): Event { +export function createEvent(params: CreateEventParams = {}): Event { return { ...params, id: params.id || createNewEventId(), invocationId: params.invocationId || '', author: params.author, - actions: params.actions || createEventActions(), + actions: createEventActions(params.actions), longRunningToolIds: params.longRunningToolIds || [], branch: params.branch, timestamp: params.timestamp || Date.now(), diff --git a/core/src/events/event_actions.ts b/core/src/events/event_actions.ts index 0316ec290..3a67cfe5b 100644 --- a/core/src/events/event_actions.ts +++ b/core/src/events/event_actions.ts @@ -56,6 +56,14 @@ export interface EventActions { * call id. */ requestedToolConfirmations: {[key: string]: ToolConfirmation}; + + /** Workflow / custom event actions */ + output?: unknown; + joinCompleted?: unknown; + toolExecution?: unknown; + requestInput?: unknown; + nodeExecutionReplay?: unknown; + [key: string]: unknown; } /** diff --git a/core/src/workflow/dynamic_node_scheduler.ts b/core/src/workflow/dynamic_node_scheduler.ts index 277430539..7d20a92ee 100644 --- a/core/src/workflow/dynamic_node_scheduler.ts +++ b/core/src/workflow/dynamic_node_scheduler.ts @@ -18,9 +18,14 @@ import {FunctionNode, FunctionNodeHandler} from './nodes/function_node.js'; /** * Type for the dynamic workflow entry point. */ +export type DynamicEntryFunction< + TInput = unknown, + TOutput = unknown, +> = FunctionNodeHandler; + export type DynamicEntry = | BaseNode - | FunctionNodeHandler; + | DynamicEntryFunction; /** * Options for the DynamicNodeScheduler. diff --git a/core/src/workflow/index.ts b/core/src/workflow/index.ts index 9121a86ba..f0ec9ff2b 100644 --- a/core/src/workflow/index.ts +++ b/core/src/workflow/index.ts @@ -8,6 +8,7 @@ export {BaseNode, type BaseNodeOptions} from './base_node.js'; export { DynamicNodeScheduler, type DynamicEntry, + type DynamicEntryFunction, type DynamicNodeSchedulerOptions, } from './dynamic_node_scheduler.js'; export { diff --git a/core/src/workflow/node_runner.ts b/core/src/workflow/node_runner.ts index 1877d633a..512ea68be 100644 --- a/core/src/workflow/node_runner.ts +++ b/core/src/workflow/node_runner.ts @@ -59,9 +59,9 @@ export async function consumeGenerator( output === undefined && lastEvent?.actions && typeof lastEvent.actions === 'object' && - 'output' in (lastEvent.actions as Record) + 'output' in (lastEvent.actions as unknown as Record) ) { - output = (lastEvent.actions as Record) + output = (lastEvent.actions as unknown as Record) .output as TOutput; } return {output, isPausedHitl, lastEvent}; diff --git a/core/src/workflow/nodes/function_node.ts b/core/src/workflow/nodes/function_node.ts index 87d4b6a18..347e046fc 100644 --- a/core/src/workflow/nodes/function_node.ts +++ b/core/src/workflow/nodes/function_node.ts @@ -74,8 +74,8 @@ export class FunctionNode extends BaseNode< res.content ?? (typeof res.actions === 'object' && res.actions !== null && - 'output' in (res.actions as Record) - ? (res.actions as Record).output + 'output' in (res.actions as unknown as Record) + ? (res.actions as unknown as Record).output : res); this.lastOutputPayload = extracted; return extracted as TOutput; diff --git a/core/src/workflow/nodes/llm_agent_wrapper.ts b/core/src/workflow/nodes/llm_agent_wrapper.ts index 919d2968c..1098ff1a5 100644 --- a/core/src/workflow/nodes/llm_agent_wrapper.ts +++ b/core/src/workflow/nodes/llm_agent_wrapper.ts @@ -4,6 +4,7 @@ * SPDX-License-Identifier: Apache-2.0 */ +import {Content} from '@google/genai'; import {BaseAgent} from '../../agents/base_agent.js'; import { InvocationContext, @@ -60,7 +61,7 @@ export class LLMAgentWrapper< 'role' in input && 'parts' in input ) { - childCtxParams.userContent = input as {role: string; parts: unknown[]}; + childCtxParams.userContent = input as unknown as Content; } } @@ -77,9 +78,10 @@ export class LLMAgentWrapper< if ( event.actions && typeof event.actions === 'object' && - 'output' in (event.actions as Record) + 'output' in (event.actions as unknown as Record) ) { - lastOutput = (event.actions as Record).output; + lastOutput = (event.actions as unknown as Record) + .output; } } diff --git a/core/src/workflow/nodes/tool_node.ts b/core/src/workflow/nodes/tool_node.ts index 8169eef80..b10cf1c85 100644 --- a/core/src/workflow/nodes/tool_node.ts +++ b/core/src/workflow/nodes/tool_node.ts @@ -4,6 +4,7 @@ * SPDX-License-Identifier: Apache-2.0 */ +import {Context} from '../../agents/context.js'; import {InvocationContext} from '../../agents/invocation_context.js'; import {createEvent, Event} from '../../events/event.js'; import {BaseTool} from '../../tools/base_tool.js'; @@ -41,10 +42,10 @@ export class ToolNode< input?: TInput, ): AsyncGenerator { const params = typeof input === 'object' && input !== null ? input : {}; - const result = await this.tool.runAsync( - {invocationContext: ctx}, - params as Record, - ); + const result = await this.tool.runAsync({ + toolContext: ctx as unknown as Context, + args: params as Record, + }); const event = createEvent({ invocationId: ctx.invocationId, diff --git a/core/src/workflow/utils/rehydration_utils.ts b/core/src/workflow/utils/rehydration_utils.ts index cf0abe6d8..0538f7c73 100644 --- a/core/src/workflow/utils/rehydration_utils.ts +++ b/core/src/workflow/utils/rehydration_utils.ts @@ -56,7 +56,7 @@ export function rehydrateAgentStates( for (const event of session.events) { if (!event || typeof event !== 'object') continue; - const eventRecord = event as Record; + const eventRecord = event as unknown as Record; const actions = eventRecord.actions as | Record | undefined; diff --git a/core/src/workflow/utils/replay_manager.ts b/core/src/workflow/utils/replay_manager.ts index d182af851..3ec210fc8 100644 --- a/core/src/workflow/utils/replay_manager.ts +++ b/core/src/workflow/utils/replay_manager.ts @@ -5,10 +5,10 @@ */ import {InvocationContext} from '../../agents/invocation_context.js'; -import {Event} from '../../events/event.js'; +import {createEvent, Event} from '../../events/event.js'; import {BaseNode} from '../base_node.js'; import {generateExecutionId, getOrInitAgentStates} from '../node_runner.js'; -import {NodeState, NodeStatus, isNodeState} from '../node_state.js'; +import {isNodeState, NodeState, NodeStatus} from '../node_state.js'; /** * Manages event replay when a workflow node is bypassed due to a completed historical checkpoint (`rerunOnResume == false`). @@ -46,7 +46,7 @@ export class ReplayManager { yield event; } } else if (existingState.outputPayload !== undefined) { - yield { + yield createEvent({ invocationId: ctx.invocationId, author: node.name, branch: ctx.branch, @@ -58,7 +58,7 @@ export class ReplayManager { timestamp: existingState.timestamp, }, }, - } as Event; + }); } return true; } From 3387e357d27df2f19c3f11d53c3ede1612cc066f Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Wed, 22 Jul 2026 14:33:44 -0700 Subject: [PATCH 04/41] feat(workflow-next): Phase 0 foundations for workflow rewrite Port the leaf data types for the new BaseNode/Context workflow architecture (parity with google/adk-python): NodeStatus (Python-aligned ordinals), Python-shaped NodeState, RetryConfig with jitter/exceptions + retry utils, and workflow errors. Extend the shared Event/EventActions model additively with first-class workflow fields (output, route, nodeInfo, isolationScope, agentState, endOfAgent) that round-trip through snake_case for cross-language session compatibility. --- core/src/events/event.ts | 46 +++++ core/src/events/event_actions.ts | 19 +++ core/src/workflow-next/errors.ts | 78 +++++++++ core/src/workflow-next/node_state.ts | 77 +++++++++ core/src/workflow-next/node_status.ts | 29 ++++ core/src/workflow-next/retry_config.ts | 85 ++++++++++ core/src/workflow-next/utils/retry_utils.ts | 112 ++++++++++++ core/test/workflow-next/event_model_test.ts | 66 ++++++++ core/test/workflow-next/foundations_test.ts | 179 ++++++++++++++++++++ 9 files changed, 691 insertions(+) create mode 100644 core/src/workflow-next/errors.ts create mode 100644 core/src/workflow-next/node_state.ts create mode 100644 core/src/workflow-next/node_status.ts create mode 100644 core/src/workflow-next/retry_config.ts create mode 100644 core/src/workflow-next/utils/retry_utils.ts create mode 100644 core/test/workflow-next/event_model_test.ts create mode 100644 core/test/workflow-next/foundations_test.ts diff --git a/core/src/events/event.ts b/core/src/events/event.ts index cddc49021..c022f14a4 100644 --- a/core/src/events/event.ts +++ b/core/src/events/event.ts @@ -11,6 +11,26 @@ import {LlmResponse} from '../models/llm_response.js'; import {toCamelCase, toSnakeCase} from '../utils/object_notation_utils.js'; import {createEventActions, EventActions} from './event_actions.js'; +/** + * Workflow-node provenance attached to an event. + * + * Mirrors `google/adk-python` `Event.node_info`. Present only on events emitted + * from within a workflow node. + */ +export interface NodeInfo { + /** The workflow node path that produced this event (e.g. `wf.child.0`). */ + path?: string; + + /** The node run id this event's output should be attributed to. */ + outputFor?: string; + + /** + * Whether the event's textual content should be promoted to the node's + * structured output. + */ + messageAsOutput?: boolean; +} + /** * Represents an event in a conversation between agents and users. @@ -62,6 +82,32 @@ export interface Event extends LlmResponse { * The timestamp of the event. */ timestamp: number; + + /** + * Workflow: the structured output produced by the emitting node, if any. + * + * First-class field mirroring `google/adk-python` `Event.output`. Used by the + * workflow engine to carry a node's return value alongside its content. + */ + output?: unknown; + + /** + * Workflow: the route key emitted by a routing node, used by the graph to + * select the matching outgoing edge. Mirrors Python `Event.route`. + */ + route?: string | number | boolean; + + /** + * Workflow: provenance of the emitting node. Mirrors Python `Event.node_info`. + */ + nodeInfo?: NodeInfo; + + /** + * Workflow: scope tag used to isolate multi-agent conversations so peer + * scopes don't see each other's events. Mirrors Python + * `Event.isolation_scope`. + */ + isolationScope?: string; } /** diff --git a/core/src/events/event_actions.ts b/core/src/events/event_actions.ts index 3a67cfe5b..47711d9c8 100644 --- a/core/src/events/event_actions.ts +++ b/core/src/events/event_actions.ts @@ -63,6 +63,25 @@ export interface EventActions { toolExecution?: unknown; requestInput?: unknown; nodeExecutionReplay?: unknown; + + /** + * Workflow: a serialized node/agent state snapshot used for resumable + * checkpointing. Mirrors Python `EventActions.agent_state`. + */ + agentState?: Record; + + /** + * Workflow: marks that the emitting agent/workflow has reached the end of its + * execution for this invocation. Mirrors Python `EventActions.end_of_agent`. + */ + endOfAgent?: boolean; + + /** + * Workflow: route key selected by a routing node (alternative carrier to the + * top-level `Event.route`, used by callbacks/tools). + */ + route?: string | number | boolean; + [key: string]: unknown; } diff --git a/core/src/workflow-next/errors.ts b/core/src/workflow-next/errors.ts new file mode 100644 index 000000000..14c73a44a --- /dev/null +++ b/core/src/workflow-next/errors.ts @@ -0,0 +1,78 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Errors raised by the workflow framework. + * + * Ported from `google/adk-python` `workflow/_errors.py`. + */ + +/** + * Internal: raised when a dynamic node interrupts (HITL). + * + * Used exclusively by `ctx.runNode()` to signal that the dynamic child has + * unresolved interrupt IDs. The parent's NodeRunner catches this and reads the + * interrupt IDs from the parent's ctx (set by `ctx.runNode()` before throwing). + * + * Internal to the framework — not part of the public API. + */ +export class NodeInterruptedError extends Error { + constructor(message = 'Node interrupted (awaiting resume input).') { + super(message); + this.name = 'NodeInterruptedError'; + // Restore prototype chain for `instanceof` across transpilation targets. + Object.setPrototypeOf(this, NodeInterruptedError.prototype); + } +} + +/** + * Raised when a node exceeds its configured timeout. + * + * This is a regular `Error` (retryable) so a timed-out node can be retried via + * `retryConfig`. + */ +export class NodeTimeoutError extends Error { + readonly nodeName: string; + readonly timeout: number; + + /** + * @param options.nodeName The name of the node that timed out. + * @param options.timeout The timeout, in seconds, that was exceeded. + */ + constructor(options: {nodeName: string; timeout: number}) { + super( + `Node '${options.nodeName}' timed out after ${options.timeout} seconds.`, + ); + this.name = 'NodeTimeoutError'; + this.nodeName = options.nodeName; + this.timeout = options.timeout; + Object.setPrototypeOf(this, NodeTimeoutError.prototype); + } +} + +/** + * Raised when a dynamic node fails. + * + * Caught by the parent node's NodeRunner to propagate the error. + * Internal to the framework — not part of the public API. + */ +export class DynamicNodeFailError extends Error { + readonly error: Error; + readonly errorNodePath: string; + + /** + * @param options.message Human-readable failure message. + * @param options.error The underlying error thrown by the dynamic node. + * @param options.errorNodePath The node path where the failure occurred. + */ + constructor(options: {message: string; error: Error; errorNodePath: string}) { + super(options.message); + this.name = 'DynamicNodeFailError'; + this.error = options.error; + this.errorNodePath = options.errorNodePath; + Object.setPrototypeOf(this, DynamicNodeFailError.prototype); + } +} diff --git a/core/src/workflow-next/node_state.ts b/core/src/workflow-next/node_state.ts new file mode 100644 index 000000000..8950bbaa5 --- /dev/null +++ b/core/src/workflow-next/node_state.ts @@ -0,0 +1,77 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {NodeStatus} from './node_status.js'; + +/** + * State of a node in the workflow. + * + * Ported from `google/adk-python` `workflow/_node_state.py`. Note that the + * node's *output* is intentionally NOT stored here — it is carried on emitted + * events / the node Context, not on the persisted node state. + */ +export interface NodeState { + /** The run status of the node. */ + status: NodeStatus; + + /** The input provided to the node. */ + input?: unknown; + + /** The attempt count for this node run (1-based). */ + attemptCount: number; + + /** The interrupt ids that are pending to be resolved. */ + interrupts: string[]; + + /** The responses for resuming the node, keyed by interrupt id. */ + resumeInputs: Record; + + /** + * Sequential counter incremented each time the node gets a fresh run. + * + * Preserving this count independently of `runId` prevents path collisions if + * a node switches between custom string IDs and auto-generated numeric IDs. + */ + runCounter: number; + + /** The run ID of this node run. */ + runId?: string; + + /** + * The run ID of the parent node which dynamically scheduled this node run. + */ + parentRunId?: string; +} + +/** + * Creates a {@link NodeState} with Python-aligned defaults, overlaying any + * provided partial values. + */ +export function createNodeState(partial?: Partial): NodeState { + return { + status: NodeStatus.INACTIVE, + attemptCount: 1, + interrupts: [], + resumeInputs: {}, + runCounter: 0, + ...partial, + }; +} + +/** + * Type guard for a {@link NodeState}-shaped object. + */ +export function isNodeState(obj: unknown): obj is NodeState { + return ( + typeof obj === 'object' && + obj !== null && + 'status' in obj && + typeof (obj as NodeState).status === 'number' && + 'attemptCount' in obj && + 'interrupts' in obj && + Array.isArray((obj as NodeState).interrupts) + ); +} diff --git a/core/src/workflow-next/node_status.ts b/core/src/workflow-next/node_status.ts new file mode 100644 index 000000000..0a6cec0eb --- /dev/null +++ b/core/src/workflow-next/node_status.ts @@ -0,0 +1,29 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * The status of a node in the workflow graph. + * + * Numeric values are aligned with `google/adk-python` + * `workflow/_node_status.py` so that persisted node state is portable across + * the Python and TypeScript runtimes. + */ +export enum NodeStatus { + /** The node is not ready to be executed. */ + INACTIVE = 0, + /** The node is ready to be executed. */ + PENDING = 1, + /** The node is being executed. */ + RUNNING = 2, + /** The node has been executed successfully. */ + COMPLETED = 3, + /** The node is waiting (e.g. for a user response or re-trigger). */ + WAITING = 4, + /** The node has failed. */ + FAILED = 5, + /** The node has been cancelled. */ + CANCELLED = 6, +} diff --git a/core/src/workflow-next/retry_config.ts b/core/src/workflow-next/retry_config.ts new file mode 100644 index 000000000..65ea24efc --- /dev/null +++ b/core/src/workflow-next/retry_config.ts @@ -0,0 +1,85 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * An error constructor usable in {@link RetryConfig.exceptions}. + */ +export type ErrorClass = new (...args: never[]) => Error; + +/** + * Configuration for retrying a node. + * + * Ported from `google/adk-python` `workflow/_retry_config.py`. Delays are + * expressed in **seconds** (fractions allowed) to match the Python semantics + * and keep configuration portable across runtimes. Unset fields fall back to + * the documented defaults inside the retry utilities. + */ +export interface RetryConfig { + /** + * Maximum number of attempts, including the original request. If 0 or 1, it + * means no retries. If not specified, defaults to 5. + */ + maxAttempts?: number; + + /** + * Initial delay before the first retry, in seconds. If not specified, + * defaults to 1.0 second. + */ + initialDelay?: number; + + /** + * Maximum delay between retries, in seconds. If not specified, defaults to + * 60.0 seconds. + */ + maxDelay?: number; + + /** + * Multiplier by which the delay increases after each attempt. If not + * specified, defaults to 2.0. + */ + backoffFactor?: number; + + /** + * Randomness factor for the delay. If not specified, defaults to 1.0. Use 0.0 + * to remove randomness. + */ + jitter?: number; + + /** + * Exceptions to retry on. Accepts error class names as strings (e.g. + * `['TypeError']`) or error classes directly (e.g. `[TypeError]`). + * `undefined`/`null` means retry on all errors. + */ + exceptions?: Array | null; +} + +/** + * Normalizes the `exceptions` field of a {@link RetryConfig} to a list of error + * class name strings, mirroring Python's `field_validator`. + * + * @returns The list of class-name strings, or `undefined` to mean "retry on all + * errors". + */ +export function normalizeRetryExceptions( + exceptions?: Array | null, +): string[] | undefined { + if (exceptions === undefined || exceptions === null) { + return undefined; + } + return exceptions.map((item) => { + if (typeof item === 'string') { + return item; + } + if (typeof item === 'function' && item.name) { + return item.name; + } + throw new Error( + `exceptions must contain error class names (string) or error classes, got: ${String( + item, + )}`, + ); + }); +} diff --git a/core/src/workflow-next/utils/retry_utils.ts b/core/src/workflow-next/utils/retry_utils.ts new file mode 100644 index 000000000..57680e801 --- /dev/null +++ b/core/src/workflow-next/utils/retry_utils.ts @@ -0,0 +1,112 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Utility functions for retrying nodes in a workflow. + * + * Ported from `google/adk-python` `workflow/utils/_retry_utils.py`. + */ + +import {NodeState} from '../node_state.js'; +import {RetryConfig, normalizeRetryExceptions} from '../retry_config.js'; + +const DEFAULT_MAX_ATTEMPTS = 5; +const DEFAULT_INITIAL_DELAY_SECONDS = 1.0; +const DEFAULT_MAX_DELAY_SECONDS = 60.0; +const DEFAULT_BACKOFF_FACTOR = 2.0; +const DEFAULT_JITTER = 1.0; + +/** + * Resolves the runtime name of a thrown value for exception-name matching. + * Mirrors Python's `type(exception).__name__`. + */ +function errorName(error: unknown): string { + if (error instanceof Error) { + // `name` is set by well-behaved Error subclasses; fall back to the + // constructor name for plain `throw new Error()` cases. + return error.name || error.constructor.name; + } + if (typeof error === 'object' && error !== null) { + return error.constructor.name; + } + return typeof error; +} + +/** + * Checks if a failed node should be retried based on its retry config. + * + * @param error The error thrown by the node. + * @param retryConfig The node's retry configuration, if any. + * @param nodeState The current node state (its `attemptCount` is 1-based). + */ +export function shouldRetryNode( + error: unknown, + retryConfig: RetryConfig | undefined, + nodeState: NodeState, +): boolean { + if (!retryConfig) { + return false; + } + + const attemptCount = nodeState.attemptCount; + const maxAttempts = retryConfig.maxAttempts ?? DEFAULT_MAX_ATTEMPTS; + + // attemptCount starts at 1 for the original request; once it reaches + // maxAttempts, the limit is exhausted. + if (attemptCount >= maxAttempts) { + return false; + } + + const exceptions = normalizeRetryExceptions(retryConfig.exceptions); + if (exceptions !== undefined) { + if (!exceptions.includes(errorName(error))) { + return false; + } + } + + return true; +} + +/** + * Calculates the delay, in seconds, before retrying a node. + * + * @param retryConfig The node's retry configuration, if any. + * @param nodeState The current node state (its `attemptCount` is the 1-based + * attempt number that just failed). + * @param randomFn Injectable uniform RNG in [0, 1) for deterministic testing. + */ +export function getRetryDelaySeconds( + retryConfig: RetryConfig | undefined, + nodeState: NodeState, + randomFn: () => number = Math.random, +): number { + if (!retryConfig) { + return DEFAULT_INITIAL_DELAY_SECONDS; + } + + const initialDelay = + retryConfig.initialDelay ?? DEFAULT_INITIAL_DELAY_SECONDS; + const maxDelay = retryConfig.maxDelay ?? DEFAULT_MAX_DELAY_SECONDS; + const backoffFactor = retryConfig.backoffFactor ?? DEFAULT_BACKOFF_FACTOR; + const jitter = retryConfig.jitter ?? DEFAULT_JITTER; + + const attemptCount = nodeState.attemptCount || 1; + // attemptCount is the attempt number that just failed (1-based); the first + // failure (attempt 1) uses exponent 0. + const attemptForCalc = Math.max(0, attemptCount - 1); + + let delay = initialDelay * Math.pow(backoffFactor, attemptForCalc); + delay = Math.min(delay, maxDelay); + + if (jitter > 0.0) { + // random.uniform(-jitter*delay, jitter*delay) + const span = jitter * delay; + const randomOffset = -span + randomFn() * (2 * span); + delay = Math.max(0.0, delay + randomOffset); + } + + return delay; +} diff --git a/core/test/workflow-next/event_model_test.ts b/core/test/workflow-next/event_model_test.ts new file mode 100644 index 000000000..b53dfe18a --- /dev/null +++ b/core/test/workflow-next/event_model_test.ts @@ -0,0 +1,66 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {describe, expect, it} from 'vitest'; +import { + createEvent, + transformToCamelCaseEvent, + transformToSnakeCaseEvent, +} from '../../src/events/event.js'; + +describe('Phase 0 — workflow event-model extensions', () => { + it('carries first-class workflow fields on Event', () => { + const ev = createEvent({ + author: 'node_a', + output: {value: 42}, + route: 'question', + nodeInfo: {path: 'wf.node_a', outputFor: 'run_1', messageAsOutput: true}, + isolationScope: 'wf:evt_123', + actions: {agentState: {status: 3}, endOfAgent: true}, + }); + expect(ev.output).toEqual({value: 42}); + expect(ev.route).toBe('question'); + expect(ev.nodeInfo?.path).toBe('wf.node_a'); + expect(ev.nodeInfo?.messageAsOutput).toBe(true); + expect(ev.isolationScope).toBe('wf:evt_123'); + expect(ev.actions.agentState).toEqual({status: 3}); + expect(ev.actions.endOfAgent).toBe(true); + }); + + it('round-trips new fields through snake_case <-> camelCase', () => { + const ev = createEvent({ + author: 'node_a', + output: {value: 42}, + route: 'question', + nodeInfo: {path: 'wf.node_a', outputFor: 'run_1', messageAsOutput: true}, + isolationScope: 'wf:evt_123', + actions: {agentState: {status: 3}, endOfAgent: true}, + }); + + const snake = transformToSnakeCaseEvent(ev); + // Verify Python-compatible key names on the wire. + expect(snake['node_info']).toBeDefined(); + expect((snake['node_info'] as Record)['output_for']).toBe( + 'run_1', + ); + expect( + (snake['node_info'] as Record)['message_as_output'], + ).toBe(true); + expect(snake['isolation_scope']).toBe('wf:evt_123'); + const snakeActions = snake['actions'] as Record; + expect(snakeActions['agent_state']).toEqual({status: 3}); + expect(snakeActions['end_of_agent']).toBe(true); + + const back = transformToCamelCaseEvent(snake); + expect(back.nodeInfo?.path).toBe('wf.node_a'); + expect(back.nodeInfo?.outputFor).toBe('run_1'); + expect(back.nodeInfo?.messageAsOutput).toBe(true); + expect(back.isolationScope).toBe('wf:evt_123'); + expect(back.route).toBe('question'); + expect(back.actions.agentState).toEqual({status: 3}); + expect(back.actions.endOfAgent).toBe(true); + }); +}); diff --git a/core/test/workflow-next/foundations_test.ts b/core/test/workflow-next/foundations_test.ts new file mode 100644 index 000000000..9f5b68c4d --- /dev/null +++ b/core/test/workflow-next/foundations_test.ts @@ -0,0 +1,179 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {describe, expect, it} from 'vitest'; +import { + DynamicNodeFailError, + NodeInterruptedError, + NodeTimeoutError, +} from '../../src/workflow-next/errors.js'; +import { + createNodeState, + isNodeState, + NodeState, +} from '../../src/workflow-next/node_state.js'; +import {NodeStatus} from '../../src/workflow-next/node_status.js'; +import {normalizeRetryExceptions} from '../../src/workflow-next/retry_config.js'; +import { + getRetryDelaySeconds, + shouldRetryNode, +} from '../../src/workflow-next/utils/retry_utils.js'; + +describe('Phase 0 — errors', () => { + it('NodeTimeoutError carries nodeName/timeout and is instanceof Error', () => { + const err = new NodeTimeoutError({nodeName: 'n1', timeout: 2.5}); + expect(err).toBeInstanceOf(Error); + expect(err).toBeInstanceOf(NodeTimeoutError); + expect(err.name).toBe('NodeTimeoutError'); + expect(err.nodeName).toBe('n1'); + expect(err.timeout).toBe(2.5); + expect(err.message).toContain("Node 'n1' timed out after 2.5 seconds"); + }); + + it('NodeInterruptedError is a distinct catchable error', () => { + const err = new NodeInterruptedError(); + expect(err).toBeInstanceOf(NodeInterruptedError); + expect(err.name).toBe('NodeInterruptedError'); + }); + + it('DynamicNodeFailError wraps the underlying error + node path', () => { + const cause = new TypeError('boom'); + const err = new DynamicNodeFailError({ + message: 'dynamic node failed', + error: cause, + errorNodePath: 'wf.child.0', + }); + expect(err).toBeInstanceOf(DynamicNodeFailError); + expect(err.error).toBe(cause); + expect(err.errorNodePath).toBe('wf.child.0'); + }); +}); + +describe('Phase 0 — NodeStatus / NodeState', () => { + it('NodeStatus values match the Python enum ordinals', () => { + expect(NodeStatus.INACTIVE).toBe(0); + expect(NodeStatus.PENDING).toBe(1); + expect(NodeStatus.RUNNING).toBe(2); + expect(NodeStatus.COMPLETED).toBe(3); + expect(NodeStatus.WAITING).toBe(4); + expect(NodeStatus.FAILED).toBe(5); + expect(NodeStatus.CANCELLED).toBe(6); + }); + + it('createNodeState applies Python-aligned defaults', () => { + const s = createNodeState(); + expect(s.status).toBe(NodeStatus.INACTIVE); + expect(s.attemptCount).toBe(1); + expect(s.interrupts).toEqual([]); + expect(s.resumeInputs).toEqual({}); + expect(s.runCounter).toBe(0); + expect(s.runId).toBeUndefined(); + expect(s.parentRunId).toBeUndefined(); + }); + + it('createNodeState overlays partial values', () => { + const s = createNodeState({ + status: NodeStatus.RUNNING, + attemptCount: 3, + runId: 'r1', + }); + expect(s.status).toBe(NodeStatus.RUNNING); + expect(s.attemptCount).toBe(3); + expect(s.runId).toBe('r1'); + }); + + it('isNodeState recognizes valid/invalid shapes', () => { + expect(isNodeState(createNodeState())).toBe(true); + expect(isNodeState({})).toBe(false); + expect(isNodeState(null)).toBe(false); + expect(isNodeState({status: 'RUNNING'})).toBe(false); + }); +}); + +describe('Phase 0 — retry config normalization', () => { + it('returns undefined (retry-all) for null/undefined', () => { + expect(normalizeRetryExceptions(undefined)).toBeUndefined(); + expect(normalizeRetryExceptions(null)).toBeUndefined(); + }); + + it('normalizes error classes and strings to class-name strings', () => { + expect(normalizeRetryExceptions([TypeError, 'RangeError'])).toEqual([ + 'TypeError', + 'RangeError', + ]); + }); +}); + +describe('Phase 0 — shouldRetryNode', () => { + const state = (attemptCount: number): NodeState => + createNodeState({attemptCount}); + + it('never retries without a config', () => { + expect(shouldRetryNode(new Error('x'), undefined, state(1))).toBe(false); + }); + + it('retries until maxAttempts is reached (default 5)', () => { + expect(shouldRetryNode(new Error('x'), {}, state(1))).toBe(true); + expect(shouldRetryNode(new Error('x'), {}, state(4))).toBe(true); + expect(shouldRetryNode(new Error('x'), {}, state(5))).toBe(false); + }); + + it('respects an explicit maxAttempts', () => { + expect(shouldRetryNode(new Error('x'), {maxAttempts: 2}, state(1))).toBe( + true, + ); + expect(shouldRetryNode(new Error('x'), {maxAttempts: 2}, state(2))).toBe( + false, + ); + }); + + it('only retries listed exception types when provided', () => { + const cfg = {exceptions: [TypeError]}; + expect(shouldRetryNode(new TypeError('x'), cfg, state(1))).toBe(true); + expect(shouldRetryNode(new RangeError('x'), cfg, state(1))).toBe(false); + }); +}); + +describe('Phase 0 — getRetryDelaySeconds', () => { + it('defaults to 1.0s with no config', () => { + expect(getRetryDelaySeconds(undefined, createNodeState())).toBe(1.0); + }); + + it('applies exponential backoff (jitter disabled)', () => { + const cfg = {initialDelay: 1, backoffFactor: 2, jitter: 0}; + // attempt 1 -> exponent 0 -> 1s + expect(getRetryDelaySeconds(cfg, createNodeState({attemptCount: 1}))).toBe( + 1, + ); + // attempt 3 -> exponent 2 -> 4s + expect(getRetryDelaySeconds(cfg, createNodeState({attemptCount: 3}))).toBe( + 4, + ); + }); + + it('caps delay at maxDelay', () => { + const cfg = {initialDelay: 10, backoffFactor: 10, maxDelay: 30, jitter: 0}; + expect(getRetryDelaySeconds(cfg, createNodeState({attemptCount: 5}))).toBe( + 30, + ); + }); + + it('applies bounded symmetric jitter using the injected RNG', () => { + const cfg = {initialDelay: 4, backoffFactor: 1, jitter: 1}; + // randomFn=0.5 -> offset 0 -> exactly base delay (4) + expect( + getRetryDelaySeconds(cfg, createNodeState({attemptCount: 1}), () => 0.5), + ).toBe(4); + // randomFn=0 -> offset -span -> max(0, 4-4)=0 + expect( + getRetryDelaySeconds(cfg, createNodeState({attemptCount: 1}), () => 0), + ).toBe(0); + // randomFn=1 -> offset +span -> 4+4=8 + expect( + getRetryDelaySeconds(cfg, createNodeState({attemptCount: 1}), () => 1), + ).toBe(8); + }); +}); From 7886de525cb5343307da74769d3c80890a6074cc Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Wed, 22 Jul 2026 14:33:58 -0700 Subject: [PATCH 05/41] feat(workflow-next): Phase 1 core execution model Add the load-bearing core: an EventChannel primitive bridging the engine's push model to the runtime's pull model; the workflow NodeContext (ctx.runNode, delta-aware ctx.state, output/route/interruptIds/isolationScope/nodePath); the new BaseNode with a run() normalizer (Event | raw | null -> Event) and START sentinel; and a per-node NodeRunner handling event enrichment, output delegation, sub-branches, retry (Python check-then-increment) and per-node timeout. --- core/src/workflow-next/base_node.ts | 222 ++++++++++++++++ core/src/workflow-next/node_context.ts | 142 ++++++++++ core/src/workflow-next/node_runner.ts | 244 ++++++++++++++++++ core/src/workflow-next/utils/event_channel.ts | 108 ++++++++ core/test/workflow-next/event_channel_test.ts | 90 +++++++ .../test/workflow-next/node_execution_test.ts | 232 +++++++++++++++++ 6 files changed, 1038 insertions(+) create mode 100644 core/src/workflow-next/base_node.ts create mode 100644 core/src/workflow-next/node_context.ts create mode 100644 core/src/workflow-next/node_runner.ts create mode 100644 core/src/workflow-next/utils/event_channel.ts create mode 100644 core/test/workflow-next/event_channel_test.ts create mode 100644 core/test/workflow-next/node_execution_test.ts diff --git a/core/src/workflow-next/base_node.ts b/core/src/workflow-next/base_node.ts new file mode 100644 index 000000000..a0e9c3c41 --- /dev/null +++ b/core/src/workflow-next/base_node.ts @@ -0,0 +1,222 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {Content} from '@google/genai'; +import type {ZodType} from 'zod'; +import {createEvent, Event, isEvent} from '../events/event.js'; +import type {NodeContext} from './node_context.js'; +import {isRequestInput} from './request_input.js'; +import {RetryConfig} from './retry_config.js'; +import {createRequestInputEvent} from './utils/hitl_utils.js'; + +/** + * Configuration shared by all workflow nodes. + * + * Mirrors the fields of `google/adk-python` `workflow/_base_node.py::BaseNode`. + */ +export interface BaseNodeConfig { + /** Canonical, unique-within-a-graph node name. */ + name: string; + + /** Human-readable description (used when a node is exposed as a tool). */ + description?: string; + + /** + * If true, the node re-executes when a workflow resumes even if it already + * completed in a prior turn. Default false. + */ + rerunOnResume?: boolean; + + /** + * If true, the node only produces its output once all of its predecessors + * have triggered it (fan-in / join semantics). Default false. + */ + waitForOutput?: boolean; + + /** Optional retry configuration for transient failures. */ + retryConfig?: RetryConfig; + + /** Maximum time, in seconds, for this node to complete. */ + timeout?: number; + + /** Optional zod schema validating the node input. */ + inputSchema?: ZodType; + + /** Optional zod schema validating the node output. */ + outputSchema?: ZodType; + + /** Optional zod schema validating relevant session state. */ + stateSchema?: ZodType; +} + +/** + * Abstract base class for all nodes in an ADK workflow. + * + * A node is a discrete unit of execution. Subclasses implement {@link runImpl}, + * which may yield {@link Event}s, raw values (boxed into an event), or + * `null`/`undefined` (skipped). {@link run} normalizes those into a stream of + * {@link Event}s consumed by the engine. + */ +export abstract class BaseNode { + readonly name: string; + readonly description: string; + readonly rerunOnResume: boolean; + readonly waitForOutput: boolean; + readonly retryConfig?: RetryConfig; + readonly timeout?: number; + readonly inputSchema?: ZodType; + readonly outputSchema?: ZodType; + readonly stateSchema?: ZodType; + + constructor(config: BaseNodeConfig) { + if ( + !config.name || + typeof config.name !== 'string' || + config.name.trim().length === 0 + ) { + throw new Error('Node name must be a non-empty string.'); + } + this.name = config.name.trim(); + this.description = config.description ?? ''; + this.rerunOnResume = config.rerunOnResume ?? false; + this.waitForOutput = config.waitForOutput ?? false; + this.retryConfig = config.retryConfig; + this.timeout = config.timeout; + this.inputSchema = config.inputSchema; + this.outputSchema = config.outputSchema; + this.stateSchema = config.stateSchema; + } + + /** + * Whether this node must wait for ALL of its predecessors to trigger before + * it runs (fan-in barrier). Overridden by `JoinNode`. + */ + get requiresAllPredecessors(): boolean { + return false; + } + + /** + * Core execution contract. Subclasses yield one of: + * - an {@link Event} (emitted as-is), + * - a raw value (boxed into an event whose `output` is that value), + * - `null`/`undefined` (skipped). + */ + protected abstract runImpl( + ctx: NodeContext, + input: TInput, + ): AsyncGenerator; + + /** + * Runs the node, normalizing every yielded item into an {@link Event}. This + * is what the engine (and `ctx.runNode()`) consumes. Validates the input + * against `inputSchema` once, up front (skipping genai `Content`, which nodes + * coerce themselves). + */ + async *run( + ctx: NodeContext, + input: TInput, + ): AsyncGenerator { + const validatedInput = this.validateInput(input); + for await (const item of this.runImpl(ctx, validatedInput)) { + if (isRequestInput(item)) { + // HITL: convert a request-for-input into an interrupt event. + yield createRequestInputEvent(item); + continue; + } + const event = this.toEvent(ctx, item); + if (event) { + yield event; + } + } + } + + /** Validates node input against `inputSchema` (Content passes through). */ + protected validateInput(input: TInput): TInput { + if (!this.inputSchema || isContent(input)) { + return input; + } + return this.inputSchema.parse(input) as TInput; + } + + /** Validates node output against `outputSchema` (Content passes through). */ + protected validateOutput(output: unknown): unknown { + if (!this.outputSchema || isContent(output)) { + return output; + } + return this.outputSchema.parse(output); + } + + /** + * Normalizes a single yielded item into an {@link Event} (or `null` to skip). + * Subclasses may override for richer coercion (e.g. `FunctionNode`). + */ + protected toEvent(ctx: NodeContext, data: unknown): Event | null { + if (data === null || data === undefined) { + return null; + } + if (isEvent(data)) { + const event = data as Event; + if (event.output !== undefined) { + event.output = this.validateOutput(event.output); + } + return event; + } + const output = this.validateOutput(data); + return createEvent({ + author: this.name, + invocationId: ctx.invocationContext.invocationId, + branch: ctx.branch, + content: toContent(output), + output, + }); + } +} + +/** Returns whether a value looks like a genai `Content` object. */ +export function isContent(value: unknown): value is Content { + return ( + typeof value === 'object' && + value !== null && + 'parts' in value && + Array.isArray((value as {parts?: unknown}).parts) + ); +} + +/** + * The sentinel node marking the entry point of a workflow graph. It is never + * executed — the orchestrator seeds triggers for its successors directly. + * + * Mirrors `google/adk-python` `START = BaseNode(name='__START__')`. + */ +class StartNode extends BaseNode { + // eslint-disable-next-line require-yield + protected async *runImpl(): AsyncGenerator { + throw new Error('START node is never executed.'); + } +} + +/** The workflow entry-point sentinel node (name `__START__`). */ +export const START: BaseNode = new StartNode({name: '__START__'}); + +/** + * Best-effort conversion of an arbitrary value to genai `Content` for display. + */ +export function toContent(val: unknown): Content | undefined { + if (val === null || val === undefined) { + return undefined; + } + if (typeof val === 'object' && 'role' in val && 'parts' in val) { + return val as Content; + } + if (typeof val === 'string') { + return {role: 'model', parts: [{text: val}]}; + } + try { + return {role: 'model', parts: [{text: JSON.stringify(val)}]}; + } catch { + return {role: 'model', parts: [{text: String(val)}]}; + } +} diff --git a/core/src/workflow-next/node_context.ts b/core/src/workflow-next/node_context.ts new file mode 100644 index 000000000..26d4f95a0 --- /dev/null +++ b/core/src/workflow-next/node_context.ts @@ -0,0 +1,142 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {InvocationContext} from '../agents/invocation_context.js'; +import {Event} from '../events/event.js'; +import {createEventActions, EventActions} from '../events/event_actions.js'; +import {State} from '../sessions/state.js'; +import type {BaseNode} from './base_node.js'; +import {executeChildNode, RunNodeOptions} from './node_runner.js'; +import type {ScheduleDynamicNode} from './schedule_dynamic_node.js'; +import {EventChannel} from './utils/event_channel.js'; + +/** + * Options for constructing a {@link NodeContext}. + */ +export interface NodeContextOptions { + invocationContext: InvocationContext; + channel: EventChannel; + /** Dotted node path of the owning node (empty string for the root). */ + nodePath: string; + /** Deterministic run id of the owning node. */ + runId: string; + /** Responses for resuming interrupted child nodes, keyed by interrupt id. */ + resumeInputs?: Record; + /** Scope tag isolating this node's events from peer scopes. */ + isolationScope?: string; + /** Accumulator for event actions (state delta, etc). */ + actions?: EventActions; +} + +/** + * The execution context for a workflow node — the TypeScript analogue of + * `google/adk-python` `agents/context.py::Context` (the workflow flavour). + * + * It exposes `ctx.runNode(...)` for programmatic child execution, `ctx.state` + * for delta-aware session state, `ctx.emit(...)` to stream an event, and the + * mutable `output`/`route`/`interruptIds` a node sets while running. + */ +export class NodeContext { + readonly invocationContext: InvocationContext; + readonly channel: EventChannel; + readonly nodePath: string; + readonly runId: string; + readonly actions: EventActions; + resumeInputs: Record; + isolationScope?: string; + + /** The structured output produced by the node during its run. */ + output: unknown = undefined; + + /** The route key emitted by the node, if any. */ + route?: string | number | boolean; + + /** Interrupt ids the node is currently blocked on (HITL). */ + interruptIds: string[] = []; + + /** + * The dynamic-node scheduler for this subtree. When set, `ctx.runNode()` + * routes through it (dedup/resume/fresh); otherwise it runs the child + * directly. Propagated to child contexts by the node runner; a nested + * Workflow overrides it with its own scheduler. + */ + scheduler?: ScheduleDynamicNode; + + private readonly _state: State; + private readonly dynamicRunCounters = new Map(); + + constructor(opts: NodeContextOptions) { + this.invocationContext = opts.invocationContext; + this.channel = opts.channel; + this.nodePath = opts.nodePath; + this.runId = opts.runId; + this.resumeInputs = opts.resumeInputs ?? {}; + this.isolationScope = opts.isolationScope; + this.actions = opts.actions ?? createEventActions(); + // Writes via `ctx.state` accumulate into `actions.stateDelta`, mirroring + // Python's `ctx.state` -> `ctx.actions.state_delta` behaviour. + this._state = new State( + opts.invocationContext.session.state, + this.actions.stateDelta, + ); + } + + /** Delta-aware session state; writes accumulate in `actions.stateDelta`. */ + get state(): State { + return this._state; + } + + /** The branch of the owning invocation context. */ + get branch(): string | undefined { + return this.invocationContext.branch; + } + + /** The current invocation id. */ + get invocationId(): string { + return this.invocationContext.invocationId; + } + + /** The current session. */ + get session() { + return this.invocationContext.session; + } + + /** Streams a single event out through the workflow's event channel. */ + emit(event: Event): void { + this.channel.push(event); + } + + /** + * 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`). + * + * When a dynamic-node {@link scheduler} is set (inside a Workflow subtree), + * the call routes through it for dedup/resume; otherwise the child runs + * directly. + */ + runNode( + node: BaseNode, + input?: unknown, + options?: RunNodeOptions, + ): Promise { + if (this.scheduler) { + const nodeName = options?.nodeName ?? node.name; + let runId = options?.runId; + if (!runId) { + const next = (this.dynamicRunCounters.get(nodeName) ?? 0) + 1; + this.dynamicRunCounters.set(nodeName, next); + runId = String(next); + } + return this.scheduler.schedule(this, node, input, { + ...options, + nodeName, + runId, + }); + } + return executeChildNode(this, node, input, options); + } +} diff --git a/core/src/workflow-next/node_runner.ts b/core/src/workflow-next/node_runner.ts new file mode 100644 index 000000000..2661484e7 --- /dev/null +++ b/core/src/workflow-next/node_runner.ts @@ -0,0 +1,244 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { + InvocationContext, + InvocationContextParams, +} from '../agents/invocation_context.js'; +import {Event} from '../events/event.js'; +import {BaseNode} from './base_node.js'; +import {BranchPath} from './branch_path.js'; +import {NodeTimeoutError} from './errors.js'; +import {NodeContext} from './node_context.js'; +import {createNodeState} from './node_state.js'; +import {NodeStatus} from './node_status.js'; +import {getRetryDelaySeconds, shouldRetryNode} from './utils/retry_utils.js'; + +/** + * Options controlling a single `ctx.runNode(...)` execution. + */ +export interface RunNodeOptions { + /** Deterministic tracking name; defaults to `node.name`. */ + nodeName?: string; + /** Unique id for this specific run; defaults to `nodeName`. */ + runId?: string; + /** If true, the child's output replaces the caller's output. */ + useAsOutput?: boolean; + /** If true, run the child in an isolated sub-branch. */ + useSubBranch?: boolean; + /** Explicit branch, overriding the default/sub-branch computation. */ + overrideBranch?: string; + /** Explicit isolation scope, overriding inheritance from the parent. */ + overrideIsolationScope?: string; +} + +/** + * Executes a child node on behalf of `parent.runNode(...)`. + * + * Responsibilities (Phase 1 scope): create the child {@link NodeContext}, + * drive `node.run()`, enrich each emitted event (author, node path, branch, + * isolation scope), track the child's `output`/`route`, apply the per-node + * `timeout`, and retry on failure per `retryConfig`. Returns the child context. + */ +export async function executeChildNode( + parent: NodeContext, + node: BaseNode, + input: unknown, + options: RunNodeOptions = {}, +): Promise { + const nodeName = options.nodeName ?? node.name; + const runId = options.runId ?? nodeName; + const nodePath = parent.nodePath + ? `${parent.nodePath}.${nodeName}` + : nodeName; + + let branch = parent.branch; + if (options.overrideBranch !== undefined) { + branch = options.overrideBranch; + } else if (options.useSubBranch) { + branch = BranchPath.createSubBranch(parent.branch, { + name: nodeName, + runId: options.runId, + }); + } + + const isolationScope = + options.overrideIsolationScope ?? parent.isolationScope; + + const childIc = + branch === parent.invocationContext.branch + ? parent.invocationContext + : withBranch(parent.invocationContext, branch); + + const child = new NodeContext({ + invocationContext: childIc, + channel: parent.channel, + nodePath, + runId, + resumeInputs: parent.resumeInputs, + isolationScope, + }); + // Propagate the dynamic scheduler down; a nested Workflow overrides it. + child.scheduler = parent.scheduler; + + const nodeState = createNodeState({ + status: NodeStatus.RUNNING, + input, + runId, + }); + + for (;;) { + // Reset per-attempt output so a retry starts clean. + child.output = undefined; + child.route = undefined; + child.interruptIds = []; + try { + await runOnce(node, child, input, nodeName, branch, isolationScope); + break; + } catch (err) { + // Check retry eligibility with the attempt that just failed, compute its + // backoff delay, THEN advance the counter (matches Python semantics). + if (shouldRetryNode(err, node.retryConfig, nodeState)) { + const delaySeconds = getRetryDelaySeconds(node.retryConfig, nodeState); + nodeState.attemptCount += 1; + await delay(delaySeconds * 1000, parent.invocationContext.abortSignal); + continue; + } + throw err; + } + } + + if (options.useAsOutput) { + parent.output = child.output; + parent.route = child.route; + } + + return child; +} + +/** + * Drives one attempt of `node.run()`, enriching and pushing each event and + * tracking the child's output/route. Wrapped in a timeout when configured. + */ +async function runOnce( + node: BaseNode, + child: NodeContext, + input: unknown, + nodeName: string, + branch: string | undefined, + isolationScope: string | undefined, +): Promise { + const body = (async () => { + for await (const event of node.run(child, input)) { + enrichEvent(event, child, nodeName, branch, isolationScope); + if (event.output !== undefined) { + child.output = event.output; + } + if (event.route !== undefined) { + child.route = event.route; + } + // HITL: an interrupt event marks its ids as long-running tool ids. + if (event.longRunningToolIds && event.longRunningToolIds.length > 0) { + for (const id of event.longRunningToolIds) { + if (!child.interruptIds.includes(id)) { + child.interruptIds.push(id); + } + } + } + child.channel.push(event); + } + })(); + + if (node.timeout && node.timeout > 0) { + await withTimeout(body, node.timeout, nodeName); + } else { + await body; + } +} + +/** + * Stamps engine-owned provenance onto an event without clobbering values the + * node explicitly set. + */ +function enrichEvent( + event: Event, + child: NodeContext, + nodeName: string, + branch: string | undefined, + isolationScope: string | undefined, +): void { + if (!event.author) { + event.author = nodeName; + } + event.nodeInfo = {...(event.nodeInfo ?? {}), path: child.nodePath}; + if (branch !== undefined && event.branch === undefined) { + event.branch = branch; + } + if (isolationScope !== undefined && event.isolationScope === undefined) { + event.isolationScope = isolationScope; + } +} + +/** + * Creates a shallow child InvocationContext with a different branch, preserving + * the shared invocation cost manager and all services/session. + */ +function withBranch( + ic: InvocationContext, + branch: string | undefined, +): InvocationContext { + return new InvocationContext({ + ...(ic as unknown as InvocationContextParams), + branch, + }); +} + +/** + * Rejects with {@link NodeTimeoutError} if `promise` does not settle within + * `timeoutSeconds`. + */ +function withTimeout( + promise: Promise, + timeoutSeconds: number, + nodeName: string, +): Promise { + return new Promise((resolve, reject) => { + const timer = setTimeout(() => { + reject(new NodeTimeoutError({nodeName, timeout: timeoutSeconds})); + }, timeoutSeconds * 1000); + promise.then( + (value) => { + clearTimeout(timer); + resolve(value); + }, + (err) => { + clearTimeout(timer); + reject(err); + }, + ); + }); +} + +/** + * Promise-based delay that rejects early if the abort signal fires. + */ +function delay(ms: number, signal?: AbortSignal): Promise { + return new Promise((resolve, reject) => { + if (signal?.aborted) { + reject(new Error('Aborted')); + return; + } + const timer = setTimeout(() => { + signal?.removeEventListener('abort', onAbort); + resolve(); + }, ms); + const onAbort = () => { + clearTimeout(timer); + reject(new Error('Aborted')); + }; + signal?.addEventListener('abort', onAbort, {once: true}); + }); +} diff --git a/core/src/workflow-next/utils/event_channel.ts b/core/src/workflow-next/utils/event_channel.ts new file mode 100644 index 000000000..48914a85d --- /dev/null +++ b/core/src/workflow-next/utils/event_channel.ts @@ -0,0 +1,108 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * A single-consumer async queue that bridges the workflow engine's *push* model + * (nodes/`ctx.runNode()` push events as they run) to the runtime's *pull* model + * (the workflow's outer async generator drains and re-yields them). + * + * Producers call {@link push} as events are produced and {@link close} (or + * {@link fail}) when finished. A single consumer drains the channel by + * `for await (const ev of channel)`. + * + * Semantics: + * - Buffered items are always delivered before an end/error signal. + * - {@link close} ends iteration cleanly (`done: true`). + * - {@link fail} surfaces the error to the consumer *after* any buffered items + * have been drained. + * - {@link push} after close/fail is ignored (the producer has already + * signalled completion). + */ +export class EventChannel implements AsyncIterable { + private readonly buffer: T[] = []; + private readonly waiters: Array<{ + resolve: (r: IteratorResult) => void; + reject: (e: unknown) => void; + }> = []; + private closed = false; + private failure?: {error: unknown}; + + /** Whether the channel has been closed or failed. */ + get isClosed(): boolean { + return this.closed; + } + + /** Number of items buffered and not yet consumed. */ + get size(): number { + return this.buffer.length; + } + + /** + * Enqueues an item. If a consumer is currently awaiting, it is resolved + * immediately; otherwise the item is buffered. No-op once closed/failed. + */ + push(item: T): void { + if (this.closed) { + return; + } + const waiter = this.waiters.shift(); + if (waiter) { + waiter.resolve({value: item, done: false}); + } else { + this.buffer.push(item); + } + } + + /** + * Signals that no more items will be produced. Any awaiting consumer receives + * `{done: true}`. Idempotent. + */ + close(): void { + if (this.closed) { + return; + } + this.closed = true; + while (this.waiters.length > 0) { + this.waiters.shift()!.resolve({value: undefined as never, done: true}); + } + } + + /** + * Signals that production failed. Buffered items are still delivered first; + * once the buffer drains, the consumer's next `next()` rejects with `error`. + * Idempotent (first failure wins). + */ + fail(error: unknown): void { + if (this.closed) { + return; + } + this.failure = {error}; + this.closed = true; + // If a consumer is awaiting, the buffer is empty, so surface the error now. + while (this.waiters.length > 0) { + this.waiters.shift()!.reject(error); + } + } + + [Symbol.asyncIterator](): AsyncIterator { + return { + next: (): Promise> => { + if (this.buffer.length > 0) { + return Promise.resolve({value: this.buffer.shift()!, done: false}); + } + if (this.failure) { + return Promise.reject(this.failure.error); + } + if (this.closed) { + return Promise.resolve({value: undefined as never, done: true}); + } + return new Promise>((resolve, reject) => { + this.waiters.push({resolve, reject}); + }); + }, + }; + } +} diff --git a/core/test/workflow-next/event_channel_test.ts b/core/test/workflow-next/event_channel_test.ts new file mode 100644 index 000000000..d9c4ea8f7 --- /dev/null +++ b/core/test/workflow-next/event_channel_test.ts @@ -0,0 +1,90 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {describe, expect, it} from 'vitest'; +import {EventChannel} from '../../src/workflow-next/utils/event_channel.js'; + +async function drain(ch: EventChannel): Promise { + const out: T[] = []; + for await (const item of ch) { + out.push(item); + } + return out; +} + +describe('Phase 1 — EventChannel', () => { + it('delivers items pushed before draining (buffered)', async () => { + const ch = new EventChannel(); + ch.push(1); + ch.push(2); + ch.push(3); + ch.close(); + expect(await drain(ch)).toEqual([1, 2, 3]); + }); + + it('delivers items pushed while a consumer is awaiting (interleaved)', async () => { + const ch = new EventChannel(); + const collected: number[] = []; + const consumer = (async () => { + for await (const item of ch) { + collected.push(item); + } + })(); + + // Push across turns of the event loop while the consumer is parked. + await Promise.resolve(); + ch.push(10); + await Promise.resolve(); + ch.push(20); + await Promise.resolve(); + ch.close(); + + await consumer; + expect(collected).toEqual([10, 20]); + }); + + it('close() terminates iteration cleanly', async () => { + const ch = new EventChannel(); + ch.push('a'); + ch.close(); + ch.push('ignored-after-close'); + expect(await drain(ch)).toEqual(['a']); + expect(ch.isClosed).toBe(true); + }); + + it('fail() surfaces the error to the consumer', async () => { + const ch = new EventChannel(); + const boom = new Error('boom'); + ch.fail(boom); + await expect(drain(ch)).rejects.toThrow('boom'); + }); + + it('fail() still drains buffered items before throwing', async () => { + const ch = new EventChannel(); + ch.push(1); + ch.push(2); + ch.fail(new Error('later')); + + const seen: number[] = []; + await expect( + (async () => { + for await (const item of ch) { + seen.push(item); + } + })(), + ).rejects.toThrow('later'); + expect(seen).toEqual([1, 2]); + }); + + it('reports buffered size and ignores push after close', async () => { + const ch = new EventChannel(); + ch.push(1); + expect(ch.size).toBe(1); + ch.close(); + ch.push(2); + expect(ch.size).toBe(1); + }); +}); diff --git a/core/test/workflow-next/node_execution_test.ts b/core/test/workflow-next/node_execution_test.ts new file mode 100644 index 000000000..0f465c896 --- /dev/null +++ b/core/test/workflow-next/node_execution_test.ts @@ -0,0 +1,232 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {describe, expect, it} from 'vitest'; +import {BaseAgent} from '../../src/agents/base_agent.js'; +import {InvocationContext} from '../../src/agents/invocation_context.js'; +import {createEvent, Event} from '../../src/events/event.js'; +import {PluginManager} from '../../src/plugins/plugin_manager.js'; +import {Session} from '../../src/sessions/session.js'; +import {BaseNode} from '../../src/workflow-next/base_node.js'; +import {NodeTimeoutError} from '../../src/workflow-next/errors.js'; +import {NodeContext} from '../../src/workflow-next/node_context.js'; +import {EventChannel} from '../../src/workflow-next/utils/event_channel.js'; + +// --- Test harness --------------------------------------------------------- + +function createIc(params?: Partial): InvocationContext { + const session: Session = { + id: 'session-123', + appName: 'test-app', + userId: 'test-user', + events: [], + state: {}, + lastUpdateTime: Date.now(), + } as unknown as Session; + + return new InvocationContext({ + invocationId: 'inv-1', + session, + agent: { + name: 'wf', + runAsync: async function* () {}, + } as unknown as BaseAgent, + pluginManager: new PluginManager(), + ...params, + }); +} + +/** + * Drives a root node the way the Phase 2 Workflow loop will: orchestration + * pushes events into the channel concurrently while the consumer drains it. + * This exercises the real push/pull bridge, not a buffer-then-read shortcut. + */ +async function driveRoot( + ic: InvocationContext, + node: BaseNode, + input?: unknown, +): Promise<{events: Event[]; output: unknown; route: unknown}> { + const channel = new EventChannel(); + const root = new NodeContext({ + invocationContext: ic, + channel, + nodePath: '', + runId: 'root', + }); + + const events: Event[] = []; + const orchestration = root.runNode(node, input, {useAsOutput: true}).then( + () => channel.close(), + (err) => channel.fail(err), + ); + + for await (const ev of channel) { + events.push(ev); + } + await orchestration; + return {events, output: root.output, route: root.route}; +} + +// A minimal node that yields whatever its function returns (value | Event). +class FnNode extends BaseNode { + constructor( + name: string, + private readonly fn: ( + ctx: NodeContext, + input: unknown, + ) => unknown | Promise, + config?: Partial< + Omit< + import('../../src/workflow-next/base_node.js').BaseNodeConfig, + 'name' + > + >, + ) { + super({name, ...config}); + } + protected async *runImpl(ctx: NodeContext, input: unknown) { + yield await this.fn(ctx, input); + } +} + +// --- Tests ---------------------------------------------------------------- + +describe('Phase 1 — node execution & the push/pull bridge', () => { + it('streams a node event and returns its output', async () => { + const node = new FnNode('greet', (_ctx, input) => `hello ${input}`); + const {events, output} = await driveRoot(createIc(), node, 'world'); + + expect(output).toBe('hello world'); + expect(events).toHaveLength(1); + expect(events[0].author).toBe('greet'); + expect(events[0].output).toBe('hello world'); + expect(events[0].nodeInfo?.path).toBe('greet'); + }); + + it('passes through an explicitly emitted Event and preserves route', async () => { + const node = new FnNode('router', () => + createEvent({route: 'question', output: 'q'}), + ); + const {events, output, route} = await driveRoot(createIc(), node, 'in'); + + expect(events).toHaveLength(1); + expect(events[0].route).toBe('question'); + expect(output).toBe('q'); + expect(route).toBe('question'); + // Engine stamps provenance without clobbering. + expect(events[0].nodeInfo?.path).toBe('router'); + expect(events[0].author).toBe('router'); + }); + + it('supports nested ctx.runNode() with correct node paths (the bridge)', async () => { + const inner = new FnNode('inner', (_ctx, input) => `inner(${input})`); + const outer = new FnNode('outer', async (ctx, input) => { + const child = await ctx.runNode(inner, input); + return `outer[${child.output}]`; + }); + + const {events, output} = await driveRoot(createIc(), outer, 'x'); + + expect(output).toBe('outer[inner(x)]'); + // Both the child and parent events streamed out, child first. + const paths = events.map((e) => e.nodeInfo?.path); + expect(paths).toContain('outer.inner'); + expect(paths).toContain('outer'); + expect(paths.indexOf('outer.inner')).toBeLessThan(paths.indexOf('outer')); + }); + + it('derives nested sub-branches as dotted paths', async () => { + let innerBranch: string | undefined = 'UNSET'; + const inner = new FnNode('inner', (ctx) => { + innerBranch = ctx.branch; + return 'ok'; + }); + const mid = new FnNode('mid', async (ctx, input) => { + await ctx.runNode(inner, input, {useSubBranch: true}); + return 'm'; + }); + const outer = new FnNode('outer', async (ctx, input) => { + await ctx.runNode(mid, input, {useSubBranch: true}); + return 'o'; + }); + + // outer runs at the root (branch undefined); mid -> 'mid'; inner -> 'mid.inner'. + await driveRoot(createIc(), outer, 'x'); + expect(innerBranch).toBe('mid.inner'); + }); + + it('accumulates ctx.state writes into the event action state delta', async () => { + const node = new FnNode('writer', (ctx) => { + ctx.state.set('counter', 7); + return 'wrote'; + }); + const channel = new EventChannel(); + const root = new NodeContext({ + invocationContext: createIc(), + channel, + nodePath: '', + runId: 'root', + }); + const child = await root.runNode(node, undefined, {useAsOutput: true}); + expect(child.state.get('counter')).toBe(7); + expect(child.actions.stateDelta['counter']).toBe(7); + }); + + it('retries a flaky node per retryConfig and then succeeds', async () => { + let attempts = 0; + const flaky = new FnNode( + 'flaky', + () => { + attempts++; + if (attempts < 3) { + throw new Error('transient'); + } + return 'ok-after-retry'; + }, + { + retryConfig: { + maxAttempts: 3, + initialDelay: 0.001, + backoffFactor: 1, + jitter: 0, + }, + }, + ); + + const {output} = await driveRoot(createIc(), flaky, 'x'); + expect(attempts).toBe(3); + expect(output).toBe('ok-after-retry'); + }); + + it('propagates the final error when retries are exhausted', async () => { + let attempts = 0; + const doomed = new FnNode( + 'doomed', + () => { + attempts++; + throw new Error('always fails'); + }, + {retryConfig: {maxAttempts: 2, initialDelay: 0.001, jitter: 0}}, + ); + + await expect(driveRoot(createIc(), doomed, 'x')).rejects.toThrow( + 'always fails', + ); + expect(attempts).toBe(2); + }); + + it('enforces a per-node timeout with NodeTimeoutError', async () => { + const slow = new FnNode( + 'slow', + () => new Promise((resolve) => setTimeout(() => resolve('late'), 200)), + {timeout: 0.02}, + ); + + await expect(driveRoot(createIc(), slow, 'x')).rejects.toBeInstanceOf( + NodeTimeoutError, + ); + }); +}); From fd18460d23c9e6bfaaed77383b1ae037fc0a2479 Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Wed, 22 Jul 2026 14:34:57 -0700 Subject: [PATCH 06/41] feat(workflow-next): Phase 2 graph orchestration engine Add the static-graph engine: Graph/Edge/RoutingMap/NodeLike with route values (bool|number|string), a graph parser (chains, fan-out, routing maps) and validator (reachability, duplicate/START rules, unconditional-cycle detection), Trigger as a data record, and Workflow(BaseNode) whose orchestration loop seeds START, schedules ready nodes, races completions, buffers downstream triggers, fans in via requiresAllPredecessors, and finalizes terminal output. Fixes the initial-input-propagation bug. --- core/src/workflow-next/graph.ts | 163 +++++++ core/src/workflow-next/trigger.ts | 26 + core/src/workflow-next/utils/graph_parser.ts | 212 +++++++++ .../workflow-next/utils/graph_validation.ts | 200 ++++++++ .../utils/workflow_graph_utils.ts | 134 ++++++ core/src/workflow-next/workflow.ts | 450 ++++++++++++++++++ core/test/workflow-next/workflow_test.ts | 190 ++++++++ 7 files changed, 1375 insertions(+) create mode 100644 core/src/workflow-next/graph.ts create mode 100644 core/src/workflow-next/trigger.ts create mode 100644 core/src/workflow-next/utils/graph_parser.ts create mode 100644 core/src/workflow-next/utils/graph_validation.ts create mode 100644 core/src/workflow-next/utils/workflow_graph_utils.ts create mode 100644 core/src/workflow-next/workflow.ts create mode 100644 core/test/workflow-next/workflow_test.ts diff --git a/core/src/workflow-next/graph.ts b/core/src/workflow-next/graph.ts new file mode 100644 index 000000000..453e7e57a --- /dev/null +++ b/core/src/workflow-next/graph.ts @@ -0,0 +1,163 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type {BaseAgent} from '../agents/base_agent.js'; +import {BaseTool} from '../tools/base_tool.js'; +import {BaseNode} from './base_node.js'; +import {parseEdgeItems} from './utils/graph_parser.js'; +import {validateGraph} from './utils/graph_validation.js'; + +/** Valid routing values used in conditional graph edges. */ +export type RouteValue = boolean | number | string; + +/** The fallback route key used when no specific route matches. */ +export const DEFAULT_ROUTE = '__DEFAULT__'; + +/** + * Any value that can be converted to a workflow node: a node, a tool, a plain + * function, or the `'START'` sentinel literal. (Agent wrapping is added in + * Phase 3 via `build_node`.) + */ +export type NodeLike = + | BaseNode + | BaseAgent + | BaseTool + | ((...args: never[]) => unknown) + | 'START'; + +/** + * A mapping from route values to destination node(s). A value may be a single + * node or an array of nodes (fan-out). + * + * @example + * {question: answerNode, statement: commentNode} + * {retry: [nodeA, nodeB]} // fan-out: both triggered + */ +export type RoutingMap = Record< + string | number, + NodeLike | readonly NodeLike[] +>; + +/** An element within a workflow chain. */ +export type ChainElement = NodeLike | readonly NodeLike[] | RoutingMap; + +/** + * An item that can be parsed into workflow edges: an explicit {@link Edge}, or a + * chain expressed as an array of {@link ChainElement}s (e.g. + * `['START', nodeA, nodeB]`). + */ +export type EdgeItem = Edge | ChainElement[]; + +/** + * A directed edge in the workflow graph. + * + * Mirrors `google/adk-python` `workflow/_graph.py::Edge`. + */ +export class Edge { + constructor( + readonly fromNode: BaseNode, + readonly toNode: BaseNode, + /** + * The route(s) this edge is associated with. `null` means unconditional + * (always triggered). A single value or a list; the edge fires when the + * emitted route matches any listed value. + */ + readonly route: RouteValue | RouteValue[] | null = null, + ) {} +} + +/** + * A compiled workflow graph. Nodes are inferred (deduped by identity) from the + * edges. + * + * Mirrors `google/adk-python` `workflow/_graph.py::Graph`. + */ +export class Graph { + readonly nodes: BaseNode[]; + readonly edges: Edge[]; + private _terminalNodeNames: ReadonlySet = new Set(); + + constructor(edges: Edge[]) { + this.edges = edges; + const seen = new Set(); + const nodes: BaseNode[] = []; + for (const edge of edges) { + for (const node of [edge.fromNode, edge.toNode]) { + if (!seen.has(node)) { + seen.add(node); + nodes.push(node); + } + } + } + this.nodes = nodes; + } + + /** Builds and returns a graph from a list of edge items. */ + static fromEdgeItems(edgeItems: EdgeItem[]): Graph { + return new Graph(parseEdgeItems(edgeItems)); + } + + /** Terminal node names (no outgoing edges); populated by {@link validate}. */ + get terminalNodeNames(): ReadonlySet { + return this._terminalNodeNames; + } + + /** + * Determines the next nodes to transition to PENDING based on the route(s) + * emitted by a completed node. Ported from Python `get_next_pending_nodes`. + */ + getNextPendingNodes( + nodeName: string, + routesToMatch: RouteValue | RouteValue[] | null | undefined, + ): string[] { + const nextPending: string[] = []; + let matchedSpecificRoute = false; + let defaultRouteNode: string | undefined; + + for (const edge of this.edges) { + if (edge.fromNode.name !== nodeName) { + continue; + } + if (edge.route === null || edge.route === undefined) { + // Unconditional edges always fire. + nextPending.push(edge.toNode.name); + continue; + } + + if (edge.route === DEFAULT_ROUTE) { + defaultRouteNode = edge.toNode.name; + continue; + } + + const edgeRoutes = new Set( + Array.isArray(edge.route) ? edge.route : [edge.route], + ); + + let edgeMatched = false; + if (Array.isArray(routesToMatch)) { + edgeMatched = routesToMatch.some((r) => edgeRoutes.has(r)); + } else if (routesToMatch !== null && routesToMatch !== undefined) { + edgeMatched = edgeRoutes.has(routesToMatch); + } + + if (edgeMatched) { + nextPending.push(edge.toNode.name); + matchedSpecificRoute = true; + } + } + + if (!matchedSpecificRoute && defaultRouteNode) { + nextPending.push(defaultRouteNode); + } + + return nextPending; + } + + /** Validates the graph and computes terminal node names. */ + validate(): void { + this._terminalNodeNames = validateGraph(this.nodes, this.edges); + } +} diff --git a/core/src/workflow-next/trigger.ts b/core/src/workflow-next/trigger.ts new file mode 100644 index 000000000..9ce2f307d --- /dev/null +++ b/core/src/workflow-next/trigger.ts @@ -0,0 +1,26 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * A buffered trigger for a downstream node. + * + * Ported from `google/adk-python` `workflow/_trigger.py`. Unlike the previous + * TypeScript `Trigger` (a route-matching predicate), this is a plain data record + * describing *how* a target node should be invoked when its turn comes. + */ +export interface Trigger { + /** The input to pass to the triggered node. */ + input?: unknown; + + /** Whether this trigger should run the node in an isolated sub-branch. */ + useSubBranch?: boolean; + + /** The branch inherited from the predecessor node. */ + branch?: string; + + /** Scope tag explicitly propagated to this trigger. */ + isolationScope?: string; +} diff --git a/core/src/workflow-next/utils/graph_parser.ts b/core/src/workflow-next/utils/graph_parser.ts new file mode 100644 index 000000000..cae1ebaeb --- /dev/null +++ b/core/src/workflow-next/utils/graph_parser.ts @@ -0,0 +1,212 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Parses workflow edge items and chains into a flat list of {@link Edge}s. + * + * Ported from `google/adk-python` `workflow/utils/_graph_parser.py`. + */ + +import {BaseNode} from '../base_node.js'; +import { + ChainElement, + Edge, + EdgeItem, + NodeLike, + RouteValue, + RoutingMap, +} from '../graph.js'; +import {buildNode, isNodeLike, isPlainObject} from './workflow_graph_utils.js'; + +function isRouteValue(value: unknown): value is RouteValue { + const t = typeof value; + return t === 'string' || t === 'number' || t === 'boolean'; +} + +/** Expands a routing map into individual (from, to, route) triples. */ +function expandRoutingMap( + fromElement: ChainElement, + routingMap: RoutingMap, +): Array<[ChainElement, NodeLike | readonly NodeLike[], RouteValue]> { + const keys = Object.keys(routingMap); + if (keys.length === 0) { + throw new Error( + 'Routing map must not be empty. Provide at least one route -> node mapping.', + ); + } + + const expanded: Array< + [ChainElement, NodeLike | readonly NodeLike[], RouteValue] + > = []; + for (const routeKey of keys) { + // Object keys are strings; numeric route keys arrive as numeric strings. + const normalizedKey: RouteValue = /^-?\d+$/.test(routeKey) + ? Number(routeKey) + : routeKey; + const target = routingMap[routeKey]; + if (Array.isArray(target)) { + for (const node of target) { + if (!isNodeLike(node)) { + throw new Error( + `Invalid node in fan-out tuple for route ${String(routeKey)}.`, + ); + } + } + } else if (!isNodeLike(target)) { + throw new Error( + `Invalid routing map value for route ${String(routeKey)}.`, + ); + } + if (!isRouteValue(normalizedKey)) { + throw new Error(`Invalid routing map key: ${String(routeKey)}.`); + } + expanded.push([ + fromElement, + target as NodeLike | readonly NodeLike[], + normalizedKey, + ]); + } + return expanded; +} + +/** Extracts all target nodes from a routing map, flattening fan-out arrays. */ +function nodesFromRoutingMap(routingMap: RoutingMap): NodeLike[] { + const nodes: NodeLike[] = []; + for (const target of Object.values(routingMap)) { + if (Array.isArray(target)) { + nodes.push(...(target as NodeLike[])); + } else { + nodes.push(target as NodeLike); + } + } + return nodes; +} + +/** Flattens a chain element into a list of individual nodes. */ +function flattenElement(element: ChainElement): NodeLike[] { + if (isPlainObject(element)) { + return nodesFromRoutingMap(element as RoutingMap); + } + if (Array.isArray(element)) { + return [...(element as readonly NodeLike[])]; + } + return [element as NodeLike]; +} + +/** Gets a node from the identity map or builds (and caches) it. */ +function getOrBuildNode( + nodeLike: NodeLike, + nodeMap: Map, +): BaseNode { + if (nodeLike === 'START') { + return buildNode('START'); + } + if (typeof nodeLike === 'object' || typeof nodeLike === 'function') { + const cached = nodeMap.get(nodeLike as object); + if (cached) { + return cached; + } + const built = buildNode(nodeLike); + // Only cache when a distinct wrapper was produced (or always, to preserve + // identity across repeated references within the same parse). + nodeMap.set(nodeLike as object, built); + return built; + } + return buildNode(nodeLike); +} + +function processExplicitEdge( + edge: Edge, + nodeMap: Map, + out: Edge[], +): void { + out.push( + new Edge( + getOrBuildNode(edge.fromNode, nodeMap), + getOrBuildNode(edge.toNode, nodeMap), + edge.route, + ), + ); +} + +function processRoutingMapEdge( + fromEl: ChainElement, + toEl: RoutingMap, + nodeMap: Map, + out: Edge[], +): void { + if (isPlainObject(fromEl)) { + throw new Error( + 'Consecutive routing maps are not allowed in a chain. Split them into separate edge items.', + ); + } + for (const [expFrom, expTo, route] of expandRoutingMap(fromEl, toEl)) { + for (const fromNode of flattenElement(expFrom)) { + for (const toNode of flattenElement(expTo as ChainElement)) { + out.push( + new Edge( + getOrBuildNode(fromNode, nodeMap), + getOrBuildNode(toNode, nodeMap), + route, + ), + ); + } + } + } +} + +function processUnconditionalEdge( + fromEl: ChainElement, + toEl: ChainElement, + nodeMap: Map, + out: Edge[], +): void { + for (const fromNode of flattenElement(fromEl)) { + for (const toNode of flattenElement(toEl)) { + out.push( + new Edge( + getOrBuildNode(fromNode, nodeMap), + getOrBuildNode(toNode, nodeMap), + null, + ), + ); + } + } +} + +function processChain( + chain: ChainElement[], + nodeMap: Map, + out: Edge[], +): void { + for (let i = 0; i < chain.length - 1; i++) { + const fromEl = chain[i]; + const toEl = chain[i + 1]; + if (isPlainObject(toEl)) { + processRoutingMapEdge(fromEl, toEl as RoutingMap, nodeMap, out); + } else { + processUnconditionalEdge(fromEl, toEl, nodeMap, out); + } + } +} + +/** Parses a list of edge items into a flat list of {@link Edge} objects. */ +export function parseEdgeItems(edgeItems: EdgeItem[]): Edge[] { + const nodeMap = new Map(); + const out: Edge[] = []; + + for (const item of edgeItems) { + if (item instanceof Edge) { + processExplicitEdge(item, nodeMap, out); + } else if (Array.isArray(item)) { + processChain(item, nodeMap, out); + } else { + throw new Error(`Invalid edge item type: ${typeof item}`); + } + } + + return out; +} diff --git a/core/src/workflow-next/utils/graph_validation.ts b/core/src/workflow-next/utils/graph_validation.ts new file mode 100644 index 000000000..29d22eae0 --- /dev/null +++ b/core/src/workflow-next/utils/graph_validation.ts @@ -0,0 +1,200 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Validates workflow graphs and computes terminal nodes. + * + * Ported from `google/adk-python` `workflow/utils/_graph_validation.py`. + * The Phase 3 static-schema check and the Phase 7 chat-agent wiring check are + * intentionally deferred to their respective phases. + */ + +import {BaseNode, START} from '../base_node.js'; +import {DEFAULT_ROUTE, Edge} from '../graph.js'; + +function validateDuplicateNodeNames(nodes: BaseNode[]): Set { + const counts = new Map(); + for (const node of nodes) { + counts.set(node.name, (counts.get(node.name) ?? 0) + 1); + } + const duplicates = [...counts.entries()] + .filter(([, c]) => c > 1) + .map(([name]) => name) + .sort(); + if (duplicates.length > 0) { + throw new Error( + `Graph validation failed. Duplicate node names found: ${JSON.stringify( + duplicates, + )}. Pass the exact same object instance to reuse a node, or give distinct nodes unique names.`, + ); + } + return new Set(counts.keys()); +} + +function validateStartNode(nodeNames: Set): void { + if (!nodeNames.has(START.name)) { + throw new Error( + `Graph validation failed. START node (name: '${START.name}') not found in graph nodes.`, + ); + } +} + +function validateStartEdges(edges: Edge[]): void { + for (const edge of edges) { + if (edge.fromNode.name === START.name && edge.route !== null) { + throw new Error( + `Graph validation failed. Edges from START must not have routes (edge to ${edge.toNode.name} has route ${String( + edge.route, + )}).`, + ); + } + } +} + +function validateConnectivity(edges: Edge[], nodeNames: Set): void { + const adj = new Map>(); + for (const name of nodeNames) { + adj.set(name, new Set()); + } + const toNodes = new Set(); + for (const edge of edges) { + adj.get(edge.fromNode.name)!.add(edge.toNode.name); + toNodes.add(edge.toNode.name); + } + + const reachable = new Set(); + const stack = [START.name]; + while (stack.length > 0) { + const node = stack.pop()!; + if (reachable.has(node)) { + continue; + } + reachable.add(node); + for (const next of adj.get(node) ?? []) { + if (!reachable.has(next)) { + stack.push(next); + } + } + } + + const unreachable = [...nodeNames].filter((n) => !reachable.has(n)).sort(); + if (unreachable.length > 0) { + throw new Error( + `Graph validation failed. The following nodes are unreachable from START: ${JSON.stringify( + unreachable, + )}`, + ); + } + if (toNodes.has(START.name)) { + throw new Error( + 'Graph validation failed. START node must not have incoming edges.', + ); + } +} + +function validateDuplicateEdges(edges: Edge[]): void { + const seen = new Set(); + for (const edge of edges) { + const key = `${edge.fromNode.name}\u0000${edge.toNode.name}`; + if (seen.has(key)) { + throw new Error( + `Graph validation failed. Duplicate edge found: from=${edge.fromNode.name}, to=${edge.toNode.name}`, + ); + } + seen.add(key); + } +} + +function validateDefaultRoutes(edges: Edge[]): void { + const defaultRouteEdges = new Map(); + for (const edge of edges) { + if (Array.isArray(edge.route) && edge.route.includes(DEFAULT_ROUTE)) { + throw new Error( + `Graph validation failed. DEFAULT_ROUTE cannot be combined with other routes in a list (edge from=${edge.fromNode.name}, to=${edge.toNode.name}). Use a separate edge for DEFAULT_ROUTE.`, + ); + } + if (edge.route === DEFAULT_ROUTE) { + const from = edge.fromNode.name; + if (defaultRouteEdges.has(from)) { + throw new Error( + `Graph validation failed. Multiple DEFAULT_ROUTE edges found from node ${from} to ${defaultRouteEdges.get( + from, + )} and ${edge.toNode.name}`, + ); + } + defaultRouteEdges.set(from, edge.toNode.name); + } + } +} + +function detectUnconditionalCycles( + edges: Edge[], + nodeNames: Set, +): void { + const adj = new Map(); + for (const name of nodeNames) { + adj.set(name, []); + } + for (const edge of edges) { + if (edge.route === null) { + adj.get(edge.fromNode.name)!.push(edge.toNode.name); + } + } + + const inStack = new Set(); + const done = new Set(); + + const dfs = (node: string, path: string[]): void => { + inStack.add(node); + path.push(node); + for (const neighbor of adj.get(node) ?? []) { + if (inStack.has(neighbor)) { + const cycleStart = path.indexOf(neighbor); + const cycle = [...path.slice(cycleStart), neighbor]; + throw new Error( + `Graph validation failed. Unconditional cycle detected: ${cycle.join( + ' -> ', + )}. Cycles must include at least one conditional (routed) edge to avoid infinite loops.`, + ); + } + if (!done.has(neighbor)) { + dfs(neighbor, path); + } + } + path.pop(); + inStack.delete(node); + done.add(node); + }; + + for (const name of nodeNames) { + if (!done.has(name)) { + dfs(name, []); + } + } +} + +function computeTerminalNodes(nodes: BaseNode[], edges: Edge[]): Set { + const fromNames = new Set(edges.map((e) => e.fromNode.name)); + return new Set( + nodes + .filter((n) => n.name !== START.name && !fromNames.has(n.name)) + .map((n) => n.name), + ); +} + +/** + * Validates the workflow graph and returns the set of terminal node names. + */ +export function validateGraph(nodes: BaseNode[], edges: Edge[]): Set { + const nodeNames = validateDuplicateNodeNames(nodes); + validateStartNode(nodeNames); + validateStartEdges(edges); + validateConnectivity(edges, nodeNames); + validateDuplicateEdges(edges); + validateDefaultRoutes(edges); + detectUnconditionalCycles(edges, nodeNames); + return computeTerminalNodes(nodes, edges); +} diff --git a/core/src/workflow-next/utils/workflow_graph_utils.ts b/core/src/workflow-next/utils/workflow_graph_utils.ts new file mode 100644 index 000000000..e9be061df --- /dev/null +++ b/core/src/workflow-next/utils/workflow_graph_utils.ts @@ -0,0 +1,134 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type {ZodType} from 'zod'; +import {BaseAgent, isBaseAgent} from '../../agents/base_agent.js'; +import {AuthConfig} from '../../auth/auth_tool.js'; +import {BaseTool, isBaseTool} from '../../tools/base_tool.js'; +import {BaseNode, START} from '../base_node.js'; +import {NodeLike} from '../graph.js'; +import {FunctionNode, FunctionNodeHandler} from '../nodes/function_node.js'; +import {LLMAgentWrapper} from '../nodes/llm_agent_wrapper.js'; +import {ParallelWorker} from '../nodes/parallel_worker.js'; +import {ToolNode} from '../nodes/tool_node.js'; +import {RetryConfig} from '../retry_config.js'; + +/** + * Property overrides applied when building a node from a {@link NodeLike}. + */ +export interface BuildNodeOptions { + name?: string; + rerunOnResume?: boolean; + retryConfig?: RetryConfig; + timeout?: number; + inputSchema?: ZodType; + outputSchema?: ZodType; + stateSchema?: ZodType; + authConfig?: AuthConfig; + /** If true, wrap the built node in a {@link ParallelWorker}. */ + parallelWorker?: boolean; + /** Concurrency limit for the parallel worker (requires `parallelWorker`). */ + maxParallelWorkers?: number; +} + +/** + * Returns whether a value is a plain object literal (a `RoutingMap`) rather than + * a class instance such as a node/tool/agent. + */ +export function isPlainObject(value: unknown): boolean { + if (typeof value !== 'object' || value === null || Array.isArray(value)) { + return false; + } + const proto = Object.getPrototypeOf(value); + return proto === Object.prototype || proto === null; +} + +/** + * Returns whether a value can be converted into a workflow node via + * {@link buildNode}. + */ +export function isNodeLike(value: unknown): value is NodeLike { + return ( + value === 'START' || + value instanceof BaseNode || + isBaseTool(value) || + typeof value === 'function' || + isAgentLike(value) + ); +} + +/** + * Converts a {@link NodeLike} into a concrete {@link BaseNode}. + * + * Supported now: the `'START'` sentinel, existing {@link BaseNode} instances, + * plain functions (→ {@link FunctionNode}), and {@link BaseTool}s (→ + * {@link ToolNode}). Wrapping agents (`LlmAgent`) lands in Phase 7. + */ +export function buildNode( + nodeLike: NodeLike, + options: BuildNodeOptions = {}, +): BaseNode { + if (options.maxParallelWorkers !== undefined && !options.parallelWorker) { + throw new Error( + 'maxParallelWorkers can only be set when parallelWorker is true.', + ); + } + + const built = buildInnerNode(nodeLike, options); + + if (options.parallelWorker) { + if (nodeLike === 'START') { + throw new Error('ParallelWorker cannot wrap a START node.'); + } + return new ParallelWorker(built, { + maxParallelWorkers: options.maxParallelWorkers, + retryConfig: options.retryConfig, + timeout: options.timeout, + }); + } + return built; +} + +function buildInnerNode( + nodeLike: NodeLike, + options: BuildNodeOptions, +): BaseNode { + if (nodeLike === 'START') { + return START; + } + if (nodeLike instanceof BaseNode) { + // TODO(phase-3+): apply property overrides via a clone when options differ. + return nodeLike; + } + if (isBaseTool(nodeLike)) { + return new ToolNode(nodeLike as BaseTool, options); + } + if (typeof nodeLike === 'function') { + const name = options.name ?? (nodeLike as {name?: string}).name; + if (!name) { + throw new Error( + 'node(): the wrapped function has no name; pass {name} explicitly.', + ); + } + return new FunctionNode(name, nodeLike as FunctionNodeHandler, options); + } + if (isBaseAgent(nodeLike) || isAgentLike(nodeLike)) { + return new LLMAgentWrapper(nodeLike as BaseAgent, options); + } + throw new Error( + `build_node: unsupported node-like value of type ${typeof nodeLike}.`, + ); +} + +/** Heuristic: an agent-like value exposes a `runAsync` generator method. */ +function isAgentLike(value: unknown): boolean { + return ( + typeof value === 'object' && + value !== null && + 'runAsync' in value && + typeof (value as {runAsync?: unknown}).runAsync === 'function' + ); +} diff --git a/core/src/workflow-next/workflow.ts b/core/src/workflow-next/workflow.ts new file mode 100644 index 000000000..e9c6ed639 --- /dev/null +++ b/core/src/workflow-next/workflow.ts @@ -0,0 +1,450 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {Event} from '../events/event.js'; +import {BaseNode, BaseNodeConfig} from './base_node.js'; +import {BranchPath} from './branch_path.js'; +import {DynamicNodeScheduler} from './dynamic_node_scheduler.js'; +import {EdgeItem, Graph, RouteValue} from './graph.js'; +import {NodeContext} 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'; + +/** + * 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}. + */ +export interface WorkflowConfig extends BaseNodeConfig { + /** Edge definitions used to build the workflow graph. */ + edges?: EdgeItem[]; + /** + * An imperative entry function driving execution via `ctx.runNode(...)`. + * Mutually exclusive with {@link edges}. + */ + dynamicEntry?: DynamicEntry; + /** + * Maximum number of graph-scheduled nodes running in parallel. `undefined` + * means unlimited. Does not throttle dynamic (`ctx.runNode`) children. + */ + maxConcurrency?: number; +} + +/** + * 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(); + errorShutDown = false; +} + +interface CompletedTask { + name: string; + childCtx?: NodeContext; + 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). + */ +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".`, + ); + } + this.maxConcurrency = config.maxConcurrency; + this.dynamicEntry = config.dynamicEntry; + if (hasEdges) { + this.graph = Graph.fromEdgeItems(config.edges!); + this.graph.validate(); + } + } + + // eslint-disable-next-line require-yield + 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(); + ctx.scheduler = new DynamicNodeScheduler(dynamicState); + + if (this.dynamicEntry) { + await this.runDynamicEntry(ctx, nodeInput, dynamicState); + return; + } + + const loop = new LoopState(); + + // --- SETUP --- + this.seedStartTriggers(loop, nodeInput); + + // --- LOOP --- + await this.runLoop(loop, ctx); + + 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; + } + } + + // --- 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): 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); + 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 && 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)!; + + let runId = nodeState.runId; + if (!runId) { + nodeState.runCounter += 1; + runId = String(nodeState.runCounter); + nodeState.runId = runId; + } + + // Static graph nodes are managed by this loop directly, bypassing the + // dynamic scheduler (which serves user-initiated ctx.runNode() calls). + const task: Promise = executeChildNode( + ctx, + node, + trigger.input, + { + 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, + ): 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 | 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 = BranchPath.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): Promise { + // Await outstanding tasks so their events flush; failures are swallowed + // because the workflow is already shutting down on error. + const outstanding = [...loop.pending.values()]; + loop.pending.clear(); + await Promise.allSettled(outstanding); + } +} diff --git a/core/test/workflow-next/workflow_test.ts b/core/test/workflow-next/workflow_test.ts new file mode 100644 index 000000000..d6c78d7d5 --- /dev/null +++ b/core/test/workflow-next/workflow_test.ts @@ -0,0 +1,190 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {describe, expect, it} from 'vitest'; +import {BaseAgent} from '../../src/agents/base_agent.js'; +import {InvocationContext} from '../../src/agents/invocation_context.js'; +import {createEvent, Event} from '../../src/events/event.js'; +import {PluginManager} from '../../src/plugins/plugin_manager.js'; +import {Session} from '../../src/sessions/session.js'; +import {BaseNode} from '../../src/workflow-next/base_node.js'; +import {DEFAULT_ROUTE} from '../../src/workflow-next/graph.js'; +import {NodeContext} from '../../src/workflow-next/node_context.js'; +import {EventChannel} from '../../src/workflow-next/utils/event_channel.js'; +import {Workflow} from '../../src/workflow-next/workflow.js'; + +function createIc(): InvocationContext { + const session = { + id: 's1', + appName: 'app', + userId: 'u', + events: [], + state: {}, + lastUpdateTime: Date.now(), + } as unknown as Session; + return new InvocationContext({ + invocationId: 'inv-1', + session, + agent: { + name: 'wf', + runAsync: async function* () {}, + } as unknown as BaseAgent, + pluginManager: new PluginManager(), + }); +} + +async function driveWorkflow( + workflow: Workflow, + input?: unknown, +): Promise<{events: Event[]; output: unknown}> { + const channel = new EventChannel(); + const root = new NodeContext({ + invocationContext: createIc(), + channel, + nodePath: '', + runId: 'root', + }); + const events: Event[] = []; + const run = root.runNode(workflow, input, {useAsOutput: true}).then( + () => channel.close(), + (err) => channel.fail(err), + ); + for await (const ev of channel) { + events.push(ev); + } + await run; + return {events, output: root.output}; +} + +// Yields whatever the function returns (value | Event). +class FnNode extends BaseNode { + constructor( + name: string, + private readonly fn: (ctx: NodeContext, input: unknown) => unknown, + ) { + super({name}); + } + protected async *runImpl(ctx: NodeContext, input: unknown) { + yield await this.fn(ctx, input); + } +} + +// 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); + }); +}); From ea2f54816cedbe16f0fe58c37d3f9efe1fca8a8d Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Wed, 22 Jul 2026 14:35:05 -0700 Subject: [PATCH 07/41] feat(workflow-next): Phase 3 user-facing node API Add the node() factory and subclassable Node, plus concrete node types: FunctionNode (idiomatic (ctx,input) handler with value/promise/generator support, Content coercion, state-delta capture, zod schema validation), JoinNode (fan-in barrier), and ToolNode (wraps a BaseTool, bridges to the tool Context, coerces args). buildNode now wraps functions and tools. --- core/src/workflow-next/node.ts | 61 +++++ core/src/workflow-next/nodes/function_node.ts | 170 +++++++++++++ core/src/workflow-next/nodes/join_node.ts | 34 +++ core/src/workflow-next/nodes/tool_node.ts | 106 ++++++++ core/test/workflow-next/node_api_test.ts | 231 ++++++++++++++++++ 5 files changed, 602 insertions(+) create mode 100644 core/src/workflow-next/node.ts create mode 100644 core/src/workflow-next/nodes/function_node.ts create mode 100644 core/src/workflow-next/nodes/join_node.ts create mode 100644 core/src/workflow-next/nodes/tool_node.ts create mode 100644 core/test/workflow-next/node_api_test.ts diff --git a/core/src/workflow-next/node.ts b/core/src/workflow-next/node.ts new file mode 100644 index 000000000..6f085838b --- /dev/null +++ b/core/src/workflow-next/node.ts @@ -0,0 +1,61 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {Event} from '../events/event.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}. + * + * Mirrors `google/adk-python` `workflow/_node.py::Node`. The `parallel_worker` + * capability is added in Phase 6. + */ +export abstract class Node< + 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-next/nodes/function_node.ts b/core/src/workflow-next/nodes/function_node.ts new file mode 100644 index 000000000..b250ab25c --- /dev/null +++ b/core/src/workflow-next/nodes/function_node.ts @@ -0,0 +1,170 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {AuthConfig} from '../../auth/auth_tool.js'; +import {createEvent, Event, isEvent} from '../../events/event.js'; +import {BaseNode, BaseNodeConfig, isContent, toContent} from '../base_node.js'; +import {NodeContext} from '../node_context.js'; + +/** + * A value a {@link FunctionNodeHandler} may return or yield. + */ +export type FunctionNodeResult = + | TOutput + | Event + | null + | undefined + | void; + +/** + * The handler wrapped by a {@link FunctionNode}. + * + * Unlike Python's `FunctionNode` (which binds named parameters from `ctx.state` + * or `node_input` via runtime signature introspection), the TypeScript form + * uses the idiomatic explicit `(ctx, input)` signature. Read `ctx.state` + * directly for state-bound values. It may return a value/`Event`, a Promise, or + * a (sync/async) generator of those. + */ +export type FunctionNodeHandler = ( + ctx: NodeContext, + input: TInput, +) => + | FunctionNodeResult + | Promise> + | Generator, void, unknown> + | AsyncGenerator, void, unknown>; + +/** + * Options for a {@link FunctionNode}. + */ +export interface FunctionNodeConfig extends Partial< + Omit +> { + /** + * If set, the framework requests user authentication before running (Phase 5 + * enables the auth gate; stored here now for API parity). + */ + authConfig?: AuthConfig; +} + +/** + * A node that wraps a plain function, async function, or (sync/async) generator. + * + * Ported (TS-idiomatic subset) from `google/adk-python` `_function_node.py`. + * Return-value handling: + * - `Event` → emitted as-is (output validated against `outputSchema`) + * - genai `Content` → emitted as the event content + * - `null`/`undefined` → skipped (unless there are pending state deltas) + * - anything else → emitted as `Event(output=value)` + * State written via `ctx.state` during execution is attached to emitted events. + */ +export class FunctionNode extends BaseNode< + TInput, + TOutput +> { + readonly authConfig?: AuthConfig; + private readonly handler: FunctionNodeHandler; + + constructor( + name: string, + handler: FunctionNodeHandler, + config: FunctionNodeConfig = {}, + ) { + if (typeof handler !== 'function') { + throw new TypeError('FunctionNode handler must be a function.'); + } + super({name, ...config}); + this.handler = handler; + this.authConfig = config.authConfig; + } + + protected async *runImpl( + ctx: NodeContext, + input: TInput, + ): AsyncGenerator { + // TODO(phase-5): auth gate (authConfig -> adk_request_credential interrupt). + const result = this.handler(ctx, input); + + if (isAsyncIterable(result)) { + for await (const item of result) { + yield item; + } + } else if (isSyncGenerator(result)) { + for (const item of result) { + yield item; + } + } else { + // Plain value or Promise of a value. + yield await (result as Promise>); + } + } + + protected override toEvent(ctx: NodeContext, data: unknown): Event | null { + const stateDelta = + Object.keys(ctx.actions.stateDelta).length > 0 + ? {...ctx.actions.stateDelta} + : undefined; + + if (data === null || data === undefined) { + return stateDelta + ? createEvent({ + author: this.name, + invocationId: ctx.invocationId, + branch: ctx.branch, + actions: {stateDelta}, + }) + : null; + } + + if (isEvent(data)) { + const event = data as Event; + if (event.output !== undefined) { + event.output = this.validateOutput(event.output); + } + if (stateDelta) { + Object.assign(event.actions.stateDelta, stateDelta); + } + return event; + } + + if (isContent(data)) { + return createEvent({ + author: this.name, + invocationId: ctx.invocationId, + branch: ctx.branch, + content: data, + actions: stateDelta ? {stateDelta} : undefined, + }); + } + + const output = this.validateOutput(data); + return createEvent({ + author: this.name, + invocationId: ctx.invocationId, + branch: ctx.branch, + content: toContent(output), + output, + actions: stateDelta ? {stateDelta} : undefined, + }); + } +} + +function isAsyncIterable(value: unknown): value is AsyncIterable { + return ( + value != null && + typeof (value as AsyncIterable)[Symbol.asyncIterator] === + 'function' + ); +} + +function isSyncGenerator(value: unknown): value is Generator { + return ( + value != null && + typeof value !== 'string' && + typeof (value as Iterable)[Symbol.iterator] === 'function' && + typeof (value as Generator).next === 'function' + ); +} diff --git a/core/src/workflow-next/nodes/join_node.ts b/core/src/workflow-next/nodes/join_node.ts new file mode 100644 index 000000000..23543d338 --- /dev/null +++ b/core/src/workflow-next/nodes/join_node.ts @@ -0,0 +1,34 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {createEvent, Event} from '../../events/event.js'; +import {BaseNode} from '../base_node.js'; +import {NodeContext} from '../node_context.js'; + +/** + * A fan-in barrier node: it waits for ALL of its predecessors to complete, then + * emits the aggregated inputs (a map of predecessor name → output) as its + * output. + * + * Ported from `google/adk-python` `workflow/_join_node.py`. + */ +export class JoinNode extends BaseNode { + override get requiresAllPredecessors(): boolean { + return true; + } + + protected async *runImpl( + ctx: NodeContext, + input: unknown, + ): AsyncGenerator { + yield createEvent({ + author: this.name, + invocationId: ctx.invocationId, + branch: ctx.branch, + output: input, + }); + } +} diff --git a/core/src/workflow-next/nodes/tool_node.ts b/core/src/workflow-next/nodes/tool_node.ts new file mode 100644 index 000000000..b70825645 --- /dev/null +++ b/core/src/workflow-next/nodes/tool_node.ts @@ -0,0 +1,106 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {Context} from '../../agents/context.js'; +import {createEvent, Event} from '../../events/event.js'; +import {BaseTool} from '../../tools/base_tool.js'; +import {randomUUID} from '../../utils/env_aware_utils.js'; +import {BaseNode, BaseNodeConfig, isContent} from '../base_node.js'; +import {NodeContext} from '../node_context.js'; +/** Options for a {@link ToolNode}. */ +export interface ToolNodeConfig extends Partial> { + /** Optional name override; defaults to the tool's name. */ + name?: string; +} + +/** + * A node that wraps an ADK {@link BaseTool} and invokes it with the node input + * as its arguments. + * + * Ported from `google/adk-python` `workflow/_tool_node.py`. The node input is + * coerced to a tool-args object: genai `Content` → its text; a JSON string → + * parsed object; `null`/empty → `{}`. + */ +export class ToolNode extends BaseNode { + readonly tool: BaseTool; + + constructor(tool: BaseTool, config: ToolNodeConfig = {}) { + super({name: config.name ?? tool.name, ...config}); + this.tool = tool; + } + + protected async *runImpl( + ctx: NodeContext, + input: unknown, + ): AsyncGenerator { + const toolContext = new Context({ + invocationContext: ctx.invocationContext, + functionCallId: randomUUID(), + }); + + const args = coerceToolArgs(input); + const response = await this.tool.runAsync({args, toolContext}); + + const stateDelta = + Object.keys(toolContext.actions.stateDelta).length > 0 + ? {...toolContext.actions.stateDelta} + : undefined; + + if (response !== undefined && response !== null) { + yield createEvent({ + author: this.name, + invocationId: ctx.invocationId, + branch: ctx.branch, + output: response, + actions: stateDelta ? {stateDelta} : undefined, + }); + } else if (stateDelta) { + yield createEvent({ + author: this.name, + invocationId: ctx.invocationId, + branch: ctx.branch, + actions: {stateDelta}, + }); + } + } +} + +/** Coerces arbitrary node input into a tool-arguments record. */ +function coerceToolArgs(input: unknown): Record { + let args: unknown = input; + + if (isContent(args)) { + args = extractText(args); + } + + if (typeof args === 'string') { + const trimmed = args.trim(); + if (!trimmed) { + args = null; + } else { + try { + args = JSON.parse(trimmed); + } catch { + // Leave as the raw string; validated below. + } + } + } + + if (args === null || args === undefined) { + return {}; + } + if (typeof args !== 'object' || Array.isArray(args)) { + throw new TypeError( + 'The input to ToolNode must be a dictionary of tool arguments or null, ' + + `but got ${typeof args}.`, + ); + } + return args as Record; +} + +function extractText(content: {parts?: Array<{text?: string}>}): string { + return (content.parts ?? []).map((p) => p.text ?? '').join(''); +} diff --git a/core/test/workflow-next/node_api_test.ts b/core/test/workflow-next/node_api_test.ts new file mode 100644 index 000000000..e478c4855 --- /dev/null +++ b/core/test/workflow-next/node_api_test.ts @@ -0,0 +1,231 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {describe, expect, it} from 'vitest'; +import {z} from 'zod'; +import {BaseAgent} from '../../src/agents/base_agent.js'; +import {InvocationContext} from '../../src/agents/invocation_context.js'; +import {Event} from '../../src/events/event.js'; +import {PluginManager} from '../../src/plugins/plugin_manager.js'; +import {Session} from '../../src/sessions/session.js'; +import {BaseTool} from '../../src/tools/base_tool.js'; +import {BaseNode} from '../../src/workflow-next/base_node.js'; +import {node, Node} from '../../src/workflow-next/node.js'; +import {NodeContext} from '../../src/workflow-next/node_context.js'; +import {FunctionNode} from '../../src/workflow-next/nodes/function_node.js'; +import {JoinNode} from '../../src/workflow-next/nodes/join_node.js'; +import {LLMAgentWrapper} from '../../src/workflow-next/nodes/llm_agent_wrapper.js'; +import {ToolNode} from '../../src/workflow-next/nodes/tool_node.js'; +import {EventChannel} from '../../src/workflow-next/utils/event_channel.js'; +import {Workflow} from '../../src/workflow-next/workflow.js'; + +function createIc(): InvocationContext { + const session = { + id: 's1', + appName: 'app', + userId: 'u', + events: [], + state: {}, + lastUpdateTime: Date.now(), + } as unknown as Session; + return new InvocationContext({ + invocationId: 'inv-1', + session, + agent: { + name: 'wf', + runAsync: async function* () {}, + } as unknown as BaseAgent, + pluginManager: new PluginManager(), + }); +} + +async function runNode( + n: BaseNode, + input?: unknown, +): Promise<{events: Event[]; output: unknown}> { + const channel = new EventChannel(); + const root = new NodeContext({ + invocationContext: createIc(), + channel, + nodePath: '', + runId: 'root', + }); + const events: Event[] = []; + const run = root.runNode(n, input, {useAsOutput: true}).then( + () => channel.close(), + (err) => channel.fail(err), + ); + for await (const ev of channel) { + events.push(ev); + } + await run; + return {events, output: root.output}; +} + +async function driveWorkflow(wf: Workflow, input?: unknown): Promise { + return (await runNode(wf, input)).output; +} + +describe('Phase 3 — FunctionNode', () => { + it('boxes a plain return value into an output event', async () => { + const n = new FunctionNode('greet', (_c, input) => `hi ${input}`); + const {output, events} = await runNode(n, 'x'); + expect(output).toBe('hi x'); + expect(events[0].output).toBe('hi x'); + }); + + it('awaits an async handler', async () => { + const n = new FunctionNode('a', async (_c, input) => { + await Promise.resolve(); + return `async:${input}`; + }); + expect((await runNode(n, 'v')).output).toBe('async:v'); + }); + + it('supports sync generators (multiple events, last output wins)', async () => { + const n = new FunctionNode('gen', function* (_c, input) { + yield `${input}-1`; + yield `${input}-2`; + }); + const {events, output} = await runNode(n, 'g'); + expect(events.map((e) => e.output)).toEqual(['g-1', 'g-2']); + expect(output).toBe('g-2'); + }); + + it('supports async generators', async () => { + const n = new FunctionNode('agen', async function* (_c) { + yield 'one'; + await Promise.resolve(); + yield 'two'; + }); + expect((await runNode(n)).events.map((e) => e.output)).toEqual([ + 'one', + 'two', + ]); + }); + + it('skips null returns but keeps state deltas', async () => { + const n = new FunctionNode('writer', (ctx) => { + ctx.state.set('flag', true); + return null; + }); + const {events} = await runNode(n); + expect(events).toHaveLength(1); + expect(events[0].output).toBeUndefined(); + expect(events[0].actions.stateDelta['flag']).toBe(true); + }); + + it('attaches state deltas to output events', async () => { + const n = new FunctionNode('w', (ctx) => { + ctx.state.set('count', 3); + return 'ok'; + }); + const {events} = await runNode(n); + expect(events[0].output).toBe('ok'); + expect(events[0].actions.stateDelta['count']).toBe(3); + }); + + it('validates output against an outputSchema', async () => { + const schema = z.object({n: z.number()}); + const good = new FunctionNode('g', () => ({n: 5}), {outputSchema: schema}); + expect((await runNode(good)).output).toEqual({n: 5}); + + const bad = new FunctionNode('b', () => ({n: 'not-a-number'}), { + outputSchema: schema, + }); + await expect(runNode(bad)).rejects.toThrow(); + }); +}); + +describe('Phase 3 — node() factory', () => { + it('wraps a function, deriving the name', () => { + function classify() { + return 'ok'; + } + const n = node(classify); + expect(n).toBeInstanceOf(FunctionNode); + expect(n.name).toBe('classify'); + }); + + it('wraps a function with an explicit name override', () => { + const n = node((_c: NodeContext, input: unknown) => input, { + name: 'passthru', + }); + expect(n.name).toBe('passthru'); + }); + + it('wraps a BaseTool into a ToolNode', () => { + const tool = new EchoTool(); + const n = node(tool); + expect(n).toBeInstanceOf(ToolNode); + expect(n.name).toBe('echo'); + }); + + it('returns an existing BaseNode unchanged', () => { + const existing = new FunctionNode('keep', () => 1); + expect(node(existing)).toBe(existing); + }); + + it('wraps an agent into an LLMAgentWrapper', () => { + const fakeAgent = {name: 'a', runAsync: async function* () {}}; + const n = node(fakeAgent as unknown as never); + expect(n).toBeInstanceOf(LLMAgentWrapper); + expect(n.name).toBe('a'); + }); +}); + +describe('Phase 3 — Node subclass', () => { + class DoubleNode extends Node { + protected async *runNodeImpl(_ctx: NodeContext, input: number) { + yield input * 2; + } + } + + it('runs a subclass via runNodeImpl', async () => { + expect((await runNode(new DoubleNode({name: 'double'}), 5)).output).toBe( + 10, + ); + }); +}); + +describe('Phase 3 — JoinNode & ToolNode in a workflow', () => { + it('fans in with the real JoinNode', async () => { + const a = node((_c: NodeContext, input: unknown) => `A(${input})`, { + name: 'A', + }); + const b = node((_c: NodeContext, input: unknown) => `B(${input})`, { + name: 'B', + }); + const join = new JoinNode({name: 'join'}); + const wf = new Workflow({name: 'fan', edges: [['START', [a, b], join]]}); + expect(await driveWorkflow(wf, 'x')).toEqual({A: 'A(x)', B: 'B(x)'}); + }); + + it('runs a ToolNode with object args', async () => { + const wf = new Workflow({ + name: 'tool_wf', + edges: [['START', new ToolNode(new EchoTool())]], + }); + expect(await driveWorkflow(wf, {msg: 'hi'})).toEqual({ + echoed: {msg: 'hi'}, + }); + }); + + it('coerces a JSON-string ToolNode input to args', async () => { + const {output} = await runNode(new ToolNode(new EchoTool()), '{"a":1}'); + expect(output).toEqual({echoed: {a: 1}}); + }); +}); + +// A trivial tool that echoes its args back. +class EchoTool extends BaseTool { + constructor() { + super({name: 'echo', description: 'Echoes the input args.'}); + } + async runAsync({args}: {args: Record}): Promise { + return {echoed: args}; + } +} From 3e16d7f9e1d8e113f817b9efd6791346e3713245 Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Wed, 22 Jul 2026 14:35:13 -0700 Subject: [PATCH 08/41] feat(workflow-next): Phase 4 dynamic node scheduling Add ScheduleDynamicNode + DynamicNodeScheduler (fresh execution, concurrent-call dedup by run path) and wire ctx.runNode() through it with per-name run-id generation. Workflow gains a dynamicEntry mode for imperative orchestration (loops, conditionals, node-as-tool), fixing the cycle-hang for the imperative path. The static loop bypasses the scheduler. --- .../workflow-next/dynamic_node_scheduler.ts | 106 ++++++++++++ .../workflow-next/schedule_dynamic_node.ts | 76 +++++++++ .../workflow-next/dynamic_workflow_test.ts | 155 ++++++++++++++++++ 3 files changed, 337 insertions(+) create mode 100644 core/src/workflow-next/dynamic_node_scheduler.ts create mode 100644 core/src/workflow-next/schedule_dynamic_node.ts create mode 100644 core/test/workflow-next/dynamic_workflow_test.ts diff --git a/core/src/workflow-next/dynamic_node_scheduler.ts b/core/src/workflow-next/dynamic_node_scheduler.ts new file mode 100644 index 000000000..c51c391da --- /dev/null +++ b/core/src/workflow-next/dynamic_node_scheduler.ts @@ -0,0 +1,106 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {BaseNode} from './base_node.js'; +import {NodeContext} 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'; + +/** + * 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 { + constructor(private readonly state: DynamicNodeState) {} + + 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}/${name}@${runId}`; + + const existing = this.state.runs.get(nodePath); + if (existing?.task) { + // Deduplicate concurrent calls: await the in-flight task. + return existing.task; + } + + // TODO(phase-5): lazy rehydration from session events + replay + // interception (dedup completed / resume waiting runs) goes here. + + 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(ctx, node, input, { + nodeName: name, + runId, + 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-next/schedule_dynamic_node.ts b/core/src/workflow-next/schedule_dynamic_node.ts new file mode 100644 index 000000000..0e0defd43 --- /dev/null +++ b/core/src/workflow-next/schedule_dynamic_node.ts @@ -0,0 +1,76 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type {BaseNode} from './base_node.js'; +import type {NodeContext} from './node_context.js'; +import {NodeState} from './node_state.js'; + +/** + * Options for scheduling a dynamic node via {@link ScheduleDynamicNode}. + */ +export interface ScheduleDynamicNodeOptions { + /** Deterministic tracking name; defaults to `node.name`. */ + nodeName?: string; + /** Unique id for this specific execution (used for path/dedup keys). */ + runId: string; + /** If true, the child's output replaces the caller's output. */ + useAsOutput?: boolean; + /** If true, run the child in an isolated sub-branch. */ + useSubBranch?: boolean; + /** Explicit branch override. */ + overrideBranch?: string; + /** Explicit isolation-scope override. */ + overrideIsolationScope?: string; +} + +/** + * Protocol for scheduling a dynamically-invoked node (via `ctx.runNode()`). + * + * Implementations handle fresh execution, deduplication of concurrent calls, + * and (Phase 5) resumption from session events. Ported from + * `google/adk-python` `workflow/_schedule_dynamic_node.py`. + */ +export interface ScheduleDynamicNode { + schedule( + ctx: NodeContext, + node: BaseNode, + input: unknown, + options: ScheduleDynamicNodeOptions, + ): Promise; +} + +/** + * Combines state, output, and the running task for a single dynamic node + * execution. + */ +export interface DynamicNodeRun { + state: NodeState; + output?: unknown; + task?: Promise; + transferToAgent?: string; +} + +/** + * State for tracking dynamic nodes scheduled via `ctx.runNode()`. + * + * Ported (Phase 4 subset) from `google/adk-python` + * `workflow/_dynamic_node_scheduler.py::DynamicNodeState`. Replay/rehydration + * fields are added in Phase 5. + */ +export class DynamicNodeState { + /** Dynamic node runs keyed by unique node path (e.g. `wf/node_a@1`). */ + readonly runs = new Map(); + + /** Union of unresolved interrupt ids across dynamic child nodes. */ + readonly interruptIds = new Set(); + + /** All in-flight dynamic node tasks. */ + getDynamicTasks(): Array> { + return [...this.runs.values()] + .map((run) => run.task) + .filter((task): task is Promise => task !== undefined); + } +} diff --git a/core/test/workflow-next/dynamic_workflow_test.ts b/core/test/workflow-next/dynamic_workflow_test.ts new file mode 100644 index 000000000..13cc5b006 --- /dev/null +++ b/core/test/workflow-next/dynamic_workflow_test.ts @@ -0,0 +1,155 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {describe, expect, it} from 'vitest'; +import {BaseAgent} from '../../src/agents/base_agent.js'; +import {InvocationContext} from '../../src/agents/invocation_context.js'; +import {Event} from '../../src/events/event.js'; +import {PluginManager} from '../../src/plugins/plugin_manager.js'; +import {Session} from '../../src/sessions/session.js'; +import {NodeContext} from '../../src/workflow-next/node_context.js'; +import {FunctionNode} from '../../src/workflow-next/nodes/function_node.js'; +import {EventChannel} from '../../src/workflow-next/utils/event_channel.js'; +import {Workflow} from '../../src/workflow-next/workflow.js'; + +function createIc(): InvocationContext { + const session = { + id: 's1', + appName: 'app', + userId: 'u', + events: [], + state: {}, + lastUpdateTime: Date.now(), + } as unknown as Session; + return new InvocationContext({ + invocationId: 'inv-1', + session, + agent: { + name: 'wf', + runAsync: async function* () {}, + } as unknown as BaseAgent, + pluginManager: new PluginManager(), + }); +} + +async function driveWorkflow( + wf: Workflow, + input?: unknown, +): Promise<{events: Event[]; output: unknown}> { + const channel = new EventChannel(); + const root = new NodeContext({ + invocationContext: createIc(), + channel, + nodePath: '', + runId: 'root', + }); + const events: Event[] = []; + const run = root.runNode(wf, input, {useAsOutput: true}).then( + () => channel.close(), + (err) => channel.fail(err), + ); + for await (const ev of channel) { + events.push(ev); + } + await run; + return {events, output: root.output}; +} + +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); + }); +}); From 31349cbfa3702a2d46dd52fe0ae5a665716b42b8 Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Wed, 22 Jul 2026 14:35:22 -0700 Subject: [PATCH 09/41] feat(workflow-next): Phase 5a human-in-the-loop core Add RequestInput and HITL utils (adk_request_input interrupt events, longRunningToolIds). BaseNode.run() normalizes a yielded/returned RequestInput into an interrupt event; the NodeRunner collects interrupt ids and the Workflow/scheduler propagate them so the workflow pauses and surfaces interrupt ids. Resume inputs flow back via ctx.resumeInputs. (Session rehydration, replay determinism, checkpoints and the auth gate remain as Phase 5b, pending Runner integration.) --- core/src/workflow-next/request_input.ts | 47 ++++++ core/src/workflow-next/utils/hitl_utils.ts | 99 ++++++++++++ core/test/workflow-next/hitl_test.ts | 173 +++++++++++++++++++++ 3 files changed, 319 insertions(+) create mode 100644 core/src/workflow-next/request_input.ts create mode 100644 core/src/workflow-next/utils/hitl_utils.ts create mode 100644 core/test/workflow-next/hitl_test.ts diff --git a/core/src/workflow-next/request_input.ts b/core/src/workflow-next/request_input.ts new file mode 100644 index 000000000..42bf8c9df --- /dev/null +++ b/core/src/workflow-next/request_input.ts @@ -0,0 +1,47 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type {ZodType} from 'zod'; +import {randomUUID} from '../utils/env_aware_utils.js'; + +/** Parameters for constructing a {@link RequestInput}. */ +export interface RequestInputParams { + /** The interrupt id (usually a function-call id). Auto-generated if omitted. */ + interruptId?: string; + /** Custom payload provided for resuming. */ + payload?: unknown; + /** A message to display to the user when requesting input. */ + message?: string; + /** The expected schema of the response. */ + responseSchema?: ZodType; +} + +/** + * A request for input from the user, yielded/returned by a node to pause a + * workflow (Human-in-the-Loop). The framework converts it to an interrupt + * event; the workflow surfaces the interrupt id to the caller, which later + * resumes by providing `resumeInputs[interruptId]`. + * + * Ported from `google/adk-python` `events/request_input.py`. + */ +export class RequestInput { + readonly interruptId: string; + readonly payload?: unknown; + readonly message?: string; + readonly responseSchema?: ZodType; + + constructor(params: RequestInputParams = {}) { + this.interruptId = params.interruptId ?? randomUUID(); + this.payload = params.payload; + this.message = params.message; + this.responseSchema = params.responseSchema; + } +} + +/** Type guard for {@link RequestInput}. */ +export function isRequestInput(value: unknown): value is RequestInput { + return value instanceof RequestInput; +} diff --git a/core/src/workflow-next/utils/hitl_utils.ts b/core/src/workflow-next/utils/hitl_utils.ts new file mode 100644 index 000000000..4f5686077 --- /dev/null +++ b/core/src/workflow-next/utils/hitl_utils.ts @@ -0,0 +1,99 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Utilities for Human-in-the-Loop (HITL) workflows. + * + * Ported (subset) from `google/adk-python` + * `workflow/utils/_workflow_hitl_utils.py`. The auth-credential helpers are + * added in the Phase 5 auth-gate follow-up. + */ + +import {Part} from '@google/genai'; +import {z} from 'zod'; +import {createEvent, Event} from '../../events/event.js'; +import {RequestInput} from '../request_input.js'; + +/** Function-call name marking a request-for-input interrupt. */ +export const REQUEST_INPUT_FUNCTION_CALL_NAME = 'adk_request_input'; + +/** Function-call name marking a request-for-credential interrupt. */ +export const REQUEST_CREDENTIAL_FUNCTION_CALL_NAME = 'adk_request_credential'; + +/** + * Creates an interrupt {@link Event} from a {@link RequestInput}. The event + * carries an `adk_request_input` function call and marks the interrupt id as a + * long-running tool id. + */ +export function createRequestInputEvent(requestInput: RequestInput): Event { + const args: Record = { + interruptId: requestInput.interruptId, + payload: requestInput.payload ?? null, + message: requestInput.message ?? null, + responseSchema: requestInput.responseSchema + ? z.toJSONSchema(requestInput.responseSchema) + : null, + }; + + return createEvent({ + content: { + role: 'model', + parts: [ + { + functionCall: { + name: REQUEST_INPUT_FUNCTION_CALL_NAME, + args, + id: requestInput.interruptId, + }, + }, + ], + }, + longRunningToolIds: [requestInput.interruptId], + }); +} + +/** Returns whether an event contains a `request_input` function call. */ +export function hasRequestInputFunctionCall(event: Event): boolean { + return (event.content?.parts ?? []).some( + (p) => p.functionCall?.name === REQUEST_INPUT_FUNCTION_CALL_NAME, + ); +} + +/** Returns whether an event contains an `adk_request_credential` function call. */ +export function hasAuthRequestFunctionCall(event: Event): boolean { + return (event.content?.parts ?? []).some( + (p) => p.functionCall?.name === REQUEST_CREDENTIAL_FUNCTION_CALL_NAME, + ); +} + +/** Extracts interrupt ids from `request_input` function calls in an event. */ +export function getRequestInputInterruptIds(event: Event): string[] { + const ids: string[] = []; + for (const part of event.content?.parts ?? []) { + const fc = part.functionCall; + if (fc && fc.name === REQUEST_INPUT_FUNCTION_CALL_NAME && fc.id) { + ids.push(fc.id); + } + } + return ids; +} + +/** + * Creates a `FunctionResponse` part answering a `request_input` interrupt, + * suitable for appending to a session as the user's resume response. + */ +export function createRequestInputResponse( + interruptId: string, + response: Record, +): Part { + return { + functionResponse: { + id: interruptId, + name: REQUEST_INPUT_FUNCTION_CALL_NAME, + response, + }, + }; +} diff --git a/core/test/workflow-next/hitl_test.ts b/core/test/workflow-next/hitl_test.ts new file mode 100644 index 000000000..097eb217e --- /dev/null +++ b/core/test/workflow-next/hitl_test.ts @@ -0,0 +1,173 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {describe, expect, it} from 'vitest'; +import {BaseAgent} from '../../src/agents/base_agent.js'; +import {InvocationContext} from '../../src/agents/invocation_context.js'; +import {Event} from '../../src/events/event.js'; +import {PluginManager} from '../../src/plugins/plugin_manager.js'; +import {Session} from '../../src/sessions/session.js'; +import {NodeContext} from '../../src/workflow-next/node_context.js'; +import {FunctionNode} from '../../src/workflow-next/nodes/function_node.js'; +import {RequestInput} from '../../src/workflow-next/request_input.js'; +import {EventChannel} from '../../src/workflow-next/utils/event_channel.js'; +import { + hasRequestInputFunctionCall, + REQUEST_INPUT_FUNCTION_CALL_NAME, +} from '../../src/workflow-next/utils/hitl_utils.js'; +import {Workflow} from '../../src/workflow-next/workflow.js'; + +function createIc(): InvocationContext { + const session = { + id: 's1', + appName: 'app', + userId: 'u', + events: [], + state: {}, + lastUpdateTime: Date.now(), + } as unknown as Session; + return new InvocationContext({ + invocationId: 'inv-1', + session, + agent: { + name: 'wf', + runAsync: async function* () {}, + } as unknown as BaseAgent, + pluginManager: new PluginManager(), + }); +} + +/** + * Drives a workflow, optionally supplying resume inputs (keyed by interrupt id). + * Returns the workflow output, its interrupt ids, and the streamed events. + */ +async function drive( + wf: Workflow, + input?: unknown, + resumeInputs: Record = {}, +): Promise<{output: unknown; interruptIds: string[]; events: Event[]}> { + const channel = new EventChannel(); + const root = new NodeContext({ + invocationContext: createIc(), + channel, + nodePath: '', + runId: 'root', + resumeInputs, + }); + const events: Event[] = []; + const wfCtxPromise = root.runNode(wf, input, {useAsOutput: true}); + const settle = wfCtxPromise.then( + () => channel.close(), + (err) => channel.fail(err), + ); + for await (const ev of channel) { + events.push(ev); + } + await settle; + const wfCtx = await wfCtxPromise; + return {output: root.output, interruptIds: wfCtx.interruptIds, events}; +} + +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'); + }); +}); From f49dd8ae6b204247ac282f74b6a2a72eb8c9eb41 Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Wed, 22 Jul 2026 14:35:29 -0700 Subject: [PATCH 10/41] feat(workflow-next): Phase 6 parallelism and branches Add BranchPath (hierarchical name@runId paths: createSubBranch, commonPrefix, isDescendantOf) and ParallelWorker (maps a list input across the inner node with bounded concurrency, order preservation, cancel-on-first-error). node()/buildNode support {parallelWorker, maxParallelWorkers}; sub-branch derivation and fan-in common-prefix now use BranchPath. --- core/src/workflow-next/branch_path.ts | 92 +++++++++++ .../workflow-next/nodes/parallel_worker.ts | 98 +++++++++++ core/test/workflow-next/parallel_test.ts | 155 ++++++++++++++++++ 3 files changed, 345 insertions(+) create mode 100644 core/src/workflow-next/branch_path.ts create mode 100644 core/src/workflow-next/nodes/parallel_worker.ts create mode 100644 core/test/workflow-next/parallel_test.ts diff --git a/core/src/workflow-next/branch_path.ts b/core/src/workflow-next/branch_path.ts new file mode 100644 index 000000000..63fd58395 --- /dev/null +++ b/core/src/workflow-next/branch_path.ts @@ -0,0 +1,92 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Hierarchical, dot-separated path for execution branches. Each segment is a + * node run, typically formatted as `name@runId` (or just `name`). + * + * Ported from `google/adk-python` `events/_branch_path.py::_BranchPath`. + * + * @example 'parent_agent@1.collect_tool@2.sub_workflow' + */ +export class BranchPath { + private readonly segments: string[]; + + constructor(segments: string[]) { + this.segments = [...segments]; + } + + /** Parses a dot-separated string into a {@link BranchPath}. */ + static fromString(path?: string | null): BranchPath { + if (!path) { + return new BranchPath([]); + } + return new BranchPath(path.split('.')); + } + + toString(): string { + return this.segments.join('.'); + } + + /** Returns a copy of the path segments. */ + getSegments(): string[] { + return [...this.segments]; + } + + /** Whether this path is a strict descendant of `ancestor`. */ + isDescendantOf(ancestor: BranchPath): boolean { + if (this.segments.length <= ancestor.segments.length) { + return false; + } + return ancestor.segments.every((seg, i) => this.segments[i] === seg); + } + + /** Returns a new path with a `name@runId` (or `name`) segment appended. */ + append(name: string, runId?: string): BranchPath { + const segment = runId !== undefined ? `${name}@${runId}` : name; + return new BranchPath([...this.segments, segment]); + } + + /** Finds the common prefix across a list of paths. */ + static commonPrefix(paths: BranchPath[]): BranchPath { + if (paths.length === 0) { + return new BranchPath([]); + } + const common: string[] = []; + const minLen = Math.min(...paths.map((p) => p.segments.length)); + for (let i = 0; i < minLen; i++) { + const seg = paths[0].segments[i]; + if (paths.every((p) => p.segments[i] === seg)) { + common.push(seg); + } else { + break; + } + } + return new BranchPath(common); + } + + /** + * Creates a new dot-separated sub-branch string by appending a segment. + * + * @example createSubBranch('parent', {name: 'child', runId: '1'}) -> 'parent.child@1' + * @example createSubBranch(undefined, {name: 'agent'}) -> 'agent' + */ + static createSubBranch( + baseBranch: string | undefined | null, + options: {name: string; runId?: string}, + ): string { + return BranchPath.fromString(baseBranch) + .append(options.name, options.runId) + .toString(); + } + + /** Finds the common prefix of a list of dot-separated branch strings. */ + static commonPrefixOf(branches: string[]): string { + return BranchPath.commonPrefix( + branches.map((b) => BranchPath.fromString(b)), + ).toString(); + } +} diff --git a/core/src/workflow-next/nodes/parallel_worker.ts b/core/src/workflow-next/nodes/parallel_worker.ts new file mode 100644 index 000000000..7d8e0e57b --- /dev/null +++ b/core/src/workflow-next/nodes/parallel_worker.ts @@ -0,0 +1,98 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {BaseNode} from '../base_node.js'; +import {NodeContext} from '../node_context.js'; +import {RetryConfig} from '../retry_config.js'; + +/** Options for a {@link ParallelWorker}. */ +export interface ParallelWorkerConfig { + /** Maximum number of items processed concurrently. `undefined` = unlimited. */ + maxParallelWorkers?: number; + retryConfig?: RetryConfig; + timeout?: number; +} + +/** + * A node that runs a wrapped node in parallel for each item of a list input, + * preserving order, bounded by `maxParallelWorkers`, cancelling on first error. + * + * Ported from `google/adk-python` `workflow/_parallel_worker.py`. A non-list + * input is treated as a single-element list. Each item runs via + * `ctx.runNode(inner, item, {useSubBranch: true})`; the node's output is the + * ordered list of the children's outputs. + */ +export class ParallelWorker extends BaseNode { + readonly maxParallelWorkers?: number; + private readonly inner: BaseNode; + + constructor(inner: BaseNode, config: ParallelWorkerConfig = {}) { + super({ + name: inner.name, + rerunOnResume: true, + retryConfig: config.retryConfig, + timeout: config.timeout, + }); + if ( + config.maxParallelWorkers !== undefined && + config.maxParallelWorkers < 1 + ) { + throw new Error('maxParallelWorkers must be greater than or equal to 1.'); + } + this.inner = inner; + this.maxParallelWorkers = config.maxParallelWorkers; + } + + protected async *runImpl( + ctx: NodeContext, + input: unknown, + ): AsyncGenerator { + const items = Array.isArray(input) ? input : [input]; + if (items.length === 0) { + yield []; + return; + } + + const results = new Array(items.length); + const poolSize = Math.min( + this.maxParallelWorkers ?? items.length, + items.length, + ); + + let nextIndex = 0; + let firstError: unknown; + + const worker = async (): Promise => { + for (;;) { + if (firstError !== undefined) { + return; + } + const i = nextIndex++; + if (i >= items.length) { + return; + } + try { + const child = await ctx.runNode(this.inner, items[i], { + useSubBranch: true, + }); + results[i] = child.output; + } catch (err) { + if (firstError === undefined) { + firstError = err; + } + return; + } + } + }; + + await Promise.all(Array.from({length: poolSize}, () => worker())); + + if (firstError !== undefined) { + throw firstError; + } + yield results; + } +} diff --git a/core/test/workflow-next/parallel_test.ts b/core/test/workflow-next/parallel_test.ts new file mode 100644 index 000000000..644bc1641 --- /dev/null +++ b/core/test/workflow-next/parallel_test.ts @@ -0,0 +1,155 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {describe, expect, it} from 'vitest'; +import {BaseAgent} from '../../src/agents/base_agent.js'; +import {InvocationContext} from '../../src/agents/invocation_context.js'; +import {Event} from '../../src/events/event.js'; +import {PluginManager} from '../../src/plugins/plugin_manager.js'; +import {Session} from '../../src/sessions/session.js'; +import {BranchPath} from '../../src/workflow-next/branch_path.js'; +import {node} from '../../src/workflow-next/node.js'; +import {NodeContext} from '../../src/workflow-next/node_context.js'; +import {FunctionNode} from '../../src/workflow-next/nodes/function_node.js'; +import {ParallelWorker} from '../../src/workflow-next/nodes/parallel_worker.js'; +import {EventChannel} from '../../src/workflow-next/utils/event_channel.js'; +import {Workflow} from '../../src/workflow-next/workflow.js'; + +function createIc(): InvocationContext { + const session = { + id: 's1', + appName: 'app', + userId: 'u', + events: [], + state: {}, + lastUpdateTime: Date.now(), + } as unknown as Session; + return new InvocationContext({ + invocationId: 'inv-1', + session, + agent: { + name: 'wf', + runAsync: async function* () {}, + } as unknown as BaseAgent, + pluginManager: new PluginManager(), + }); +} + +async function driveWorkflow(wf: Workflow, input?: unknown): Promise { + const channel = new EventChannel(); + const root = new NodeContext({ + invocationContext: createIc(), + channel, + nodePath: '', + runId: 'root', + }); + const run = root.runNode(wf, input, {useAsOutput: true}).then( + () => channel.close(), + (err) => channel.fail(err), + ); + // drain + for await (const _ of channel) { + void _; + } + await run; + return root.output; +} + +describe('Phase 6 — BranchPath', () => { + it('creates sub-branches with and without run ids', () => { + expect( + BranchPath.createSubBranch('parent', {name: 'child', runId: '1'}), + ).toBe('parent.child@1'); + expect(BranchPath.createSubBranch(undefined, {name: 'agent'})).toBe( + 'agent', + ); + }); + + it('computes the common prefix of branches', () => { + expect(BranchPath.commonPrefixOf(['a@1.b@2', 'a@1.c@3'])).toBe('a@1'); + expect(BranchPath.commonPrefixOf(['a@1', 'b@1'])).toBe(''); + expect(BranchPath.commonPrefixOf([])).toBe(''); + }); + + it('detects descendants', () => { + const parent = BranchPath.fromString('a@1'); + expect(BranchPath.fromString('a@1.b@2').isDescendantOf(parent)).toBe(true); + expect(BranchPath.fromString('a@1').isDescendantOf(parent)).toBe(false); + expect(BranchPath.fromString('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])).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')).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, [])).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 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])).toEqual([11, 21, 31]); + }); + + it('rejects maxParallelWorkers without parallelWorker', () => { + expect(() => + node((_c: NodeContext, x: unknown) => x, { + name: 'x', + maxParallelWorkers: 2, + }), + ).toThrow(/maxParallelWorkers/); + }); +}); From 42c18e1a37be236aa0b1d6fd9f57f287be8f42f0 Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Wed, 22 Jul 2026 14:35:37 -0700 Subject: [PATCH 11/41] feat(workflow-next): Phase 7a run LlmAgent as a node (single_turn) Add LLMAgentWrapper (single_turn): append the node input as a user turn, run the agent under the node's invocation context, and promote its final model text to the node output. buildNode now wraps BaseAgent (and agent-like objects), so node(agent) works and agents can be used directly in edges. (task/chat modes with FinishTaskTool, delegation, transfer and isolation scopes remain as Phase 7b.) --- .../workflow-next/nodes/llm_agent_wrapper.ts | 104 +++++++++++++ core/test/workflow-next/llm_agent_test.ts | 147 ++++++++++++++++++ 2 files changed, 251 insertions(+) create mode 100644 core/src/workflow-next/nodes/llm_agent_wrapper.ts create mode 100644 core/test/workflow-next/llm_agent_test.ts diff --git a/core/src/workflow-next/nodes/llm_agent_wrapper.ts b/core/src/workflow-next/nodes/llm_agent_wrapper.ts new file mode 100644 index 000000000..4af1aba00 --- /dev/null +++ b/core/src/workflow-next/nodes/llm_agent_wrapper.ts @@ -0,0 +1,104 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {Content} from '@google/genai'; +import {BaseAgent} from '../../agents/base_agent.js'; +import {createEvent, Event} from '../../events/event.js'; +import {BaseNode, BaseNodeConfig, isContent} from '../base_node.js'; +import {NodeContext} from '../node_context.js'; + +/** Options for an {@link LLMAgentWrapper}. */ +export interface LLMAgentWrapperConfig extends Partial< + Omit +> { + name?: string; +} + +/** + * Runs a {@link BaseAgent} (typically an `LlmAgent`) as a workflow node in + * `single_turn` mode: the node input is appended as a user turn, the agent runs + * once, and its final model text becomes the node output. + * + * Ported (single_turn subset) from `google/adk-python` + * `workflow/_llm_agent_wrapper.py`. The `task` and `chat` modes (FinishTaskTool, + * task delegation, transfer, isolation scopes) are a Phase 7b continuation. + */ +export class LLMAgentWrapper extends BaseNode { + readonly agent: BaseAgent; + + constructor(agent: BaseAgent, config: LLMAgentWrapperConfig = {}) { + super({ + name: config.name ?? agent.name, + description: agent.description, + ...config, + }); + this.agent = agent; + } + + protected async *runImpl( + ctx: NodeContext, + input: unknown, + ): AsyncGenerator { + // Append the node input as a user turn so the agent responds to it. + if (input !== undefined && input !== null) { + const userEvent = createEvent({ + author: 'user', + invocationId: ctx.invocationId, + branch: ctx.branch, + content: toUserContent(input), + }); + if (ctx.isolationScope) { + userEvent.isolationScope = ctx.isolationScope; + } + ctx.session.events.push(userEvent); + } + + // Run the agent under the node's invocation context (it sets agent=itself). + for await (const event of this.agent.runAsync(ctx.invocationContext)) { + this.maybeSetOutput(event); + yield event; + } + } + + /** + * Promotes the final model text of an event to the node output (mirroring + * Python `process_llm_agent_output`). + */ + private maybeSetOutput(event: Event): void { + if (event.partial) { + return; + } + if (hasFunctionCalls(event)) { + return; + } + const content = event.content; + if (!content || content.role !== 'model' || !content.parts) { + return; + } + const text = content.parts + .filter((p) => p.text && !p.thought) + .map((p) => p.text) + .join(''); + + event.output = text; + event.nodeInfo = {...(event.nodeInfo ?? {}), messageAsOutput: true}; + } +} + +function hasFunctionCalls(event: Event): boolean { + return (event.content?.parts ?? []).some((p) => p.functionCall); +} + +/** Converts an arbitrary node input into a user-role `Content`. */ +function toUserContent(input: unknown): Content { + if (isContent(input)) { + return {...input, role: 'user'}; + } + if (typeof input === 'string') { + return {role: 'user', parts: [{text: input}]}; + } + return {role: 'user', parts: [{text: JSON.stringify(input)}]}; +} diff --git a/core/test/workflow-next/llm_agent_test.ts b/core/test/workflow-next/llm_agent_test.ts new file mode 100644 index 000000000..e880c2e80 --- /dev/null +++ b/core/test/workflow-next/llm_agent_test.ts @@ -0,0 +1,147 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {describe, expect, it} from 'vitest'; +import {BaseAgent} from '../../src/agents/base_agent.js'; +import {InvocationContext} from '../../src/agents/invocation_context.js'; +import {createEvent, Event} from '../../src/events/event.js'; +import {PluginManager} from '../../src/plugins/plugin_manager.js'; +import {Session} from '../../src/sessions/session.js'; +import {node} from '../../src/workflow-next/node.js'; +import {NodeContext} from '../../src/workflow-next/node_context.js'; +import {LLMAgentWrapper} from '../../src/workflow-next/nodes/llm_agent_wrapper.js'; +import {EventChannel} from '../../src/workflow-next/utils/event_channel.js'; +import {Workflow} from '../../src/workflow-next/workflow.js'; + +function createIc(): InvocationContext { + const session = { + id: 's1', + appName: 'app', + userId: 'u', + events: [], + state: {}, + lastUpdateTime: Date.now(), + } as unknown as Session; + return new InvocationContext({ + invocationId: 'inv-1', + session, + agent: { + name: 'wf', + runAsync: async function* () {}, + } as unknown as BaseAgent, + pluginManager: new PluginManager(), + }); +} + +async function driveWorkflow( + wf: Workflow, + input?: unknown, +): Promise<{output: unknown; events: Event[]}> { + const channel = new EventChannel(); + const root = new NodeContext({ + invocationContext: createIc(), + channel, + nodePath: '', + runId: 'root', + }); + const events: Event[] = []; + const run = root.runNode(wf, input, {useAsOutput: true}).then( + () => channel.close(), + (err) => channel.fail(err), + ); + for await (const ev of channel) { + events.push(ev); + } + await run; + return {output: root.output, events}; +} + +/** + * A fake agent that echoes the most recent user turn as a model response — a + * stand-in for a real LlmAgent so the wrapper can be tested without a model. + */ +class EchoAgent extends BaseAgent { + constructor(name = 'echo') { + super({name}); + } + protected async *runAsyncImpl( + ctx: InvocationContext, + ): AsyncGenerator { + const lastUser = [...ctx.session.events] + .reverse() + .find((e) => e.author === 'user'); + const text = (lastUser?.content?.parts ?? []) + .map((p) => p.text ?? '') + .join(''); + yield createEvent({ + author: this.name, + invocationId: ctx.invocationId, + branch: ctx.branch, + content: {role: 'model', parts: [{text: `echo:${text}`}]}, + }); + } + // eslint-disable-next-line require-yield + protected async *runLiveImpl(): AsyncGenerator { + return; + } +} + +describe('Phase 7 — LlmAgent as a node (single_turn)', () => { + it('runs an agent as a node and extracts its text output', async () => { + const wf = new Workflow({ + name: 'agent_wf', + edges: [['START', new EchoAgent()]], + }); + const {output, events} = await driveWorkflow(wf, 'hello'); + expect(output).toBe('echo:hello'); + // The agent's model event streamed through, authored by the agent. + expect(events.some((e) => e.author === 'echo')).toBe(true); + }); + + it('lets an agent be used directly in edges, feeding a downstream node (baseline bug #3)', async () => { + const upper = node( + (_c: NodeContext, input: string) => input.toUpperCase(), + { + name: 'upper', + }, + ); + const wf = new Workflow({ + name: 'agent_then_fn', + edges: [['START', new EchoAgent(), upper]], + }); + expect((await driveWorkflow(wf, 'hi')).output).toBe('ECHO:HI'); + }); + + it('node(agent) produces an LLMAgentWrapper carrying the agent name', () => { + const wrapped = node(new EchoAgent('assistant')); + expect(wrapped).toBeInstanceOf(LLMAgentWrapper); + expect(wrapped.name).toBe('assistant'); + }); + + it('routes on an agent-produced value', async () => { + // The classifier agent echoes; a function maps it to a route. + const classify = node( + (_c: NodeContext, input: string) => + createEvent({route: input.includes('?') ? 'q' : 's', output: input}), + {name: 'route_fn'}, + ); + const answer = node((_c: NodeContext, i: string) => `A:${i}`, { + name: 'answer', + }); + const comment = node((_c: NodeContext, i: string) => `C:${i}`, { + name: 'comment', + }); + const wf = new Workflow({ + name: 'agent_route', + edges: [ + ['START', new EchoAgent(), classify], + [classify, {q: answer, s: comment}], + ], + }); + // echo:'what?' contains '?', so route 'q'. + expect((await driveWorkflow(wf, 'what?')).output).toBe('A:echo:what?'); + }); +}); From 61db53db3c4d73b394327bb505d3fe7c044b6b3b Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Wed, 22 Jul 2026 14:42:04 -0700 Subject: [PATCH 12/41] feat(workflow-next): Phase 8 Runner integration and public barrel Add WorkflowAgent, a BaseAgent adapter that runs a Workflow (a BaseNode) via the standard ADK Runner: it sets up the event-channel bridge and maps the user message to the workflow input. Add the module barrel (index.ts) exposing the public surface (Workflow, WorkflowAgent, BaseNode, Node, node, FunctionNode, JoinNode, ToolNode, ParallelWorker, LLMAgentWrapper, Edge, START, DEFAULT_ROUTE, RequestInput, RetryConfig, NodeTimeoutError, ...). Includes an end-to-end integration test proving sequence, routing and DEFAULT_ROUTE workflows run through the real Runner + InMemorySessionService. --- core/src/workflow-next/index.ts | 64 ++++++++++ core/src/workflow-next/workflow_agent.ts | 87 +++++++++++++ .../workflow-next/runner_integration_test.ts | 115 ++++++++++++++++++ 3 files changed, 266 insertions(+) create mode 100644 core/src/workflow-next/index.ts create mode 100644 core/src/workflow-next/workflow_agent.ts create mode 100644 core/test/workflow-next/runner_integration_test.ts diff --git a/core/src/workflow-next/index.ts b/core/src/workflow-next/index.ts new file mode 100644 index 000000000..680753936 --- /dev/null +++ b/core/src/workflow-next/index.ts @@ -0,0 +1,64 @@ +/** + * @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 {Node, node} from './node.js'; +export type {NodeOptions} from './node.js'; +export {FunctionNode} from './nodes/function_node.js'; +export type { + FunctionNodeConfig, + FunctionNodeHandler, +} from './nodes/function_node.js'; +export {JoinNode} from './nodes/join_node.js'; +export {LLMAgentWrapper} from './nodes/llm_agent_wrapper.js'; +export type {LLMAgentWrapperConfig} from './nodes/llm_agent_wrapper.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'; + +// --- 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} from './branch_path.js'; +export {NodeContext} from './node_context.js'; +export {createNodeState, isNodeState} from './node_state.js'; +export type {NodeState} from './node_state.js'; +export {NodeStatus} from './node_status.js'; + +// --- HITL --- +export {RequestInput, isRequestInput} from './request_input.js'; +export type {RequestInputParams} from './request_input.js'; + +// --- Retry --- +export {normalizeRetryExceptions} from './retry_config.js'; +export type {ErrorClass, RetryConfig} from './retry_config.js'; + +// --- Errors --- +export {NodeTimeoutError} from './errors.js'; diff --git a/core/src/workflow-next/workflow_agent.ts b/core/src/workflow-next/workflow_agent.ts new file mode 100644 index 000000000..c5b5ea035 --- /dev/null +++ b/core/src/workflow-next/workflow_agent.ts @@ -0,0 +1,87 @@ +/** + * @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 {Event} from '../events/event.js'; +import {NodeContext} from './node_context.js'; +import {EventChannel} from './utils/event_channel.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. + */ +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 EventChannel(); + const root = new NodeContext({ + invocationContext: ic, + channel, + nodePath: '', + runId: this.name, + // TODO(phase-5b): reconstruct resumeInputs from session function + // responses so an interrupted workflow can resume via the Runner. + resumeInputs: {}, + }); + + const input = extractWorkflowInput(ic.userContent); + + const settle = root.runNode(this.workflow, input, {useAsOutput: true}).then( + () => channel.close(), + (err) => channel.fail(err), + ); + + for await (const event of channel) { + yield event; + } + await settle; + } + + // eslint-disable-next-line require-yield + protected async *runLiveImpl(): AsyncGenerator { + throw new Error('WorkflowAgent does not support live mode.'); + } +} + +/** + * 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-next/runner_integration_test.ts b/core/test/workflow-next/runner_integration_test.ts new file mode 100644 index 000000000..a1fa2a4b7 --- /dev/null +++ b/core/test/workflow-next/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-next/graph.js'; +import {node} from '../../src/workflow-next/node.js'; +import {NodeContext} from '../../src/workflow-next/node_context.js'; +import {Workflow} from '../../src/workflow-next/workflow.js'; +import {WorkflowAgent} from '../../src/workflow-next/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); + }); +}); From 7ecf669e0e59633ed7a6d4a3e578f310c8414109 Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Wed, 22 Jul 2026 14:49:56 -0700 Subject: [PATCH 13/41] refactor(workflow)!: cut over to the new BaseNode/Context workflow engine Replace the old BaseAgent/InvocationContext-based workflow module with the new Python-parity BaseNode/Context implementation: delete core/src/workflow, rename workflow-next -> workflow (the @google/adk barrel already points at ./workflow/index.js), and move the test suite to core/test/workflow. Remove the old-module unit, sample, and integration tests (superseded by the 82 workflow tests, including end-to-end Runner integration). BREAKING CHANGE: the @google/adk workflow API is now the new node-based model (Workflow is a BaseNode; run it via WorkflowAgent). --- core/src/workflow-next/base_node.ts | 222 -------- .../workflow-next/dynamic_node_scheduler.ts | 106 ---- core/src/workflow-next/index.ts | 64 --- core/src/workflow-next/node_runner.ts | 244 --------- core/src/workflow-next/node_state.ts | 77 --- core/src/workflow-next/nodes/function_node.ts | 170 ------ core/src/workflow-next/nodes/join_node.ts | 34 -- .../workflow-next/nodes/llm_agent_wrapper.ts | 104 ---- core/src/workflow-next/nodes/tool_node.ts | 106 ---- core/src/workflow-next/retry_config.ts | 85 --- core/src/workflow-next/trigger.ts | 26 - core/src/workflow-next/utils/graph_parser.ts | 212 -------- .../workflow-next/utils/graph_validation.ts | 200 ------- core/src/workflow-next/utils/hitl_utils.ts | 99 ---- core/src/workflow-next/utils/retry_utils.ts | 112 ---- core/src/workflow-next/workflow.ts | 450 ---------------- core/src/workflow/base_node.ts | 220 ++++++-- .../branch_path.ts | 0 core/src/workflow/dynamic_node_scheduler.ts | 221 +++----- .../src/{workflow-next => workflow}/errors.ts | 0 core/src/{workflow-next => workflow}/graph.ts | 0 core/src/workflow/index.ts | 96 ++-- core/src/{workflow-next => workflow}/node.ts | 0 .../node_context.ts | 0 core/src/workflow/node_runner.ts | 456 +++++++--------- core/src/workflow/node_state.ts | 105 ++-- .../node_status.ts | 0 core/src/workflow/nodes/function_node.ts | 206 +++++--- core/src/workflow/nodes/join_node.ts | 127 +---- core/src/workflow/nodes/llm_agent_wrapper.ts | 146 ++--- .../nodes/parallel_worker.ts | 0 core/src/workflow/nodes/tool_node.ts | 129 +++-- core/src/workflow/parallel_worker.ts | 132 ----- .../request_input.ts | 0 core/src/workflow/retry_config.ts | 86 +-- core/src/workflow/run_node.ts | 96 ---- .../schedule_dynamic_node.ts | 0 core/src/workflow/trigger.ts | 121 +---- .../utils/event_channel.ts | 0 core/src/workflow/utils/graph_parser.ts | 363 ++++++------- core/src/workflow/utils/graph_validation.ts | 232 +++++--- core/src/workflow/utils/hitl_utils.ts | 144 ++--- core/src/workflow/utils/rehydration_utils.ts | 114 ---- core/src/workflow/utils/replay_manager.ts | 90 ---- core/src/workflow/utils/retry_utils.ts | 164 +++--- .../utils/workflow_graph_utils.ts | 0 core/src/workflow/workflow.ts | 499 ++++++++++++++---- .../workflow_agent.ts | 0 .../workflow-next/dynamic_workflow_test.ts | 155 ------ core/test/workflow/dynamic_workflow_test.ts | 218 +++++--- .../event_channel_test.ts | 2 +- .../event_model_test.ts | 0 .../foundations_test.ts | 10 +- core/test/workflow/graph_parser_test.ts | 105 ---- .../workflow/hitl_and_rehydration_test.ts | 189 ------- .../{workflow-next => workflow}/hitl_test.ts | 12 +- .../workflow/join_node_and_parallel_test.ts | 130 ----- .../llm_agent_test.ts | 10 +- .../node_api_test.ts | 18 +- .../node_execution_test.ts | 13 +- core/test/workflow/node_runner_test.ts | 177 ------- .../parallel_test.ts | 14 +- .../runner_integration_test.ts | 10 +- core/test/workflow/workflow_agent_test.ts | 141 ----- .../workflow_test.ts | 10 +- .../workflows/dynamic_nodes_workflow_test.ts | 105 ---- .../workflows/fan_out_fan_in_workflow_test.ts | 77 --- .../workflows/loop_workflow_test.ts | 121 ----- .../workflows/nested_workflow_test.ts | 88 --- .../workflows/node_as_tool_workflow_test.ts | 116 ---- .../workflows/route_workflow_test.ts | 207 -------- .../workflows/sequence_workflow_test.ts | 157 ------ 72 files changed, 2021 insertions(+), 6122 deletions(-) delete mode 100644 core/src/workflow-next/base_node.ts delete mode 100644 core/src/workflow-next/dynamic_node_scheduler.ts delete mode 100644 core/src/workflow-next/index.ts delete mode 100644 core/src/workflow-next/node_runner.ts delete mode 100644 core/src/workflow-next/node_state.ts delete mode 100644 core/src/workflow-next/nodes/function_node.ts delete mode 100644 core/src/workflow-next/nodes/join_node.ts delete mode 100644 core/src/workflow-next/nodes/llm_agent_wrapper.ts delete mode 100644 core/src/workflow-next/nodes/tool_node.ts delete mode 100644 core/src/workflow-next/retry_config.ts delete mode 100644 core/src/workflow-next/trigger.ts delete mode 100644 core/src/workflow-next/utils/graph_parser.ts delete mode 100644 core/src/workflow-next/utils/graph_validation.ts delete mode 100644 core/src/workflow-next/utils/hitl_utils.ts delete mode 100644 core/src/workflow-next/utils/retry_utils.ts delete mode 100644 core/src/workflow-next/workflow.ts rename core/src/{workflow-next => workflow}/branch_path.ts (100%) rename core/src/{workflow-next => workflow}/errors.ts (100%) rename core/src/{workflow-next => workflow}/graph.ts (100%) rename core/src/{workflow-next => workflow}/node.ts (100%) rename core/src/{workflow-next => workflow}/node_context.ts (100%) rename core/src/{workflow-next => workflow}/node_status.ts (100%) rename core/src/{workflow-next => workflow}/nodes/parallel_worker.ts (100%) delete mode 100644 core/src/workflow/parallel_worker.ts rename core/src/{workflow-next => workflow}/request_input.ts (100%) delete mode 100644 core/src/workflow/run_node.ts rename core/src/{workflow-next => workflow}/schedule_dynamic_node.ts (100%) rename core/src/{workflow-next => workflow}/utils/event_channel.ts (100%) delete mode 100644 core/src/workflow/utils/rehydration_utils.ts delete mode 100644 core/src/workflow/utils/replay_manager.ts rename core/src/{workflow-next => workflow}/utils/workflow_graph_utils.ts (100%) rename core/src/{workflow-next => workflow}/workflow_agent.ts (100%) delete mode 100644 core/test/workflow-next/dynamic_workflow_test.ts rename core/test/{workflow-next => workflow}/event_channel_test.ts (96%) rename core/test/{workflow-next => workflow}/event_model_test.ts (100%) rename core/test/{workflow-next => workflow}/foundations_test.ts (95%) delete mode 100644 core/test/workflow/graph_parser_test.ts delete mode 100644 core/test/workflow/hitl_and_rehydration_test.ts rename core/test/{workflow-next => workflow}/hitl_test.ts (93%) delete mode 100644 core/test/workflow/join_node_and_parallel_test.ts rename core/test/{workflow-next => workflow}/llm_agent_test.ts (92%) rename core/test/{workflow-next => workflow}/node_api_test.ts (91%) rename core/test/{workflow-next => workflow}/node_execution_test.ts (94%) delete mode 100644 core/test/workflow/node_runner_test.ts rename core/test/{workflow-next => workflow}/parallel_test.ts (91%) rename core/test/{workflow-next => workflow}/runner_integration_test.ts (91%) delete mode 100644 core/test/workflow/workflow_agent_test.ts rename core/test/{workflow-next => workflow}/workflow_test.ts (94%) delete mode 100644 tests/integration/workflows/dynamic_nodes_workflow_test.ts delete mode 100644 tests/integration/workflows/fan_out_fan_in_workflow_test.ts delete mode 100644 tests/integration/workflows/loop_workflow_test.ts delete mode 100644 tests/integration/workflows/nested_workflow_test.ts delete mode 100644 tests/integration/workflows/node_as_tool_workflow_test.ts delete mode 100644 tests/integration/workflows/route_workflow_test.ts delete mode 100644 tests/integration/workflows/sequence_workflow_test.ts diff --git a/core/src/workflow-next/base_node.ts b/core/src/workflow-next/base_node.ts deleted file mode 100644 index a0e9c3c41..000000000 --- a/core/src/workflow-next/base_node.ts +++ /dev/null @@ -1,222 +0,0 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import {Content} from '@google/genai'; -import type {ZodType} from 'zod'; -import {createEvent, Event, isEvent} from '../events/event.js'; -import type {NodeContext} from './node_context.js'; -import {isRequestInput} from './request_input.js'; -import {RetryConfig} from './retry_config.js'; -import {createRequestInputEvent} from './utils/hitl_utils.js'; - -/** - * Configuration shared by all workflow nodes. - * - * Mirrors the fields of `google/adk-python` `workflow/_base_node.py::BaseNode`. - */ -export interface BaseNodeConfig { - /** Canonical, unique-within-a-graph node name. */ - name: string; - - /** Human-readable description (used when a node is exposed as a tool). */ - description?: string; - - /** - * If true, the node re-executes when a workflow resumes even if it already - * completed in a prior turn. Default false. - */ - rerunOnResume?: boolean; - - /** - * If true, the node only produces its output once all of its predecessors - * have triggered it (fan-in / join semantics). Default false. - */ - waitForOutput?: boolean; - - /** Optional retry configuration for transient failures. */ - retryConfig?: RetryConfig; - - /** Maximum time, in seconds, for this node to complete. */ - timeout?: number; - - /** Optional zod schema validating the node input. */ - inputSchema?: ZodType; - - /** Optional zod schema validating the node output. */ - outputSchema?: ZodType; - - /** Optional zod schema validating relevant session state. */ - stateSchema?: ZodType; -} - -/** - * Abstract base class for all nodes in an ADK workflow. - * - * A node is a discrete unit of execution. Subclasses implement {@link runImpl}, - * which may yield {@link Event}s, raw values (boxed into an event), or - * `null`/`undefined` (skipped). {@link run} normalizes those into a stream of - * {@link Event}s consumed by the engine. - */ -export abstract class BaseNode { - readonly name: string; - readonly description: string; - readonly rerunOnResume: boolean; - readonly waitForOutput: boolean; - readonly retryConfig?: RetryConfig; - readonly timeout?: number; - readonly inputSchema?: ZodType; - readonly outputSchema?: ZodType; - readonly stateSchema?: ZodType; - - constructor(config: BaseNodeConfig) { - if ( - !config.name || - typeof config.name !== 'string' || - config.name.trim().length === 0 - ) { - throw new Error('Node name must be a non-empty string.'); - } - this.name = config.name.trim(); - this.description = config.description ?? ''; - this.rerunOnResume = config.rerunOnResume ?? false; - this.waitForOutput = config.waitForOutput ?? false; - this.retryConfig = config.retryConfig; - this.timeout = config.timeout; - this.inputSchema = config.inputSchema; - this.outputSchema = config.outputSchema; - this.stateSchema = config.stateSchema; - } - - /** - * Whether this node must wait for ALL of its predecessors to trigger before - * it runs (fan-in barrier). Overridden by `JoinNode`. - */ - get requiresAllPredecessors(): boolean { - return false; - } - - /** - * Core execution contract. Subclasses yield one of: - * - an {@link Event} (emitted as-is), - * - a raw value (boxed into an event whose `output` is that value), - * - `null`/`undefined` (skipped). - */ - protected abstract runImpl( - ctx: NodeContext, - input: TInput, - ): AsyncGenerator; - - /** - * Runs the node, normalizing every yielded item into an {@link Event}. This - * is what the engine (and `ctx.runNode()`) consumes. Validates the input - * against `inputSchema` once, up front (skipping genai `Content`, which nodes - * coerce themselves). - */ - async *run( - ctx: NodeContext, - input: TInput, - ): AsyncGenerator { - const validatedInput = this.validateInput(input); - for await (const item of this.runImpl(ctx, validatedInput)) { - if (isRequestInput(item)) { - // HITL: convert a request-for-input into an interrupt event. - yield createRequestInputEvent(item); - continue; - } - const event = this.toEvent(ctx, item); - if (event) { - yield event; - } - } - } - - /** Validates node input against `inputSchema` (Content passes through). */ - protected validateInput(input: TInput): TInput { - if (!this.inputSchema || isContent(input)) { - return input; - } - return this.inputSchema.parse(input) as TInput; - } - - /** Validates node output against `outputSchema` (Content passes through). */ - protected validateOutput(output: unknown): unknown { - if (!this.outputSchema || isContent(output)) { - return output; - } - return this.outputSchema.parse(output); - } - - /** - * Normalizes a single yielded item into an {@link Event} (or `null` to skip). - * Subclasses may override for richer coercion (e.g. `FunctionNode`). - */ - protected toEvent(ctx: NodeContext, data: unknown): Event | null { - if (data === null || data === undefined) { - return null; - } - if (isEvent(data)) { - const event = data as Event; - if (event.output !== undefined) { - event.output = this.validateOutput(event.output); - } - return event; - } - const output = this.validateOutput(data); - return createEvent({ - author: this.name, - invocationId: ctx.invocationContext.invocationId, - branch: ctx.branch, - content: toContent(output), - output, - }); - } -} - -/** Returns whether a value looks like a genai `Content` object. */ -export function isContent(value: unknown): value is Content { - return ( - typeof value === 'object' && - value !== null && - 'parts' in value && - Array.isArray((value as {parts?: unknown}).parts) - ); -} - -/** - * The sentinel node marking the entry point of a workflow graph. It is never - * executed — the orchestrator seeds triggers for its successors directly. - * - * Mirrors `google/adk-python` `START = BaseNode(name='__START__')`. - */ -class StartNode extends BaseNode { - // eslint-disable-next-line require-yield - protected async *runImpl(): AsyncGenerator { - throw new Error('START node is never executed.'); - } -} - -/** The workflow entry-point sentinel node (name `__START__`). */ -export const START: BaseNode = new StartNode({name: '__START__'}); - -/** - * Best-effort conversion of an arbitrary value to genai `Content` for display. - */ -export function toContent(val: unknown): Content | undefined { - if (val === null || val === undefined) { - return undefined; - } - if (typeof val === 'object' && 'role' in val && 'parts' in val) { - return val as Content; - } - if (typeof val === 'string') { - return {role: 'model', parts: [{text: val}]}; - } - try { - return {role: 'model', parts: [{text: JSON.stringify(val)}]}; - } catch { - return {role: 'model', parts: [{text: String(val)}]}; - } -} diff --git a/core/src/workflow-next/dynamic_node_scheduler.ts b/core/src/workflow-next/dynamic_node_scheduler.ts deleted file mode 100644 index c51c391da..000000000 --- a/core/src/workflow-next/dynamic_node_scheduler.ts +++ /dev/null @@ -1,106 +0,0 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import {BaseNode} from './base_node.js'; -import {NodeContext} 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'; - -/** - * 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 { - constructor(private readonly state: DynamicNodeState) {} - - 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}/${name}@${runId}`; - - const existing = this.state.runs.get(nodePath); - if (existing?.task) { - // Deduplicate concurrent calls: await the in-flight task. - return existing.task; - } - - // TODO(phase-5): lazy rehydration from session events + replay - // interception (dedup completed / resume waiting runs) goes here. - - 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(ctx, node, input, { - nodeName: name, - runId, - 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-next/index.ts b/core/src/workflow-next/index.ts deleted file mode 100644 index 680753936..000000000 --- a/core/src/workflow-next/index.ts +++ /dev/null @@ -1,64 +0,0 @@ -/** - * @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 {Node, node} from './node.js'; -export type {NodeOptions} from './node.js'; -export {FunctionNode} from './nodes/function_node.js'; -export type { - FunctionNodeConfig, - FunctionNodeHandler, -} from './nodes/function_node.js'; -export {JoinNode} from './nodes/join_node.js'; -export {LLMAgentWrapper} from './nodes/llm_agent_wrapper.js'; -export type {LLMAgentWrapperConfig} from './nodes/llm_agent_wrapper.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'; - -// --- 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} from './branch_path.js'; -export {NodeContext} from './node_context.js'; -export {createNodeState, isNodeState} from './node_state.js'; -export type {NodeState} from './node_state.js'; -export {NodeStatus} from './node_status.js'; - -// --- HITL --- -export {RequestInput, isRequestInput} from './request_input.js'; -export type {RequestInputParams} from './request_input.js'; - -// --- Retry --- -export {normalizeRetryExceptions} from './retry_config.js'; -export type {ErrorClass, RetryConfig} from './retry_config.js'; - -// --- Errors --- -export {NodeTimeoutError} from './errors.js'; diff --git a/core/src/workflow-next/node_runner.ts b/core/src/workflow-next/node_runner.ts deleted file mode 100644 index 2661484e7..000000000 --- a/core/src/workflow-next/node_runner.ts +++ /dev/null @@ -1,244 +0,0 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import { - InvocationContext, - InvocationContextParams, -} from '../agents/invocation_context.js'; -import {Event} from '../events/event.js'; -import {BaseNode} from './base_node.js'; -import {BranchPath} from './branch_path.js'; -import {NodeTimeoutError} from './errors.js'; -import {NodeContext} from './node_context.js'; -import {createNodeState} from './node_state.js'; -import {NodeStatus} from './node_status.js'; -import {getRetryDelaySeconds, shouldRetryNode} from './utils/retry_utils.js'; - -/** - * Options controlling a single `ctx.runNode(...)` execution. - */ -export interface RunNodeOptions { - /** Deterministic tracking name; defaults to `node.name`. */ - nodeName?: string; - /** Unique id for this specific run; defaults to `nodeName`. */ - runId?: string; - /** If true, the child's output replaces the caller's output. */ - useAsOutput?: boolean; - /** If true, run the child in an isolated sub-branch. */ - useSubBranch?: boolean; - /** Explicit branch, overriding the default/sub-branch computation. */ - overrideBranch?: string; - /** Explicit isolation scope, overriding inheritance from the parent. */ - overrideIsolationScope?: string; -} - -/** - * Executes a child node on behalf of `parent.runNode(...)`. - * - * Responsibilities (Phase 1 scope): create the child {@link NodeContext}, - * drive `node.run()`, enrich each emitted event (author, node path, branch, - * isolation scope), track the child's `output`/`route`, apply the per-node - * `timeout`, and retry on failure per `retryConfig`. Returns the child context. - */ -export async function executeChildNode( - parent: NodeContext, - node: BaseNode, - input: unknown, - options: RunNodeOptions = {}, -): Promise { - const nodeName = options.nodeName ?? node.name; - const runId = options.runId ?? nodeName; - const nodePath = parent.nodePath - ? `${parent.nodePath}.${nodeName}` - : nodeName; - - let branch = parent.branch; - if (options.overrideBranch !== undefined) { - branch = options.overrideBranch; - } else if (options.useSubBranch) { - branch = BranchPath.createSubBranch(parent.branch, { - name: nodeName, - runId: options.runId, - }); - } - - const isolationScope = - options.overrideIsolationScope ?? parent.isolationScope; - - const childIc = - branch === parent.invocationContext.branch - ? parent.invocationContext - : withBranch(parent.invocationContext, branch); - - const child = new NodeContext({ - invocationContext: childIc, - channel: parent.channel, - nodePath, - runId, - resumeInputs: parent.resumeInputs, - isolationScope, - }); - // Propagate the dynamic scheduler down; a nested Workflow overrides it. - child.scheduler = parent.scheduler; - - const nodeState = createNodeState({ - status: NodeStatus.RUNNING, - input, - runId, - }); - - for (;;) { - // Reset per-attempt output so a retry starts clean. - child.output = undefined; - child.route = undefined; - child.interruptIds = []; - try { - await runOnce(node, child, input, nodeName, branch, isolationScope); - break; - } catch (err) { - // Check retry eligibility with the attempt that just failed, compute its - // backoff delay, THEN advance the counter (matches Python semantics). - if (shouldRetryNode(err, node.retryConfig, nodeState)) { - const delaySeconds = getRetryDelaySeconds(node.retryConfig, nodeState); - nodeState.attemptCount += 1; - await delay(delaySeconds * 1000, parent.invocationContext.abortSignal); - continue; - } - throw err; - } - } - - if (options.useAsOutput) { - parent.output = child.output; - parent.route = child.route; - } - - return child; -} - -/** - * Drives one attempt of `node.run()`, enriching and pushing each event and - * tracking the child's output/route. Wrapped in a timeout when configured. - */ -async function runOnce( - node: BaseNode, - child: NodeContext, - input: unknown, - nodeName: string, - branch: string | undefined, - isolationScope: string | undefined, -): Promise { - const body = (async () => { - for await (const event of node.run(child, input)) { - enrichEvent(event, child, nodeName, branch, isolationScope); - if (event.output !== undefined) { - child.output = event.output; - } - if (event.route !== undefined) { - child.route = event.route; - } - // HITL: an interrupt event marks its ids as long-running tool ids. - if (event.longRunningToolIds && event.longRunningToolIds.length > 0) { - for (const id of event.longRunningToolIds) { - if (!child.interruptIds.includes(id)) { - child.interruptIds.push(id); - } - } - } - child.channel.push(event); - } - })(); - - if (node.timeout && node.timeout > 0) { - await withTimeout(body, node.timeout, nodeName); - } else { - await body; - } -} - -/** - * Stamps engine-owned provenance onto an event without clobbering values the - * node explicitly set. - */ -function enrichEvent( - event: Event, - child: NodeContext, - nodeName: string, - branch: string | undefined, - isolationScope: string | undefined, -): void { - if (!event.author) { - event.author = nodeName; - } - event.nodeInfo = {...(event.nodeInfo ?? {}), path: child.nodePath}; - if (branch !== undefined && event.branch === undefined) { - event.branch = branch; - } - if (isolationScope !== undefined && event.isolationScope === undefined) { - event.isolationScope = isolationScope; - } -} - -/** - * Creates a shallow child InvocationContext with a different branch, preserving - * the shared invocation cost manager and all services/session. - */ -function withBranch( - ic: InvocationContext, - branch: string | undefined, -): InvocationContext { - return new InvocationContext({ - ...(ic as unknown as InvocationContextParams), - branch, - }); -} - -/** - * Rejects with {@link NodeTimeoutError} if `promise` does not settle within - * `timeoutSeconds`. - */ -function withTimeout( - promise: Promise, - timeoutSeconds: number, - nodeName: string, -): Promise { - return new Promise((resolve, reject) => { - const timer = setTimeout(() => { - reject(new NodeTimeoutError({nodeName, timeout: timeoutSeconds})); - }, timeoutSeconds * 1000); - promise.then( - (value) => { - clearTimeout(timer); - resolve(value); - }, - (err) => { - clearTimeout(timer); - reject(err); - }, - ); - }); -} - -/** - * Promise-based delay that rejects early if the abort signal fires. - */ -function delay(ms: number, signal?: AbortSignal): Promise { - return new Promise((resolve, reject) => { - if (signal?.aborted) { - reject(new Error('Aborted')); - return; - } - const timer = setTimeout(() => { - signal?.removeEventListener('abort', onAbort); - resolve(); - }, ms); - const onAbort = () => { - clearTimeout(timer); - reject(new Error('Aborted')); - }; - signal?.addEventListener('abort', onAbort, {once: true}); - }); -} diff --git a/core/src/workflow-next/node_state.ts b/core/src/workflow-next/node_state.ts deleted file mode 100644 index 8950bbaa5..000000000 --- a/core/src/workflow-next/node_state.ts +++ /dev/null @@ -1,77 +0,0 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import {NodeStatus} from './node_status.js'; - -/** - * State of a node in the workflow. - * - * Ported from `google/adk-python` `workflow/_node_state.py`. Note that the - * node's *output* is intentionally NOT stored here — it is carried on emitted - * events / the node Context, not on the persisted node state. - */ -export interface NodeState { - /** The run status of the node. */ - status: NodeStatus; - - /** The input provided to the node. */ - input?: unknown; - - /** The attempt count for this node run (1-based). */ - attemptCount: number; - - /** The interrupt ids that are pending to be resolved. */ - interrupts: string[]; - - /** The responses for resuming the node, keyed by interrupt id. */ - resumeInputs: Record; - - /** - * Sequential counter incremented each time the node gets a fresh run. - * - * Preserving this count independently of `runId` prevents path collisions if - * a node switches between custom string IDs and auto-generated numeric IDs. - */ - runCounter: number; - - /** The run ID of this node run. */ - runId?: string; - - /** - * The run ID of the parent node which dynamically scheduled this node run. - */ - parentRunId?: string; -} - -/** - * Creates a {@link NodeState} with Python-aligned defaults, overlaying any - * provided partial values. - */ -export function createNodeState(partial?: Partial): NodeState { - return { - status: NodeStatus.INACTIVE, - attemptCount: 1, - interrupts: [], - resumeInputs: {}, - runCounter: 0, - ...partial, - }; -} - -/** - * Type guard for a {@link NodeState}-shaped object. - */ -export function isNodeState(obj: unknown): obj is NodeState { - return ( - typeof obj === 'object' && - obj !== null && - 'status' in obj && - typeof (obj as NodeState).status === 'number' && - 'attemptCount' in obj && - 'interrupts' in obj && - Array.isArray((obj as NodeState).interrupts) - ); -} diff --git a/core/src/workflow-next/nodes/function_node.ts b/core/src/workflow-next/nodes/function_node.ts deleted file mode 100644 index b250ab25c..000000000 --- a/core/src/workflow-next/nodes/function_node.ts +++ /dev/null @@ -1,170 +0,0 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import {AuthConfig} from '../../auth/auth_tool.js'; -import {createEvent, Event, isEvent} from '../../events/event.js'; -import {BaseNode, BaseNodeConfig, isContent, toContent} from '../base_node.js'; -import {NodeContext} from '../node_context.js'; - -/** - * A value a {@link FunctionNodeHandler} may return or yield. - */ -export type FunctionNodeResult = - | TOutput - | Event - | null - | undefined - | void; - -/** - * The handler wrapped by a {@link FunctionNode}. - * - * Unlike Python's `FunctionNode` (which binds named parameters from `ctx.state` - * or `node_input` via runtime signature introspection), the TypeScript form - * uses the idiomatic explicit `(ctx, input)` signature. Read `ctx.state` - * directly for state-bound values. It may return a value/`Event`, a Promise, or - * a (sync/async) generator of those. - */ -export type FunctionNodeHandler = ( - ctx: NodeContext, - input: TInput, -) => - | FunctionNodeResult - | Promise> - | Generator, void, unknown> - | AsyncGenerator, void, unknown>; - -/** - * Options for a {@link FunctionNode}. - */ -export interface FunctionNodeConfig extends Partial< - Omit -> { - /** - * If set, the framework requests user authentication before running (Phase 5 - * enables the auth gate; stored here now for API parity). - */ - authConfig?: AuthConfig; -} - -/** - * A node that wraps a plain function, async function, or (sync/async) generator. - * - * Ported (TS-idiomatic subset) from `google/adk-python` `_function_node.py`. - * Return-value handling: - * - `Event` → emitted as-is (output validated against `outputSchema`) - * - genai `Content` → emitted as the event content - * - `null`/`undefined` → skipped (unless there are pending state deltas) - * - anything else → emitted as `Event(output=value)` - * State written via `ctx.state` during execution is attached to emitted events. - */ -export class FunctionNode extends BaseNode< - TInput, - TOutput -> { - readonly authConfig?: AuthConfig; - private readonly handler: FunctionNodeHandler; - - constructor( - name: string, - handler: FunctionNodeHandler, - config: FunctionNodeConfig = {}, - ) { - if (typeof handler !== 'function') { - throw new TypeError('FunctionNode handler must be a function.'); - } - super({name, ...config}); - this.handler = handler; - this.authConfig = config.authConfig; - } - - protected async *runImpl( - ctx: NodeContext, - input: TInput, - ): AsyncGenerator { - // TODO(phase-5): auth gate (authConfig -> adk_request_credential interrupt). - const result = this.handler(ctx, input); - - if (isAsyncIterable(result)) { - for await (const item of result) { - yield item; - } - } else if (isSyncGenerator(result)) { - for (const item of result) { - yield item; - } - } else { - // Plain value or Promise of a value. - yield await (result as Promise>); - } - } - - protected override toEvent(ctx: NodeContext, data: unknown): Event | null { - const stateDelta = - Object.keys(ctx.actions.stateDelta).length > 0 - ? {...ctx.actions.stateDelta} - : undefined; - - if (data === null || data === undefined) { - return stateDelta - ? createEvent({ - author: this.name, - invocationId: ctx.invocationId, - branch: ctx.branch, - actions: {stateDelta}, - }) - : null; - } - - if (isEvent(data)) { - const event = data as Event; - if (event.output !== undefined) { - event.output = this.validateOutput(event.output); - } - if (stateDelta) { - Object.assign(event.actions.stateDelta, stateDelta); - } - return event; - } - - if (isContent(data)) { - return createEvent({ - author: this.name, - invocationId: ctx.invocationId, - branch: ctx.branch, - content: data, - actions: stateDelta ? {stateDelta} : undefined, - }); - } - - const output = this.validateOutput(data); - return createEvent({ - author: this.name, - invocationId: ctx.invocationId, - branch: ctx.branch, - content: toContent(output), - output, - actions: stateDelta ? {stateDelta} : undefined, - }); - } -} - -function isAsyncIterable(value: unknown): value is AsyncIterable { - return ( - value != null && - typeof (value as AsyncIterable)[Symbol.asyncIterator] === - 'function' - ); -} - -function isSyncGenerator(value: unknown): value is Generator { - return ( - value != null && - typeof value !== 'string' && - typeof (value as Iterable)[Symbol.iterator] === 'function' && - typeof (value as Generator).next === 'function' - ); -} diff --git a/core/src/workflow-next/nodes/join_node.ts b/core/src/workflow-next/nodes/join_node.ts deleted file mode 100644 index 23543d338..000000000 --- a/core/src/workflow-next/nodes/join_node.ts +++ /dev/null @@ -1,34 +0,0 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import {createEvent, Event} from '../../events/event.js'; -import {BaseNode} from '../base_node.js'; -import {NodeContext} from '../node_context.js'; - -/** - * A fan-in barrier node: it waits for ALL of its predecessors to complete, then - * emits the aggregated inputs (a map of predecessor name → output) as its - * output. - * - * Ported from `google/adk-python` `workflow/_join_node.py`. - */ -export class JoinNode extends BaseNode { - override get requiresAllPredecessors(): boolean { - return true; - } - - protected async *runImpl( - ctx: NodeContext, - input: unknown, - ): AsyncGenerator { - yield createEvent({ - author: this.name, - invocationId: ctx.invocationId, - branch: ctx.branch, - output: input, - }); - } -} diff --git a/core/src/workflow-next/nodes/llm_agent_wrapper.ts b/core/src/workflow-next/nodes/llm_agent_wrapper.ts deleted file mode 100644 index 4af1aba00..000000000 --- a/core/src/workflow-next/nodes/llm_agent_wrapper.ts +++ /dev/null @@ -1,104 +0,0 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import {Content} from '@google/genai'; -import {BaseAgent} from '../../agents/base_agent.js'; -import {createEvent, Event} from '../../events/event.js'; -import {BaseNode, BaseNodeConfig, isContent} from '../base_node.js'; -import {NodeContext} from '../node_context.js'; - -/** Options for an {@link LLMAgentWrapper}. */ -export interface LLMAgentWrapperConfig extends Partial< - Omit -> { - name?: string; -} - -/** - * Runs a {@link BaseAgent} (typically an `LlmAgent`) as a workflow node in - * `single_turn` mode: the node input is appended as a user turn, the agent runs - * once, and its final model text becomes the node output. - * - * Ported (single_turn subset) from `google/adk-python` - * `workflow/_llm_agent_wrapper.py`. The `task` and `chat` modes (FinishTaskTool, - * task delegation, transfer, isolation scopes) are a Phase 7b continuation. - */ -export class LLMAgentWrapper extends BaseNode { - readonly agent: BaseAgent; - - constructor(agent: BaseAgent, config: LLMAgentWrapperConfig = {}) { - super({ - name: config.name ?? agent.name, - description: agent.description, - ...config, - }); - this.agent = agent; - } - - protected async *runImpl( - ctx: NodeContext, - input: unknown, - ): AsyncGenerator { - // Append the node input as a user turn so the agent responds to it. - if (input !== undefined && input !== null) { - const userEvent = createEvent({ - author: 'user', - invocationId: ctx.invocationId, - branch: ctx.branch, - content: toUserContent(input), - }); - if (ctx.isolationScope) { - userEvent.isolationScope = ctx.isolationScope; - } - ctx.session.events.push(userEvent); - } - - // Run the agent under the node's invocation context (it sets agent=itself). - for await (const event of this.agent.runAsync(ctx.invocationContext)) { - this.maybeSetOutput(event); - yield event; - } - } - - /** - * Promotes the final model text of an event to the node output (mirroring - * Python `process_llm_agent_output`). - */ - private maybeSetOutput(event: Event): void { - if (event.partial) { - return; - } - if (hasFunctionCalls(event)) { - return; - } - const content = event.content; - if (!content || content.role !== 'model' || !content.parts) { - return; - } - const text = content.parts - .filter((p) => p.text && !p.thought) - .map((p) => p.text) - .join(''); - - event.output = text; - event.nodeInfo = {...(event.nodeInfo ?? {}), messageAsOutput: true}; - } -} - -function hasFunctionCalls(event: Event): boolean { - return (event.content?.parts ?? []).some((p) => p.functionCall); -} - -/** Converts an arbitrary node input into a user-role `Content`. */ -function toUserContent(input: unknown): Content { - if (isContent(input)) { - return {...input, role: 'user'}; - } - if (typeof input === 'string') { - return {role: 'user', parts: [{text: input}]}; - } - return {role: 'user', parts: [{text: JSON.stringify(input)}]}; -} diff --git a/core/src/workflow-next/nodes/tool_node.ts b/core/src/workflow-next/nodes/tool_node.ts deleted file mode 100644 index b70825645..000000000 --- a/core/src/workflow-next/nodes/tool_node.ts +++ /dev/null @@ -1,106 +0,0 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import {Context} from '../../agents/context.js'; -import {createEvent, Event} from '../../events/event.js'; -import {BaseTool} from '../../tools/base_tool.js'; -import {randomUUID} from '../../utils/env_aware_utils.js'; -import {BaseNode, BaseNodeConfig, isContent} from '../base_node.js'; -import {NodeContext} from '../node_context.js'; -/** Options for a {@link ToolNode}. */ -export interface ToolNodeConfig extends Partial> { - /** Optional name override; defaults to the tool's name. */ - name?: string; -} - -/** - * A node that wraps an ADK {@link BaseTool} and invokes it with the node input - * as its arguments. - * - * Ported from `google/adk-python` `workflow/_tool_node.py`. The node input is - * coerced to a tool-args object: genai `Content` → its text; a JSON string → - * parsed object; `null`/empty → `{}`. - */ -export class ToolNode extends BaseNode { - readonly tool: BaseTool; - - constructor(tool: BaseTool, config: ToolNodeConfig = {}) { - super({name: config.name ?? tool.name, ...config}); - this.tool = tool; - } - - protected async *runImpl( - ctx: NodeContext, - input: unknown, - ): AsyncGenerator { - const toolContext = new Context({ - invocationContext: ctx.invocationContext, - functionCallId: randomUUID(), - }); - - const args = coerceToolArgs(input); - const response = await this.tool.runAsync({args, toolContext}); - - const stateDelta = - Object.keys(toolContext.actions.stateDelta).length > 0 - ? {...toolContext.actions.stateDelta} - : undefined; - - if (response !== undefined && response !== null) { - yield createEvent({ - author: this.name, - invocationId: ctx.invocationId, - branch: ctx.branch, - output: response, - actions: stateDelta ? {stateDelta} : undefined, - }); - } else if (stateDelta) { - yield createEvent({ - author: this.name, - invocationId: ctx.invocationId, - branch: ctx.branch, - actions: {stateDelta}, - }); - } - } -} - -/** Coerces arbitrary node input into a tool-arguments record. */ -function coerceToolArgs(input: unknown): Record { - let args: unknown = input; - - if (isContent(args)) { - args = extractText(args); - } - - if (typeof args === 'string') { - const trimmed = args.trim(); - if (!trimmed) { - args = null; - } else { - try { - args = JSON.parse(trimmed); - } catch { - // Leave as the raw string; validated below. - } - } - } - - if (args === null || args === undefined) { - return {}; - } - if (typeof args !== 'object' || Array.isArray(args)) { - throw new TypeError( - 'The input to ToolNode must be a dictionary of tool arguments or null, ' + - `but got ${typeof args}.`, - ); - } - return args as Record; -} - -function extractText(content: {parts?: Array<{text?: string}>}): string { - return (content.parts ?? []).map((p) => p.text ?? '').join(''); -} diff --git a/core/src/workflow-next/retry_config.ts b/core/src/workflow-next/retry_config.ts deleted file mode 100644 index 65ea24efc..000000000 --- a/core/src/workflow-next/retry_config.ts +++ /dev/null @@ -1,85 +0,0 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -/** - * An error constructor usable in {@link RetryConfig.exceptions}. - */ -export type ErrorClass = new (...args: never[]) => Error; - -/** - * Configuration for retrying a node. - * - * Ported from `google/adk-python` `workflow/_retry_config.py`. Delays are - * expressed in **seconds** (fractions allowed) to match the Python semantics - * and keep configuration portable across runtimes. Unset fields fall back to - * the documented defaults inside the retry utilities. - */ -export interface RetryConfig { - /** - * Maximum number of attempts, including the original request. If 0 or 1, it - * means no retries. If not specified, defaults to 5. - */ - maxAttempts?: number; - - /** - * Initial delay before the first retry, in seconds. If not specified, - * defaults to 1.0 second. - */ - initialDelay?: number; - - /** - * Maximum delay between retries, in seconds. If not specified, defaults to - * 60.0 seconds. - */ - maxDelay?: number; - - /** - * Multiplier by which the delay increases after each attempt. If not - * specified, defaults to 2.0. - */ - backoffFactor?: number; - - /** - * Randomness factor for the delay. If not specified, defaults to 1.0. Use 0.0 - * to remove randomness. - */ - jitter?: number; - - /** - * Exceptions to retry on. Accepts error class names as strings (e.g. - * `['TypeError']`) or error classes directly (e.g. `[TypeError]`). - * `undefined`/`null` means retry on all errors. - */ - exceptions?: Array | null; -} - -/** - * Normalizes the `exceptions` field of a {@link RetryConfig} to a list of error - * class name strings, mirroring Python's `field_validator`. - * - * @returns The list of class-name strings, or `undefined` to mean "retry on all - * errors". - */ -export function normalizeRetryExceptions( - exceptions?: Array | null, -): string[] | undefined { - if (exceptions === undefined || exceptions === null) { - return undefined; - } - return exceptions.map((item) => { - if (typeof item === 'string') { - return item; - } - if (typeof item === 'function' && item.name) { - return item.name; - } - throw new Error( - `exceptions must contain error class names (string) or error classes, got: ${String( - item, - )}`, - ); - }); -} diff --git a/core/src/workflow-next/trigger.ts b/core/src/workflow-next/trigger.ts deleted file mode 100644 index 9ce2f307d..000000000 --- a/core/src/workflow-next/trigger.ts +++ /dev/null @@ -1,26 +0,0 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -/** - * A buffered trigger for a downstream node. - * - * Ported from `google/adk-python` `workflow/_trigger.py`. Unlike the previous - * TypeScript `Trigger` (a route-matching predicate), this is a plain data record - * describing *how* a target node should be invoked when its turn comes. - */ -export interface Trigger { - /** The input to pass to the triggered node. */ - input?: unknown; - - /** Whether this trigger should run the node in an isolated sub-branch. */ - useSubBranch?: boolean; - - /** The branch inherited from the predecessor node. */ - branch?: string; - - /** Scope tag explicitly propagated to this trigger. */ - isolationScope?: string; -} diff --git a/core/src/workflow-next/utils/graph_parser.ts b/core/src/workflow-next/utils/graph_parser.ts deleted file mode 100644 index cae1ebaeb..000000000 --- a/core/src/workflow-next/utils/graph_parser.ts +++ /dev/null @@ -1,212 +0,0 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -/** - * Parses workflow edge items and chains into a flat list of {@link Edge}s. - * - * Ported from `google/adk-python` `workflow/utils/_graph_parser.py`. - */ - -import {BaseNode} from '../base_node.js'; -import { - ChainElement, - Edge, - EdgeItem, - NodeLike, - RouteValue, - RoutingMap, -} from '../graph.js'; -import {buildNode, isNodeLike, isPlainObject} from './workflow_graph_utils.js'; - -function isRouteValue(value: unknown): value is RouteValue { - const t = typeof value; - return t === 'string' || t === 'number' || t === 'boolean'; -} - -/** Expands a routing map into individual (from, to, route) triples. */ -function expandRoutingMap( - fromElement: ChainElement, - routingMap: RoutingMap, -): Array<[ChainElement, NodeLike | readonly NodeLike[], RouteValue]> { - const keys = Object.keys(routingMap); - if (keys.length === 0) { - throw new Error( - 'Routing map must not be empty. Provide at least one route -> node mapping.', - ); - } - - const expanded: Array< - [ChainElement, NodeLike | readonly NodeLike[], RouteValue] - > = []; - for (const routeKey of keys) { - // Object keys are strings; numeric route keys arrive as numeric strings. - const normalizedKey: RouteValue = /^-?\d+$/.test(routeKey) - ? Number(routeKey) - : routeKey; - const target = routingMap[routeKey]; - if (Array.isArray(target)) { - for (const node of target) { - if (!isNodeLike(node)) { - throw new Error( - `Invalid node in fan-out tuple for route ${String(routeKey)}.`, - ); - } - } - } else if (!isNodeLike(target)) { - throw new Error( - `Invalid routing map value for route ${String(routeKey)}.`, - ); - } - if (!isRouteValue(normalizedKey)) { - throw new Error(`Invalid routing map key: ${String(routeKey)}.`); - } - expanded.push([ - fromElement, - target as NodeLike | readonly NodeLike[], - normalizedKey, - ]); - } - return expanded; -} - -/** Extracts all target nodes from a routing map, flattening fan-out arrays. */ -function nodesFromRoutingMap(routingMap: RoutingMap): NodeLike[] { - const nodes: NodeLike[] = []; - for (const target of Object.values(routingMap)) { - if (Array.isArray(target)) { - nodes.push(...(target as NodeLike[])); - } else { - nodes.push(target as NodeLike); - } - } - return nodes; -} - -/** Flattens a chain element into a list of individual nodes. */ -function flattenElement(element: ChainElement): NodeLike[] { - if (isPlainObject(element)) { - return nodesFromRoutingMap(element as RoutingMap); - } - if (Array.isArray(element)) { - return [...(element as readonly NodeLike[])]; - } - return [element as NodeLike]; -} - -/** Gets a node from the identity map or builds (and caches) it. */ -function getOrBuildNode( - nodeLike: NodeLike, - nodeMap: Map, -): BaseNode { - if (nodeLike === 'START') { - return buildNode('START'); - } - if (typeof nodeLike === 'object' || typeof nodeLike === 'function') { - const cached = nodeMap.get(nodeLike as object); - if (cached) { - return cached; - } - const built = buildNode(nodeLike); - // Only cache when a distinct wrapper was produced (or always, to preserve - // identity across repeated references within the same parse). - nodeMap.set(nodeLike as object, built); - return built; - } - return buildNode(nodeLike); -} - -function processExplicitEdge( - edge: Edge, - nodeMap: Map, - out: Edge[], -): void { - out.push( - new Edge( - getOrBuildNode(edge.fromNode, nodeMap), - getOrBuildNode(edge.toNode, nodeMap), - edge.route, - ), - ); -} - -function processRoutingMapEdge( - fromEl: ChainElement, - toEl: RoutingMap, - nodeMap: Map, - out: Edge[], -): void { - if (isPlainObject(fromEl)) { - throw new Error( - 'Consecutive routing maps are not allowed in a chain. Split them into separate edge items.', - ); - } - for (const [expFrom, expTo, route] of expandRoutingMap(fromEl, toEl)) { - for (const fromNode of flattenElement(expFrom)) { - for (const toNode of flattenElement(expTo as ChainElement)) { - out.push( - new Edge( - getOrBuildNode(fromNode, nodeMap), - getOrBuildNode(toNode, nodeMap), - route, - ), - ); - } - } - } -} - -function processUnconditionalEdge( - fromEl: ChainElement, - toEl: ChainElement, - nodeMap: Map, - out: Edge[], -): void { - for (const fromNode of flattenElement(fromEl)) { - for (const toNode of flattenElement(toEl)) { - out.push( - new Edge( - getOrBuildNode(fromNode, nodeMap), - getOrBuildNode(toNode, nodeMap), - null, - ), - ); - } - } -} - -function processChain( - chain: ChainElement[], - nodeMap: Map, - out: Edge[], -): void { - for (let i = 0; i < chain.length - 1; i++) { - const fromEl = chain[i]; - const toEl = chain[i + 1]; - if (isPlainObject(toEl)) { - processRoutingMapEdge(fromEl, toEl as RoutingMap, nodeMap, out); - } else { - processUnconditionalEdge(fromEl, toEl, nodeMap, out); - } - } -} - -/** Parses a list of edge items into a flat list of {@link Edge} objects. */ -export function parseEdgeItems(edgeItems: EdgeItem[]): Edge[] { - const nodeMap = new Map(); - const out: Edge[] = []; - - for (const item of edgeItems) { - if (item instanceof Edge) { - processExplicitEdge(item, nodeMap, out); - } else if (Array.isArray(item)) { - processChain(item, nodeMap, out); - } else { - throw new Error(`Invalid edge item type: ${typeof item}`); - } - } - - return out; -} diff --git a/core/src/workflow-next/utils/graph_validation.ts b/core/src/workflow-next/utils/graph_validation.ts deleted file mode 100644 index 29d22eae0..000000000 --- a/core/src/workflow-next/utils/graph_validation.ts +++ /dev/null @@ -1,200 +0,0 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -/** - * Validates workflow graphs and computes terminal nodes. - * - * Ported from `google/adk-python` `workflow/utils/_graph_validation.py`. - * The Phase 3 static-schema check and the Phase 7 chat-agent wiring check are - * intentionally deferred to their respective phases. - */ - -import {BaseNode, START} from '../base_node.js'; -import {DEFAULT_ROUTE, Edge} from '../graph.js'; - -function validateDuplicateNodeNames(nodes: BaseNode[]): Set { - const counts = new Map(); - for (const node of nodes) { - counts.set(node.name, (counts.get(node.name) ?? 0) + 1); - } - const duplicates = [...counts.entries()] - .filter(([, c]) => c > 1) - .map(([name]) => name) - .sort(); - if (duplicates.length > 0) { - throw new Error( - `Graph validation failed. Duplicate node names found: ${JSON.stringify( - duplicates, - )}. Pass the exact same object instance to reuse a node, or give distinct nodes unique names.`, - ); - } - return new Set(counts.keys()); -} - -function validateStartNode(nodeNames: Set): void { - if (!nodeNames.has(START.name)) { - throw new Error( - `Graph validation failed. START node (name: '${START.name}') not found in graph nodes.`, - ); - } -} - -function validateStartEdges(edges: Edge[]): void { - for (const edge of edges) { - if (edge.fromNode.name === START.name && edge.route !== null) { - throw new Error( - `Graph validation failed. Edges from START must not have routes (edge to ${edge.toNode.name} has route ${String( - edge.route, - )}).`, - ); - } - } -} - -function validateConnectivity(edges: Edge[], nodeNames: Set): void { - const adj = new Map>(); - for (const name of nodeNames) { - adj.set(name, new Set()); - } - const toNodes = new Set(); - for (const edge of edges) { - adj.get(edge.fromNode.name)!.add(edge.toNode.name); - toNodes.add(edge.toNode.name); - } - - const reachable = new Set(); - const stack = [START.name]; - while (stack.length > 0) { - const node = stack.pop()!; - if (reachable.has(node)) { - continue; - } - reachable.add(node); - for (const next of adj.get(node) ?? []) { - if (!reachable.has(next)) { - stack.push(next); - } - } - } - - const unreachable = [...nodeNames].filter((n) => !reachable.has(n)).sort(); - if (unreachable.length > 0) { - throw new Error( - `Graph validation failed. The following nodes are unreachable from START: ${JSON.stringify( - unreachable, - )}`, - ); - } - if (toNodes.has(START.name)) { - throw new Error( - 'Graph validation failed. START node must not have incoming edges.', - ); - } -} - -function validateDuplicateEdges(edges: Edge[]): void { - const seen = new Set(); - for (const edge of edges) { - const key = `${edge.fromNode.name}\u0000${edge.toNode.name}`; - if (seen.has(key)) { - throw new Error( - `Graph validation failed. Duplicate edge found: from=${edge.fromNode.name}, to=${edge.toNode.name}`, - ); - } - seen.add(key); - } -} - -function validateDefaultRoutes(edges: Edge[]): void { - const defaultRouteEdges = new Map(); - for (const edge of edges) { - if (Array.isArray(edge.route) && edge.route.includes(DEFAULT_ROUTE)) { - throw new Error( - `Graph validation failed. DEFAULT_ROUTE cannot be combined with other routes in a list (edge from=${edge.fromNode.name}, to=${edge.toNode.name}). Use a separate edge for DEFAULT_ROUTE.`, - ); - } - if (edge.route === DEFAULT_ROUTE) { - const from = edge.fromNode.name; - if (defaultRouteEdges.has(from)) { - throw new Error( - `Graph validation failed. Multiple DEFAULT_ROUTE edges found from node ${from} to ${defaultRouteEdges.get( - from, - )} and ${edge.toNode.name}`, - ); - } - defaultRouteEdges.set(from, edge.toNode.name); - } - } -} - -function detectUnconditionalCycles( - edges: Edge[], - nodeNames: Set, -): void { - const adj = new Map(); - for (const name of nodeNames) { - adj.set(name, []); - } - for (const edge of edges) { - if (edge.route === null) { - adj.get(edge.fromNode.name)!.push(edge.toNode.name); - } - } - - const inStack = new Set(); - const done = new Set(); - - const dfs = (node: string, path: string[]): void => { - inStack.add(node); - path.push(node); - for (const neighbor of adj.get(node) ?? []) { - if (inStack.has(neighbor)) { - const cycleStart = path.indexOf(neighbor); - const cycle = [...path.slice(cycleStart), neighbor]; - throw new Error( - `Graph validation failed. Unconditional cycle detected: ${cycle.join( - ' -> ', - )}. Cycles must include at least one conditional (routed) edge to avoid infinite loops.`, - ); - } - if (!done.has(neighbor)) { - dfs(neighbor, path); - } - } - path.pop(); - inStack.delete(node); - done.add(node); - }; - - for (const name of nodeNames) { - if (!done.has(name)) { - dfs(name, []); - } - } -} - -function computeTerminalNodes(nodes: BaseNode[], edges: Edge[]): Set { - const fromNames = new Set(edges.map((e) => e.fromNode.name)); - return new Set( - nodes - .filter((n) => n.name !== START.name && !fromNames.has(n.name)) - .map((n) => n.name), - ); -} - -/** - * Validates the workflow graph and returns the set of terminal node names. - */ -export function validateGraph(nodes: BaseNode[], edges: Edge[]): Set { - const nodeNames = validateDuplicateNodeNames(nodes); - validateStartNode(nodeNames); - validateStartEdges(edges); - validateConnectivity(edges, nodeNames); - validateDuplicateEdges(edges); - validateDefaultRoutes(edges); - detectUnconditionalCycles(edges, nodeNames); - return computeTerminalNodes(nodes, edges); -} diff --git a/core/src/workflow-next/utils/hitl_utils.ts b/core/src/workflow-next/utils/hitl_utils.ts deleted file mode 100644 index 4f5686077..000000000 --- a/core/src/workflow-next/utils/hitl_utils.ts +++ /dev/null @@ -1,99 +0,0 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -/** - * Utilities for Human-in-the-Loop (HITL) workflows. - * - * Ported (subset) from `google/adk-python` - * `workflow/utils/_workflow_hitl_utils.py`. The auth-credential helpers are - * added in the Phase 5 auth-gate follow-up. - */ - -import {Part} from '@google/genai'; -import {z} from 'zod'; -import {createEvent, Event} from '../../events/event.js'; -import {RequestInput} from '../request_input.js'; - -/** Function-call name marking a request-for-input interrupt. */ -export const REQUEST_INPUT_FUNCTION_CALL_NAME = 'adk_request_input'; - -/** Function-call name marking a request-for-credential interrupt. */ -export const REQUEST_CREDENTIAL_FUNCTION_CALL_NAME = 'adk_request_credential'; - -/** - * Creates an interrupt {@link Event} from a {@link RequestInput}. The event - * carries an `adk_request_input` function call and marks the interrupt id as a - * long-running tool id. - */ -export function createRequestInputEvent(requestInput: RequestInput): Event { - const args: Record = { - interruptId: requestInput.interruptId, - payload: requestInput.payload ?? null, - message: requestInput.message ?? null, - responseSchema: requestInput.responseSchema - ? z.toJSONSchema(requestInput.responseSchema) - : null, - }; - - return createEvent({ - content: { - role: 'model', - parts: [ - { - functionCall: { - name: REQUEST_INPUT_FUNCTION_CALL_NAME, - args, - id: requestInput.interruptId, - }, - }, - ], - }, - longRunningToolIds: [requestInput.interruptId], - }); -} - -/** Returns whether an event contains a `request_input` function call. */ -export function hasRequestInputFunctionCall(event: Event): boolean { - return (event.content?.parts ?? []).some( - (p) => p.functionCall?.name === REQUEST_INPUT_FUNCTION_CALL_NAME, - ); -} - -/** Returns whether an event contains an `adk_request_credential` function call. */ -export function hasAuthRequestFunctionCall(event: Event): boolean { - return (event.content?.parts ?? []).some( - (p) => p.functionCall?.name === REQUEST_CREDENTIAL_FUNCTION_CALL_NAME, - ); -} - -/** Extracts interrupt ids from `request_input` function calls in an event. */ -export function getRequestInputInterruptIds(event: Event): string[] { - const ids: string[] = []; - for (const part of event.content?.parts ?? []) { - const fc = part.functionCall; - if (fc && fc.name === REQUEST_INPUT_FUNCTION_CALL_NAME && fc.id) { - ids.push(fc.id); - } - } - return ids; -} - -/** - * Creates a `FunctionResponse` part answering a `request_input` interrupt, - * suitable for appending to a session as the user's resume response. - */ -export function createRequestInputResponse( - interruptId: string, - response: Record, -): Part { - return { - functionResponse: { - id: interruptId, - name: REQUEST_INPUT_FUNCTION_CALL_NAME, - response, - }, - }; -} diff --git a/core/src/workflow-next/utils/retry_utils.ts b/core/src/workflow-next/utils/retry_utils.ts deleted file mode 100644 index 57680e801..000000000 --- a/core/src/workflow-next/utils/retry_utils.ts +++ /dev/null @@ -1,112 +0,0 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -/** - * Utility functions for retrying nodes in a workflow. - * - * Ported from `google/adk-python` `workflow/utils/_retry_utils.py`. - */ - -import {NodeState} from '../node_state.js'; -import {RetryConfig, normalizeRetryExceptions} from '../retry_config.js'; - -const DEFAULT_MAX_ATTEMPTS = 5; -const DEFAULT_INITIAL_DELAY_SECONDS = 1.0; -const DEFAULT_MAX_DELAY_SECONDS = 60.0; -const DEFAULT_BACKOFF_FACTOR = 2.0; -const DEFAULT_JITTER = 1.0; - -/** - * Resolves the runtime name of a thrown value for exception-name matching. - * Mirrors Python's `type(exception).__name__`. - */ -function errorName(error: unknown): string { - if (error instanceof Error) { - // `name` is set by well-behaved Error subclasses; fall back to the - // constructor name for plain `throw new Error()` cases. - return error.name || error.constructor.name; - } - if (typeof error === 'object' && error !== null) { - return error.constructor.name; - } - return typeof error; -} - -/** - * Checks if a failed node should be retried based on its retry config. - * - * @param error The error thrown by the node. - * @param retryConfig The node's retry configuration, if any. - * @param nodeState The current node state (its `attemptCount` is 1-based). - */ -export function shouldRetryNode( - error: unknown, - retryConfig: RetryConfig | undefined, - nodeState: NodeState, -): boolean { - if (!retryConfig) { - return false; - } - - const attemptCount = nodeState.attemptCount; - const maxAttempts = retryConfig.maxAttempts ?? DEFAULT_MAX_ATTEMPTS; - - // attemptCount starts at 1 for the original request; once it reaches - // maxAttempts, the limit is exhausted. - if (attemptCount >= maxAttempts) { - return false; - } - - const exceptions = normalizeRetryExceptions(retryConfig.exceptions); - if (exceptions !== undefined) { - if (!exceptions.includes(errorName(error))) { - return false; - } - } - - return true; -} - -/** - * Calculates the delay, in seconds, before retrying a node. - * - * @param retryConfig The node's retry configuration, if any. - * @param nodeState The current node state (its `attemptCount` is the 1-based - * attempt number that just failed). - * @param randomFn Injectable uniform RNG in [0, 1) for deterministic testing. - */ -export function getRetryDelaySeconds( - retryConfig: RetryConfig | undefined, - nodeState: NodeState, - randomFn: () => number = Math.random, -): number { - if (!retryConfig) { - return DEFAULT_INITIAL_DELAY_SECONDS; - } - - const initialDelay = - retryConfig.initialDelay ?? DEFAULT_INITIAL_DELAY_SECONDS; - const maxDelay = retryConfig.maxDelay ?? DEFAULT_MAX_DELAY_SECONDS; - const backoffFactor = retryConfig.backoffFactor ?? DEFAULT_BACKOFF_FACTOR; - const jitter = retryConfig.jitter ?? DEFAULT_JITTER; - - const attemptCount = nodeState.attemptCount || 1; - // attemptCount is the attempt number that just failed (1-based); the first - // failure (attempt 1) uses exponent 0. - const attemptForCalc = Math.max(0, attemptCount - 1); - - let delay = initialDelay * Math.pow(backoffFactor, attemptForCalc); - delay = Math.min(delay, maxDelay); - - if (jitter > 0.0) { - // random.uniform(-jitter*delay, jitter*delay) - const span = jitter * delay; - const randomOffset = -span + randomFn() * (2 * span); - delay = Math.max(0.0, delay + randomOffset); - } - - return delay; -} diff --git a/core/src/workflow-next/workflow.ts b/core/src/workflow-next/workflow.ts deleted file mode 100644 index e9c6ed639..000000000 --- a/core/src/workflow-next/workflow.ts +++ /dev/null @@ -1,450 +0,0 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import {Event} from '../events/event.js'; -import {BaseNode, BaseNodeConfig} from './base_node.js'; -import {BranchPath} from './branch_path.js'; -import {DynamicNodeScheduler} from './dynamic_node_scheduler.js'; -import {EdgeItem, Graph, RouteValue} from './graph.js'; -import {NodeContext} 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'; - -/** - * 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}. - */ -export interface WorkflowConfig extends BaseNodeConfig { - /** Edge definitions used to build the workflow graph. */ - edges?: EdgeItem[]; - /** - * An imperative entry function driving execution via `ctx.runNode(...)`. - * Mutually exclusive with {@link edges}. - */ - dynamicEntry?: DynamicEntry; - /** - * Maximum number of graph-scheduled nodes running in parallel. `undefined` - * means unlimited. Does not throttle dynamic (`ctx.runNode`) children. - */ - maxConcurrency?: number; -} - -/** - * 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(); - errorShutDown = false; -} - -interface CompletedTask { - name: string; - childCtx?: NodeContext; - 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). - */ -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".`, - ); - } - this.maxConcurrency = config.maxConcurrency; - this.dynamicEntry = config.dynamicEntry; - if (hasEdges) { - this.graph = Graph.fromEdgeItems(config.edges!); - this.graph.validate(); - } - } - - // eslint-disable-next-line require-yield - 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(); - ctx.scheduler = new DynamicNodeScheduler(dynamicState); - - if (this.dynamicEntry) { - await this.runDynamicEntry(ctx, nodeInput, dynamicState); - return; - } - - const loop = new LoopState(); - - // --- SETUP --- - this.seedStartTriggers(loop, nodeInput); - - // --- LOOP --- - await this.runLoop(loop, ctx); - - 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; - } - } - - // --- 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): 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); - 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 && 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)!; - - let runId = nodeState.runId; - if (!runId) { - nodeState.runCounter += 1; - runId = String(nodeState.runCounter); - nodeState.runId = runId; - } - - // Static graph nodes are managed by this loop directly, bypassing the - // dynamic scheduler (which serves user-initiated ctx.runNode() calls). - const task: Promise = executeChildNode( - ctx, - node, - trigger.input, - { - 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, - ): 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 | 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 = BranchPath.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): Promise { - // Await outstanding tasks so their events flush; failures are swallowed - // because the workflow is already shutting down on error. - const outstanding = [...loop.pending.values()]; - loop.pending.clear(); - await Promise.allSettled(outstanding); - } -} diff --git a/core/src/workflow/base_node.ts b/core/src/workflow/base_node.ts index e762eb9f6..a0e9c3c41 100644 --- a/core/src/workflow/base_node.ts +++ b/core/src/workflow/base_node.ts @@ -4,71 +4,219 @@ * SPDX-License-Identifier: Apache-2.0 */ -import {InvocationContext} from '../agents/invocation_context.js'; -import {Event} from '../events/event.js'; -import {RetryConfig, normalizeRetryConfig} from './retry_config.js'; +import {Content} from '@google/genai'; +import type {ZodType} from 'zod'; +import {createEvent, Event, isEvent} from '../events/event.js'; +import type {NodeContext} from './node_context.js'; +import {isRequestInput} from './request_input.js'; +import {RetryConfig} from './retry_config.js'; +import {createRequestInputEvent} from './utils/hitl_utils.js'; /** - * Options for configuring a BaseNode. + * Configuration shared by all workflow nodes. + * + * Mirrors the fields of `google/adk-python` `workflow/_base_node.py::BaseNode`. */ -export interface BaseNodeOptions { +export interface BaseNodeConfig { + /** Canonical, unique-within-a-graph node name. */ + name: string; + + /** Human-readable description (used when a node is exposed as a tool). */ + description?: string; + /** - * If true, the node will re-execute when a workflow is resumed even if - * historical completed state exists in `InvocationContext.agentStates`. - * Default is false. + * If true, the node re-executes when a workflow resumes even if it already + * completed in a prior turn. Default false. */ rerunOnResume?: boolean; /** - * Optional retry configuration for handling transient errors during execution. + * If true, the node only produces its output once all of its predecessors + * have triggered it (fan-in / join semantics). Default false. */ + waitForOutput?: boolean; + + /** Optional retry configuration for transient failures. */ retryConfig?: RetryConfig; + + /** Maximum time, in seconds, for this node to complete. */ + timeout?: number; + + /** Optional zod schema validating the node input. */ + inputSchema?: ZodType; + + /** Optional zod schema validating the node output. */ + outputSchema?: ZodType; + + /** Optional zod schema validating relevant session state. */ + stateSchema?: ZodType; } /** - * Abstract base class for all nodes in an ADK Workflow. - * A node represents a discrete unit of execution within a static graph or dynamic chain. + * Abstract base class for all nodes in an ADK workflow. + * + * A node is a discrete unit of execution. Subclasses implement {@link runImpl}, + * which may yield {@link Event}s, raw values (boxed into an event), or + * `null`/`undefined` (skipped). {@link run} normalizes those into a stream of + * {@link Event}s consumed by the engine. */ export abstract class BaseNode { - /** - * The canonical name of the node. Must be unique within a workflow graph. - */ readonly name: string; + readonly description: string; + readonly rerunOnResume: boolean; + readonly waitForOutput: boolean; + readonly retryConfig?: RetryConfig; + readonly timeout?: number; + readonly inputSchema?: ZodType; + readonly outputSchema?: ZodType; + readonly stateSchema?: ZodType; + + constructor(config: BaseNodeConfig) { + if ( + !config.name || + typeof config.name !== 'string' || + config.name.trim().length === 0 + ) { + throw new Error('Node name must be a non-empty string.'); + } + this.name = config.name.trim(); + this.description = config.description ?? ''; + this.rerunOnResume = config.rerunOnResume ?? false; + this.waitForOutput = config.waitForOutput ?? false; + this.retryConfig = config.retryConfig; + this.timeout = config.timeout; + this.inputSchema = config.inputSchema; + this.outputSchema = config.outputSchema; + this.stateSchema = config.stateSchema; + } /** - * Whether this node should re-execute when resuming a paused or rehydrated workflow. + * Whether this node must wait for ALL of its predecessors to trigger before + * it runs (fan-in barrier). Overridden by `JoinNode`. */ - readonly rerunOnResume: boolean; + get requiresAllPredecessors(): boolean { + return false; + } /** - * The normalized retry configuration for this node, if any. + * Core execution contract. Subclasses yield one of: + * - an {@link Event} (emitted as-is), + * - a raw value (boxed into an event whose `output` is that value), + * - `null`/`undefined` (skipped). */ - readonly retryConfig?: Required; + protected abstract runImpl( + ctx: NodeContext, + input: TInput, + ): AsyncGenerator; /** - * Optional cached output payload stored on the instance during generator execution. + * Runs the node, normalizing every yielded item into an {@link Event}. This + * is what the engine (and `ctx.runNode()`) consumes. Validates the input + * against `inputSchema` once, up front (skipping genai `Content`, which nodes + * coerce themselves). */ - lastOutputPayload?: unknown; + async *run( + ctx: NodeContext, + input: TInput, + ): AsyncGenerator { + const validatedInput = this.validateInput(input); + for await (const item of this.runImpl(ctx, validatedInput)) { + if (isRequestInput(item)) { + // HITL: convert a request-for-input into an interrupt event. + yield createRequestInputEvent(item); + continue; + } + const event = this.toEvent(ctx, item); + if (event) { + yield event; + } + } + } - constructor(name: string, options?: BaseNodeOptions) { - if (!name || typeof name !== 'string' || name.trim().length === 0) { - throw new Error('Node name must be a non-empty string.'); + /** Validates node input against `inputSchema` (Content passes through). */ + protected validateInput(input: TInput): TInput { + if (!this.inputSchema || isContent(input)) { + return input; } - this.name = name.trim(); - this.rerunOnResume = options?.rerunOnResume ?? false; - this.retryConfig = normalizeRetryConfig(options?.retryConfig); + return this.inputSchema.parse(input) as TInput; + } + + /** Validates node output against `outputSchema` (Content passes through). */ + protected validateOutput(output: unknown): unknown { + if (!this.outputSchema || isContent(output)) { + return output; + } + return this.outputSchema.parse(output); } /** - * Core execution contract for a node. - * - * @param ctx The invocation context of the current workflow execution. - * @param input Optional input payload passed from upstream nodes or dynamic scheduler. - * @yields Events generated during node execution (including partial output or route events). - * @returns The final output payload of this node. + * Normalizes a single yielded item into an {@link Event} (or `null` to skip). + * Subclasses may override for richer coercion (e.g. `FunctionNode`). */ - abstract run( - ctx: InvocationContext, - input?: TInput, - ): AsyncGenerator; + protected toEvent(ctx: NodeContext, data: unknown): Event | null { + if (data === null || data === undefined) { + return null; + } + if (isEvent(data)) { + const event = data as Event; + if (event.output !== undefined) { + event.output = this.validateOutput(event.output); + } + return event; + } + const output = this.validateOutput(data); + return createEvent({ + author: this.name, + invocationId: ctx.invocationContext.invocationId, + branch: ctx.branch, + content: toContent(output), + output, + }); + } +} + +/** Returns whether a value looks like a genai `Content` object. */ +export function isContent(value: unknown): value is Content { + return ( + typeof value === 'object' && + value !== null && + 'parts' in value && + Array.isArray((value as {parts?: unknown}).parts) + ); +} + +/** + * The sentinel node marking the entry point of a workflow graph. It is never + * executed — the orchestrator seeds triggers for its successors directly. + * + * Mirrors `google/adk-python` `START = BaseNode(name='__START__')`. + */ +class StartNode extends BaseNode { + // eslint-disable-next-line require-yield + protected async *runImpl(): AsyncGenerator { + throw new Error('START node is never executed.'); + } +} + +/** The workflow entry-point sentinel node (name `__START__`). */ +export const START: BaseNode = new StartNode({name: '__START__'}); + +/** + * Best-effort conversion of an arbitrary value to genai `Content` for display. + */ +export function toContent(val: unknown): Content | undefined { + if (val === null || val === undefined) { + return undefined; + } + if (typeof val === 'object' && 'role' in val && 'parts' in val) { + return val as Content; + } + if (typeof val === 'string') { + return {role: 'model', parts: [{text: val}]}; + } + try { + return {role: 'model', parts: [{text: JSON.stringify(val)}]}; + } catch { + return {role: 'model', parts: [{text: String(val)}]}; + } } diff --git a/core/src/workflow-next/branch_path.ts b/core/src/workflow/branch_path.ts similarity index 100% rename from core/src/workflow-next/branch_path.ts rename to core/src/workflow/branch_path.ts diff --git a/core/src/workflow/dynamic_node_scheduler.ts b/core/src/workflow/dynamic_node_scheduler.ts index 7d20a92ee..c51c391da 100644 --- a/core/src/workflow/dynamic_node_scheduler.ts +++ b/core/src/workflow/dynamic_node_scheduler.ts @@ -4,152 +4,103 @@ * SPDX-License-Identifier: Apache-2.0 */ -import {InvocationContext} from '../agents/invocation_context.js'; -import {Event, isEvent} from '../events/event.js'; import {BaseNode} from './base_node.js'; +import {NodeContext} from './node_context.js'; +import {executeChildNode} from './node_runner.js'; +import {createNodeState} from './node_state.js'; +import {NodeStatus} from './node_status.js'; import { - consumeGenerator, - generateExecutionId, - getOrInitAgentStates, -} from './node_runner.js'; -import {NodeState, NodeStatus, isNodeState} from './node_state.js'; -import {FunctionNode, FunctionNodeHandler} from './nodes/function_node.js'; + DynamicNodeRun, + DynamicNodeState, + ScheduleDynamicNode, + ScheduleDynamicNodeOptions, +} from './schedule_dynamic_node.js'; /** - * Type for the dynamic workflow entry point. + * 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 type DynamicEntryFunction< - TInput = unknown, - TOutput = unknown, -> = FunctionNodeHandler; - -export type DynamicEntry = - | BaseNode - | DynamicEntryFunction; - -/** - * Options for the DynamicNodeScheduler. - */ -export interface DynamicNodeSchedulerOptions { - /** - * Key inside `InvocationContext.agentStates` where the final output of the - * dynamic entry point should be saved upon completion. - */ - outputKey?: string; -} - -/** - * Coordinates and executes a dynamic workflow where control flow (`async/await`, loops, conditionals) - * is driven programmatically by Python/TS code calling `ctx.runNode(...)`. - * Manages deterministic ID counters (`exec_node__`) and checkpoint skip-on-resume. - */ -export class DynamicNodeScheduler { - readonly entryNode: BaseNode; - readonly options: DynamicNodeSchedulerOptions; - - /** - * @param entry A BaseNode instance or a function handler to serve as the root of the dynamic workflow. - * @param options Optional configuration (outputKey). - */ - constructor(entry: DynamicEntry, options?: DynamicNodeSchedulerOptions) { - if (typeof entry === 'function') { - this.entryNode = new FunctionNode('dynamic_entry_node', entry); - } else if (isBaseNode(entry)) { - this.entryNode = entry; - } else { - throw new Error( - 'DynamicNodeScheduler requires a valid BaseNode instance or function handler.', - ); +export class DynamicNodeScheduler implements ScheduleDynamicNode { + constructor(private readonly state: DynamicNodeState) {} + + 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}/${name}@${runId}`; + + const existing = this.state.runs.get(nodePath); + if (existing?.task) { + // Deduplicate concurrent calls: await the in-flight task. + return existing.task; } - this.options = options || {}; - } - /** - * Runs the dynamic workflow entry node, intercepting events and handling checkpointing. - */ - async *runAsync( - ctx: InvocationContext, - initialInput?: unknown, - ): AsyncGenerator { - const agentStates = getOrInitAgentStates(ctx); - const execId = generateExecutionId(ctx, this.entryNode.name); + // TODO(phase-5): lazy rehydration from session events + replay + // interception (dedup completed / resume waiting runs) goes here. - const existingState = agentStates[execId] as NodeState | undefined; - if ( - existingState && - isNodeState(existingState) && - existingState.status === NodeStatus.COMPLETED && - !this.entryNode.rerunOnResume - ) { - if (this.options.outputKey) { - agentStates[this.options.outputKey] = existingState.outputPayload; - } - return; - } + return this.runFresh(ctx, node, input, name, runId, nodePath, options); + } - const stateRecord: NodeState = { - executionId: execId, - nodeName: this.entryNode.name, - status: NodeStatus.RUNNING, - inputPayload: initialInput, - timestamp: Date.now(), + 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, + }), }; - agentStates[execId] = stateRecord; - - try { - const generator = this.entryNode.run(ctx, initialInput); - const yieldedEvents: Event[] = []; - - const {output, isPausedHitl} = await consumeGenerator( - generator, - async (ev) => { - if (isEvent(ev)) { - yieldedEvents.push(ev); - } - }, - ); - - for (const ev of yieldedEvents) { - yield ev; - } - - if (isPausedHitl || ctx.endInvocation || ctx.abortSignal?.aborted) { - stateRecord.status = NodeStatus.PAUSED_HITL; - stateRecord.timestamp = Date.now(); - ctx.endInvocation = true; - return; - } - - const finalResult = - output !== undefined - ? output - : (this.entryNode.lastOutputPayload ?? - stateRecord.lastOutputPayload ?? - initialInput); - stateRecord.status = NodeStatus.COMPLETED; - stateRecord.outputPayload = finalResult; - stateRecord.timestamp = Date.now(); + this.state.runs.set(nodePath, run); + + run.task = executeChildNode(ctx, node, input, { + nodeName: name, + runId, + useAsOutput: options.useAsOutput, + useSubBranch: options.useSubBranch, + overrideBranch: options.overrideBranch, + overrideIsolationScope: options.overrideIsolationScope, + }); + + const childCtx = await run.task; + this.recordResult(run, childCtx, node); + return childCtx; + } - if (this.options.outputKey) { - agentStates[this.options.outputKey] = finalResult; - } - } catch (error: unknown) { - stateRecord.status = NodeStatus.FAILED; - stateRecord.errorMessage = - error instanceof Error ? error.message : String(error); - stateRecord.timestamp = Date.now(); - throw error; + 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; } } } - -function isBaseNode(obj: unknown): obj is BaseNode { - return ( - typeof obj === 'object' && - obj !== null && - 'name' in obj && - typeof (obj as BaseNode).name === 'string' && - 'run' in obj && - typeof (obj as BaseNode).run === 'function' - ); -} diff --git a/core/src/workflow-next/errors.ts b/core/src/workflow/errors.ts similarity index 100% rename from core/src/workflow-next/errors.ts rename to core/src/workflow/errors.ts diff --git a/core/src/workflow-next/graph.ts b/core/src/workflow/graph.ts similarity index 100% rename from core/src/workflow-next/graph.ts rename to core/src/workflow/graph.ts diff --git a/core/src/workflow/index.ts b/core/src/workflow/index.ts index f0ec9ff2b..680753936 100644 --- a/core/src/workflow/index.ts +++ b/core/src/workflow/index.ts @@ -4,45 +4,61 @@ * SPDX-License-Identifier: Apache-2.0 */ -export {BaseNode, type BaseNodeOptions} from './base_node.js'; -export { - DynamicNodeScheduler, - type DynamicEntry, - type DynamicEntryFunction, - type DynamicNodeSchedulerOptions, -} from './dynamic_node_scheduler.js'; -export { - NodeRunner, - generateExecutionId, - getOrInitAgentStates, - type NodeRunnerOptions, -} from './node_runner.js'; -export {NodeStatus, isNodeState, type NodeState} from './node_state.js'; -export {FunctionNode, type FunctionNodeHandler} from './nodes/function_node.js'; -export {JoinNode, type JoinNodeOptions} from './nodes/join_node.js'; +/** + * 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 {Node, node} from './node.js'; +export type {NodeOptions} from './node.js'; +export {FunctionNode} from './nodes/function_node.js'; +export type { + FunctionNodeConfig, + FunctionNodeHandler, +} from './nodes/function_node.js'; +export {JoinNode} from './nodes/join_node.js'; export {LLMAgentWrapper} from './nodes/llm_agent_wrapper.js'; +export type {LLMAgentWrapperConfig} from './nodes/llm_agent_wrapper.js'; +export {ParallelWorker} from './nodes/parallel_worker.js'; +export type {ParallelWorkerConfig} from './nodes/parallel_worker.js'; export {ToolNode} from './nodes/tool_node.js'; -export {runInParallel, type ParallelRunOptions} from './parallel_worker.js'; -export {normalizeRetryConfig, type RetryConfig} from './retry_config.js'; -export {runNode, type RunNodeOptions} from './run_node.js'; -export {DEFAULT_ROUTE, Trigger, type TriggerPredicate} from './trigger.js'; -export { - ParsedGraph, - parseGraph, - type AdjacencyEdge, - type EdgeElement, - type GraphEdge, -} from './utils/graph_parser.js'; -export {validateGraph} from './utils/graph_validation.js'; -export { - createRequestInputEvent, - injectHitlResumptionInput, - type RequestInputOptions, -} from './utils/hitl_utils.js'; -export { - persistAgentStatesToSession, - rehydrateAgentStates, -} from './utils/rehydration_utils.js'; -export {ReplayManager} from './utils/replay_manager.js'; -export {runWithRetry} from './utils/retry_utils.js'; -export {Workflow, isWorkflow, type WorkflowConfig} from './workflow.js'; +export type {ToolNodeConfig} from './nodes/tool_node.js'; + +// --- 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} from './branch_path.js'; +export {NodeContext} from './node_context.js'; +export {createNodeState, isNodeState} from './node_state.js'; +export type {NodeState} from './node_state.js'; +export {NodeStatus} from './node_status.js'; + +// --- HITL --- +export {RequestInput, isRequestInput} from './request_input.js'; +export type {RequestInputParams} from './request_input.js'; + +// --- Retry --- +export {normalizeRetryExceptions} from './retry_config.js'; +export type {ErrorClass, RetryConfig} from './retry_config.js'; + +// --- Errors --- +export {NodeTimeoutError} from './errors.js'; diff --git a/core/src/workflow-next/node.ts b/core/src/workflow/node.ts similarity index 100% rename from core/src/workflow-next/node.ts rename to core/src/workflow/node.ts diff --git a/core/src/workflow-next/node_context.ts b/core/src/workflow/node_context.ts similarity index 100% rename from core/src/workflow-next/node_context.ts rename to core/src/workflow/node_context.ts diff --git a/core/src/workflow/node_runner.ts b/core/src/workflow/node_runner.ts index 512ea68be..2661484e7 100644 --- a/core/src/workflow/node_runner.ts +++ b/core/src/workflow/node_runner.ts @@ -4,303 +4,241 @@ * SPDX-License-Identifier: Apache-2.0 */ -import {InvocationContext} from '../agents/invocation_context.js'; +import { + InvocationContext, + InvocationContextParams, +} from '../agents/invocation_context.js'; import {Event} from '../events/event.js'; import {BaseNode} from './base_node.js'; -import {NodeState, NodeStatus, isNodeState} from './node_state.js'; -import {GraphEdge, ParsedGraph, parseGraph} from './utils/graph_parser.js'; -import {validateGraph} from './utils/graph_validation.js'; -import {runWithRetry} from './utils/retry_utils.js'; +import {BranchPath} from './branch_path.js'; +import {NodeTimeoutError} from './errors.js'; +import {NodeContext} from './node_context.js'; +import {createNodeState} from './node_state.js'; +import {NodeStatus} from './node_status.js'; +import {getRetryDelaySeconds, shouldRetryNode} from './utils/retry_utils.js'; /** - * Options for configuring the NodeRunner. + * Options controlling a single `ctx.runNode(...)` execution. */ -export interface NodeRunnerOptions { - /** - * Whether to allow cycles in the graph during validation. - * Default is false. - */ - allowCycles?: boolean; - - /** - * Key inside `InvocationContext.agentStates` where the final leaf node outputs - * should also be written or aggregated, if requested by the workflow. - */ - outputKey?: string; -} - -interface QueueItem { - readonly node: BaseNode; - readonly inputPayload?: unknown; - readonly sourceNodeName: string; +export interface RunNodeOptions { + /** Deterministic tracking name; defaults to `node.name`. */ + nodeName?: string; + /** Unique id for this specific run; defaults to `nodeName`. */ + runId?: string; + /** If true, the child's output replaces the caller's output. */ + useAsOutput?: boolean; + /** If true, run the child in an isolated sub-branch. */ + useSubBranch?: boolean; + /** Explicit branch, overriding the default/sub-branch computation. */ + overrideBranch?: string; + /** Explicit isolation scope, overriding inheritance from the parent. */ + overrideIsolationScope?: string; } /** - * Consumes an AsyncGenerator to completion, capturing all yielded Events via onEvent - * and extracting the final return value when `done: true`. - * Also detects if any yielded Event signals a Human-in-the-Loop (`RequestInput`) pause condition. + * Executes a child node on behalf of `parent.runNode(...)`. + * + * Responsibilities (Phase 1 scope): create the child {@link NodeContext}, + * drive `node.run()`, enrich each emitted event (author, node path, branch, + * isolation scope), track the child's `output`/`route`, apply the per-node + * `timeout`, and retry on failure per `retryConfig`. Returns the child context. */ -export async function consumeGenerator( - generator: AsyncGenerator, - onEvent?: (event: Event) => void | Promise, -): Promise<{ - output: TOutput | undefined; - isPausedHitl: boolean; - lastEvent?: Event; -}> { - let isPausedHitl = false; - let lastEvent: Event | undefined; +export async function executeChildNode( + parent: NodeContext, + node: BaseNode, + input: unknown, + options: RunNodeOptions = {}, +): Promise { + const nodeName = options.nodeName ?? node.name; + const runId = options.runId ?? nodeName; + const nodePath = parent.nodePath + ? `${parent.nodePath}.${nodeName}` + : nodeName; + + let branch = parent.branch; + if (options.overrideBranch !== undefined) { + branch = options.overrideBranch; + } else if (options.useSubBranch) { + branch = BranchPath.createSubBranch(parent.branch, { + name: nodeName, + runId: options.runId, + }); + } - while (true) { - const {value, done} = await generator.next(); - if (done) { - let output = value as TOutput | undefined; - if ( - output === undefined && - lastEvent?.actions && - typeof lastEvent.actions === 'object' && - 'output' in (lastEvent.actions as unknown as Record) - ) { - output = (lastEvent.actions as unknown as Record) - .output as TOutput; + const isolationScope = + options.overrideIsolationScope ?? parent.isolationScope; + + const childIc = + branch === parent.invocationContext.branch + ? parent.invocationContext + : withBranch(parent.invocationContext, branch); + + const child = new NodeContext({ + invocationContext: childIc, + channel: parent.channel, + nodePath, + runId, + resumeInputs: parent.resumeInputs, + isolationScope, + }); + // Propagate the dynamic scheduler down; a nested Workflow overrides it. + child.scheduler = parent.scheduler; + + const nodeState = createNodeState({ + status: NodeStatus.RUNNING, + input, + runId, + }); + + for (;;) { + // Reset per-attempt output so a retry starts clean. + child.output = undefined; + child.route = undefined; + child.interruptIds = []; + try { + await runOnce(node, child, input, nodeName, branch, isolationScope); + break; + } catch (err) { + // Check retry eligibility with the attempt that just failed, compute its + // backoff delay, THEN advance the counter (matches Python semantics). + if (shouldRetryNode(err, node.retryConfig, nodeState)) { + const delaySeconds = getRetryDelaySeconds(node.retryConfig, nodeState); + nodeState.attemptCount += 1; + await delay(delaySeconds * 1000, parent.invocationContext.abortSignal); + continue; } - return {output, isPausedHitl, lastEvent}; + throw err; } + } - const event = value as Event; - lastEvent = event; - if (onEvent) { - await onEvent(event); - } - if (isHitlPauseEvent(event)) { - isPausedHitl = true; - return {output: undefined, isPausedHitl, lastEvent: event}; - } + if (options.useAsOutput) { + parent.output = child.output; + parent.route = child.route; } + + return child; } /** - * Executes a static graph workflow (`edges`) using topological queue-based scheduling, - * evaluating edge triggers upon node completion, checkpointing state in `InvocationContext.agentStates`, - * and handling Human-in-the-Loop (`RequestInput` / `PAUSED_HITL`) interruptions cleanly. + * Drives one attempt of `node.run()`, enriching and pushing each event and + * tracking the child's output/route. Wrapped in a timeout when configured. */ -export class NodeRunner { - readonly graph: ParsedGraph; - readonly options: NodeRunnerOptions; - - /** - * @param edgesOrGraph Array of GraphEdge sequences or a pre-parsed ParsedGraph. - * @param options Optional configuration for the runner. - */ - constructor( - edgesOrGraph: GraphEdge[] | ParsedGraph, - options?: NodeRunnerOptions, - ) { - if (edgesOrGraph instanceof ParsedGraph) { - this.graph = edgesOrGraph; - } else { - this.graph = parseGraph(edgesOrGraph); - } - this.options = options || {}; - validateGraph(this.graph, {allowCycles: this.options.allowCycles}); - } - - /** - * Executes the workflow graph from "START" (or from paused/rehydrated checkpoints). - * @param ctx The invocation context for the workflow run. - * @param initialInput Optional initial input payload passed to START nodes. - * @yields All events generated during node execution. - */ - async *runAsync( - ctx: InvocationContext, - initialInput?: unknown, - ): AsyncGenerator { - const agentStates = getOrInitAgentStates(ctx); - const queue: QueueItem[] = []; - - // 1. Initialize queue with edges originating from "START" - const startEdges = this.graph.adjacencyList.get('START') || []; - for (const edge of startEdges) { - queue.push({ - node: edge.target, - inputPayload: initialInput, - sourceNodeName: 'START', - }); - } - - // 2. Queue processing loop - while (queue.length > 0) { - if (ctx.endInvocation || ctx.abortSignal?.aborted) { - break; +async function runOnce( + node: BaseNode, + child: NodeContext, + input: unknown, + nodeName: string, + branch: string | undefined, + isolationScope: string | undefined, +): Promise { + const body = (async () => { + for await (const event of node.run(child, input)) { + enrichEvent(event, child, nodeName, branch, isolationScope); + if (event.output !== undefined) { + child.output = event.output; } - - const item = queue.shift()!; - const execId = generateExecutionId(ctx, item.node.name); - - const existingState = agentStates[execId] as NodeState | undefined; - let nodeOutput: unknown = undefined; - - if ( - existingState && - isNodeState(existingState) && - existingState.status === NodeStatus.COMPLETED && - !item.node.rerunOnResume - ) { - nodeOutput = existingState.outputPayload; - } else { - const effectiveInput = - existingState?.inputPayload !== undefined - ? existingState.inputPayload - : item.inputPayload; - const stateRecord: NodeState = { - executionId: execId, - nodeName: item.node.name, - status: NodeStatus.RUNNING, - inputPayload: effectiveInput, - timestamp: Date.now(), - }; - agentStates[execId] = stateRecord; - - try { - const generator = runWithRetry( - () => item.node.run(ctx, effectiveInput), - item.node.retryConfig, - ctx.abortSignal, - ); - - const yieldedEvents: Event[] = []; - const {output, isPausedHitl} = await consumeGenerator( - generator, - async (event) => { - yieldedEvents.push(event); - }, - ); - - for (const ev of yieldedEvents) { - yield ev; - } - - if (isPausedHitl || ctx.endInvocation || ctx.abortSignal?.aborted) { - stateRecord.status = NodeStatus.PAUSED_HITL; - stateRecord.timestamp = Date.now(); - ctx.endInvocation = true; - break; - } - - nodeOutput = - output !== undefined - ? output - : (item.node.lastOutputPayload ?? - stateRecord.lastOutputPayload ?? - item.inputPayload); - stateRecord.status = NodeStatus.COMPLETED; - stateRecord.outputPayload = nodeOutput; - stateRecord.timestamp = Date.now(); - } catch (error: unknown) { - stateRecord.status = NodeStatus.FAILED; - stateRecord.errorMessage = - error instanceof Error ? error.message : String(error); - stateRecord.timestamp = Date.now(); - throw error; - } + if (event.route !== undefined) { + child.route = event.route; } - - // 3. Evaluate outgoing edges and enqueue successors whose triggers are satisfied - const outgoingEdges = this.graph.adjacencyList.get(item.node.name) || []; - let hasRoutingEdges = false; - let matchedSpecificRoute = false; - let defaultRouteEdge: (typeof outgoingEdges)[0] | undefined; - - for (const edge of outgoingEdges) { - if (!edge.trigger) { - queue.push({ - node: edge.target, - inputPayload: nodeOutput, - sourceNodeName: item.node.name, - }); - continue; - } - - hasRoutingEdges = true; - if (edge.trigger.isDefaultRoute()) { - defaultRouteEdge = edge; - continue; - } - - const triggerSatisfied = await edge.trigger.evaluate(ctx, nodeOutput); - if (triggerSatisfied) { - matchedSpecificRoute = true; - queue.push({ - node: edge.target, - inputPayload: nodeOutput, - sourceNodeName: item.node.name, - }); + // HITL: an interrupt event marks its ids as long-running tool ids. + if (event.longRunningToolIds && event.longRunningToolIds.length > 0) { + for (const id of event.longRunningToolIds) { + if (!child.interruptIds.includes(id)) { + child.interruptIds.push(id); + } } } - - if (hasRoutingEdges && !matchedSpecificRoute && defaultRouteEdge) { - queue.push({ - node: defaultRouteEdge.target, - inputPayload: nodeOutput, - sourceNodeName: item.node.name, - }); - } + child.channel.push(event); } + })(); - if (this.options.outputKey) { - const finalStates: Record = {}; - for (const state of Object.values(agentStates)) { - if (isNodeState(state) && state.status === NodeStatus.COMPLETED) { - finalStates[state.nodeName] = state.outputPayload; - } - } - agentStates[this.options.outputKey] = finalStates; - } + if (node.timeout && node.timeout > 0) { + await withTimeout(body, node.timeout, nodeName); + } else { + await body; } } /** - * Gets or initializes the `agentStates` record on the invocation context. + * Stamps engine-owned provenance onto an event without clobbering values the + * node explicitly set. */ -export function getOrInitAgentStates( - ctx: InvocationContext, -): Record { - const unknownCtx = ctx as unknown as Record; - if (!unknownCtx.agentStates || typeof unknownCtx.agentStates !== 'object') { - unknownCtx.agentStates = {}; +function enrichEvent( + event: Event, + child: NodeContext, + nodeName: string, + branch: string | undefined, + isolationScope: string | undefined, +): void { + if (!event.author) { + event.author = nodeName; + } + event.nodeInfo = {...(event.nodeInfo ?? {}), path: child.nodePath}; + if (branch !== undefined && event.branch === undefined) { + event.branch = branch; + } + if (isolationScope !== undefined && event.isolationScope === undefined) { + event.isolationScope = isolationScope; } - return unknownCtx.agentStates as Record; } /** - * Generates a deterministic execution ID for a node based on the context branch and node name. + * Creates a shallow child InvocationContext with a different branch, preserving + * the shared invocation cost manager and all services/session. */ -export function generateExecutionId( - ctx: InvocationContext, +function withBranch( + ic: InvocationContext, + branch: string | undefined, +): InvocationContext { + return new InvocationContext({ + ...(ic as unknown as InvocationContextParams), + branch, + }); +} + +/** + * Rejects with {@link NodeTimeoutError} if `promise` does not settle within + * `timeoutSeconds`. + */ +function withTimeout( + promise: Promise, + timeoutSeconds: number, nodeName: string, -): string { - const branchPrefix = ctx.branch ? `${ctx.branch}.` : ''; - return `exec_node_${branchPrefix}${nodeName}`; +): Promise { + return new Promise((resolve, reject) => { + const timer = setTimeout(() => { + reject(new NodeTimeoutError({nodeName, timeout: timeoutSeconds})); + }, timeoutSeconds * 1000); + promise.then( + (value) => { + clearTimeout(timer); + resolve(value); + }, + (err) => { + clearTimeout(timer); + reject(err); + }, + ); + }); } /** - * Checks whether an event signals a Human-in-the-Loop pause (`RequestInput`). + * Promise-based delay that rejects early if the abort signal fires. */ -export function isHitlPauseEvent(event: Event): boolean { - if (!event) return false; - if (event.actions && typeof event.actions === 'object') { - if ( - 'requestInput' in event.actions && - Boolean((event.actions as Record).requestInput) - ) { - return true; +function delay(ms: number, signal?: AbortSignal): Promise { + return new Promise((resolve, reject) => { + if (signal?.aborted) { + reject(new Error('Aborted')); + return; } - } - if ( - 'requestInput' in event && - Boolean((event as Record).requestInput) - ) { - return true; - } - return false; + const timer = setTimeout(() => { + signal?.removeEventListener('abort', onAbort); + resolve(); + }, ms); + const onAbort = () => { + clearTimeout(timer); + reject(new Error('Aborted')); + }; + signal?.addEventListener('abort', onAbort, {once: true}); + }); } diff --git a/core/src/workflow/node_state.ts b/core/src/workflow/node_state.ts index 4edcd4dfb..8950bbaa5 100644 --- a/core/src/workflow/node_state.ts +++ b/core/src/workflow/node_state.ts @@ -4,89 +4,74 @@ * SPDX-License-Identifier: Apache-2.0 */ -import {Event} from '../events/event.js'; +import {NodeStatus} from './node_status.js'; /** - * Represents the execution status of a node in a workflow graph or dynamic chain. + * State of a node in the workflow. + * + * Ported from `google/adk-python` `workflow/_node_state.py`. Note that the + * node's *output* is intentionally NOT stored here — it is carried on emitted + * events / the node Context, not on the persisted node state. */ -export enum NodeStatus { - PENDING = 'PENDING', - RUNNING = 'RUNNING', - COMPLETED = 'COMPLETED', - PAUSED_HITL = 'PAUSED_HITL', - FAILED = 'FAILED', -} - -/** - * Checkpointed state for a specific node execution. - * Stored inside `InvocationContext.agentStates[executionId]`. - */ -export interface NodeState { - /** - * The deterministic execution ID assigned to this node execution. - */ - executionId: string; - - /** - * The canonical name of the node. - */ - nodeName: string; - - /** - * The current status of the node execution. - */ +export interface NodeState { + /** The run status of the node. */ status: NodeStatus; - /** - * The input payload passed into the node during execution. - */ - inputPayload?: TInput; + /** The input provided to the node. */ + input?: unknown; - /** - * The final output payload yielded or returned by the node upon completion. - */ - outputPayload?: TOutput; + /** The attempt count for this node run (1-based). */ + attemptCount: number; - /** - * Error message if the node execution failed (`status === FAILED`). - */ - errorMessage?: string; + /** The interrupt ids that are pending to be resolved. */ + interrupts: string[]; - /** - * Timestamp in milliseconds when this state record was last updated. - */ - timestamp: number; + /** The responses for resuming the node, keyed by interrupt id. */ + resumeInputs: Record; /** - * Events emitted by the node during live execution, cached for replaying on resumption. + * Sequential counter incremented each time the node gets a fresh run. + * + * Preserving this count independently of `runId` prevents path collisions if + * a node switches between custom string IDs and auto-generated numeric IDs. */ - cachedEvents?: Event[]; + runCounter: number; - /** - * Indicates if this node previously paused for Human-in-the-Loop (`RequestInput`). - */ - wasPausedHitl?: boolean; + /** The run ID of this node run. */ + runId?: string; /** - * Stores intermediate or final payload before completion status transition. + * The run ID of the parent node which dynamically scheduled this node run. */ - lastOutputPayload?: unknown; + parentRunId?: string; +} + +/** + * Creates a {@link NodeState} with Python-aligned defaults, overlaying any + * provided partial values. + */ +export function createNodeState(partial?: Partial): NodeState { + return { + status: NodeStatus.INACTIVE, + attemptCount: 1, + interrupts: [], + resumeInputs: {}, + runCounter: 0, + ...partial, + }; } /** - * Type guard to check if an object is a valid NodeState instance. - * @param obj The object to check. - * @returns True if the object matches the NodeState structure. + * Type guard for a {@link NodeState}-shaped object. */ export function isNodeState(obj: unknown): obj is NodeState { return ( typeof obj === 'object' && obj !== null && - 'executionId' in obj && - typeof (obj as NodeState).executionId === 'string' && - 'nodeName' in obj && - typeof (obj as NodeState).nodeName === 'string' && 'status' in obj && - Object.values(NodeStatus).includes((obj as NodeState).status) + typeof (obj as NodeState).status === 'number' && + 'attemptCount' in obj && + 'interrupts' in obj && + Array.isArray((obj as NodeState).interrupts) ); } diff --git a/core/src/workflow-next/node_status.ts b/core/src/workflow/node_status.ts similarity index 100% rename from core/src/workflow-next/node_status.ts rename to core/src/workflow/node_status.ts diff --git a/core/src/workflow/nodes/function_node.ts b/core/src/workflow/nodes/function_node.ts index 347e046fc..b250ab25c 100644 --- a/core/src/workflow/nodes/function_node.ts +++ b/core/src/workflow/nodes/function_node.ts @@ -4,131 +4,167 @@ * SPDX-License-Identifier: Apache-2.0 */ -import {Content} from '@google/genai'; -import {InvocationContext} from '../../agents/invocation_context.js'; +import {AuthConfig} from '../../auth/auth_tool.js'; import {createEvent, Event, isEvent} from '../../events/event.js'; -import {BaseNode, BaseNodeOptions} from '../base_node.js'; +import {BaseNode, BaseNodeConfig, isContent, toContent} from '../base_node.js'; +import {NodeContext} from '../node_context.js'; /** - * Type for the function wrapped by a FunctionNode. - * Can return a direct value, a Promise of a value, an Event, or an AsyncGenerator of Events. + * A value a {@link FunctionNodeHandler} may return or yield. + */ +export type FunctionNodeResult = + | TOutput + | Event + | null + | undefined + | void; + +/** + * The handler wrapped by a {@link FunctionNode}. + * + * Unlike Python's `FunctionNode` (which binds named parameters from `ctx.state` + * or `node_input` via runtime signature introspection), the TypeScript form + * uses the idiomatic explicit `(ctx, input)` signature. Read `ctx.state` + * directly for state-bound values. It may return a value/`Event`, a Promise, or + * a (sync/async) generator of those. */ export type FunctionNodeHandler = ( - ctx: InvocationContext, - input?: TInput, + ctx: NodeContext, + input: TInput, ) => - | AsyncGenerator - | Promise - | TOutput - | Event; + | FunctionNodeResult + | Promise> + | Generator, void, unknown> + | AsyncGenerator, void, unknown>; /** - * A concrete node that wraps a deterministic JavaScript/TypeScript function. - * Automatically handles generator streams, Event returns, or boxes plain return values into Event outputs. + * Options for a {@link FunctionNode}. + */ +export interface FunctionNodeConfig extends Partial< + Omit +> { + /** + * If set, the framework requests user authentication before running (Phase 5 + * enables the auth gate; stored here now for API parity). + */ + authConfig?: AuthConfig; +} + +/** + * A node that wraps a plain function, async function, or (sync/async) generator. + * + * Ported (TS-idiomatic subset) from `google/adk-python` `_function_node.py`. + * Return-value handling: + * - `Event` → emitted as-is (output validated against `outputSchema`) + * - genai `Content` → emitted as the event content + * - `null`/`undefined` → skipped (unless there are pending state deltas) + * - anything else → emitted as `Event(output=value)` + * State written via `ctx.state` during execution is attached to emitted events. */ export class FunctionNode extends BaseNode< TInput, TOutput > { + readonly authConfig?: AuthConfig; private readonly handler: FunctionNodeHandler; - /** - * @param name Unique name for this function node. - * @param handler The execution logic function. - * @param options Optional BaseNode configuration (rerunOnResume, retryConfig). - */ constructor( name: string, handler: FunctionNodeHandler, - options?: BaseNodeOptions, + config: FunctionNodeConfig = {}, ) { - super(name, options); if (typeof handler !== 'function') { - throw new Error( - `FunctionNode "${name}" requires a valid function handler.`, - ); + throw new TypeError('FunctionNode handler must be a function.'); } + super({name, ...config}); this.handler = handler; + this.authConfig = config.authConfig; } - /** - * Executes the wrapped handler function inside the workflow context. - */ - async *run( - ctx: InvocationContext, - input?: TInput, - ): AsyncGenerator { - const resultOrGen = this.handler(ctx, input); + protected async *runImpl( + ctx: NodeContext, + input: TInput, + ): AsyncGenerator { + // TODO(phase-5): auth gate (authConfig -> adk_request_credential interrupt). + const result = this.handler(ctx, input); - if (isAsyncGenerator(resultOrGen)) { - const finalVal = yield* resultOrGen; - this.lastOutputPayload = finalVal; - return finalVal; + if (isAsyncIterable(result)) { + for await (const item of result) { + yield item; + } + } else if (isSyncGenerator(result)) { + for (const item of result) { + yield item; + } + } else { + // Plain value or Promise of a value. + yield await (result as Promise>); } + } - const res = await Promise.resolve(resultOrGen); + protected override toEvent(ctx: NodeContext, data: unknown): Event | null { + const stateDelta = + Object.keys(ctx.actions.stateDelta).length > 0 + ? {...ctx.actions.stateDelta} + : undefined; - if (isEvent(res)) { - yield res; - const extracted = - res.content ?? - (typeof res.actions === 'object' && - res.actions !== null && - 'output' in (res.actions as unknown as Record) - ? (res.actions as unknown as Record).output - : res); - this.lastOutputPayload = extracted; - return extracted as TOutput; + if (data === null || data === undefined) { + return stateDelta + ? createEvent({ + author: this.name, + invocationId: ctx.invocationId, + branch: ctx.branch, + actions: {stateDelta}, + }) + : null; } - if (res !== undefined && res !== null) { - const boxedEvent = createEvent({ - invocationId: ctx.invocationId, + if (isEvent(data)) { + const event = data as Event; + if (event.output !== undefined) { + event.output = this.validateOutput(event.output); + } + if (stateDelta) { + Object.assign(event.actions.stateDelta, stateDelta); + } + return event; + } + + if (isContent(data)) { + return createEvent({ author: this.name, + invocationId: ctx.invocationId, branch: ctx.branch, - content: toContent(res), - actions: {output: res}, + content: data, + actions: stateDelta ? {stateDelta} : undefined, }); - yield boxedEvent; } - this.lastOutputPayload = res; - return res as TOutput; + const output = this.validateOutput(data); + return createEvent({ + author: this.name, + invocationId: ctx.invocationId, + branch: ctx.branch, + content: toContent(output), + output, + actions: stateDelta ? {stateDelta} : undefined, + }); } } -function isAsyncGenerator( - obj: unknown, -): obj is AsyncGenerator { +function isAsyncIterable(value: unknown): value is AsyncIterable { return ( - typeof obj === 'object' && - obj !== null && - Symbol.asyncIterator in obj && - 'next' in obj && - typeof (obj as Record).next === 'function' + value != null && + typeof (value as AsyncIterable)[Symbol.asyncIterator] === + 'function' ); } -function toContent(val: unknown): Content | undefined { - if (!val) return undefined; - if (typeof val === 'object' && 'role' in val && 'parts' in val) { - return val as Content; - } - if (typeof val === 'string') { - return { - role: 'model', - parts: [{text: val}], - }; - } - try { - return { - role: 'model', - parts: [{text: JSON.stringify(val)}], - }; - } catch { - return { - role: 'model', - parts: [{text: String(val)}], - }; - } +function isSyncGenerator(value: unknown): value is Generator { + return ( + value != null && + typeof value !== 'string' && + typeof (value as Iterable)[Symbol.iterator] === 'function' && + typeof (value as Generator).next === 'function' + ); } diff --git a/core/src/workflow/nodes/join_node.ts b/core/src/workflow/nodes/join_node.ts index ddcfd772f..23543d338 100644 --- a/core/src/workflow/nodes/join_node.ts +++ b/core/src/workflow/nodes/join_node.ts @@ -4,124 +4,31 @@ * SPDX-License-Identifier: Apache-2.0 */ -import {InvocationContext} from '../../agents/invocation_context.js'; import {createEvent, Event} from '../../events/event.js'; -import {BaseNode, BaseNodeOptions} from '../base_node.js'; -import {getOrInitAgentStates} from '../node_runner.js'; -import {isNodeState, NodeStatus} from '../node_state.js'; +import {BaseNode} from '../base_node.js'; +import {NodeContext} from '../node_context.js'; /** - * Options for configuring a JoinNode. + * A fan-in barrier node: it waits for ALL of its predecessors to complete, then + * emits the aggregated inputs (a map of predecessor name → output) as its + * output. + * + * Ported from `google/adk-python` `workflow/_join_node.py`. */ -export interface JoinNodeOptions extends BaseNodeOptions { - /** - * The number of distinct upstream predecessor nodes that must complete - * before this join node unblocks and emits a combined output. - * Must be >= 1. - */ - upstreamCount: number; - - /** - * Optional array of explicit predecessor node names to wait on. - * If provided, `upstreamCount` must match `predecessors.length`. - */ - predecessors?: string[]; -} - -/** - * A synchronization barrier node used in fan-out/fan-in parallel workflows. - * Waits until `upstreamCount` predecessor branches have reached completion (`COMPLETED`), - * then aggregates their outputs into a dictionary and yields a single combined event. - */ -export class JoinNode< - TInput = unknown, - TOutput = Record, -> extends BaseNode { - readonly upstreamCount: number; - readonly predecessors?: string[]; - - constructor(name: string, options?: Partial) { - const count = options?.upstreamCount ?? 0; - if (count < 0) { - throw new Error( - `JoinNode "${name}" requires a valid upstreamCount >= 0.`, - ); - } - if ( - options?.predecessors && - options.predecessors.length !== count && - count > 0 - ) { - throw new Error( - `JoinNode "${name}" upstreamCount (${count}) does not match predecessors.length (${options.predecessors.length}).`, - ); - } - super(name, options); - this.upstreamCount = count; - this.predecessors = options?.predecessors; +export class JoinNode extends BaseNode { + override get requiresAllPredecessors(): boolean { + return true; } - /** - * Evaluates the completion status of upstream predecessor nodes in `InvocationContext.agentStates`. - * Only unblocks and yields combined output when all required predecessors have completed. - */ - async *run( - ctx: InvocationContext, - _input?: TInput, - ): AsyncGenerator { - const agentStates = getOrInitAgentStates(ctx); - - const completedPredecessors: Record = {}; - let count = 0; - - if (this.predecessors && this.predecessors.length > 0) { - for (const predName of this.predecessors) { - for (const state of Object.values(agentStates)) { - if ( - isNodeState(state) && - state.nodeName === predName && - state.status === NodeStatus.COMPLETED - ) { - completedPredecessors[predName] = state.outputPayload; - count++; - break; - } - } - } - } else { - for (const state of Object.values(agentStates)) { - if ( - isNodeState(state) && - state.nodeName !== this.name && - state.status === NodeStatus.COMPLETED && - !(state.nodeName in completedPredecessors) - ) { - completedPredecessors[state.nodeName] = state.outputPayload; - count++; - } - } - } - - if (count < this.upstreamCount) { - return completedPredecessors as unknown as TOutput; - } - - const joinEvent = createEvent({ - invocationId: ctx.invocationId, + protected async *runImpl( + ctx: NodeContext, + input: unknown, + ): AsyncGenerator { + yield createEvent({ author: this.name, + invocationId: ctx.invocationId, branch: ctx.branch, - actions: { - joinCompleted: { - node: this.name, - upstreamCount: this.upstreamCount, - predecessors: Object.keys(completedPredecessors), - outputs: completedPredecessors, - }, - }, + output: input, }); - - yield joinEvent; - this.lastOutputPayload = completedPredecessors; - return completedPredecessors as unknown as TOutput; } } diff --git a/core/src/workflow/nodes/llm_agent_wrapper.ts b/core/src/workflow/nodes/llm_agent_wrapper.ts index 1098ff1a5..4af1aba00 100644 --- a/core/src/workflow/nodes/llm_agent_wrapper.ts +++ b/core/src/workflow/nodes/llm_agent_wrapper.ts @@ -6,87 +6,99 @@ import {Content} from '@google/genai'; import {BaseAgent} from '../../agents/base_agent.js'; -import { - InvocationContext, - InvocationContextParams, -} from '../../agents/invocation_context.js'; -import {Event} from '../../events/event.js'; -import {BaseNode, BaseNodeOptions} from '../base_node.js'; +import {createEvent, Event} from '../../events/event.js'; +import {BaseNode, BaseNodeConfig, isContent} from '../base_node.js'; +import {NodeContext} from '../node_context.js'; + +/** Options for an {@link LLMAgentWrapper}. */ +export interface LLMAgentWrapperConfig extends Partial< + Omit +> { + name?: string; +} /** - * A concrete node that wraps any ADK BaseAgent (e.g., LlmAgent, SequentialAgent) - * so it can participate as a node inside a workflow graph. - * Enforces single-turn task execution mode and relays generated events. + * Runs a {@link BaseAgent} (typically an `LlmAgent`) as a workflow node in + * `single_turn` mode: the node input is appended as a user turn, the agent runs + * once, and its final model text becomes the node output. + * + * Ported (single_turn subset) from `google/adk-python` + * `workflow/_llm_agent_wrapper.py`. The `task` and `chat` modes (FinishTaskTool, + * task delegation, transfer, isolation scopes) are a Phase 7b continuation. */ -export class LLMAgentWrapper< - TInput = unknown, - TOutput = unknown, -> extends BaseNode { +export class LLMAgentWrapper extends BaseNode { readonly agent: BaseAgent; - /** - * @param agent The BaseAgent instance to wrap. - * @param options Optional BaseNode configuration (name override, rerunOnResume, retryConfig). - */ - constructor(agent: BaseAgent, options?: BaseNodeOptions & {name?: string}) { - if (!agent || typeof agent.runAsync !== 'function') { - throw new Error('LLMAgentWrapper requires a valid BaseAgent instance.'); - } - super(options?.name || agent.name || 'llm_agent_wrapper', options); + constructor(agent: BaseAgent, config: LLMAgentWrapperConfig = {}) { + super({ + name: config.name ?? agent.name, + description: agent.description, + ...config, + }); this.agent = agent; } - /** - * Invokes the wrapped agent via runAsync and relays all produced events. - */ - async *run( - ctx: InvocationContext, - input?: TInput, - ): AsyncGenerator { - let lastOutput: unknown = undefined; - - const childCtxParams: InvocationContextParams & Record = { - ...ctx, - agent: this.agent, - }; - + protected async *runImpl( + ctx: NodeContext, + input: unknown, + ): AsyncGenerator { + // Append the node input as a user turn so the agent responds to it. if (input !== undefined && input !== null) { - if (typeof input === 'string') { - childCtxParams.userContent = { - role: 'user', - parts: [{text: input}], - }; - } else if ( - typeof input === 'object' && - 'role' in input && - 'parts' in input - ) { - childCtxParams.userContent = input as unknown as Content; + const userEvent = createEvent({ + author: 'user', + invocationId: ctx.invocationId, + branch: ctx.branch, + content: toUserContent(input), + }); + if (ctx.isolationScope) { + userEvent.isolationScope = ctx.isolationScope; } + ctx.session.events.push(userEvent); } - const childCtx = new InvocationContext(childCtxParams); - - for await (const event of this.agent.runAsync(childCtx)) { + // Run the agent under the node's invocation context (it sets agent=itself). + for await (const event of this.agent.runAsync(ctx.invocationContext)) { + this.maybeSetOutput(event); yield event; - if (event.content?.parts?.length) { - const texts = event.content.parts.map((p) => p.text).filter(Boolean); - if (texts.length > 0) { - lastOutput = texts.join('\n'); - } - } - if ( - event.actions && - typeof event.actions === 'object' && - 'output' in (event.actions as unknown as Record) - ) { - lastOutput = (event.actions as unknown as Record) - .output; - } } + } - const finalVal = lastOutput ?? input; - this.lastOutputPayload = finalVal; - return finalVal as TOutput; + /** + * Promotes the final model text of an event to the node output (mirroring + * Python `process_llm_agent_output`). + */ + private maybeSetOutput(event: Event): void { + if (event.partial) { + return; + } + if (hasFunctionCalls(event)) { + return; + } + const content = event.content; + if (!content || content.role !== 'model' || !content.parts) { + return; + } + const text = content.parts + .filter((p) => p.text && !p.thought) + .map((p) => p.text) + .join(''); + + event.output = text; + event.nodeInfo = {...(event.nodeInfo ?? {}), messageAsOutput: true}; + } +} + +function hasFunctionCalls(event: Event): boolean { + return (event.content?.parts ?? []).some((p) => p.functionCall); +} + +/** Converts an arbitrary node input into a user-role `Content`. */ +function toUserContent(input: unknown): Content { + if (isContent(input)) { + return {...input, role: 'user'}; + } + if (typeof input === 'string') { + return {role: 'user', parts: [{text: input}]}; } + return {role: 'user', parts: [{text: JSON.stringify(input)}]}; } diff --git a/core/src/workflow-next/nodes/parallel_worker.ts b/core/src/workflow/nodes/parallel_worker.ts similarity index 100% rename from core/src/workflow-next/nodes/parallel_worker.ts rename to core/src/workflow/nodes/parallel_worker.ts diff --git a/core/src/workflow/nodes/tool_node.ts b/core/src/workflow/nodes/tool_node.ts index b10cf1c85..b70825645 100644 --- a/core/src/workflow/nodes/tool_node.ts +++ b/core/src/workflow/nodes/tool_node.ts @@ -5,63 +5,102 @@ */ import {Context} from '../../agents/context.js'; -import {InvocationContext} from '../../agents/invocation_context.js'; import {createEvent, Event} from '../../events/event.js'; import {BaseTool} from '../../tools/base_tool.js'; -import {BaseNode, BaseNodeOptions} from '../base_node.js'; +import {randomUUID} from '../../utils/env_aware_utils.js'; +import {BaseNode, BaseNodeConfig, isContent} from '../base_node.js'; +import {NodeContext} from '../node_context.js'; +/** Options for a {@link ToolNode}. */ +export interface ToolNodeConfig extends Partial> { + /** Optional name override; defaults to the tool's name. */ + name?: string; +} /** - * A concrete node that wraps an ADK BaseTool (or FunctionTool) so it can be executed - * directly as a step in a workflow graph without requiring an LLM wrapper. + * A node that wraps an ADK {@link BaseTool} and invokes it with the node input + * as its arguments. + * + * Ported from `google/adk-python` `workflow/_tool_node.py`. The node input is + * coerced to a tool-args object: genai `Content` → its text; a JSON string → + * parsed object; `null`/empty → `{}`. */ -export class ToolNode< - TInput = Record, - TOutput = unknown, -> extends BaseNode { +export class ToolNode extends BaseNode { readonly tool: BaseTool; - /** - * @param tool The BaseTool instance to execute when this node runs. - * @param options Optional BaseNode configuration (name override, rerunOnResume, retryConfig). - */ - constructor(tool: BaseTool, options?: BaseNodeOptions & {name?: string}) { - if (!tool || typeof tool.runAsync !== 'function') { - throw new Error( - 'ToolNode requires a valid BaseTool instance with runAsync().', - ); - } - super(options?.name || tool.name || 'tool_node', options); + constructor(tool: BaseTool, config: ToolNodeConfig = {}) { + super({name: config.name ?? tool.name, ...config}); this.tool = tool; } - /** - * Executes the wrapped tool using parameters from the input payload. - */ - async *run( - ctx: InvocationContext, - input?: TInput, - ): AsyncGenerator { - const params = typeof input === 'object' && input !== null ? input : {}; - const result = await this.tool.runAsync({ - toolContext: ctx as unknown as Context, - args: params as Record, + protected async *runImpl( + ctx: NodeContext, + input: unknown, + ): AsyncGenerator { + const toolContext = new Context({ + invocationContext: ctx.invocationContext, + functionCallId: randomUUID(), }); - const event = createEvent({ - invocationId: ctx.invocationId, - author: this.name, - branch: ctx.branch, - actions: { - toolExecution: { - name: this.tool.name, - input: params, - output: result, - }, - }, - }); + const args = coerceToolArgs(input); + const response = await this.tool.runAsync({args, toolContext}); + + const stateDelta = + Object.keys(toolContext.actions.stateDelta).length > 0 + ? {...toolContext.actions.stateDelta} + : undefined; + + if (response !== undefined && response !== null) { + yield createEvent({ + author: this.name, + invocationId: ctx.invocationId, + branch: ctx.branch, + output: response, + actions: stateDelta ? {stateDelta} : undefined, + }); + } else if (stateDelta) { + yield createEvent({ + author: this.name, + invocationId: ctx.invocationId, + branch: ctx.branch, + actions: {stateDelta}, + }); + } + } +} + +/** Coerces arbitrary node input into a tool-arguments record. */ +function coerceToolArgs(input: unknown): Record { + let args: unknown = input; - yield event; - this.lastOutputPayload = result; - return result as TOutput; + if (isContent(args)) { + args = extractText(args); } + + if (typeof args === 'string') { + const trimmed = args.trim(); + if (!trimmed) { + args = null; + } else { + try { + args = JSON.parse(trimmed); + } catch { + // Leave as the raw string; validated below. + } + } + } + + if (args === null || args === undefined) { + return {}; + } + if (typeof args !== 'object' || Array.isArray(args)) { + throw new TypeError( + 'The input to ToolNode must be a dictionary of tool arguments or null, ' + + `but got ${typeof args}.`, + ); + } + return args as Record; +} + +function extractText(content: {parts?: Array<{text?: string}>}): string { + return (content.parts ?? []).map((p) => p.text ?? '').join(''); } diff --git a/core/src/workflow/parallel_worker.ts b/core/src/workflow/parallel_worker.ts deleted file mode 100644 index 3d3a83ead..000000000 --- a/core/src/workflow/parallel_worker.ts +++ /dev/null @@ -1,132 +0,0 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import {InvocationContext} from '../agents/invocation_context.js'; -import {BaseNode} from './base_node.js'; -import {consumeGenerator, getOrInitAgentStates} from './node_runner.js'; -import {NodeState, NodeStatus, isNodeState} from './node_state.js'; -import {FunctionNode, FunctionNodeHandler} from './nodes/function_node.js'; - -/** - * Options for configuring parallel branch execution (`runInParallel`). - */ -export interface ParallelRunOptions { - /** - * Optional prefix for naming the child branches in `InvocationContext.branch`. - * Defaults to the target node name. - */ - branchPrefix?: string; - - /** - * If true, throws an error immediately if any parallel worker branch fails. - * If false, catches branch errors and returns undefined/error markers for failed items. - * Default is true. - */ - stopOnError?: boolean; -} - -/** - * Executes a target node (or function) concurrently across an array of input items using isolated - * child `InvocationContext` branches. Prevents concurrent async tasks from corrupting shared - * session event histories or racing on shared node state checkpoints. - * - * @param ctx The parent invocation context. - * @param nodeOrFunc The BaseNode or function handler to execute for each item. - * @param items Array of input items to process in parallel. - * @param options Optional settings (branchPrefix, stopOnError). - * @returns Promise resolving to an array of output payloads corresponding 1-to-1 with the items array. - */ -export async function runInParallel( - ctx: InvocationContext, - nodeOrFunc: BaseNode | FunctionNodeHandler, - items: TInput[], - options?: ParallelRunOptions, -): Promise { - if (!Array.isArray(items) || items.length === 0) { - return []; - } - - const node = - typeof nodeOrFunc === 'function' - ? new FunctionNode('parallel_worker_node', nodeOrFunc) - : nodeOrFunc; - - const prefix = options?.branchPrefix || node.name; - const stopOnError = options?.stopOnError ?? true; - const parentStates = getOrInitAgentStates(ctx); - - const tasks = items.map(async (item, index) => { - const branchName = ctx.branch - ? `${ctx.branch}.${prefix}_${index}` - : `${prefix}_${index}`; - - const childCtx = new InvocationContext({ - ...ctx, - branch: branchName, - }); - - const childStates: Record = {...parentStates}; - (childCtx as unknown as Record).agentStates = childStates; - - const execId = `exec_node_${branchName}.${node.name}`; - - const existingState = childStates[execId] as NodeState | undefined; - if ( - existingState && - isNodeState(existingState) && - existingState.status === NodeStatus.COMPLETED && - !node.rerunOnResume - ) { - return existingState.outputPayload as TOutput; - } - - const stateRecord: NodeState = { - executionId: execId, - nodeName: `${node.name}_${index}`, - status: NodeStatus.RUNNING, - inputPayload: item, - timestamp: Date.now(), - }; - childStates[execId] = stateRecord; - - try { - const generator = node.run(childCtx, item); - const {output, isPausedHitl} = await consumeGenerator(generator); - - if (isPausedHitl || ctx.endInvocation || ctx.abortSignal?.aborted) { - stateRecord.status = NodeStatus.PAUSED_HITL; - stateRecord.timestamp = Date.now(); - ctx.endInvocation = true; - throw new Error( - `Parallel worker branch "${branchName}" requested HITL pause.`, - ); - } - - const finalVal = - output !== undefined - ? output - : (node.lastOutputPayload ?? stateRecord.lastOutputPayload ?? item); - stateRecord.status = NodeStatus.COMPLETED; - stateRecord.outputPayload = finalVal; - stateRecord.timestamp = Date.now(); - - parentStates[execId] = stateRecord; - return finalVal as TOutput; - } catch (err: unknown) { - stateRecord.status = NodeStatus.FAILED; - stateRecord.errorMessage = - err instanceof Error ? err.message : String(err); - stateRecord.timestamp = Date.now(); - parentStates[execId] = stateRecord; - if (stopOnError) { - throw err; - } - return undefined as unknown as TOutput; - } - }); - - return await Promise.all(tasks); -} diff --git a/core/src/workflow-next/request_input.ts b/core/src/workflow/request_input.ts similarity index 100% rename from core/src/workflow-next/request_input.ts rename to core/src/workflow/request_input.ts diff --git a/core/src/workflow/retry_config.ts b/core/src/workflow/retry_config.ts index 0b6f678e0..65ea24efc 100644 --- a/core/src/workflow/retry_config.ts +++ b/core/src/workflow/retry_config.ts @@ -5,63 +5,81 @@ */ /** - * Configuration options for node execution retries upon transient failures. + * An error constructor usable in {@link RetryConfig.exceptions}. + */ +export type ErrorClass = new (...args: never[]) => Error; + +/** + * Configuration for retrying a node. + * + * Ported from `google/adk-python` `workflow/_retry_config.py`. Delays are + * expressed in **seconds** (fractions allowed) to match the Python semantics + * and keep configuration portable across runtimes. Unset fields fall back to + * the documented defaults inside the retry utilities. */ export interface RetryConfig { /** - * Maximum number of execution attempts (including the initial attempt). - * Must be >= 1. + * Maximum number of attempts, including the original request. If 0 or 1, it + * means no retries. If not specified, defaults to 5. */ - maxAttempts: number; + maxAttempts?: number; /** - * Initial delay in milliseconds before the first retry. - * Default is 1000ms (1 second). + * Initial delay before the first retry, in seconds. If not specified, + * defaults to 1.0 second. */ - initialDelayMs?: number; + initialDelay?: number; /** - * Maximum delay in milliseconds between retries. - * Default is 30000ms (30 seconds). + * Maximum delay between retries, in seconds. If not specified, defaults to + * 60.0 seconds. */ - maxDelayMs?: number; + maxDelay?: number; /** - * Multiplier applied to the delay after each retry attempt (exponential backoff). - * Default is 2.0. + * Multiplier by which the delay increases after each attempt. If not + * specified, defaults to 2.0. */ backoffFactor?: number; /** - * Optional array of Error constructors or error message patterns that should trigger a retry. - * If not specified, all errors are considered retryable up to `maxAttempts`. + * Randomness factor for the delay. If not specified, defaults to 1.0. Use 0.0 + * to remove randomness. */ - retryableErrors?: Array Error | string | RegExp>; + jitter?: number; + + /** + * Exceptions to retry on. Accepts error class names as strings (e.g. + * `['TypeError']`) or error classes directly (e.g. `[TypeError]`). + * `undefined`/`null` means retry on all errors. + */ + exceptions?: Array | null; } /** - * Validates and normalizes a RetryConfig into canonical defaults. - * @param config Optional user-provided RetryConfig. - * @returns Normalized RetryConfig or undefined if not provided. + * Normalizes the `exceptions` field of a {@link RetryConfig} to a list of error + * class name strings, mirroring Python's `field_validator`. + * + * @returns The list of class-name strings, or `undefined` to mean "retry on all + * errors". */ -export function normalizeRetryConfig( - config?: RetryConfig, -): Required | undefined { - if (!config) { +export function normalizeRetryExceptions( + exceptions?: Array | null, +): string[] | undefined { + if (exceptions === undefined || exceptions === null) { return undefined; } - - if (config.maxAttempts < 1) { + return exceptions.map((item) => { + if (typeof item === 'string') { + return item; + } + if (typeof item === 'function' && item.name) { + return item.name; + } throw new Error( - `RetryConfig.maxAttempts must be at least 1, received: ${config.maxAttempts}`, + `exceptions must contain error class names (string) or error classes, got: ${String( + item, + )}`, ); - } - - return { - maxAttempts: config.maxAttempts, - initialDelayMs: config.initialDelayMs ?? 1000, - maxDelayMs: config.maxDelayMs ?? 30000, - backoffFactor: config.backoffFactor ?? 2.0, - retryableErrors: config.retryableErrors ?? [], - }; + }); } diff --git a/core/src/workflow/run_node.ts b/core/src/workflow/run_node.ts deleted file mode 100644 index 2d8e7a046..000000000 --- a/core/src/workflow/run_node.ts +++ /dev/null @@ -1,96 +0,0 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import {InvocationContext} from '../agents/invocation_context.js'; -import {BaseNode} from './base_node.js'; -import {consumeGenerator, getOrInitAgentStates} from './node_runner.js'; -import {NodeState, NodeStatus, isNodeState} from './node_state.js'; -import {FunctionNode, FunctionNodeHandler} from './nodes/function_node.js'; - -/** - * Options for `runNode`. - */ -export interface RunNodeOptions { - /** - * Custom execution ID override for this node execution. - * If not provided, a deterministic execution ID based on branch and node name is used. - */ - customExecutionId?: string; -} - -/** - * Programmatically runs a target node (or function) inside the invocation context, - * checking and persisting state checkpoints (`InvocationContext.agentStates[execId]`). - * Essential for dynamic workflows where execution flow is controlled by TS code. - * - * @param ctx The current invocation context. - * @param nodeOrFunc A BaseNode instance or function handler to execute. - * @param input Optional input payload. - * @param options Optional settings (customExecutionId). - * @returns Promise resolving to the final output payload of the node. - */ -export async function runNode( - ctx: InvocationContext, - nodeOrFunc: BaseNode | FunctionNodeHandler, - input?: TInput, - options?: RunNodeOptions, -): Promise { - const node = - typeof nodeOrFunc === 'function' - ? new FunctionNode('run_node_dynamic', nodeOrFunc) - : nodeOrFunc; - - const agentStates = getOrInitAgentStates(ctx); - const execId = - options?.customExecutionId ?? - `exec_node_${ctx.branch ? ctx.branch + '.' : ''}${node.name}`; - - const existingState = agentStates[execId] as NodeState | undefined; - if ( - existingState && - isNodeState(existingState) && - existingState.status === NodeStatus.COMPLETED && - !node.rerunOnResume - ) { - return existingState.outputPayload as TOutput; - } - - const stateRecord: NodeState = { - executionId: execId, - nodeName: node.name, - status: NodeStatus.RUNNING, - inputPayload: input, - timestamp: Date.now(), - }; - agentStates[execId] = stateRecord; - - try { - const generator = node.run(ctx, input); - const {output, isPausedHitl} = await consumeGenerator(generator); - - if (isPausedHitl || ctx.endInvocation || ctx.abortSignal?.aborted) { - stateRecord.status = NodeStatus.PAUSED_HITL; - stateRecord.timestamp = Date.now(); - ctx.endInvocation = true; - throw new Error(`Node "${node.name}" requested HITL pause.`); - } - - const finalVal = - output !== undefined - ? output - : (node.lastOutputPayload ?? stateRecord.lastOutputPayload ?? input); - stateRecord.status = NodeStatus.COMPLETED; - stateRecord.outputPayload = finalVal; - stateRecord.timestamp = Date.now(); - return finalVal as TOutput; - } catch (error: unknown) { - stateRecord.status = NodeStatus.FAILED; - stateRecord.errorMessage = - error instanceof Error ? error.message : String(error); - stateRecord.timestamp = Date.now(); - throw error; - } -} diff --git a/core/src/workflow-next/schedule_dynamic_node.ts b/core/src/workflow/schedule_dynamic_node.ts similarity index 100% rename from core/src/workflow-next/schedule_dynamic_node.ts rename to core/src/workflow/schedule_dynamic_node.ts diff --git a/core/src/workflow/trigger.ts b/core/src/workflow/trigger.ts index f0da3eb9a..9ce2f307d 100644 --- a/core/src/workflow/trigger.ts +++ b/core/src/workflow/trigger.ts @@ -4,116 +4,23 @@ * SPDX-License-Identifier: Apache-2.0 */ -import {InvocationContext} from '../agents/invocation_context.js'; -import {Event} from '../events/event.js'; - -/** - * The default route key used as fallback when no conditional routes match. - */ -export const DEFAULT_ROUTE = '__DEFAULT__'; - /** - * A predicate function evaluated against the event or output yielded by an upstream node. - * @param ctx The current invocation context. - * @param eventOrOutput The event or output payload produced by the source node. - * @returns True if the transition condition is satisfied, false otherwise. + * A buffered trigger for a downstream node. + * + * Ported from `google/adk-python` `workflow/_trigger.py`. Unlike the previous + * TypeScript `Trigger` (a route-matching predicate), this is a plain data record + * describing *how* a target node should be invoked when its turn comes. */ -export type TriggerPredicate = ( - ctx: InvocationContext, - eventOrOutput: Event | unknown, -) => boolean | Promise; - -/** - * Represents a conditional trigger attached to a workflow graph edge. - * Determines whether a specific transition route should be taken upon upstream node completion. - */ -export class Trigger { - private readonly routeKey?: string; - private readonly predicate?: TriggerPredicate; - - private constructor(options: { - routeKey?: string; - predicate?: TriggerPredicate; - }) { - this.routeKey = options.routeKey; - this.predicate = options.predicate; - } - - /** - * Whether this trigger represents the fallback default route (`DEFAULT_ROUTE`). - */ - isDefaultRoute(): boolean { - return this.routeKey === DEFAULT_ROUTE; - } - - /** - * Creates a route matching trigger. - * The trigger evaluates to true if the emitted `Event.actions.route` (or event payload route) matches `routeKey`. - * @param routeKey The exact string route key to match. - */ - static fromRoute(routeKey: string | symbol): Trigger { - if (!routeKey && typeof routeKey !== 'symbol') { - throw new Error( - 'Trigger.fromRoute requires a non-empty string or symbol routeKey.', - ); - } - const keyStr = - typeof routeKey === 'symbol' - ? routeKey === Symbol.for('DEFAULT_ROUTE') - ? DEFAULT_ROUTE - : routeKey.description || String(routeKey) - : routeKey; - return new Trigger({routeKey: keyStr}); - } - - /** - * Creates a predicate-based trigger. - * The trigger evaluates to true if the provided predicate function returns true. - * @param predicate A boolean function evaluated against the context and node output/event. - */ - static fromPredicate(predicate: TriggerPredicate): Trigger { - if (typeof predicate !== 'function') { - throw new Error('Trigger.fromPredicate requires a function predicate.'); - } - return new Trigger({predicate}); - } +export interface Trigger { + /** The input to pass to the triggered node. */ + input?: unknown; - /** - * Evaluates whether this trigger is satisfied by the given event or output payload. - * @param ctx The current invocation context. - * @param eventOrOutput The event or output payload produced by the source node. - * @returns Promise resolving to true if the transition should occur, false otherwise. - */ - async evaluate( - ctx: InvocationContext, - eventOrOutput: Event | unknown, - ): Promise { - if (this.routeKey) { - if (eventOrOutput && typeof eventOrOutput === 'object') { - const obj = eventOrOutput as Record; - if ( - 'actions' in obj && - obj.actions && - typeof obj.actions === 'object' && - 'route' in (obj.actions as Record) && - (obj.actions as Record).route === this.routeKey - ) { - return true; - } - if ('route' in obj && obj.route === this.routeKey) { - return true; - } - } - if (eventOrOutput === this.routeKey) { - return true; - } - return false; - } + /** Whether this trigger should run the node in an isolated sub-branch. */ + useSubBranch?: boolean; - if (this.predicate) { - return await this.predicate(ctx, eventOrOutput); - } + /** The branch inherited from the predecessor node. */ + branch?: string; - return true; - } + /** Scope tag explicitly propagated to this trigger. */ + isolationScope?: string; } diff --git a/core/src/workflow-next/utils/event_channel.ts b/core/src/workflow/utils/event_channel.ts similarity index 100% rename from core/src/workflow-next/utils/event_channel.ts rename to core/src/workflow/utils/event_channel.ts diff --git a/core/src/workflow/utils/graph_parser.ts b/core/src/workflow/utils/graph_parser.ts index 6835f0ba1..cae1ebaeb 100644 --- a/core/src/workflow/utils/graph_parser.ts +++ b/core/src/workflow/utils/graph_parser.ts @@ -4,226 +4,209 @@ * SPDX-License-Identifier: Apache-2.0 */ -import {BaseNode} from '../base_node.js'; -import {Trigger} from '../trigger.js'; - -/** - * A single element inside a GraphEdge array. - */ -export type EdgeElement = - | string // e.g., "START" - | BaseNode // concrete node instance - | Record // route map: { ROUTE_A: nodeA, ROUTE_B: nodeB } - | [Trigger, BaseNode] // conditional tuple: [trigger, targetNode] - | BaseNode[] // fan-out node array: [nodeA, nodeB] - | readonly BaseNode[]; // fan-out node array - -/** - * A workflow graph edge definition. - * Examples: - * ["START", nodeA, nodeB, nodeC] - * ["START", [nodeA, nodeB], joinNode] - * [routerNode, { ROUTE_X: nodeC, ROUTE_Y: nodeD }] - * [nodeA, [Trigger.fromPredicate(...), nodeB]] - */ -export type GraphEdge = EdgeElement[]; - /** - * Internal representation of a directed edge between two nodes. + * Parses workflow edge items and chains into a flat list of {@link Edge}s. + * + * Ported from `google/adk-python` `workflow/utils/_graph_parser.py`. */ -export interface AdjacencyEdge { - readonly source: string; // source node name or "START" - readonly target: BaseNode; - readonly trigger?: Trigger; -} - -/** - * A parsed and structured workflow graph ready for execution or validation. - */ -export class ParsedGraph { - readonly nodes = new Map(); - readonly adjacencyList = new Map(); - readonly inboundCounts = new Map(); - constructor() { - this.adjacencyList.set('START', []); - this.inboundCounts.set('START', 0); - } - - /** - * Registers a node in the graph and initializes its adjacency and inbound counters if new. - * @param node The node to add. - */ - addNode(node: BaseNode): void { - if (!this.nodes.has(node.name)) { - this.nodes.set(node.name, node); - this.adjacencyList.set(node.name, []); - this.inboundCounts.set(node.name, 0); - } - } - - /** - * Adds a directed edge from `source` to `target` with an optional `trigger`. - * @param source Source node name (or "START"). - * @param target Target node instance. - * @param trigger Optional trigger condition. - */ - addEdge(source: string, target: BaseNode, trigger?: Trigger): void { - this.addNode(target); - if (source !== 'START' && !this.nodes.has(source)) { - throw new Error( - `Source node "${source}" must be added to the graph or referenced before defining an edge from it.`, - ); - } - - const edges = this.adjacencyList.get(source) || []; - edges.push({source, target, trigger}); - this.adjacencyList.set(source, edges); - - const currentInbound = this.inboundCounts.get(target.name) || 0; - this.inboundCounts.set(target.name, currentInbound + 1); - } +import {BaseNode} from '../base_node.js'; +import { + ChainElement, + Edge, + EdgeItem, + NodeLike, + RouteValue, + RoutingMap, +} from '../graph.js'; +import {buildNode, isNodeLike, isPlainObject} from './workflow_graph_utils.js'; + +function isRouteValue(value: unknown): value is RouteValue { + const t = typeof value; + return t === 'string' || t === 'number' || t === 'boolean'; } -/** - * Parses an array of user-defined GraphEdge structures into an internal ParsedGraph. - * @param edges Array of GraphEdge sequences or branch definitions. - * @returns A structured ParsedGraph. - */ -export function parseGraph(edges: GraphEdge[]): ParsedGraph { - if (!Array.isArray(edges) || edges.length === 0) { +/** Expands a routing map into individual (from, to, route) triples. */ +function expandRoutingMap( + fromElement: ChainElement, + routingMap: RoutingMap, +): Array<[ChainElement, NodeLike | readonly NodeLike[], RouteValue]> { + const keys = Object.keys(routingMap); + if (keys.length === 0) { throw new Error( - 'parseGraph requires a non-empty array of GraphEdge definitions.', + 'Routing map must not be empty. Provide at least one route -> node mapping.', ); } - const graph = new ParsedGraph(); - - for (const edgeSeq of edges) { - if (!Array.isArray(edgeSeq) || edgeSeq.length < 2) { - throw new Error( - 'Each GraphEdge definition must be an array with at least 2 elements (e.g., ["START", nodeA]).', - ); - } - - for (let i = 0; i < edgeSeq.length - 1; i++) { - const current = edgeSeq[i]; - const next = edgeSeq[i + 1]; - - const sources = flattenSource(current, graph, i); - for (const sourceName of sources) { - if (isBaseNode(next)) { - graph.addEdge(sourceName, next); - } else if (Array.isArray(next) && !isTriggerTuple(next)) { - for (const targetNode of next) { - if (!isBaseNode(targetNode)) { - throw new Error( - `All elements in target fan-out array from source "${sourceName}" must be BaseNode instances.`, - ); - } - graph.addEdge(sourceName, targetNode); - } - } else if (isRouteMap(next)) { - const entries: [string | symbol, BaseNode][] = [ - ...Object.entries(next), - ...Object.getOwnPropertySymbols(next).map( - (sym) => - [sym, (next as Record)[sym]] as [ - string | symbol, - BaseNode, - ], - ), - ]; - for (const [routeKey, targetNode] of entries) { - if (!isBaseNode(targetNode)) { - throw new Error( - `Target for route "${String(routeKey)}" from source "${sourceName}" must be a BaseNode instance.`, - ); - } - graph.addEdge(sourceName, targetNode, Trigger.fromRoute(routeKey)); - } - } else if (isTriggerTuple(next)) { - const [trigger, targetNode] = next; - graph.addEdge(sourceName, targetNode, trigger); - } else { + const expanded: Array< + [ChainElement, NodeLike | readonly NodeLike[], RouteValue] + > = []; + for (const routeKey of keys) { + // Object keys are strings; numeric route keys arrive as numeric strings. + const normalizedKey: RouteValue = /^-?\d+$/.test(routeKey) + ? Number(routeKey) + : routeKey; + const target = routingMap[routeKey]; + if (Array.isArray(target)) { + for (const node of target) { + if (!isNodeLike(node)) { throw new Error( - `Invalid target element at index ${i + 1} from source "${sourceName}". Must be a BaseNode, BaseNode array, route dictionary, or [Trigger, BaseNode] tuple.`, + `Invalid node in fan-out tuple for route ${String(routeKey)}.`, ); } } + } else if (!isNodeLike(target)) { + throw new Error( + `Invalid routing map value for route ${String(routeKey)}.`, + ); + } + if (!isRouteValue(normalizedKey)) { + throw new Error(`Invalid routing map key: ${String(routeKey)}.`); } + expanded.push([ + fromElement, + target as NodeLike | readonly NodeLike[], + normalizedKey, + ]); } + return expanded; +} - for (const [nodeName, node] of graph.nodes.entries()) { - if ( - node && - node.constructor && - node.constructor.name === 'JoinNode' && - (node as unknown as {upstreamCount: number}).upstreamCount === 0 - ) { - const count = graph.inboundCounts.get(nodeName) || 0; - if (count >= 1) { - (node as unknown as {upstreamCount: number}).upstreamCount = count; - } +/** Extracts all target nodes from a routing map, flattening fan-out arrays. */ +function nodesFromRoutingMap(routingMap: RoutingMap): NodeLike[] { + const nodes: NodeLike[] = []; + for (const target of Object.values(routingMap)) { + if (Array.isArray(target)) { + nodes.push(...(target as NodeLike[])); + } else { + nodes.push(target as NodeLike); } } + return nodes; +} - return graph; +/** Flattens a chain element into a list of individual nodes. */ +function flattenElement(element: ChainElement): NodeLike[] { + if (isPlainObject(element)) { + return nodesFromRoutingMap(element as RoutingMap); + } + if (Array.isArray(element)) { + return [...(element as readonly NodeLike[])]; + } + return [element as NodeLike]; } -function flattenSource( - element: EdgeElement, - graph: ParsedGraph, - index: number, -): string[] { - if (typeof element === 'string' && element === 'START') { - return ['START']; +/** Gets a node from the identity map or builds (and caches) it. */ +function getOrBuildNode( + nodeLike: NodeLike, + nodeMap: Map, +): BaseNode { + if (nodeLike === 'START') { + return buildNode('START'); } - if (isBaseNode(element)) { - graph.addNode(element); - return [element.name]; + if (typeof nodeLike === 'object' || typeof nodeLike === 'function') { + const cached = nodeMap.get(nodeLike as object); + if (cached) { + return cached; + } + const built = buildNode(nodeLike); + // Only cache when a distinct wrapper was produced (or always, to preserve + // identity across repeated references within the same parse). + nodeMap.set(nodeLike as object, built); + return built; } - if (Array.isArray(element) && !isTriggerTuple(element)) { - return element.map((node, idx) => { - if (!isBaseNode(node)) { - throw new Error( - `Invalid source element inside array at index ${index}[${idx}]. Must be a BaseNode instance.`, + return buildNode(nodeLike); +} + +function processExplicitEdge( + edge: Edge, + nodeMap: Map, + out: Edge[], +): void { + out.push( + new Edge( + getOrBuildNode(edge.fromNode, nodeMap), + getOrBuildNode(edge.toNode, nodeMap), + edge.route, + ), + ); +} + +function processRoutingMapEdge( + fromEl: ChainElement, + toEl: RoutingMap, + nodeMap: Map, + out: Edge[], +): void { + if (isPlainObject(fromEl)) { + throw new Error( + 'Consecutive routing maps are not allowed in a chain. Split them into separate edge items.', + ); + } + for (const [expFrom, expTo, route] of expandRoutingMap(fromEl, toEl)) { + for (const fromNode of flattenElement(expFrom)) { + for (const toNode of flattenElement(expTo as ChainElement)) { + out.push( + new Edge( + getOrBuildNode(fromNode, nodeMap), + getOrBuildNode(toNode, nodeMap), + route, + ), ); } - graph.addNode(node); - return node.name; - }); + } } - throw new Error( - `Invalid source element at index ${index} in edge sequence. Must be "START", a BaseNode, or an array of BaseNode instances.`, - ); } -function isBaseNode(obj: unknown): obj is BaseNode { - return ( - typeof obj === 'object' && - obj !== null && - 'name' in obj && - typeof (obj as BaseNode).name === 'string' && - 'run' in obj && - typeof (obj as BaseNode).run === 'function' - ); +function processUnconditionalEdge( + fromEl: ChainElement, + toEl: ChainElement, + nodeMap: Map, + out: Edge[], +): void { + for (const fromNode of flattenElement(fromEl)) { + for (const toNode of flattenElement(toEl)) { + out.push( + new Edge( + getOrBuildNode(fromNode, nodeMap), + getOrBuildNode(toNode, nodeMap), + null, + ), + ); + } + } } -function isRouteMap(obj: unknown): obj is Record { - return ( - typeof obj === 'object' && - obj !== null && - !isBaseNode(obj) && - !Array.isArray(obj) - ); +function processChain( + chain: ChainElement[], + nodeMap: Map, + out: Edge[], +): void { + for (let i = 0; i < chain.length - 1; i++) { + const fromEl = chain[i]; + const toEl = chain[i + 1]; + if (isPlainObject(toEl)) { + processRoutingMapEdge(fromEl, toEl as RoutingMap, nodeMap, out); + } else { + processUnconditionalEdge(fromEl, toEl, nodeMap, out); + } + } } -function isTriggerTuple(obj: unknown): obj is [Trigger, BaseNode] { - return ( - Array.isArray(obj) && - obj.length === 2 && - obj[0] instanceof Trigger && - isBaseNode(obj[1]) - ); +/** Parses a list of edge items into a flat list of {@link Edge} objects. */ +export function parseEdgeItems(edgeItems: EdgeItem[]): Edge[] { + const nodeMap = new Map(); + const out: Edge[] = []; + + for (const item of edgeItems) { + if (item instanceof Edge) { + processExplicitEdge(item, nodeMap, out); + } else if (Array.isArray(item)) { + processChain(item, nodeMap, out); + } else { + throw new Error(`Invalid edge item type: ${typeof item}`); + } + } + + return out; } diff --git a/core/src/workflow/utils/graph_validation.ts b/core/src/workflow/utils/graph_validation.ts index e8f75ef17..29d22eae0 100644 --- a/core/src/workflow/utils/graph_validation.ts +++ b/core/src/workflow/utils/graph_validation.ts @@ -4,89 +4,197 @@ * SPDX-License-Identifier: Apache-2.0 */ -import {ParsedGraph} from './graph_parser.js'; - /** - * Performs structural validation on a ParsedGraph before execution begins. - * Verifies reachability, checks for unintended cycles in DAG mode, and validates JoinNode upstream counts. + * Validates workflow graphs and computes terminal nodes. * - * @param graph The ParsedGraph to validate. - * @param options Optional validation settings (e.g. `allowCycles`). - * @throws Error if the graph structure is invalid or malformed. + * Ported from `google/adk-python` `workflow/utils/_graph_validation.py`. + * The Phase 3 static-schema check and the Phase 7 chat-agent wiring check are + * intentionally deferred to their respective phases. */ -export function validateGraph( - graph: ParsedGraph, - options?: {allowCycles?: boolean}, -): void { - // 1. Reachability check from "START" - const visited = new Set(['START']); - const queue: string[] = ['START']; - - while (queue.length > 0) { - const current = queue.shift()!; - const edges = graph.adjacencyList.get(current) || []; - for (const edge of edges) { - if (!visited.has(edge.target.name)) { - visited.add(edge.target.name); - queue.push(edge.target.name); - } - } + +import {BaseNode, START} from '../base_node.js'; +import {DEFAULT_ROUTE, Edge} from '../graph.js'; + +function validateDuplicateNodeNames(nodes: BaseNode[]): Set { + const counts = new Map(); + for (const node of nodes) { + counts.set(node.name, (counts.get(node.name) ?? 0) + 1); } + const duplicates = [...counts.entries()] + .filter(([, c]) => c > 1) + .map(([name]) => name) + .sort(); + if (duplicates.length > 0) { + throw new Error( + `Graph validation failed. Duplicate node names found: ${JSON.stringify( + duplicates, + )}. Pass the exact same object instance to reuse a node, or give distinct nodes unique names.`, + ); + } + return new Set(counts.keys()); +} + +function validateStartNode(nodeNames: Set): void { + if (!nodeNames.has(START.name)) { + throw new Error( + `Graph validation failed. START node (name: '${START.name}') not found in graph nodes.`, + ); + } +} - for (const [nodeName] of graph.nodes) { - if (!visited.has(nodeName)) { +function validateStartEdges(edges: Edge[]): void { + for (const edge of edges) { + if (edge.fromNode.name === START.name && edge.route !== null) { throw new Error( - `Graph validation failed: Node "${nodeName}" is unreachable from "START". Check your edge definitions.`, + `Graph validation failed. Edges from START must not have routes (edge to ${edge.toNode.name} has route ${String( + edge.route, + )}).`, ); } } +} + +function validateConnectivity(edges: Edge[], nodeNames: Set): void { + const adj = new Map>(); + for (const name of nodeNames) { + adj.set(name, new Set()); + } + const toNodes = new Set(); + for (const edge of edges) { + adj.get(edge.fromNode.name)!.add(edge.toNode.name); + toNodes.add(edge.toNode.name); + } - // 2. Cycle detection (DFS via recursion stack) if !allowCycles - if (!options?.allowCycles) { - const recursionStack = new Set(); - const dfsVisited = new Set(); - - const checkCycles = (nodeName: string): void => { - dfsVisited.add(nodeName); - recursionStack.add(nodeName); - - const edges = graph.adjacencyList.get(nodeName) || []; - for (const edge of edges) { - const targetName = edge.target.name; - if (!dfsVisited.has(targetName)) { - checkCycles(targetName); - } else if (recursionStack.has(targetName)) { - throw new Error( - `Graph validation failed: Cycle detected involving node "${targetName}". If your workflow intentionally contains loops, enable cycle support or use dynamic routing.`, - ); - } + const reachable = new Set(); + const stack = [START.name]; + while (stack.length > 0) { + const node = stack.pop()!; + if (reachable.has(node)) { + continue; + } + reachable.add(node); + for (const next of adj.get(node) ?? []) { + if (!reachable.has(next)) { + stack.push(next); } + } + } - recursionStack.delete(nodeName); - }; + const unreachable = [...nodeNames].filter((n) => !reachable.has(n)).sort(); + if (unreachable.length > 0) { + throw new Error( + `Graph validation failed. The following nodes are unreachable from START: ${JSON.stringify( + unreachable, + )}`, + ); + } + if (toNodes.has(START.name)) { + throw new Error( + 'Graph validation failed. START node must not have incoming edges.', + ); + } +} - checkCycles('START'); +function validateDuplicateEdges(edges: Edge[]): void { + const seen = new Set(); + for (const edge of edges) { + const key = `${edge.fromNode.name}\u0000${edge.toNode.name}`; + if (seen.has(key)) { + throw new Error( + `Graph validation failed. Duplicate edge found: from=${edge.fromNode.name}, to=${edge.toNode.name}`, + ); + } + seen.add(key); } +} - // 3. JoinNode upstream predecessor validation - for (const [nodeName, node] of graph.nodes) { - const nodeObj = node as unknown as Record; - if ( - 'upstreamCount' in nodeObj && - typeof nodeObj.upstreamCount === 'number' - ) { - const upstreamCount = nodeObj.upstreamCount as number; - const actualInbound = graph.inboundCounts.get(nodeName) || 0; - if (upstreamCount < 1) { +function validateDefaultRoutes(edges: Edge[]): void { + const defaultRouteEdges = new Map(); + for (const edge of edges) { + if (Array.isArray(edge.route) && edge.route.includes(DEFAULT_ROUTE)) { + throw new Error( + `Graph validation failed. DEFAULT_ROUTE cannot be combined with other routes in a list (edge from=${edge.fromNode.name}, to=${edge.toNode.name}). Use a separate edge for DEFAULT_ROUTE.`, + ); + } + if (edge.route === DEFAULT_ROUTE) { + const from = edge.fromNode.name; + if (defaultRouteEdges.has(from)) { throw new Error( - `JoinNode "${nodeName}" has invalid upstreamCount: ${upstreamCount}. Must be >= 1.`, + `Graph validation failed. Multiple DEFAULT_ROUTE edges found from node ${from} to ${defaultRouteEdges.get( + from, + )} and ${edge.toNode.name}`, ); } - if (upstreamCount > actualInbound) { + defaultRouteEdges.set(from, edge.toNode.name); + } + } +} + +function detectUnconditionalCycles( + edges: Edge[], + nodeNames: Set, +): void { + const adj = new Map(); + for (const name of nodeNames) { + adj.set(name, []); + } + for (const edge of edges) { + if (edge.route === null) { + adj.get(edge.fromNode.name)!.push(edge.toNode.name); + } + } + + const inStack = new Set(); + const done = new Set(); + + const dfs = (node: string, path: string[]): void => { + inStack.add(node); + path.push(node); + for (const neighbor of adj.get(node) ?? []) { + if (inStack.has(neighbor)) { + const cycleStart = path.indexOf(neighbor); + const cycle = [...path.slice(cycleStart), neighbor]; throw new Error( - `JoinNode "${nodeName}" expects ${upstreamCount} upstream predecessors, but only has ${actualInbound} inbound edges defined in the graph.`, + `Graph validation failed. Unconditional cycle detected: ${cycle.join( + ' -> ', + )}. Cycles must include at least one conditional (routed) edge to avoid infinite loops.`, ); } + if (!done.has(neighbor)) { + dfs(neighbor, path); + } + } + path.pop(); + inStack.delete(node); + done.add(node); + }; + + for (const name of nodeNames) { + if (!done.has(name)) { + dfs(name, []); } } } + +function computeTerminalNodes(nodes: BaseNode[], edges: Edge[]): Set { + const fromNames = new Set(edges.map((e) => e.fromNode.name)); + return new Set( + nodes + .filter((n) => n.name !== START.name && !fromNames.has(n.name)) + .map((n) => n.name), + ); +} + +/** + * Validates the workflow graph and returns the set of terminal node names. + */ +export function validateGraph(nodes: BaseNode[], edges: Edge[]): Set { + const nodeNames = validateDuplicateNodeNames(nodes); + validateStartNode(nodeNames); + validateStartEdges(edges); + validateConnectivity(edges, nodeNames); + validateDuplicateEdges(edges); + validateDefaultRoutes(edges); + detectUnconditionalCycles(edges, nodeNames); + return computeTerminalNodes(nodes, edges); +} diff --git a/core/src/workflow/utils/hitl_utils.ts b/core/src/workflow/utils/hitl_utils.ts index a19515eec..4f5686077 100644 --- a/core/src/workflow/utils/hitl_utils.ts +++ b/core/src/workflow/utils/hitl_utils.ts @@ -4,86 +4,96 @@ * SPDX-License-Identifier: Apache-2.0 */ -import {InvocationContext} from '../../agents/invocation_context.js'; -import {createEvent, Event} from '../../events/event.js'; -import {isNodeState, NodeStatus} from '../node_state.js'; - -import {getOrInitAgentStates} from '../node_runner.js'; - /** - * Options when creating a HITL input request. + * Utilities for Human-in-the-Loop (HITL) workflows. + * + * Ported (subset) from `google/adk-python` + * `workflow/utils/_workflow_hitl_utils.py`. The auth-credential helpers are + * added in the Phase 5 auth-gate follow-up. */ -export interface RequestInputOptions { - /** - * Optional custom prompt or question to display to the user. - */ - prompt?: string; - /** - * Optional structured schema or options describing what input is required. - */ - schema?: Record; -} +import {Part} from '@google/genai'; +import {z} from 'zod'; +import {createEvent, Event} from '../../events/event.js'; +import {RequestInput} from '../request_input.js'; + +/** Function-call name marking a request-for-input interrupt. */ +export const REQUEST_INPUT_FUNCTION_CALL_NAME = 'adk_request_input'; + +/** Function-call name marking a request-for-credential interrupt. */ +export const REQUEST_CREDENTIAL_FUNCTION_CALL_NAME = 'adk_request_credential'; /** - * Creates an Event that signals a Human-in-the-Loop (`RequestInput`) pause condition to the workflow engine. - * - * @param ctx The current invocation context. - * @param nodeName Name of the node requesting input. - * @param options Optional prompt and schema describing the required input. + * Creates an interrupt {@link Event} from a {@link RequestInput}. The event + * carries an `adk_request_input` function call and marks the interrupt id as a + * long-running tool id. */ -export function createRequestInputEvent( - ctx: InvocationContext, - nodeName: string, - options?: RequestInputOptions, -): Event { +export function createRequestInputEvent(requestInput: RequestInput): Event { + const args: Record = { + interruptId: requestInput.interruptId, + payload: requestInput.payload ?? null, + message: requestInput.message ?? null, + responseSchema: requestInput.responseSchema + ? z.toJSONSchema(requestInput.responseSchema) + : null, + }; + return createEvent({ - invocationId: ctx.invocationId, - author: nodeName, - branch: ctx.branch, - content: options?.prompt - ? {role: 'model', parts: [{text: options.prompt}]} - : undefined, - actions: { - requestInput: { - nodeName, - prompt: options?.prompt, - schema: options?.schema, - }, + content: { + role: 'model', + parts: [ + { + functionCall: { + name: REQUEST_INPUT_FUNCTION_CALL_NAME, + args, + id: requestInput.interruptId, + }, + }, + ], }, + longRunningToolIds: [requestInput.interruptId], }); } -/** - * Locates any node inside `InvocationContext.agentStates` whose status is `PAUSED_HITL`, - * and injects the resumption input payload so that subsequent workflow execution can proceed from that node. - * - * @param ctx The invocation context being resumed. - * @param resumptionInput The user's input payload provided upon resumption. - * @returns The name and execution ID of the resumed node, or undefined if no paused node was found. - */ -export function injectHitlResumptionInput( - ctx: InvocationContext, - resumptionInput: unknown, -): {nodeName: string; executionId: string} | undefined { - const agentStates = getOrInitAgentStates(ctx); +/** Returns whether an event contains a `request_input` function call. */ +export function hasRequestInputFunctionCall(event: Event): boolean { + return (event.content?.parts ?? []).some( + (p) => p.functionCall?.name === REQUEST_INPUT_FUNCTION_CALL_NAME, + ); +} - for (const [execId, state] of Object.entries(agentStates)) { - if ( - isNodeState(state) && - state.status === NodeStatus.COMPLETED && - state.wasPausedHitl - ) { - continue; - } - if (isNodeState(state) && state.status === NodeStatus.PAUSED_HITL) { - state.status = NodeStatus.RUNNING; - state.inputPayload = resumptionInput; - state.wasPausedHitl = true; - state.timestamp = Date.now(); - return {nodeName: state.nodeName, executionId: execId}; +/** Returns whether an event contains an `adk_request_credential` function call. */ +export function hasAuthRequestFunctionCall(event: Event): boolean { + return (event.content?.parts ?? []).some( + (p) => p.functionCall?.name === REQUEST_CREDENTIAL_FUNCTION_CALL_NAME, + ); +} + +/** Extracts interrupt ids from `request_input` function calls in an event. */ +export function getRequestInputInterruptIds(event: Event): string[] { + const ids: string[] = []; + for (const part of event.content?.parts ?? []) { + const fc = part.functionCall; + if (fc && fc.name === REQUEST_INPUT_FUNCTION_CALL_NAME && fc.id) { + ids.push(fc.id); } } + return ids; +} - return undefined; +/** + * Creates a `FunctionResponse` part answering a `request_input` interrupt, + * suitable for appending to a session as the user's resume response. + */ +export function createRequestInputResponse( + interruptId: string, + response: Record, +): Part { + return { + functionResponse: { + id: interruptId, + name: REQUEST_INPUT_FUNCTION_CALL_NAME, + response, + }, + }; } diff --git a/core/src/workflow/utils/rehydration_utils.ts b/core/src/workflow/utils/rehydration_utils.ts deleted file mode 100644 index 0538f7c73..000000000 --- a/core/src/workflow/utils/rehydration_utils.ts +++ /dev/null @@ -1,114 +0,0 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import {InvocationContext} from '../../agents/invocation_context.js'; -import {Session} from '../../sessions/session.js'; -import {getOrInitAgentStates} from '../node_runner.js'; -import {NodeStatus, isNodeState} from '../node_state.js'; - -/** - * Scans historical session events and state metadata to reconstruct `InvocationContext.agentStates` - * and `InvocationContext.endOfAgents` when resuming a durable session from storage (such as database or GCS). - * - * @param session The session loaded from storage containing historical events and state. - * @param ctx The invocation context to populate with rehydrated checkpoints. - */ -export function rehydrateAgentStates( - session: Session, - ctx: InvocationContext, -): void { - const agentStates = getOrInitAgentStates(ctx); - - if (session.state && typeof session.state === 'object') { - const sessionState = session.state as Record; - if ( - 'agentStates' in sessionState && - sessionState.agentStates && - typeof sessionState.agentStates === 'object' - ) { - for (const [key, val] of Object.entries( - sessionState.agentStates as Record, - )) { - if (!agentStates[key] && isNodeState(val)) { - agentStates[key] = val; - } - } - } - if ( - 'endOfAgents' in sessionState && - sessionState.endOfAgents && - typeof sessionState.endOfAgents === 'object' - ) { - for (const [key, val] of Object.entries( - sessionState.endOfAgents as Record, - )) { - if (typeof val === 'boolean') { - ctx.endOfAgents[key] = val; - } - } - } - } - - if (Array.isArray(session.events)) { - for (const event of session.events) { - if (!event || typeof event !== 'object') continue; - - const eventRecord = event as unknown as Record; - const actions = eventRecord.actions as - | Record - | undefined; - if (actions && typeof actions === 'object') { - if ( - 'nodeExecution' in actions && - actions.nodeExecution && - typeof actions.nodeExecution === 'object' - ) { - const {executionId, nodeName, status, outputPayload} = - actions.nodeExecution as Record; - if ( - executionId && - typeof executionId === 'string' && - !agentStates[executionId] - ) { - agentStates[executionId] = { - executionId, - nodeName: - typeof nodeName === 'string' ? nodeName : 'unknown_node', - status: - status === 'PAUSED_HITL' - ? NodeStatus.PAUSED_HITL - : NodeStatus.COMPLETED, - outputPayload, - timestamp: - typeof eventRecord.timestamp === 'number' - ? eventRecord.timestamp - : Date.now(), - }; - } - } - } - } - } -} - -/** - * Persists current `InvocationContext.agentStates` and `InvocationContext.endOfAgents` snapshots - * onto the session's state dictionary so they can be securely serialized by session services. - * - * @param ctx The invocation context whose states should be saved. - * @param session The session object to update. - */ -export function persistAgentStatesToSession( - ctx: InvocationContext, - session: Session, -): void { - if (!session.state || typeof session.state !== 'object') { - session.state = {}; - } - const sessionState = session.state as Record; - sessionState.agentStates = {...getOrInitAgentStates(ctx)}; - sessionState.endOfAgents = {...ctx.endOfAgents}; -} diff --git a/core/src/workflow/utils/replay_manager.ts b/core/src/workflow/utils/replay_manager.ts deleted file mode 100644 index 3ec210fc8..000000000 --- a/core/src/workflow/utils/replay_manager.ts +++ /dev/null @@ -1,90 +0,0 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import {InvocationContext} from '../../agents/invocation_context.js'; -import {createEvent, Event} from '../../events/event.js'; -import {BaseNode} from '../base_node.js'; -import {generateExecutionId, getOrInitAgentStates} from '../node_runner.js'; -import {isNodeState, NodeState, NodeStatus} from '../node_state.js'; - -/** - * Manages event replay when a workflow node is bypassed due to a completed historical checkpoint (`rerunOnResume == false`). - * Yields historical events associated with that node execution so client UIs and event subscribers can reconstruct the full trajectory. - */ -export class ReplayManager { - /** - * Checks if the given node has a historical COMPLETED checkpoint in `InvocationContext.agentStates` - * and yields its historical events (or a synthetic replay event) if `!node.rerunOnResume`. - * - * @param ctx The current invocation context. - * @param node The node being checked for replay. - * @yields Historical or synthetic replay events if checkpoint exists. - * @returns True if the node was successfully replayed (and execution should be skipped), false otherwise. - */ - static async *replayIfCompleted( - ctx: InvocationContext, - node: BaseNode, - ): AsyncGenerator { - if (node.rerunOnResume) { - return false; - } - - const agentStates = getOrInitAgentStates(ctx); - const execId = generateExecutionId(ctx, node.name); - const existingState = agentStates[execId] as NodeState | undefined; - - if ( - existingState && - isNodeState(existingState) && - existingState.status === NodeStatus.COMPLETED - ) { - if (Array.isArray(existingState.cachedEvents)) { - for (const event of existingState.cachedEvents) { - yield event; - } - } else if (existingState.outputPayload !== undefined) { - yield createEvent({ - invocationId: ctx.invocationId, - author: node.name, - branch: ctx.branch, - actions: { - nodeExecutionReplay: { - executionId: execId, - nodeName: node.name, - outputPayload: existingState.outputPayload, - timestamp: existingState.timestamp, - }, - }, - }); - } - return true; - } - - return false; - } - - /** - * Caches emitted events onto a node's state record during live execution so they can be replayed on subsequent resumptions. - * - * @param ctx The invocation context. - * @param execId The execution ID of the running node. - * @param event The emitted event to cache. - */ - static cacheEventForReplay( - ctx: InvocationContext, - execId: string, - event: Event, - ): void { - const agentStates = getOrInitAgentStates(ctx); - const state = agentStates[execId] as NodeState | undefined; - if (state && isNodeState(state)) { - if (!Array.isArray(state.cachedEvents)) { - state.cachedEvents = []; - } - state.cachedEvents.push(event); - } - } -} diff --git a/core/src/workflow/utils/retry_utils.ts b/core/src/workflow/utils/retry_utils.ts index 104cebd48..57680e801 100644 --- a/core/src/workflow/utils/retry_utils.ts +++ b/core/src/workflow/utils/retry_utils.ts @@ -4,111 +4,109 @@ * SPDX-License-Identifier: Apache-2.0 */ -import {Event} from '../../events/event.js'; -import {RetryConfig, normalizeRetryConfig} from '../retry_config.js'; +/** + * Utility functions for retrying nodes in a workflow. + * + * Ported from `google/adk-python` `workflow/utils/_retry_utils.py`. + */ + +import {NodeState} from '../node_state.js'; +import {RetryConfig, normalizeRetryExceptions} from '../retry_config.js'; + +const DEFAULT_MAX_ATTEMPTS = 5; +const DEFAULT_INITIAL_DELAY_SECONDS = 1.0; +const DEFAULT_MAX_DELAY_SECONDS = 60.0; +const DEFAULT_BACKOFF_FACTOR = 2.0; +const DEFAULT_JITTER = 1.0; /** - * Sleeps for a specified number of milliseconds unless the abort signal fires. - * @param ms Delay in milliseconds. - * @param abortSignal Optional AbortSignal to cancel sleeping early. + * Resolves the runtime name of a thrown value for exception-name matching. + * Mirrors Python's `type(exception).__name__`. */ -async function sleepWithSignal( - ms: number, - abortSignal?: AbortSignal, -): Promise { - if (abortSignal?.aborted) { - throw new Error('Aborted before retry delay.'); +function errorName(error: unknown): string { + if (error instanceof Error) { + // `name` is set by well-behaved Error subclasses; fall back to the + // constructor name for plain `throw new Error()` cases. + return error.name || error.constructor.name; } - return new Promise((resolve, reject) => { - const timer = setTimeout(() => { - abortSignal?.removeEventListener('abort', onAbort); - resolve(); - }, ms); - const onAbort = () => { - clearTimeout(timer); - reject(new Error('Aborted during retry delay.')); - }; - abortSignal?.addEventListener('abort', onAbort); - }); + if (typeof error === 'object' && error !== null) { + return error.constructor.name; + } + return typeof error; } /** - * Checks whether an error is retryable according to the RetryConfig. - * @param error The error thrown during execution. - * @param retryableErrors Array of Error constructors, strings, or RegExps. + * Checks if a failed node should be retried based on its retry config. + * + * @param error The error thrown by the node. + * @param retryConfig The node's retry configuration, if any. + * @param nodeState The current node state (its `attemptCount` is 1-based). */ -function isErrorRetryable( +export function shouldRetryNode( error: unknown, - retryableErrors: Required['retryableErrors'], + retryConfig: RetryConfig | undefined, + nodeState: NodeState, ): boolean { - if (!retryableErrors || retryableErrors.length === 0) { - return true; + if (!retryConfig) { + return false; } - const errObj = - typeof error === 'object' && error !== null - ? (error as Record) - : undefined; - const errMsg = - errObj && typeof errObj.message === 'string' ? errObj.message : undefined; + const attemptCount = nodeState.attemptCount; + const maxAttempts = retryConfig.maxAttempts ?? DEFAULT_MAX_ATTEMPTS; - return retryableErrors.some((matcher) => { - if (typeof matcher === 'string') { - return errMsg && errMsg.includes(matcher); - } - if (matcher instanceof RegExp) { - return errMsg && matcher.test(errMsg); - } - if (typeof matcher === 'function' && error instanceof matcher) { - return true; - } + // attemptCount starts at 1 for the original request; once it reaches + // maxAttempts, the limit is exhausted. + if (attemptCount >= maxAttempts) { return false; - }); + } + + const exceptions = normalizeRetryExceptions(retryConfig.exceptions); + if (exceptions !== undefined) { + if (!exceptions.includes(errorName(error))) { + return false; + } + } + + return true; } /** - * Wraps an async generator with retry logic according to the provided RetryConfig. - * If the generator throws a retryable error mid-stream or during start, it will back off and retry from the beginning. + * Calculates the delay, in seconds, before retrying a node. * - * @param generatorFactory A factory function that creates a fresh AsyncGenerator for each attempt. - * @param retryConfig Optional RetryConfig or undefined (if undefined, runs once without retrying). - * @param abortSignal Optional AbortSignal to halt retries upon cancellation. + * @param retryConfig The node's retry configuration, if any. + * @param nodeState The current node state (its `attemptCount` is the 1-based + * attempt number that just failed). + * @param randomFn Injectable uniform RNG in [0, 1) for deterministic testing. */ -export async function* runWithRetry( - generatorFactory: () => AsyncGenerator, - retryConfig?: RetryConfig, - abortSignal?: AbortSignal, -): AsyncGenerator { - const config = normalizeRetryConfig(retryConfig); - if (!config || config.maxAttempts <= 1) { - return yield* generatorFactory(); +export function getRetryDelaySeconds( + retryConfig: RetryConfig | undefined, + nodeState: NodeState, + randomFn: () => number = Math.random, +): number { + if (!retryConfig) { + return DEFAULT_INITIAL_DELAY_SECONDS; } - let attempt = 1; - while (true) { - if (abortSignal?.aborted) { - throw new Error('Execution aborted before attempt.'); - } + const initialDelay = + retryConfig.initialDelay ?? DEFAULT_INITIAL_DELAY_SECONDS; + const maxDelay = retryConfig.maxDelay ?? DEFAULT_MAX_DELAY_SECONDS; + const backoffFactor = retryConfig.backoffFactor ?? DEFAULT_BACKOFF_FACTOR; + const jitter = retryConfig.jitter ?? DEFAULT_JITTER; - const generator = generatorFactory(); - try { - const result = yield* generator; - return result; - } catch (error: unknown) { - if ( - attempt >= config.maxAttempts || - !isErrorRetryable(error, config.retryableErrors) || - abortSignal?.aborted - ) { - throw error; - } + const attemptCount = nodeState.attemptCount || 1; + // attemptCount is the attempt number that just failed (1-based); the first + // failure (attempt 1) uses exponent 0. + const attemptForCalc = Math.max(0, attemptCount - 1); - const delayMs = Math.min( - config.initialDelayMs * Math.pow(config.backoffFactor, attempt - 1), - config.maxDelayMs, - ); - await sleepWithSignal(delayMs, abortSignal); - attempt++; - } + let delay = initialDelay * Math.pow(backoffFactor, attemptForCalc); + delay = Math.min(delay, maxDelay); + + if (jitter > 0.0) { + // random.uniform(-jitter*delay, jitter*delay) + const span = jitter * delay; + const randomOffset = -span + randomFn() * (2 * span); + delay = Math.max(0.0, delay + randomOffset); } + + return delay; } diff --git a/core/src/workflow-next/utils/workflow_graph_utils.ts b/core/src/workflow/utils/workflow_graph_utils.ts similarity index 100% rename from core/src/workflow-next/utils/workflow_graph_utils.ts rename to core/src/workflow/utils/workflow_graph_utils.ts diff --git a/core/src/workflow/workflow.ts b/core/src/workflow/workflow.ts index fe21f8779..e9c6ed639 100644 --- a/core/src/workflow/workflow.ts +++ b/core/src/workflow/workflow.ts @@ -4,138 +4,447 @@ * SPDX-License-Identifier: Apache-2.0 */ -import {BaseAgent, BaseAgentConfig} from '../agents/base_agent.js'; -import {InvocationContext} from '../agents/invocation_context.js'; import {Event} from '../events/event.js'; -import {BaseNode} from './base_node.js'; -import { - DynamicEntryFunction, - DynamicNodeScheduler, -} from './dynamic_node_scheduler.js'; -import {NodeRunner} from './node_runner.js'; -import {GraphEdge} from './utils/graph_parser.js'; +import {BaseNode, BaseNodeConfig} from './base_node.js'; +import {BranchPath} from './branch_path.js'; +import {DynamicNodeScheduler} from './dynamic_node_scheduler.js'; +import {EdgeItem, Graph, RouteValue} from './graph.js'; +import {NodeContext} 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'; /** - * A unique symbol to identify ADK Workflow agent instances. + * 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`. */ -const WORKFLOW_SIGNATURE_SYMBOL = Symbol.for('google.adk.workflow'); +export type DynamicEntry = ( + ctx: NodeContext, + input: unknown, +) => unknown | Promise; /** - * Type guard to check if an object is an instance of Workflow. - * @param obj The object to check. - * @returns True if the object is an instance of Workflow, false otherwise. + * Configuration for a {@link Workflow}. */ -export function isWorkflow(obj: unknown): obj is Workflow { - return ( - typeof obj === 'object' && - obj !== null && - WORKFLOW_SIGNATURE_SYMBOL in obj && - (obj as Record)[WORKFLOW_SIGNATURE_SYMBOL] === true - ); -} - -/** - * Configuration options for creating a Workflow agent. - * Workflows must define exactly one of `edges` (for static DAG execution) or `dynamicEntry` (for programmatic execution). - */ -export interface WorkflowConfig extends BaseAgentConfig { +export interface WorkflowConfig extends BaseNodeConfig { + /** Edge definitions used to build the workflow graph. */ + edges?: EdgeItem[]; /** - * Static graph edge definitions (e.g., `["START", nodeA, nodeB]` or `[routerNode, { ROUTE_A: nodeC }]`). - * Mutually exclusive with `dynamicEntry`. + * An imperative entry function driving execution via `ctx.runNode(...)`. + * Mutually exclusive with {@link edges}. */ - edges?: GraphEdge[]; - + dynamicEntry?: DynamicEntry; /** - * Programmatic entry node or async function handler (`async (ctx, input) => ...`) that coordinates - * child nodes using `ctx.runNode(...)`. Mutually exclusive with `edges`. + * Maximum number of graph-scheduled nodes running in parallel. `undefined` + * means unlimited. Does not throttle dynamic (`ctx.runNode`) children. */ - dynamicEntry?: BaseNode | DynamicEntryFunction; - - /** - * Optional key inside `InvocationContext.agentStates` where the final output of the workflow - * should be stored upon successful completion. - */ - outputKey?: string; + maxConcurrency?: number; +} - /** - * If true, the workflow will force re-execution on resumption even if historical outputs exist. - * Default is false. - */ - rerunOnResume?: boolean; +/** + * 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(); + errorShutDown = false; +} - /** - * If true, allows directed cycles inside static `edges` graph validation. - * Default is false. - */ - allowCycles?: boolean; +interface CompletedTask { + name: string; + childCtx?: NodeContext; + error?: unknown; } /** - * The top-level Workflow agent in ADK-JS. - * Inherits from `BaseAgent` and orchestrates multi-step node execution using either a static graph DAG (`NodeRunner`) - * or dynamic programmatic scheduling (`DynamicNodeScheduler`). + * 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). */ -export class Workflow extends BaseAgent { - readonly [WORKFLOW_SIGNATURE_SYMBOL] = true; - - readonly edges?: GraphEdge[]; - readonly dynamicEntry?: BaseNode | DynamicEntryFunction; - readonly outputKey?: string; - readonly rerunOnResume: boolean; - readonly allowCycles: boolean; +export class Workflow extends BaseNode { + readonly graph?: Graph; + readonly dynamicEntry?: DynamicEntry; + readonly maxConcurrency?: number; constructor(config: WorkflowConfig) { - super(config); - if (config.edges && config.dynamicEntry) { + super({...config, rerunOnResume: config.rerunOnResume ?? true}); + const hasEdges = !!config.edges && config.edges.length > 0; + if (hasEdges && config.dynamicEntry) { throw new Error( - `Workflow "${this.name}" cannot have both "edges" and "dynamicEntry" defined. They are mutually exclusive.`, + `Workflow "${this.name}": "edges" and "dynamicEntry" are mutually exclusive.`, ); } - if (!config.edges && !config.dynamicEntry) { + if (!hasEdges && !config.dynamicEntry) { throw new Error( - `Workflow "${this.name}" must define either "edges" (for static graphs) or "dynamicEntry" (for dynamic code execution).`, + `Workflow "${this.name}" requires either "edges" or "dynamicEntry".`, ); } - - this.edges = config.edges; + this.maxConcurrency = config.maxConcurrency; this.dynamicEntry = config.dynamicEntry; - this.outputKey = config.outputKey; - this.rerunOnResume = config.rerunOnResume ?? false; - this.allowCycles = config.allowCycles ?? false; + if (hasEdges) { + this.graph = Graph.fromEdgeItems(config.edges!); + this.graph.validate(); + } } - /** - * Executes the workflow via text-based or programmatic invocation. - */ - protected async *runAsyncImpl( - context: InvocationContext, + // eslint-disable-next-line require-yield + protected async *runImpl( + ctx: NodeContext, + nodeInput: unknown, ): AsyncGenerator { - if (context.endOfAgents[this.name]) { + // Child events are streamed through ctx.channel by ctx.runNode(), so this + // orchestration generator itself yields nothing. + const dynamicState = new DynamicNodeState(); + ctx.scheduler = new DynamicNodeScheduler(dynamicState); + + if (this.dynamicEntry) { + await this.runDynamicEntry(ctx, nodeInput, dynamicState); return; } - if (this.edges) { - const runner = new NodeRunner(this.edges, { - outputKey: this.outputKey, - allowCycles: this.allowCycles, - }); - yield* runner.runAsync(context, context.userContent); - } else if (this.dynamicEntry) { - const scheduler = new DynamicNodeScheduler(this.dynamicEntry, { - outputKey: this.outputKey, - }); - yield* scheduler.runAsync(context, context.userContent); + const loop = new LoopState(); + + // --- SETUP --- + this.seedStartTriggers(loop, nodeInput); + + // --- LOOP --- + await this.runLoop(loop, ctx); + + 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); } - context.endOfAgents[this.name] = true; + // --- FINALIZE --- + this.finalize(loop, ctx); } /** - * Executes the workflow via audio/video live streaming invocation. + * Runs an imperative `dynamicEntry` workflow. The entry drives execution via + * `ctx.runNode(...)` (routed through the scheduler) and returns the output. */ - protected async *runLiveImpl( - context: InvocationContext, - ): AsyncGenerator { - yield* this.runAsyncImpl(context); + 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; + } + } + + // --- 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): 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); + 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 && 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)!; + + let runId = nodeState.runId; + if (!runId) { + nodeState.runCounter += 1; + runId = String(nodeState.runCounter); + nodeState.runId = runId; + } + + // Static graph nodes are managed by this loop directly, bypassing the + // dynamic scheduler (which serves user-initiated ctx.runNode() calls). + const task: Promise = executeChildNode( + ctx, + node, + trigger.input, + { + 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, + ): 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 | 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 = BranchPath.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): Promise { + // Await outstanding tasks so their events flush; failures are swallowed + // because the workflow is already shutting down on error. + const outstanding = [...loop.pending.values()]; + loop.pending.clear(); + await Promise.allSettled(outstanding); } } diff --git a/core/src/workflow-next/workflow_agent.ts b/core/src/workflow/workflow_agent.ts similarity index 100% rename from core/src/workflow-next/workflow_agent.ts rename to core/src/workflow/workflow_agent.ts diff --git a/core/test/workflow-next/dynamic_workflow_test.ts b/core/test/workflow-next/dynamic_workflow_test.ts deleted file mode 100644 index 13cc5b006..000000000 --- a/core/test/workflow-next/dynamic_workflow_test.ts +++ /dev/null @@ -1,155 +0,0 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import {describe, expect, it} from 'vitest'; -import {BaseAgent} from '../../src/agents/base_agent.js'; -import {InvocationContext} from '../../src/agents/invocation_context.js'; -import {Event} from '../../src/events/event.js'; -import {PluginManager} from '../../src/plugins/plugin_manager.js'; -import {Session} from '../../src/sessions/session.js'; -import {NodeContext} from '../../src/workflow-next/node_context.js'; -import {FunctionNode} from '../../src/workflow-next/nodes/function_node.js'; -import {EventChannel} from '../../src/workflow-next/utils/event_channel.js'; -import {Workflow} from '../../src/workflow-next/workflow.js'; - -function createIc(): InvocationContext { - const session = { - id: 's1', - appName: 'app', - userId: 'u', - events: [], - state: {}, - lastUpdateTime: Date.now(), - } as unknown as Session; - return new InvocationContext({ - invocationId: 'inv-1', - session, - agent: { - name: 'wf', - runAsync: async function* () {}, - } as unknown as BaseAgent, - pluginManager: new PluginManager(), - }); -} - -async function driveWorkflow( - wf: Workflow, - input?: unknown, -): Promise<{events: Event[]; output: unknown}> { - const channel = new EventChannel(); - const root = new NodeContext({ - invocationContext: createIc(), - channel, - nodePath: '', - runId: 'root', - }); - const events: Event[] = []; - const run = root.runNode(wf, input, {useAsOutput: true}).then( - () => channel.close(), - (err) => channel.fail(err), - ); - for await (const ev of channel) { - events.push(ev); - } - await run; - return {events, output: root.output}; -} - -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/dynamic_workflow_test.ts b/core/test/workflow/dynamic_workflow_test.ts index 43b219b18..e8d4476b7 100644 --- a/core/test/workflow/dynamic_workflow_test.ts +++ b/core/test/workflow/dynamic_workflow_test.ts @@ -4,104 +4,152 @@ * SPDX-License-Identifier: Apache-2.0 */ -import {describe, expect, it, vi} from 'vitest'; +import {describe, expect, it} from 'vitest'; import {BaseAgent} from '../../src/agents/base_agent.js'; import {InvocationContext} from '../../src/agents/invocation_context.js'; +import {Event} from '../../src/events/event.js'; import {PluginManager} from '../../src/plugins/plugin_manager.js'; import {Session} from '../../src/sessions/session.js'; -import { - DynamicNodeScheduler, - FunctionNode, - NodeStatus, - runNode, -} from '../../src/workflow/index.js'; +import {NodeContext} from '../../src/workflow/node_context.js'; +import {FunctionNode} from '../../src/workflow/nodes/function_node.js'; +import {EventChannel} from '../../src/workflow/utils/event_channel.js'; +import {Workflow} from '../../src/workflow/workflow.js'; -describe('Workflow DynamicNodeScheduler & runNode', () => { - function createTestContext( - params?: Partial, - ): InvocationContext { - const session: Session = { - id: 'session-dyn', - appName: 'test-app', - userId: 'test-user', - events: [], - state: {}, - }; +function createIc(): InvocationContext { + const session = { + id: 's1', + appName: 'app', + userId: 'u', + events: [], + state: {}, + lastUpdateTime: Date.now(), + } as unknown as Session; + return new InvocationContext({ + invocationId: 'inv-1', + session, + agent: { + name: 'wf', + runAsync: async function* () {}, + } as unknown as BaseAgent, + pluginManager: new PluginManager(), + }); +} - return new InvocationContext({ - invocationId: 'inv-dyn', - session, - agent: { - name: 'mock_agent', - runAsync: async function* () {}, - } as unknown as BaseAgent, - pluginManager: new PluginManager(), - ...params, - }); +async function driveWorkflow( + wf: Workflow, + input?: unknown, +): Promise<{events: Event[]; output: unknown}> { + const channel = new EventChannel(); + const root = new NodeContext({ + invocationContext: createIc(), + channel, + nodePath: '', + runId: 'root', + }); + const events: Event[] = []; + const run = root.runNode(wf, input, {useAsOutput: true}).then( + () => channel.close(), + (err) => channel.fail(err), + ); + for await (const ev of channel) { + events.push(ev); } + await run; + return {events, output: root.output}; +} - it('should run dynamic workflows and track individual node checkpoints via runNode', async () => { - const ctx = createTestContext(); - const nodeA = new FunctionNode('calc_a', (_ctx, num: number) => num * 2); - const nodeB = new FunctionNode('calc_b', (_ctx, num: number) => num + 10); - - const dynamicEntry = async (context: InvocationContext, input?: number) => { - const resA = await runNode(context, nodeA, input || 5); - if (resA > 5) { - const resB = await runNode(context, nodeB, resA); - return resB; - } - return resA; - }; - - const scheduler = new DynamicNodeScheduler(dynamicEntry, { - outputKey: 'dynResult', +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}]`; + }, }); - for await (const _ of scheduler.runAsync(ctx, 4)) { - /* consume events */ - } - - expect(ctx.agentStates['exec_node_calc_a'].status).toBe( - NodeStatus.COMPLETED, - ); - expect(ctx.agentStates['exec_node_calc_a'].outputPayload).toBe(8); - expect(ctx.agentStates['exec_node_calc_b'].status).toBe( - NodeStatus.COMPLETED, - ); - expect(ctx.agentStates['exec_node_calc_b'].outputPayload).toBe(18); - expect(ctx.agentStates['dynResult']).toBe(18); + expect((await driveWorkflow(wf, 'x')).output).toBe('wrapped[step(x)]'); }); - it('should skip completed nodes inside dynamic execution on resume', async () => { - const ctx = createTestContext(); - const spyA = vi.fn((_ctx, num: number) => num * 100); - const spyB = vi.fn((_ctx, num: number) => num + 50); - const nodeA = new FunctionNode('node_dyn_a', spyA, {rerunOnResume: false}); - const nodeB = new FunctionNode('node_dyn_b', spyB, {rerunOnResume: true}); - - ctx.agentStates['exec_node_node_dyn_a'] = { - executionId: 'exec_node_node_dyn_a', - nodeName: 'node_dyn_a', - status: NodeStatus.COMPLETED, - outputPayload: 999, - timestamp: Date.now(), - }; + 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, + }); + }); - const dynamicEntry = async (context: InvocationContext, input?: number) => { - const resA = await runNode(context, nodeA, input || 1); - const resB = await runNode(context, nodeB, resA); - return resB; - }; + 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); + }); - const scheduler = new DynamicNodeScheduler(dynamicEntry); - for await (const _ of scheduler.runAsync(ctx, 2)) { - /* consume events */ - } + 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, + }); + }); - // spyA skipped, reused 999 - expect(spyA).not.toHaveBeenCalled(); - // spyB executed with 999 - expect(spyB).toHaveBeenCalledTimes(1); - expect(spyB).toHaveBeenCalledWith(ctx, 999); + 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-next/event_channel_test.ts b/core/test/workflow/event_channel_test.ts similarity index 96% rename from core/test/workflow-next/event_channel_test.ts rename to core/test/workflow/event_channel_test.ts index d9c4ea8f7..e2b1a446e 100644 --- a/core/test/workflow-next/event_channel_test.ts +++ b/core/test/workflow/event_channel_test.ts @@ -5,7 +5,7 @@ */ import {describe, expect, it} from 'vitest'; -import {EventChannel} from '../../src/workflow-next/utils/event_channel.js'; +import {EventChannel} from '../../src/workflow/utils/event_channel.js'; async function drain(ch: EventChannel): Promise { const out: T[] = []; diff --git a/core/test/workflow-next/event_model_test.ts b/core/test/workflow/event_model_test.ts similarity index 100% rename from core/test/workflow-next/event_model_test.ts rename to core/test/workflow/event_model_test.ts diff --git a/core/test/workflow-next/foundations_test.ts b/core/test/workflow/foundations_test.ts similarity index 95% rename from core/test/workflow-next/foundations_test.ts rename to core/test/workflow/foundations_test.ts index 9f5b68c4d..6d69e10bd 100644 --- a/core/test/workflow-next/foundations_test.ts +++ b/core/test/workflow/foundations_test.ts @@ -9,18 +9,18 @@ import { DynamicNodeFailError, NodeInterruptedError, NodeTimeoutError, -} from '../../src/workflow-next/errors.js'; +} from '../../src/workflow/errors.js'; import { createNodeState, isNodeState, NodeState, -} from '../../src/workflow-next/node_state.js'; -import {NodeStatus} from '../../src/workflow-next/node_status.js'; -import {normalizeRetryExceptions} from '../../src/workflow-next/retry_config.js'; +} from '../../src/workflow/node_state.js'; +import {NodeStatus} from '../../src/workflow/node_status.js'; +import {normalizeRetryExceptions} from '../../src/workflow/retry_config.js'; import { getRetryDelaySeconds, shouldRetryNode, -} from '../../src/workflow-next/utils/retry_utils.js'; +} from '../../src/workflow/utils/retry_utils.js'; describe('Phase 0 — errors', () => { it('NodeTimeoutError carries nodeName/timeout and is instanceof Error', () => { diff --git a/core/test/workflow/graph_parser_test.ts b/core/test/workflow/graph_parser_test.ts deleted file mode 100644 index e49557045..000000000 --- a/core/test/workflow/graph_parser_test.ts +++ /dev/null @@ -1,105 +0,0 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import {describe, expect, it} from 'vitest'; -import { - FunctionNode, - JoinNode, - parseGraph, - Trigger, - validateGraph, -} from '../../src/workflow/index.js'; - -describe('Workflow Graph Parser & Validation', () => { - const nodeA = new FunctionNode('node_a', () => 'A'); - const nodeB = new FunctionNode('node_b', () => 'B'); - const nodeC = new FunctionNode('node_c', () => 'C'); - const router = new FunctionNode('router', () => 'router_res'); - - it('should parse sequential edges accurately', () => { - const graph = parseGraph([['START', nodeA, nodeB, nodeC]]); - expect(graph.nodes.size).toBe(3); - expect(graph.nodes.get('node_a')).toBe(nodeA); - expect(graph.nodes.get('node_b')).toBe(nodeB); - expect(graph.nodes.get('node_c')).toBe(nodeC); - - const startEdges = graph.adjacencyList.get('START') || []; - expect(startEdges.length).toBe(1); - expect(startEdges[0].target).toBe(nodeA); - - const aEdges = graph.adjacencyList.get('node_a') || []; - expect(aEdges.length).toBe(1); - expect(aEdges[0].target).toBe(nodeB); - - const bEdges = graph.adjacencyList.get('node_b') || []; - expect(bEdges.length).toBe(1); - expect(bEdges[0].target).toBe(nodeC); - - expect(graph.inboundCounts.get('node_a')).toBe(1); - expect(graph.inboundCounts.get('node_b')).toBe(1); - expect(graph.inboundCounts.get('node_c')).toBe(1); - }); - - it('should parse route maps and trigger tuples', () => { - const customTrigger = Trigger.fromPredicate(() => true); - const graph = parseGraph([ - ['START', router], - [router, {ROUTE_X: nodeA, ROUTE_Y: nodeB}], - [nodeA, [customTrigger, nodeC]], - ]); - - expect(graph.nodes.size).toBe(4); - const routerEdges = graph.adjacencyList.get('router') || []; - expect(routerEdges.length).toBe(2); - expect(routerEdges[0].target.name).toBe('node_a'); - expect(routerEdges[1].target.name).toBe('node_b'); - - const aEdges = graph.adjacencyList.get('node_a') || []; - expect(aEdges.length).toBe(1); - expect(aEdges[0].trigger).toBe(customTrigger); - }); - - it('should throw during validation when a node is unreachable from START', () => { - const graph = parseGraph([ - ['START', nodeA], - [nodeB, nodeC], // nodeB and nodeC have no path from START - ]); - - expect(() => validateGraph(graph)).toThrowError( - /unreachable from "START"/i, - ); - }); - - it('should detect cycles and throw unless allowCycles is true', () => { - const graph = parseGraph([ - ['START', nodeA, nodeB], - [nodeB, nodeA], // cycle nodeB -> nodeA - ]); - - expect(() => validateGraph(graph, {allowCycles: false})).toThrowError( - /Cycle detected/i, - ); - - expect(() => validateGraph(graph, {allowCycles: true})).not.toThrow(); - }); - - it('should validate JoinNode upstreamCount integrity', () => { - const joinNode = new JoinNode('join_node', {upstreamCount: 2}); - const validGraph = parseGraph([ - ['START', nodeA, joinNode], - ['START', nodeB, joinNode], - ]); - - expect(() => validateGraph(validGraph)).not.toThrow(); - - const invalidJoin = new JoinNode('bad_join', {upstreamCount: 5}); - const invalidGraph = parseGraph([['START', nodeA, invalidJoin]]); - - expect(() => validateGraph(invalidGraph)).toThrowError( - /expects 5 upstream predecessors, but only has 1 inbound edges/i, - ); - }); -}); diff --git a/core/test/workflow/hitl_and_rehydration_test.ts b/core/test/workflow/hitl_and_rehydration_test.ts deleted file mode 100644 index 1801b05f7..000000000 --- a/core/test/workflow/hitl_and_rehydration_test.ts +++ /dev/null @@ -1,189 +0,0 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import {describe, expect, it} from 'vitest'; -import {BaseAgent} from '../../src/agents/base_agent.js'; -import {InvocationContext} from '../../src/agents/invocation_context.js'; -import {PluginManager} from '../../src/plugins/plugin_manager.js'; -import {Session} from '../../src/sessions/session.js'; -import { - createRequestInputEvent, - FunctionNode, - injectHitlResumptionInput, - NodeRunner, - NodeStatus, - persistAgentStatesToSession, - rehydrateAgentStates, - ReplayManager, -} from '../../src/workflow/index.js'; - -describe('Workflow HITL & State Rehydration', () => { - function createTestContext(session?: Session): InvocationContext { - const s: Session = session || { - id: 'session-hitl', - appName: 'test-app', - userId: 'test-user', - events: [], - state: {}, - }; - - return new InvocationContext({ - invocationId: 'inv-hitl', - session: s, - agent: { - name: 'mock_agent', - runAsync: async function* () {}, - } as unknown as BaseAgent, - pluginManager: new PluginManager(), - }); - } - - it('should pause workflow execution when a node yields a RequestInput event', async () => { - const ctx = createTestContext(); - const hitlNode = new FunctionNode('approval_step', (context) => { - return createRequestInputEvent(context, 'approval_step', { - prompt: 'Please approve this transaction (yes/no)', - }); - }); - const downstreamNode = new FunctionNode('post_approval', () => 'EXECUTED'); - - const runner = new NodeRunner([['START', hitlNode, downstreamNode]]); - const events: unknown[] = []; - for await (const ev of runner.runAsync(ctx)) { - events.push(ev); - } - - expect(events.length).toBe(1); - expect( - ( - events[0] as unknown as Record< - string, - Record - > - ).actions.requestInput.nodeName, - ).toBe('approval_step'); - expect(ctx.agentStates['exec_node_approval_step'].status).toBe( - NodeStatus.PAUSED_HITL, - ); - - expect(ctx.endInvocation).toBe(true); - // downstreamNode must NOT have run while paused - expect(ctx.agentStates['exec_node_post_approval']).toBeUndefined(); - }); - - it('should inject resumption input and complete the paused node on subsequent turn', async () => { - const ctx = createTestContext(); - const hitlNode = new FunctionNode( - 'approval_step', - (_context, input?: string) => { - if (!input) { - return createRequestInputEvent(_context, 'approval_step', { - prompt: 'Approve?', - }); - } - return `APPROVED_WITH_${input}`; - }, - ); - const downstreamNode = new FunctionNode( - 'post_approval', - (_context, input: string) => `DONE_${input}`, - ); - - const runner = new NodeRunner([['START', hitlNode, downstreamNode]]); - - // Turn 1: Pauses - for await (const _ of runner.runAsync(ctx)) { - /* consume events */ - } - expect(ctx.agentStates['exec_node_approval_step'].status).toBe( - NodeStatus.PAUSED_HITL, - ); - - // Turn 2: User provides input 'YES'. Inject it and run again. - ctx.endInvocation = false; - const resumedInfo = injectHitlResumptionInput(ctx, 'YES'); - expect(resumedInfo?.nodeName).toBe('approval_step'); - expect(ctx.agentStates['exec_node_approval_step'].status).toBe( - NodeStatus.RUNNING, - ); - expect(ctx.agentStates['exec_node_approval_step'].inputPayload).toBe('YES'); - - for await (const _ of runner.runAsync(ctx)) { - /* consume events */ - } - - expect(ctx.agentStates['exec_node_approval_step'].status).toBe( - NodeStatus.COMPLETED, - ); - expect(ctx.agentStates['exec_node_approval_step'].outputPayload).toBe( - 'APPROVED_WITH_YES', - ); - expect(ctx.agentStates['exec_node_post_approval'].status).toBe( - NodeStatus.COMPLETED, - ); - expect(ctx.agentStates['exec_node_post_approval'].outputPayload).toBe( - 'DONE_APPROVED_WITH_YES', - ); - }); - - it('should persist and rehydrate checkpoints from session state accurately', () => { - const session: Session = { - id: 's-rehydrate', - appName: 'app', - userId: 'user', - events: [], - state: {}, - }; - const ctx1 = createTestContext(session); - ctx1.agentStates['exec_node_saved'] = { - executionId: 'exec_node_saved', - nodeName: 'saved', - status: NodeStatus.COMPLETED, - outputPayload: {foo: 'bar'}, - timestamp: 12345, - }; - ctx1.endOfAgents['my_wf'] = true; - - persistAgentStatesToSession(ctx1, session); - - // Now simulate a fresh context loading from the same session - const ctx2 = createTestContext(session); - rehydrateAgentStates(session, ctx2); - - expect(ctx2.agentStates['exec_node_saved'].status).toBe( - NodeStatus.COMPLETED, - ); - expect(ctx2.agentStates['exec_node_saved'].outputPayload).toEqual({ - foo: 'bar', - }); - expect(ctx2.endOfAgents['my_wf']).toBe(true); - }); - - it('should yield historical replay events when ReplayManager inspects completed checkpoints', async () => { - const ctx = createTestContext(); - const node = new FunctionNode('past_node', () => 'historical_res', { - rerunOnResume: false, - }); - ctx.agentStates['exec_node_past_node'] = { - executionId: 'exec_node_past_node', - nodeName: 'past_node', - status: NodeStatus.COMPLETED, - outputPayload: 'historical_res', - timestamp: 99999, - }; - - const gen = ReplayManager.replayIfCompleted(ctx, node); - const ev = await gen.next(); - expect(ev.done).toBe(false); - expect(ev.value.actions.nodeExecutionReplay.outputPayload).toBe( - 'historical_res', - ); - - const res = await gen.next(); - expect(res.done).toBe(true); - expect(res.value).toBe(true); // Replayed successfully, skip real run - }); -}); diff --git a/core/test/workflow-next/hitl_test.ts b/core/test/workflow/hitl_test.ts similarity index 93% rename from core/test/workflow-next/hitl_test.ts rename to core/test/workflow/hitl_test.ts index 097eb217e..69d41f1e0 100644 --- a/core/test/workflow-next/hitl_test.ts +++ b/core/test/workflow/hitl_test.ts @@ -10,15 +10,15 @@ import {InvocationContext} from '../../src/agents/invocation_context.js'; import {Event} from '../../src/events/event.js'; import {PluginManager} from '../../src/plugins/plugin_manager.js'; import {Session} from '../../src/sessions/session.js'; -import {NodeContext} from '../../src/workflow-next/node_context.js'; -import {FunctionNode} from '../../src/workflow-next/nodes/function_node.js'; -import {RequestInput} from '../../src/workflow-next/request_input.js'; -import {EventChannel} from '../../src/workflow-next/utils/event_channel.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 {EventChannel} from '../../src/workflow/utils/event_channel.js'; import { hasRequestInputFunctionCall, REQUEST_INPUT_FUNCTION_CALL_NAME, -} from '../../src/workflow-next/utils/hitl_utils.js'; -import {Workflow} from '../../src/workflow-next/workflow.js'; +} from '../../src/workflow/utils/hitl_utils.js'; +import {Workflow} from '../../src/workflow/workflow.js'; function createIc(): InvocationContext { const session = { diff --git a/core/test/workflow/join_node_and_parallel_test.ts b/core/test/workflow/join_node_and_parallel_test.ts deleted file mode 100644 index a78db2ccd..000000000 --- a/core/test/workflow/join_node_and_parallel_test.ts +++ /dev/null @@ -1,130 +0,0 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import {describe, expect, it} from 'vitest'; -import {BaseAgent} from '../../src/agents/base_agent.js'; -import {InvocationContext} from '../../src/agents/invocation_context.js'; -import {PluginManager} from '../../src/plugins/plugin_manager.js'; -import {Session} from '../../src/sessions/session.js'; -import { - FunctionNode, - JoinNode, - NodeStatus, - runInParallel, -} from '../../src/workflow/index.js'; - -describe('Workflow ParallelWorker & JoinNode', () => { - function createTestContext( - params?: Partial, - ): InvocationContext { - const session: Session = { - id: 'session-par', - appName: 'test-app', - userId: 'test-user', - events: [], - state: {}, - }; - - return new InvocationContext({ - invocationId: 'inv-par', - session, - agent: { - name: 'mock_agent', - runAsync: async function* () {}, - } as unknown as BaseAgent, - pluginManager: new PluginManager(), - ...params, - }); - } - - it('should run items in parallel and merge child checkpoints back into parent context', async () => { - const ctx = createTestContext(); - const workerNode = new FunctionNode( - 'worker', - async (_ctx, item: string) => { - return `processed_${item}`; - }, - ); - - const results = await runInParallel(ctx, workerNode, [ - 'alpha', - 'beta', - 'gamma', - ]); - - expect(results).toEqual([ - 'processed_alpha', - 'processed_beta', - 'processed_gamma', - ]); - expect(ctx.agentStates['exec_node_worker_0.worker'].status).toBe( - NodeStatus.COMPLETED, - ); - expect(ctx.agentStates['exec_node_worker_0.worker'].outputPayload).toBe( - 'processed_alpha', - ); - expect(ctx.agentStates['exec_node_worker_1.worker'].status).toBe( - NodeStatus.COMPLETED, - ); - expect(ctx.agentStates['exec_node_worker_1.worker'].outputPayload).toBe( - 'processed_beta', - ); - expect(ctx.agentStates['exec_node_worker_2.worker'].status).toBe( - NodeStatus.COMPLETED, - ); - expect(ctx.agentStates['exec_node_worker_2.worker'].outputPayload).toBe( - 'processed_gamma', - ); - }); - - it('should synchronize at a JoinNode when all upstream predecessors reach COMPLETED', async () => { - const ctx = createTestContext(); - const joinNode = new JoinNode('join_sync', { - upstreamCount: 2, - predecessors: ['branch_1', 'branch_2'], - }); - - // 1. First branch finishes, second has not started - ctx.agentStates['exec_node_branch_1'] = { - executionId: 'exec_node_branch_1', - nodeName: 'branch_1', - status: NodeStatus.COMPLETED, - outputPayload: {data: 100}, - timestamp: Date.now(), - }; - - const gen1 = joinNode.run(ctx); - const res1 = await gen1.next(); - // Since only 1 of 2 predecessors completed, joinNode returns partial state and does NOT yield a joinCompleted event - expect(res1.done).toBe(true); - expect(res1.value).toEqual({branch_1: {data: 100}}); - - // 2. Second branch now finishes - ctx.agentStates['exec_node_branch_2'] = { - executionId: 'exec_node_branch_2', - nodeName: 'branch_2', - status: NodeStatus.COMPLETED, - outputPayload: {data: 200}, - timestamp: Date.now(), - }; - - const gen2 = joinNode.run(ctx); - const ev = await gen2.next(); // Should yield joinCompleted event - expect(ev.done).toBe(false); - expect(ev.value.actions.joinCompleted).toBeDefined(); - expect(ev.value.actions.joinCompleted.outputs).toEqual({ - branch_1: {data: 100}, - branch_2: {data: 200}, - }); - - const finalRes = await gen2.next(); - expect(finalRes.done).toBe(true); - expect(finalRes.value).toEqual({ - branch_1: {data: 100}, - branch_2: {data: 200}, - }); - }); -}); diff --git a/core/test/workflow-next/llm_agent_test.ts b/core/test/workflow/llm_agent_test.ts similarity index 92% rename from core/test/workflow-next/llm_agent_test.ts rename to core/test/workflow/llm_agent_test.ts index e880c2e80..a53796093 100644 --- a/core/test/workflow-next/llm_agent_test.ts +++ b/core/test/workflow/llm_agent_test.ts @@ -10,11 +10,11 @@ import {InvocationContext} from '../../src/agents/invocation_context.js'; import {createEvent, Event} from '../../src/events/event.js'; import {PluginManager} from '../../src/plugins/plugin_manager.js'; import {Session} from '../../src/sessions/session.js'; -import {node} from '../../src/workflow-next/node.js'; -import {NodeContext} from '../../src/workflow-next/node_context.js'; -import {LLMAgentWrapper} from '../../src/workflow-next/nodes/llm_agent_wrapper.js'; -import {EventChannel} from '../../src/workflow-next/utils/event_channel.js'; -import {Workflow} from '../../src/workflow-next/workflow.js'; +import {node} from '../../src/workflow/node.js'; +import {NodeContext} from '../../src/workflow/node_context.js'; +import {LLMAgentWrapper} from '../../src/workflow/nodes/llm_agent_wrapper.js'; +import {EventChannel} from '../../src/workflow/utils/event_channel.js'; +import {Workflow} from '../../src/workflow/workflow.js'; function createIc(): InvocationContext { const session = { diff --git a/core/test/workflow-next/node_api_test.ts b/core/test/workflow/node_api_test.ts similarity index 91% rename from core/test/workflow-next/node_api_test.ts rename to core/test/workflow/node_api_test.ts index e478c4855..39ad99405 100644 --- a/core/test/workflow-next/node_api_test.ts +++ b/core/test/workflow/node_api_test.ts @@ -12,15 +12,15 @@ import {Event} from '../../src/events/event.js'; import {PluginManager} from '../../src/plugins/plugin_manager.js'; import {Session} from '../../src/sessions/session.js'; import {BaseTool} from '../../src/tools/base_tool.js'; -import {BaseNode} from '../../src/workflow-next/base_node.js'; -import {node, Node} from '../../src/workflow-next/node.js'; -import {NodeContext} from '../../src/workflow-next/node_context.js'; -import {FunctionNode} from '../../src/workflow-next/nodes/function_node.js'; -import {JoinNode} from '../../src/workflow-next/nodes/join_node.js'; -import {LLMAgentWrapper} from '../../src/workflow-next/nodes/llm_agent_wrapper.js'; -import {ToolNode} from '../../src/workflow-next/nodes/tool_node.js'; -import {EventChannel} from '../../src/workflow-next/utils/event_channel.js'; -import {Workflow} from '../../src/workflow-next/workflow.js'; +import {BaseNode} from '../../src/workflow/base_node.js'; +import {node, Node} from '../../src/workflow/node.js'; +import {NodeContext} from '../../src/workflow/node_context.js'; +import {FunctionNode} from '../../src/workflow/nodes/function_node.js'; +import {JoinNode} from '../../src/workflow/nodes/join_node.js'; +import {LLMAgentWrapper} from '../../src/workflow/nodes/llm_agent_wrapper.js'; +import {ToolNode} from '../../src/workflow/nodes/tool_node.js'; +import {EventChannel} from '../../src/workflow/utils/event_channel.js'; +import {Workflow} from '../../src/workflow/workflow.js'; function createIc(): InvocationContext { const session = { diff --git a/core/test/workflow-next/node_execution_test.ts b/core/test/workflow/node_execution_test.ts similarity index 94% rename from core/test/workflow-next/node_execution_test.ts rename to core/test/workflow/node_execution_test.ts index 0f465c896..ec588ff63 100644 --- a/core/test/workflow-next/node_execution_test.ts +++ b/core/test/workflow/node_execution_test.ts @@ -10,10 +10,10 @@ import {InvocationContext} from '../../src/agents/invocation_context.js'; import {createEvent, Event} from '../../src/events/event.js'; import {PluginManager} from '../../src/plugins/plugin_manager.js'; import {Session} from '../../src/sessions/session.js'; -import {BaseNode} from '../../src/workflow-next/base_node.js'; -import {NodeTimeoutError} from '../../src/workflow-next/errors.js'; -import {NodeContext} from '../../src/workflow-next/node_context.js'; -import {EventChannel} from '../../src/workflow-next/utils/event_channel.js'; +import {BaseNode} from '../../src/workflow/base_node.js'; +import {NodeTimeoutError} from '../../src/workflow/errors.js'; +import {NodeContext} from '../../src/workflow/node_context.js'; +import {EventChannel} from '../../src/workflow/utils/event_channel.js'; // --- Test harness --------------------------------------------------------- @@ -79,10 +79,7 @@ class FnNode extends BaseNode { input: unknown, ) => unknown | Promise, config?: Partial< - Omit< - import('../../src/workflow-next/base_node.js').BaseNodeConfig, - 'name' - > + Omit >, ) { super({name, ...config}); diff --git a/core/test/workflow/node_runner_test.ts b/core/test/workflow/node_runner_test.ts deleted file mode 100644 index 868a5f833..000000000 --- a/core/test/workflow/node_runner_test.ts +++ /dev/null @@ -1,177 +0,0 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import {describe, expect, it, vi} from 'vitest'; -import {BaseAgent} from '../../src/agents/base_agent.js'; -import {InvocationContext} from '../../src/agents/invocation_context.js'; -import {PluginManager} from '../../src/plugins/plugin_manager.js'; -import {Session} from '../../src/sessions/session.js'; -import { - FunctionNode, - NodeRunner, - NodeStatus, -} from '../../src/workflow/index.js'; - -describe('Workflow NodeRunner & Checkpointing', () => { - function createTestContext( - params?: Partial, - ): InvocationContext { - const session: Session = { - id: 'session-123', - appName: 'test-app', - userId: 'test-user', - events: [], - state: {}, - }; - - return new InvocationContext({ - invocationId: 'inv-456', - session, - agent: { - name: 'mock_agent', - runAsync: async function* () {}, - } as unknown as BaseAgent, - pluginManager: new PluginManager(), - ...params, - }); - } - - it('should run sequential nodes and pass input payloads downstream', async () => { - const ctx = createTestContext(); - const nodeA = new FunctionNode( - 'step_1', - (_ctx, input: string) => `${input}_A`, - ); - const nodeB = new FunctionNode( - 'step_2', - (_ctx, input: string) => `${input}_B`, - ); - - const runner = new NodeRunner([['START', nodeA, nodeB]], { - outputKey: 'finalOutput', - }); - const events: unknown[] = []; - for await (const event of runner.runAsync(ctx, 'INITIAL')) { - events.push(event); - } - - expect(ctx.agentStates['exec_node_step_1'].status).toBe( - NodeStatus.COMPLETED, - ); - expect(ctx.agentStates['exec_node_step_1'].outputPayload).toBe('INITIAL_A'); - expect(ctx.agentStates['exec_node_step_2'].status).toBe( - NodeStatus.COMPLETED, - ); - expect(ctx.agentStates['exec_node_step_2'].outputPayload).toBe( - 'INITIAL_A_B', - ); - expect(ctx.agentStates['finalOutput']).toEqual({ - step_1: 'INITIAL_A', - step_2: 'INITIAL_A_B', - }); - }); - - it('should route conditionally when a router node emits a route action or string', async () => { - const ctx = createTestContext(); - const router = new FunctionNode('router', () => 'ROUTE_Y'); - const nodeX = new FunctionNode( - 'node_x', - vi.fn(() => 'X'), - ); - const nodeY = new FunctionNode( - 'node_y', - vi.fn(() => 'Y'), - ); - - const runner = new NodeRunner([ - ['START', router], - [router, {ROUTE_X: nodeX, ROUTE_Y: nodeY}], - ]); - - for await (const _ of runner.runAsync(ctx)) { - /* consume events */ - } - - expect(ctx.agentStates['exec_node_router'].status).toBe( - NodeStatus.COMPLETED, - ); - expect(ctx.agentStates['exec_node_node_y'].status).toBe( - NodeStatus.COMPLETED, - ); - expect(ctx.agentStates['exec_node_node_x']).toBeUndefined(); // ROUTE_X never enqueued - }); - - it('should skip completed nodes on resume unless rerunOnResume is true', async () => { - const ctx = createTestContext(); - const spyA = vi.fn(() => 'fresh_A'); - const spyB = vi.fn(() => 'fresh_B'); - const nodeA = new FunctionNode('node_a', spyA, {rerunOnResume: false}); - const nodeB = new FunctionNode('node_b', spyB, {rerunOnResume: true}); - - // Pre-populate agentStates as if nodeA and nodeB completed in a previous run - ctx.agentStates['exec_node_node_a'] = { - executionId: 'exec_node_node_a', - nodeName: 'node_a', - status: NodeStatus.COMPLETED, - outputPayload: 'cached_A', - timestamp: Date.now() - 10000, - }; - ctx.agentStates['exec_node_node_b'] = { - executionId: 'exec_node_node_b', - nodeName: 'node_b', - status: NodeStatus.COMPLETED, - outputPayload: 'cached_B', - timestamp: Date.now() - 10000, - }; - - const runner = new NodeRunner([['START', nodeA, nodeB]]); - for await (const _ of runner.runAsync(ctx)) { - /* consume events */ - } - - // nodeA should have been skipped (spyA not called), and cached_A passed to nodeB - expect(spyA).not.toHaveBeenCalled(); - // nodeB has rerunOnResume: true, so spyB MUST be called with cached_A - expect(spyB).toHaveBeenCalledTimes(1); - expect(spyB).toHaveBeenCalledWith(ctx, 'cached_A'); - }); - - it('should retry node execution upon transient errors according to retryConfig', async () => { - const ctx = createTestContext(); - let attempts = 0; - const flakyNode = new FunctionNode( - 'flaky_node', - () => { - attempts++; - if (attempts < 3) { - throw new Error('Transient timeout error'); - } - return 'SUCCESS_AFTER_RETRY'; - }, - { - retryConfig: { - maxAttempts: 3, - initialDelayMs: 10, - maxDelayMs: 50, - backoffFactor: 1.5, - }, - }, - ); - - const runner = new NodeRunner([['START', flakyNode]]); - for await (const _ of runner.runAsync(ctx)) { - /* consume events */ - } - - expect(attempts).toBe(3); - expect(ctx.agentStates['exec_node_flaky_node'].status).toBe( - NodeStatus.COMPLETED, - ); - expect(ctx.agentStates['exec_node_flaky_node'].outputPayload).toBe( - 'SUCCESS_AFTER_RETRY', - ); - }); -}); diff --git a/core/test/workflow-next/parallel_test.ts b/core/test/workflow/parallel_test.ts similarity index 91% rename from core/test/workflow-next/parallel_test.ts rename to core/test/workflow/parallel_test.ts index 644bc1641..5cfdebbe2 100644 --- a/core/test/workflow-next/parallel_test.ts +++ b/core/test/workflow/parallel_test.ts @@ -10,13 +10,13 @@ import {InvocationContext} from '../../src/agents/invocation_context.js'; import {Event} from '../../src/events/event.js'; import {PluginManager} from '../../src/plugins/plugin_manager.js'; import {Session} from '../../src/sessions/session.js'; -import {BranchPath} from '../../src/workflow-next/branch_path.js'; -import {node} from '../../src/workflow-next/node.js'; -import {NodeContext} from '../../src/workflow-next/node_context.js'; -import {FunctionNode} from '../../src/workflow-next/nodes/function_node.js'; -import {ParallelWorker} from '../../src/workflow-next/nodes/parallel_worker.js'; -import {EventChannel} from '../../src/workflow-next/utils/event_channel.js'; -import {Workflow} from '../../src/workflow-next/workflow.js'; +import {BranchPath} 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 {EventChannel} from '../../src/workflow/utils/event_channel.js'; +import {Workflow} from '../../src/workflow/workflow.js'; function createIc(): InvocationContext { const session = { diff --git a/core/test/workflow-next/runner_integration_test.ts b/core/test/workflow/runner_integration_test.ts similarity index 91% rename from core/test/workflow-next/runner_integration_test.ts rename to core/test/workflow/runner_integration_test.ts index a1fa2a4b7..fa8742fb3 100644 --- a/core/test/workflow-next/runner_integration_test.ts +++ b/core/test/workflow/runner_integration_test.ts @@ -8,11 +8,11 @@ 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-next/graph.js'; -import {node} from '../../src/workflow-next/node.js'; -import {NodeContext} from '../../src/workflow-next/node_context.js'; -import {Workflow} from '../../src/workflow-next/workflow.js'; -import {WorkflowAgent} from '../../src/workflow-next/workflow_agent.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, diff --git a/core/test/workflow/workflow_agent_test.ts b/core/test/workflow/workflow_agent_test.ts deleted file mode 100644 index a3a786c2d..000000000 --- a/core/test/workflow/workflow_agent_test.ts +++ /dev/null @@ -1,141 +0,0 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import {describe, expect, it, vi} from 'vitest'; -import {BaseAgent} from '../../src/agents/base_agent.js'; -import {InvocationContext} from '../../src/agents/invocation_context.js'; - -import {PluginManager} from '../../src/plugins/plugin_manager.js'; -import {Session} from '../../src/sessions/session.js'; -import { - FunctionNode, - isWorkflow, - NodeStatus, - Workflow, -} from '../../src/workflow/index.js'; - -describe('Workflow Agent Orchestrator (`Workflow`)', () => { - function createTestContext( - params?: Partial, - ): InvocationContext { - const session: Session = { - id: 'session-wf', - appName: 'test-app', - userId: 'test-user', - events: [], - state: {}, - }; - - return new InvocationContext({ - invocationId: 'inv-wf', - session, - agent: { - name: 'mock_parent', - runAsync: async function* () {}, - } as unknown as BaseAgent, - pluginManager: new PluginManager(), - ...params, - }); - } - - it('should identify Workflow instances accurately via isWorkflow type guard', () => { - const wf = new Workflow({ - name: 'test_wf', - edges: [['START', new FunctionNode('a', () => 'a')]], - }); - - expect(isWorkflow(wf)).toBe(true); - expect(isWorkflow({name: 'not_wf'})).toBe(false); - }); - - it('should throw an error if both or neither of edges and dynamicEntry are defined', () => { - expect(() => new Workflow({name: 'empty_wf'})).toThrowError( - /must define either "edges"/i, - ); - - expect( - () => - new Workflow({ - name: 'conflicted_wf', - edges: [['START', new FunctionNode('a', () => 'a')]], - dynamicEntry: async () => 'b', - }), - ).toThrowError(/cannot have both "edges" and "dynamicEntry" defined/i); - }); - - it('should run a static graph Workflow from runAsync and mark endOfAgents upon completion', async () => { - const ctx = createTestContext(); - const nodeA = new FunctionNode('step_first', () => 'First'); - const nodeB = new FunctionNode( - 'step_second', - (_ctx, input: string) => `${input}_Second`, - ); - - const wf = new Workflow({ - name: 'static_wf', - edges: [['START', nodeA, nodeB]], - outputKey: 'wfResult', - }); - - for await (const _ of wf.runAsync(ctx)) { - /* consume events */ - } - - expect(ctx.agentStates['exec_node_step_first'].status).toBe( - NodeStatus.COMPLETED, - ); - expect(ctx.agentStates['exec_node_step_second'].status).toBe( - NodeStatus.COMPLETED, - ); - expect(ctx.agentStates['wfResult']).toEqual({ - step_first: 'First', - step_second: 'First_Second', - }); - expect(ctx.endOfAgents['static_wf']).toBe(true); - }); - - it('should run a dynamic Workflow from runAsync and mark endOfAgents upon completion', async () => { - const ctx = createTestContext(); - const spy = vi.fn(async (_ctx, input: number) => input * 5); - const dynNode = new FunctionNode('dyn_mul', spy); - - const wf = new Workflow({ - name: 'dynamic_wf', - dynamicEntry: dynNode, - outputKey: 'dynOut', - }); - - // Provide initial input via userContent text - ctx.userContent = {role: 'user', parts: [{text: '10'}]}; - - for await (const _ of wf.runAsync(ctx)) { - /* consume events */ - } - - expect(spy).toHaveBeenCalledTimes(1); - expect(spy.mock.calls[0][1]).toEqual({role: 'user', parts: [{text: '10'}]}); - expect(ctx.agentStates['dynOut']).toBeDefined(); - expect(ctx.endOfAgents['dynamic_wf']).toBe(true); - }); - - it('should skip execution if context.endOfAgents already marks the workflow as true', async () => { - const ctx = createTestContext(); - const spy = vi.fn(() => 'should_not_run'); - const wf = new Workflow({ - name: 'already_done_wf', - edges: [['START', new FunctionNode('step_never', spy)]], - }); - - ctx.endOfAgents['already_done_wf'] = true; - - for await (const _ of wf.runAsync(ctx)) { - /* consume events */ - } - - expect(spy).not.toHaveBeenCalled(); - expect(ctx.agentStates['exec_node_step_never']).toBeUndefined(); - }); -}); diff --git a/core/test/workflow-next/workflow_test.ts b/core/test/workflow/workflow_test.ts similarity index 94% rename from core/test/workflow-next/workflow_test.ts rename to core/test/workflow/workflow_test.ts index d6c78d7d5..81f1b611a 100644 --- a/core/test/workflow-next/workflow_test.ts +++ b/core/test/workflow/workflow_test.ts @@ -10,11 +10,11 @@ import {InvocationContext} from '../../src/agents/invocation_context.js'; import {createEvent, Event} from '../../src/events/event.js'; import {PluginManager} from '../../src/plugins/plugin_manager.js'; import {Session} from '../../src/sessions/session.js'; -import {BaseNode} from '../../src/workflow-next/base_node.js'; -import {DEFAULT_ROUTE} from '../../src/workflow-next/graph.js'; -import {NodeContext} from '../../src/workflow-next/node_context.js'; -import {EventChannel} from '../../src/workflow-next/utils/event_channel.js'; -import {Workflow} from '../../src/workflow-next/workflow.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 {EventChannel} from '../../src/workflow/utils/event_channel.js'; +import {Workflow} from '../../src/workflow/workflow.js'; function createIc(): InvocationContext { const session = { diff --git a/tests/integration/workflows/dynamic_nodes_workflow_test.ts b/tests/integration/workflows/dynamic_nodes_workflow_test.ts deleted file mode 100644 index 99b4b775e..000000000 --- a/tests/integration/workflows/dynamic_nodes_workflow_test.ts +++ /dev/null @@ -1,105 +0,0 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import { - createEvent, - Event, - FunctionNode, - InMemoryRunner, - InvocationContext, - Workflow, -} from '@google/adk'; -import {describe, expect, it} from 'vitest'; - -describe('Workflow Samples: Dynamic Nodes & Dynamic Fan-Out', () => { - it('should run a dynamic entry workflow scheduling downstream nodes via ctx.runNode (dynamic_nodes sample parity)', async () => { - const generatorNode = new FunctionNode( - 'generator', - (_ctx, topic: string) => `Catchy headline about ${topic}`, - ); - - const evaluatorNode = new FunctionNode( - 'evaluator', - (_ctx, headline: string) => `Evaluated: [${headline}] - Grade A`, - ); - - const dynamicEntryNode = new FunctionNode( - 'orchestrate', - async (ctx: InvocationContext, topic: string) => { - // Execute generator using ctx.runNode - const genOutput = await ctx.runNode(generatorNode, topic); - // Pass output to evaluator using ctx.runNode - const evalOutput = await ctx.runNode(evaluatorNode, genOutput); - return createEvent({message: `Final report: ${evalOutput}`}); - }, - ); - - const rootAgent = new Workflow({ - name: 'dynamic_nodes_workflow', - edges: [['START', dynamicEntryNode]], - }); - - const runner = new InMemoryRunner({agent: rootAgent}); - const events: Event[] = []; - - for await (const event of runner.runEphemeral({ - userId: 'test_user', - newMessage: {role: 'user', parts: [{text: 'AI Innovations'}]}, - })) { - events.push(event); - } - - const messages = events - .flatMap((e) => e.content?.parts?.map((p) => p.text) ?? []) - .join(''); - expect(messages).toContain( - 'Final report: Evaluated: [Catchy headline about AI Innovations] - Grade A', - ); - }); - - it('should perform dynamic fan-out and fan-in across items (dynamic_fan_out_fan_in sample parity)', async () => { - const processTopicNode = new FunctionNode( - 'process_topic', - (_ctx, topic: string) => `Processed: ${topic.trim().toUpperCase()}`, - ); - - const dynamicOrchestrator = new FunctionNode( - 'orchestrate_fan_out', - async (ctx: InvocationContext, input: string) => { - const topics = input.split(',').map((t) => t.trim()); - // Dynamic fan-out executing multiple nodes in parallel via Promise.all with ctx.runNode - const results = await Promise.all( - topics.map((topic) => ctx.runNode(processTopicNode, topic)), - ); - return createEvent({ - message: `Aggregated Topics: ${results.join(' | ')}`, - }); - }, - ); - - const rootAgent = new Workflow({ - name: 'dynamic_fanout_workflow', - edges: [['START', dynamicOrchestrator]], - }); - - const runner = new InMemoryRunner({agent: rootAgent}); - const events: Event[] = []; - - for await (const event of runner.runEphemeral({ - userId: 'test_user', - newMessage: {role: 'user', parts: [{text: 'apple, banana, cherry'}]}, - })) { - events.push(event); - } - - const messages = events - .flatMap((e) => e.content?.parts?.map((p) => p.text) ?? []) - .join(''); - expect(messages).toContain( - 'Aggregated Topics: Processed: APPLE | Processed: BANANA | Processed: CHERRY', - ); - }); -}); diff --git a/tests/integration/workflows/fan_out_fan_in_workflow_test.ts b/tests/integration/workflows/fan_out_fan_in_workflow_test.ts deleted file mode 100644 index 37b8ac0c7..000000000 --- a/tests/integration/workflows/fan_out_fan_in_workflow_test.ts +++ /dev/null @@ -1,77 +0,0 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import { - createEvent, - Event, - FunctionNode, - InMemoryRunner, - JoinNode, - Workflow, -} from '@google/adk'; -import {describe, expect, it} from 'vitest'; - -describe('Workflow Samples: Fan Out Fan In with JoinNode', () => { - it('should run parallel nodes and aggregate outputs at JoinNode (fan_out_fan_in sample parity)', async () => { - const makeUppercase = new FunctionNode( - 'make_uppercase', - (_ctx, input: string) => input.toUpperCase(), - ); - - const countCharacters = new FunctionNode( - 'count_characters', - (_ctx, input: string) => input.length, - ); - - const reverseString = new FunctionNode( - 'reverse_string', - (_ctx, input: string) => input.split('').reverse().join(''), - ); - - const joinNode = new JoinNode('join_for_results'); - - const aggregateNode = new FunctionNode( - 'aggregate', - (_ctx, input: Record) => { - return createEvent({ - message: - `Uppercase: ${input['make_uppercase']}\n\n` + - `Character Count: ${input['count_characters']}\n\n` + - `Reversed: ${input['reverse_string']}\n\n`, - }); - }, - ); - - const rootAgent = new Workflow({ - name: 'fan_out_fan_in_workflow', - edges: [ - [ - 'START', - [makeUppercase, countCharacters, reverseString], - joinNode, - aggregateNode, - ], - ], - }); - - const runner = new InMemoryRunner({agent: rootAgent}); - const events: Event[] = []; - - for await (const event of runner.runEphemeral({ - userId: 'test_user', - newMessage: {role: 'user', parts: [{text: 'adk workflow'}]}, - })) { - events.push(event); - } - - const messages = events - .flatMap((e) => e.content?.parts?.map((p) => p.text) ?? []) - .join(''); - expect(messages).toContain('Uppercase: ADK WORKFLOW'); - expect(messages).toContain('Character Count: 12'); - expect(messages).toContain('Reversed: wolfkrow kda'); - }); -}); diff --git a/tests/integration/workflows/loop_workflow_test.ts b/tests/integration/workflows/loop_workflow_test.ts deleted file mode 100644 index 90a6131fe..000000000 --- a/tests/integration/workflows/loop_workflow_test.ts +++ /dev/null @@ -1,121 +0,0 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import { - createEvent, - Event, - FunctionNode, - InMemoryRunner, - Workflow, -} from '@google/adk'; -import {describe, expect, it} from 'vitest'; - -describe('Workflow Samples: Loop & Loop Self', () => { - it('should support multi-node looping with conditional exit (loop sample parity)', async () => { - const processInput = new FunctionNode( - 'process_input', - (_ctx, input: string) => - createEvent({state: {topic: input, attempts: 0}}), - ); - - const generateHeadline = new FunctionNode('generate_headline', (ctx) => { - const attempts = ((ctx.session.state['attempts'] as number) || 0) + 1; - ctx.session.state['attempts'] = attempts; - const topic = ctx.session.state['topic'] as string; - const headline = - attempts === 1 - ? `General news about ${topic}` - : `Tech Breakthrough in ${topic}`; - return createEvent({state: {currentHeadline: headline}}); - }); - - const evaluateHeadline = new FunctionNode('evaluate_headline', (ctx) => { - const headline = ctx.session.state['currentHeadline'] as string; - const isTech = headline.includes('Tech'); - return {grade: isTech ? 'tech-related' : 'unrelated', headline}; - }); - - const routeHeadline = new FunctionNode( - 'route_headline', - (_ctx, input: {grade: string}) => createEvent({route: input.grade}), - ); - - const rootAgent = new Workflow({ - name: 'loop_workflow', - edges: [ - [ - 'START', - processInput, - generateHeadline, - evaluateHeadline, - routeHeadline, - ], - [routeHeadline, {unrelated: generateHeadline}], - ], - outputKey: 'loopResult', - allowCycles: true, - }); - - const runner = new InMemoryRunner({agent: rootAgent}); - const events: Event[] = []; - - for await (const event of runner.runEphemeral({ - userId: 'test_user', - newMessage: {role: 'user', parts: [{text: 'Software Engineering'}]}, - })) { - events.push(event); - } - - // It should have cycled: attempts should be 2 when it exits via "tech-related" (no route handler -> end) - expect(events.length).toBeGreaterThanOrEqual(1); - const finalEvent = events[events.length - 1]; - expect(finalEvent).toBeDefined(); - }); - - it('should support a node looping back to itself (loop_self sample parity)', async () => { - let guessCount = 0; - const guessNode = new FunctionNode('guess_node', () => { - guessCount++; - if (guessCount < 3) { - return createEvent({ - message: `Guess ${guessCount}: wrong`, - route: 'guessed_wrong', - }); - } - return createEvent({ - message: `Guess ${guessCount}: correct!`, - route: 'guessed_right', - }); - }); - - const rootAgent = new Workflow({ - name: 'loop_self_workflow', - edges: [ - ['START', guessNode], - [guessNode, {guessed_wrong: guessNode}], - ], - allowCycles: true, - }); - - const runner = new InMemoryRunner({agent: rootAgent}); - const events: Event[] = []; - - for await (const event of runner.runEphemeral({ - userId: 'test_user', - newMessage: {role: 'user', parts: [{text: 'Guess a number'}]}, - })) { - events.push(event); - } - - expect(guessCount).toBe(3); - const messages = events - .flatMap((e) => e.content?.parts?.map((p) => p.text) ?? []) - .filter(Boolean); - expect(messages).toContain('Guess 1: wrong'); - expect(messages).toContain('Guess 2: wrong'); - expect(messages).toContain('Guess 3: correct!'); - }); -}); diff --git a/tests/integration/workflows/nested_workflow_test.ts b/tests/integration/workflows/nested_workflow_test.ts deleted file mode 100644 index ff54e20e5..000000000 --- a/tests/integration/workflows/nested_workflow_test.ts +++ /dev/null @@ -1,88 +0,0 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import { - createEvent, - Event, - FunctionNode, - InMemoryRunner, - JoinNode, - Workflow, -} from '@google/adk'; -import {describe, expect, it} from 'vitest'; - -describe('Workflow Samples: Nested Workflow Composition', () => { - it('should compose a sub-Workflow inside a parent Workflow (nested_workflow sample parity)', async () => { - const findNameNode = new FunctionNode( - 'find_name', - (_ctx, input: string) => `Person_${input}`, - ); - - const generateBioNode = new FunctionNode( - 'generate_bio', - (_ctx, name: string) => `Bio for ${name}: Famous historical author.`, - ); - - // Sub-workflow wraps two sequential steps - const findFamousPersonWorkflow = new Workflow({ - name: 'find_famous_person_workflow', - edges: [['START', findNameNode, generateBioNode]], - outputKey: 'famousPersonResult', - }); - - const findHistoricalEventNode = new FunctionNode( - 'find_historical_event', - (_ctx, input: string) => - `Historical event in year ${input}: First publication of landmark novel.`, - ); - - const joinNode = new JoinNode('join_for_aggregation'); - - const formatOutputNode = new FunctionNode( - 'format_output', - (_ctx, input: Record) => { - // When a Workflow node completes inside a parent graph, its final output payload is passed to the join node - return createEvent({ - message: - `Person Bio: ${JSON.stringify(input['find_famous_person_workflow'])}\n` + - `Event: ${input['find_historical_event']}`, - }); - }, - ); - - const rootAgent = new Workflow({ - name: 'nested_root_workflow', - edges: [ - [ - 'START', - [findFamousPersonWorkflow, findHistoricalEventNode], - joinNode, - formatOutputNode, - ], - ], - }); - - const runner = new InMemoryRunner({agent: rootAgent}); - const events: Event[] = []; - - for await (const event of runner.runEphemeral({ - userId: 'test_user', - newMessage: {role: 'user', parts: [{text: '1984'}]}, - })) { - events.push(event); - } - - const messages = events - .flatMap((e) => e.content?.parts?.map((p) => p.text) ?? []) - .join(''); - expect(messages).toContain( - 'Bio for Person_1984: Famous historical author.', - ); - expect(messages).toContain( - 'Historical event in year 1984: First publication of landmark novel.', - ); - }); -}); diff --git a/tests/integration/workflows/node_as_tool_workflow_test.ts b/tests/integration/workflows/node_as_tool_workflow_test.ts deleted file mode 100644 index 2d7e4b716..000000000 --- a/tests/integration/workflows/node_as_tool_workflow_test.ts +++ /dev/null @@ -1,116 +0,0 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import { - AgentTool, - createEvent, - FunctionNode, - LlmAgent, - Workflow, -} from '@google/adk'; -import {describe, it} from 'vitest'; -import {runTestCase} from '../test_case_utils.js'; - -describe('Workflow Samples: Node / Workflow as Tool', () => { - it('should allow an LlmAgent to call a Workflow wrapped via AgentTool (node_as_tool sample parity)', async () => { - const lookupDbNode = new FunctionNode( - 'lookup_db', - (_ctx, customerId: string) => { - if (customerId === 'C123') { - return JSON.stringify({name: 'Jane Doe', tier: 'Gold', balance: 450}); - } - return JSON.stringify({error: 'Customer not found'}); - }, - ); - - const customerLookupWorkflow = new Workflow({ - name: 'customer_lookup_workflow', - edges: [['START', lookupDbNode]], - }); - - const workflowTool = new AgentTool({ - agent: customerLookupWorkflow, - }); - - const rootAgent = new LlmAgent({ - name: 'customer_service_agent', - instruction: - 'Use the lookup_customer tool when asked for customer details.', - tools: [workflowTool], - }); - - await runTestCase({ - agent: rootAgent, - turns: [ - { - userPrompt: 'Can you check details for customer C123?', - expectedEvents: [ - createEvent({ - author: 'customer_service_agent', - content: { - role: 'model', - parts: [ - { - functionCall: { - id: 'call_lookup_1', - name: 'lookup_customer', - args: {input: 'C123'}, - }, - }, - ], - }, - }), - createEvent({ - author: 'customer_service_agent', - content: { - role: 'model', - parts: [ - { - text: 'Customer C123 is Jane Doe with Gold tier and a balance of $450.', - }, - ], - }, - }), - ], - }, - ], - modelResponses: [ - { - candidates: [ - { - content: { - role: 'model', - parts: [ - { - functionCall: { - id: 'call_lookup_1', - name: 'lookup_customer', - args: {input: 'C123'}, - }, - }, - ], - }, - }, - ], - }, - { - candidates: [ - { - content: { - role: 'model', - parts: [ - { - text: 'Customer C123 is Jane Doe with Gold tier and a balance of $450.', - }, - ], - }, - }, - ], - }, - ], - }); - }); -}); diff --git a/tests/integration/workflows/route_workflow_test.ts b/tests/integration/workflows/route_workflow_test.ts deleted file mode 100644 index 6a92f20c9..000000000 --- a/tests/integration/workflows/route_workflow_test.ts +++ /dev/null @@ -1,207 +0,0 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import { - createEvent, - DEFAULT_ROUTE, - Event, - FunctionNode, - InMemoryRunner, - LlmAgent, - Workflow, -} from '@google/adk'; -import {describe, expect, it} from 'vitest'; -import {runTestCase} from '../test_case_utils.js'; - -describe('Workflow Samples: Route & Agent in Workflow with DEFAULT_ROUTE', () => { - it('should route conditionally based on event route (route sample parity)', async () => { - const classifyInputNode = new FunctionNode( - 'classify_input', - (_ctx, input: string) => { - const isQuestion = input.endsWith('?'); - return createEvent({ - route: isQuestion ? 'question' : 'statement', - }); - }, - ); - - const answerQuestionAgent = new LlmAgent({ - name: 'answer_question', - instruction: 'Answer the user question concisely.', - }); - - const commentOnStatementAgent = new LlmAgent({ - name: 'comment_on_statement', - instruction: 'Provide a brief comment on the statement.', - }); - - const rootAgent = new Workflow({ - name: 'router_workflow', - edges: [ - ['START', classifyInputNode], - [ - classifyInputNode, - { - question: answerQuestionAgent, - statement: commentOnStatementAgent, - }, - ], - ], - }); - - // Test question route - await runTestCase({ - agent: rootAgent, - turns: [ - { - userPrompt: 'What is ADK?', - expectedEvents: [ - createEvent({ - author: 'answer_question', - content: { - role: 'model', - parts: [{text: 'ADK is the Agent Development Kit.'}], - }, - }), - ], - }, - ], - modelResponses: [ - { - candidates: [ - { - content: { - role: 'model', - parts: [{text: 'ADK is the Agent Development Kit.'}], - }, - }, - ], - }, - ], - }); - - // Test statement route - await runTestCase({ - agent: rootAgent, - turns: [ - { - userPrompt: 'ADK supports workflows and routing.', - expectedEvents: [ - createEvent({ - author: 'comment_on_statement', - content: { - role: 'model', - parts: [{text: 'That is correct and very powerful!'}], - }, - }), - ], - }, - ], - modelResponses: [ - { - candidates: [ - { - content: { - role: 'model', - parts: [{text: 'That is correct and very powerful!'}], - }, - }, - ], - }, - ], - }); - }); - - it('should support DEFAULT_ROUTE fallback when specific route does not match (agent_in_workflow sample parity)', async () => { - const checkIdentityNode = new FunctionNode( - 'check_identity', - (_ctx, name: string) => { - if (name.toLowerCase() !== 'jane doe') { - return createEvent({ - message: `Could not find matching records for ${name}. Let's try again.`, - route: 'retry', - }); - } - return createEvent({ - message: `Hello ${name}! Let me look up your orders.`, - }); - }, - ); - - const retryHandlerNode = new FunctionNode('retry_handler', () => - createEvent({message: 'Retrying identity check...'}), - ); - - const generateInstructionAgent = new LlmAgent({ - name: 'generate_instruction', - instruction: 'Generate preparation instruction for Jane Doe.', - }); - - const rootAgent = new Workflow({ - name: 'agent_in_workflow', - edges: [ - ['START', checkIdentityNode], - [ - checkIdentityNode, - { - retry: retryHandlerNode, - [DEFAULT_ROUTE]: generateInstructionAgent, - }, - ], - ], - }); - - // Test when route="retry" matches specific route in routing table - const runnerRetry = new InMemoryRunner({agent: rootAgent}); - const retryEvents: Event[] = []; - for await (const event of runnerRetry.runEphemeral({ - userId: 'user1', - newMessage: {role: 'user', parts: [{text: 'John Smith'}]}, - })) { - retryEvents.push(event); - } - expect( - retryEvents.some( - (e) => e.content?.parts?.[0].text === 'Retrying identity check...', - ), - ).toBe(true); - - // Test when no route is yielded, taking DEFAULT_ROUTE fallback - await runTestCase({ - agent: rootAgent, - turns: [ - { - userPrompt: 'Jane Doe', - expectedEvents: [ - createEvent({ - author: 'generate_instruction', - content: { - role: 'model', - parts: [ - {text: 'Please fast for 12 hours before your lipid panel.'}, - ], - }, - }), - ], - }, - ], - modelResponses: [ - { - candidates: [ - { - content: { - role: 'model', - parts: [ - {text: 'Please fast for 12 hours before your lipid panel.'}, - ], - }, - }, - ], - }, - ], - }); - }); -}); diff --git a/tests/integration/workflows/sequence_workflow_test.ts b/tests/integration/workflows/sequence_workflow_test.ts deleted file mode 100644 index 0d06b84e1..000000000 --- a/tests/integration/workflows/sequence_workflow_test.ts +++ /dev/null @@ -1,157 +0,0 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import { - createEvent, - Event, - FunctionNode, - InMemoryRunner, - LlmAgent, - Workflow, -} from '@google/adk'; -import {describe, expect, it} from 'vitest'; -import {runTestCase} from '../test_case_utils.js'; - -describe('Workflow Samples: Sequence, Message & State', () => { - it('should run sequential workflow with LLM agents (sequence sample parity)', async () => { - const generateFruitAgent = new LlmAgent({ - name: 'generate_fruit_agent', - instruction: - 'Return the name of a random fruit. Return only the name, nothing else.', - }); - - const generateBenefitAgent = new LlmAgent({ - name: 'generate_benefit_agent', - instruction: 'Tell me a health benefit about the specified fruit.', - }); - - const rootAgent = new Workflow({ - name: 'root_agent', - edges: [['START', generateFruitAgent, generateBenefitAgent]], - }); - - await runTestCase({ - agent: rootAgent, - turns: [ - { - userPrompt: 'Tell me about a fruit.', - expectedEvents: [ - createEvent({ - author: 'generate_fruit_agent', - content: {role: 'model', parts: [{text: 'Apple'}]}, - }), - createEvent({ - author: 'generate_benefit_agent', - content: { - role: 'model', - parts: [{text: 'Apples are rich in fiber and vitamin C.'}], - }, - }), - ], - }, - ], - modelResponses: [ - {candidates: [{content: {role: 'model', parts: [{text: 'Apple'}]}}]}, - { - candidates: [ - { - content: { - role: 'model', - parts: [{text: 'Apples are rich in fiber and vitamin C.'}], - }, - }, - ], - }, - ], - }); - }); - - it('should emit event message from a FunctionNode (message sample parity)', async () => { - const messageNode = new FunctionNode('emit_message', () => - createEvent({ - content: { - role: 'model', - parts: [{text: 'Hello from workflow function node!'}], - }, - }), - ); - - const rootAgent = new Workflow({ - name: 'message_workflow', - edges: [['START', messageNode]], - }); - - const runner = new InMemoryRunner({agent: rootAgent}); - const events: Event[] = []; - - for await (const event of runner.runEphemeral({ - userId: 'test_user', - newMessage: {role: 'user', parts: [{text: 'Start'}]}, - })) { - events.push(event); - } - - expect(events.length).toBeGreaterThanOrEqual(1); - const msgEvents = events.filter((e) => - e.content?.parts?.some( - (p) => p.text === 'Hello from workflow function node!', - ), - ); - expect(msgEvents.length).toBeGreaterThanOrEqual(1); - }); - - it('should read and update session state across FunctionNodes (state sample parity)', async () => { - const initNode = new FunctionNode('init_state', (ctx, input: string) => { - ctx.session.state['topic'] = input; - ctx.session.state['count'] = 1; - return createEvent({ - actions: {stateDelta: {topic: input, count: 1}}, - content: { - role: 'model', - parts: [{text: `Initialized ${input}`}], - }, - }); - }); - - const updateNode = new FunctionNode('update_state', (ctx) => { - const currentCount = (ctx.session.state['count'] as number) || 0; - const topic = (ctx.session.state['topic'] as string) || ''; - const newCount = currentCount + 1; - ctx.session.state['count'] = newCount; - ctx.session.state['lastProcessed'] = `${topic}_processed`; - return createEvent({ - actions: { - stateDelta: {count: newCount, lastProcessed: `${topic}_processed`}, - }, - content: { - role: 'model', - parts: [{text: `Processed ${topic} with count ${newCount}`}], - }, - }); - }); - - const rootAgent = new Workflow({ - name: 'state_workflow', - edges: [['START', initNode, updateNode]], - }); - - const runner = new InMemoryRunner({agent: rootAgent}); - const events: Event[] = []; - - for await (const event of runner.runEphemeral({ - userId: 'test_user', - newMessage: {role: 'user', parts: [{text: 'AI Workflows'}]}, - })) { - events.push(event); - } - - expect(events.length).toBeGreaterThanOrEqual(2); - const lastEvent = events[events.length - 1]; - expect(lastEvent.content?.parts?.[0].text).toContain( - 'Processed AI Workflows with count 2', - ); - }); -}); From b7c488755e502c036c3ca0ddd773e8be0d18e4b3 Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Wed, 22 Jul 2026 14:57:58 -0700 Subject: [PATCH 14/41] feat(workflow): Phase 5b static-graph resume via session rehydration Reconstruct node state from prior session events (reconstructNodeStates): completed outputs/routes, raised interrupts, and interrupt responses resolved from user FunctionResponses. On resume, the Workflow merges resolved responses into ctx.resumeInputs (so waiting nodes resume) and fast-forwards already-completed nodes (cached output, no re-execution) while continuing the graph. Verified end-to-end via the Runner: an interrupted node resumes and upstream nodes are not re-run. (Replay sequence barrier for parallel/dynamic determinism, agent_state checkpoints, and the auth gate remain a 5b continuation.) --- core/src/workflow/utils/rehydration_utils.ts | 129 ++++++++++++++++ core/src/workflow/workflow.ts | 61 ++++++++ core/test/workflow/resume_test.ts | 151 +++++++++++++++++++ 3 files changed, 341 insertions(+) create mode 100644 core/src/workflow/utils/rehydration_utils.ts create mode 100644 core/test/workflow/resume_test.ts diff --git a/core/src/workflow/utils/rehydration_utils.ts b/core/src/workflow/utils/rehydration_utils.ts new file mode 100644 index 000000000..4a38de3a6 --- /dev/null +++ b/core/src/workflow/utils/rehydration_utils.ts @@ -0,0 +1,129 @@ +/** + * @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'; + +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 the node emitted, if any. */ + route?: RouteValue; + /** The branch the node ran on. */ + branch?: string; + /** 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). + */ +export function reconstructNodeStates( + events: Event[], +): 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. Key by the node path leaf (robust if the runtime + // rewrites the author), falling back to the author. + const nodeName = event.nodeInfo?.path + ? nodeNameFromPath(event.nodeInfo.path) + : event.author; + if (!nodeName) { + continue; + } + const node = getNode(nodeName); + if (event.output !== undefined) { + node.output = event.output; + node.branch = event.branch; + } + if (event.route !== undefined) { + node.route = event.route as RouteValue; + } + for (const id of event.longRunningToolIds ?? []) { + node.interruptIds.add(id); + interruptOwner.set(id, nodeName); + } + } + + 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; +} + +/** 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]; +} + +/** 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 index e9c6ed639..84bc63f2d 100644 --- a/core/src/workflow/workflow.ts +++ b/core/src/workflow/workflow.ts @@ -15,6 +15,11 @@ 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, + reconstructNodeStates, + RehydratedNode, +} from './utils/rehydration_utils.js'; /** * An imperative workflow entry point. Receives the workflow's node context and @@ -56,6 +61,8 @@ class LoopState { 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; } @@ -110,12 +117,19 @@ export class Workflow extends BaseNode { const dynamicState = new DynamicNodeState(); ctx.scheduler = new DynamicNodeScheduler(dynamicState); + // --- REHYDRATE (resume) --- + // Reconstruct node state from prior session events and surface resolved + // interrupt responses so waiting nodes can resume. + const rehydrated = reconstructNodeStates(ctx.session?.events ?? []); + this.applyResumeInputs(ctx, rehydrated); + if (this.dynamicEntry) { await this.runDynamicEntry(ctx, nodeInput, dynamicState); return; } const loop = new LoopState(); + loop.rehydrated = rehydrated; // --- SETUP --- this.seedStartTriggers(loop, nodeInput); @@ -156,6 +170,22 @@ export class Workflow extends BaseNode { } } + /** + * 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 { @@ -258,6 +288,20 @@ export class Workflow extends BaseNode { 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: makeFastForwardContext(ctx, prior), + }), + ); + return; + } + let runId = nodeState.runId; if (!runId) { nodeState.runCounter += 1; @@ -448,3 +492,20 @@ export class Workflow extends BaseNode { await Promise.allSettled(outstanding); } } + +/** + * Builds a minimal completion result for a fast-forwarded (cached) node on + * resume. Only the fields read by `handleCompletion` are populated; the node's + * events are NOT re-emitted (they already exist in the session). + */ +function makeFastForwardContext( + parent: NodeContext, + prior: RehydratedNode, +): NodeContext { + return { + output: prior.output, + route: prior.route, + branch: prior.branch ?? parent.branch, + interruptIds: [], + } as unknown as NodeContext; +} diff --git a/core/test/workflow/resume_test.ts b/core/test/workflow/resume_test.ts new file mode 100644 index 000000000..00cf17710 --- /dev/null +++ b/core/test/workflow/resume_test.ts @@ -0,0 +1,151 @@ +/** + * @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 {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', + ); + }); +}); + +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'}, + ); + 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'}, + ); + 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); + }); +}); From 457202059ddf13abf4736082bb9e57a3d1d063a0 Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Thu, 23 Jul 2026 12:20:11 -0700 Subject: [PATCH 15/41] feat(workflow): Phase 5b-cont dynamic (ctx.runNode) resume + dedup Extend resume to dynamically-scheduled nodes. The DynamicNodeScheduler now assigns run-id-qualified node paths (wf.node@runId, via a new overrideNodePath option on executeChildNode) so distinct ctx.runNode iterations are distinguishable across turns. On resume it rehydrates each run from session events (reconstructNodeStatesByPath): completed runs are fast-forwarded with their cached output (no re-execution), while waiting runs fall through to a fresh run that reads the resolved resumeInputs merged by the Workflow. Verified end-to-end via the Runner: a completed pre-interrupt dynamic node is not re-run and the waiting node resumes. (Replay sequence barrier for concurrent determinism, agent_state checkpoints, and the auth gate remain.) --- core/src/workflow/dynamic_node_scheduler.ts | 36 ++++++- core/src/workflow/node_runner.ts | 12 ++- core/src/workflow/utils/rehydration_utils.ts | 53 ++++++++-- core/src/workflow/workflow.ts | 18 +--- core/test/workflow/dynamic_resume_test.ts | 100 +++++++++++++++++++ 5 files changed, 188 insertions(+), 31 deletions(-) create mode 100644 core/test/workflow/dynamic_resume_test.ts diff --git a/core/src/workflow/dynamic_node_scheduler.ts b/core/src/workflow/dynamic_node_scheduler.ts index c51c391da..415d9662a 100644 --- a/core/src/workflow/dynamic_node_scheduler.ts +++ b/core/src/workflow/dynamic_node_scheduler.ts @@ -15,6 +15,11 @@ import { ScheduleDynamicNode, ScheduleDynamicNodeOptions, } from './schedule_dynamic_node.js'; +import { + isFastForwardable, + makeFastForwardContext, + reconstructNodeStatesByPath, +} from './utils/rehydration_utils.js'; /** * Handles `ctx.runNode()` calls for a {@link Workflow} subtree. @@ -36,7 +41,9 @@ export class DynamicNodeScheduler implements ScheduleDynamicNode { ): Promise { const name = options.nodeName ?? node.name; const runId = options.runId; - const nodePath = `${ctx.nodePath}/${name}@${runId}`; + const nodePath = ctx.nodePath + ? `${ctx.nodePath}.${name}@${runId}` + : `${name}@${runId}`; const existing = this.state.runs.get(nodePath); if (existing?.task) { @@ -44,8 +51,30 @@ export class DynamicNodeScheduler implements ScheduleDynamicNode { return existing.task; } - // TODO(phase-5): lazy rehydration from session events + replay - // interception (dedup completed / resume waiting runs) goes here. + // 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 makeFastForwardContext(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); } @@ -72,6 +101,7 @@ export class DynamicNodeScheduler implements ScheduleDynamicNode { run.task = executeChildNode(ctx, node, input, { nodeName: name, runId, + overrideNodePath: nodePath, useAsOutput: options.useAsOutput, useSubBranch: options.useSubBranch, overrideBranch: options.overrideBranch, diff --git a/core/src/workflow/node_runner.ts b/core/src/workflow/node_runner.ts index 2661484e7..528091eb2 100644 --- a/core/src/workflow/node_runner.ts +++ b/core/src/workflow/node_runner.ts @@ -33,6 +33,12 @@ export interface RunNodeOptions { overrideBranch?: string; /** Explicit isolation scope, overriding inheritance from the parent. */ overrideIsolationScope?: string; + /** + * Explicit node path for the child (used by the dynamic scheduler to embed + * the run id, e.g. `wf.node@1`, so distinct runs are distinguishable on + * resume). Defaults to `${parent.nodePath}.${nodeName}`. + */ + overrideNodePath?: string; } /** @@ -51,9 +57,9 @@ export async function executeChildNode( ): Promise { const nodeName = options.nodeName ?? node.name; const runId = options.runId ?? nodeName; - const nodePath = parent.nodePath - ? `${parent.nodePath}.${nodeName}` - : nodeName; + const nodePath = + options.overrideNodePath ?? + (parent.nodePath ? `${parent.nodePath}.${nodeName}` : nodeName); let branch = parent.branch; if (options.overrideBranch !== undefined) { diff --git a/core/src/workflow/utils/rehydration_utils.ts b/core/src/workflow/utils/rehydration_utils.ts index 4a38de3a6..0ee0b307d 100644 --- a/core/src/workflow/utils/rehydration_utils.ts +++ b/core/src/workflow/utils/rehydration_utils.ts @@ -15,6 +15,7 @@ import {Event} from '../../events/event.js'; import {RouteValue} from '../graph.js'; +import type {NodeContext} from '../node_context.js'; const RESULT_KEY = 'result'; @@ -38,6 +39,28 @@ export interface RehydratedNode { */ export function reconstructNodeStates( events: Event[], +): Map { + // Key static-graph nodes by their name (path leaf), robust to author rewrite. + 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(); @@ -67,15 +90,12 @@ export function reconstructNodeStates( continue; } - // 2. Node events. Key by the node path leaf (robust if the runtime - // rewrites the author), falling back to the author. - const nodeName = event.nodeInfo?.path - ? nodeNameFromPath(event.nodeInfo.path) - : event.author; - if (!nodeName) { + // 2. Node events. + const key = keyFor(event); + if (!key) { continue; } - const node = getNode(nodeName); + const node = getNode(key); if (event.output !== undefined) { node.output = event.output; node.branch = event.branch; @@ -85,7 +105,7 @@ export function reconstructNodeStates( } for (const id of event.longRunningToolIds ?? []) { node.interruptIds.add(id); - interruptOwner.set(id, nodeName); + interruptOwner.set(id, key); } } @@ -108,6 +128,23 @@ export function isFastForwardable(node: RehydratedNode): boolean { return true; } +/** + * Builds a minimal completion result for a fast-forwarded (cached) node on + * resume. Only the fields read by completion handling are populated; the node's + * events are NOT re-emitted (they already exist in the session). + */ +export function makeFastForwardContext( + parent: NodeContext, + prior: RehydratedNode, +): NodeContext { + return { + output: prior.output, + route: prior.route, + branch: prior.branch ?? parent.branch, + interruptIds: [], + } as unknown as NodeContext; +} + /** 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; diff --git a/core/src/workflow/workflow.ts b/core/src/workflow/workflow.ts index 84bc63f2d..b7e4c9aa6 100644 --- a/core/src/workflow/workflow.ts +++ b/core/src/workflow/workflow.ts @@ -17,6 +17,7 @@ import {DynamicNodeState} from './schedule_dynamic_node.js'; import {Trigger} from './trigger.js'; import { isFastForwardable, + makeFastForwardContext, reconstructNodeStates, RehydratedNode, } from './utils/rehydration_utils.js'; @@ -492,20 +493,3 @@ export class Workflow extends BaseNode { await Promise.allSettled(outstanding); } } - -/** - * Builds a minimal completion result for a fast-forwarded (cached) node on - * resume. Only the fields read by `handleCompletion` are populated; the node's - * events are NOT re-emitted (they already exist in the session). - */ -function makeFastForwardContext( - parent: NodeContext, - prior: RehydratedNode, -): NodeContext { - return { - output: prior.output, - route: prior.route, - branch: prior.branch ?? parent.branch, - interruptIds: [], - } as unknown as NodeContext; -} 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); + }); +}); From 4645c11d7e131c4eba085752b9e77715d88704be Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Thu, 23 Jul 2026 12:28:55 -0700 Subject: [PATCH 16/41] feat(workflow): Phase 5b-cont auth gate on FunctionNode Add an authentication gate to FunctionNode. When a node has an authConfig and no credential is present, it emits an adk_request_credential interrupt (createAuthRequestEvent) and pauses. On resume, the credential supplied via ctx.resumeInputs is stored into session state (processAuthResume + AuthHandler.parseAndStoreAuthResponse) and the node proceeds. hasAuthCredential short-circuits when a credential already exists. The credentialKey doubles as the deterministic interrupt id so the resume response matches across turns. Verified end-to-end via the Runner with an API-key scheme (request -> resume -> run) and directly for the pre-existing-credential case. --- core/src/workflow/nodes/function_node.ts | 40 ++++- core/src/workflow/utils/hitl_utils.ts | 116 ++++++++++++++- core/test/workflow/auth_gate_test.ts | 180 +++++++++++++++++++++++ 3 files changed, 333 insertions(+), 3 deletions(-) create mode 100644 core/test/workflow/auth_gate_test.ts diff --git a/core/src/workflow/nodes/function_node.ts b/core/src/workflow/nodes/function_node.ts index b250ab25c..40596f899 100644 --- a/core/src/workflow/nodes/function_node.ts +++ b/core/src/workflow/nodes/function_node.ts @@ -8,6 +8,11 @@ import {AuthConfig} from '../../auth/auth_tool.js'; import {createEvent, Event, isEvent} from '../../events/event.js'; import {BaseNode, BaseNodeConfig, isContent, toContent} from '../base_node.js'; import {NodeContext} from '../node_context.js'; +import { + createAuthRequestEvent, + hasAuthCredential, + processAuthResume, +} from '../utils/hitl_utils.js'; /** * A value a {@link FunctionNodeHandler} may return or yield. @@ -85,7 +90,15 @@ export class FunctionNode extends BaseNode< ctx: NodeContext, input: TInput, ): AsyncGenerator { - // TODO(phase-5): auth gate (authConfig -> adk_request_credential interrupt). + // Auth gate: request credentials (and interrupt) if not yet available. + if (this.authConfig) { + const authRequest = await this.runAuthGate(ctx); + if (authRequest) { + yield authRequest; + return; + } + } + const result = this.handler(ctx, input); if (isAsyncIterable(result)) { @@ -102,6 +115,31 @@ export class FunctionNode extends BaseNode< } } + /** + * Ensures a credential for `authConfig` is available. Returns an + * `adk_request_credential` interrupt event if the credential must be + * requested from the user, or `undefined` if the node may proceed. + * + * On resume, a credential provided via `ctx.resumeInputs[credentialKey]` is + * stored into state before re-checking. + */ + private async runAuthGate(ctx: NodeContext): Promise { + const authConfig = this.authConfig!; + if (hasAuthCredential(authConfig, ctx.state)) { + return undefined; + } + const resumeResponse = ctx.resumeInputs[authConfig.credentialKey]; + if (resumeResponse !== undefined) { + await processAuthResume(resumeResponse, authConfig, ctx.state); + if (hasAuthCredential(authConfig, ctx.state)) { + return undefined; + } + } + // The credential key doubles as a deterministic interrupt id so the resume + // response matches across turns. + return createAuthRequestEvent(authConfig, authConfig.credentialKey); + } + protected override toEvent(ctx: NodeContext, data: unknown): Event | null { const stateDelta = Object.keys(ctx.actions.stateDelta).length > 0 diff --git a/core/src/workflow/utils/hitl_utils.ts b/core/src/workflow/utils/hitl_utils.ts index 4f5686077..938c321d4 100644 --- a/core/src/workflow/utils/hitl_utils.ts +++ b/core/src/workflow/utils/hitl_utils.ts @@ -8,13 +8,19 @@ * Utilities for Human-in-the-Loop (HITL) workflows. * * Ported (subset) from `google/adk-python` - * `workflow/utils/_workflow_hitl_utils.py`. The auth-credential helpers are - * added in the Phase 5 auth-gate follow-up. + * `workflow/utils/_workflow_hitl_utils.py`. */ import {Part} from '@google/genai'; import {z} from 'zod'; +import { + AuthCredential, + AuthCredentialTypes, +} from '../../auth/auth_credential.js'; +import {AuthHandler} from '../../auth/auth_handler.js'; +import {AuthConfig} from '../../auth/auth_tool.js'; import {createEvent, Event} from '../../events/event.js'; +import {State} from '../../sessions/state.js'; import {RequestInput} from '../request_input.js'; /** Function-call name marking a request-for-input interrupt. */ @@ -97,3 +103,109 @@ export function createRequestInputResponse( }, }; } + +// --------------------------------------------------------------------------- +// Auth-credential utilities (auth gate) +// --------------------------------------------------------------------------- + +/** Whether a credential for the given auth config already exists in state. */ +export function hasAuthCredential( + authConfig: AuthConfig, + state: State, +): boolean { + return new AuthHandler(authConfig).getAuthResponse(state) !== undefined; +} + +/** + * Creates an event requesting user authentication credentials + * (`adk_request_credential`), marking the interrupt id as a long-running tool. + * + * Ported from `google/adk-python` `create_auth_request_event`. + */ +export function createAuthRequestEvent( + authConfig: AuthConfig, + interruptId: string, +): Event { + const authRequest = new AuthHandler(authConfig).generateAuthRequest(); + const args: Record = { + functionCallId: interruptId, + authConfig: authRequest, + message: buildAuthMessage(authConfig), + }; + return createEvent({ + content: { + role: 'model', + parts: [ + { + functionCall: { + name: REQUEST_CREDENTIAL_FUNCTION_CALL_NAME, + id: interruptId, + args, + }, + }, + ], + }, + longRunningToolIds: [interruptId], + }); +} + +/** + * Stores credentials from an auth resume response into session state. Accepts a + * full {@link AuthConfig} (web UI flow) or a plain value (e.g. an API key + * string), mirroring `google/adk-python` `process_auth_resume`. + */ +export async function processAuthResume( + responseData: unknown, + authConfig: AuthConfig, + state: State, +): Promise { + let responseConfig: AuthConfig; + if (isAuthConfigLike(responseData)) { + responseConfig = { + ...(responseData as AuthConfig), + credentialKey: authConfig.credentialKey, + }; + } else { + responseConfig = { + ...authConfig, + exchangedAuthCredential: buildCredentialFromValue( + authConfig, + responseData, + ), + }; + } + await new AuthHandler(responseConfig).parseAndStoreAuthResponse(state); +} + +function isAuthConfigLike(value: unknown): value is AuthConfig { + return ( + typeof value === 'object' && + value !== null && + 'authScheme' in value && + 'credentialKey' in value + ); +} + +function buildCredentialFromValue( + authConfig: AuthConfig, + value: unknown, +): AuthCredential { + if (authConfig.rawAuthCredential?.authType === AuthCredentialTypes.API_KEY) { + return {authType: AuthCredentialTypes.API_KEY, apiKey: String(value)}; + } + return value as AuthCredential; +} + +function buildAuthMessage(authConfig: AuthConfig): string { + const authType = authConfig.rawAuthCredential?.authType; + if (authType === AuthCredentialTypes.API_KEY) { + return 'Please provide your API key.'; + } + if ( + authType === AuthCredentialTypes.OAUTH2 || + authType === AuthCredentialTypes.OPEN_ID_CONNECT + ) { + return 'Please complete the authentication flow.'; + } + return 'Please provide your authentication credentials.'; +} diff --git a/core/test/workflow/auth_gate_test.ts b/core/test/workflow/auth_gate_test.ts new file mode 100644 index 000000000..32878d015 --- /dev/null +++ b/core/test/workflow/auth_gate_test.ts @@ -0,0 +1,180 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {describe, expect, it} from 'vitest'; +import {BaseAgent} from '../../src/agents/base_agent.js'; +import {InvocationContext} from '../../src/agents/invocation_context.js'; +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 {PluginManager} from '../../src/plugins/plugin_manager.js'; +import {Runner} from '../../src/runner/runner.js'; +import {InMemorySessionService} from '../../src/sessions/in_memory_session_service.js'; +import {Session} from '../../src/sessions/session.js'; +import {NodeContext} from '../../src/workflow/node_context.js'; +import {FunctionNode} from '../../src/workflow/nodes/function_node.js'; +import {EventChannel} from '../../src/workflow/utils/event_channel.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'; + +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; + + 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()}, + ); + + 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 session = { + id: 's1', + appName: 'app', + userId: 'u', + events: [], + state: { + ['temp:' + CREDENTIAL_KEY]: { + authType: AuthCredentialTypes.API_KEY, + apiKey: 'pre-existing', + }, + }, + lastUpdateTime: Date.now(), + } as unknown as Session; + const ic = new InvocationContext({ + invocationId: 'inv-1', + session, + agent: { + name: 'wf', + runAsync: async function* () {}, + } as unknown as BaseAgent, + pluginManager: new PluginManager(), + }); + + const channel = new EventChannel(); + const root = new NodeContext({ + invocationContext: ic, + channel, + nodePath: '', + runId: 'root', + }); + const events: Event[] = []; + const settle = root.runNode(wf, 'go', {useAsOutput: true}).then( + () => channel.close(), + (err) => channel.fail(err), + ); + for await (const e of channel) { + events.push(e); + } + await settle; + + expect(runs).toBe(1); + expect(root.output).toBe('ok'); + expect(events.some(hasAuthRequestFunctionCall)).toBe(false); + }); +}); From 8908f12bd1a4db4c4ff3dd0fc0d2d744908ae126 Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Thu, 23 Jul 2026 12:32:31 -0700 Subject: [PATCH 17/41] feat(workflow): Phase 7b multi-agent hand-off (transfer_to_agent) LLMAgentWrapper now follows transfer_to_agent hand-offs: when a wrapped agent emits a transfer action, the target is resolved in the agent tree (rootAgent.findAgent) and execution continues with it, chaining until an agent responds without transferring (depth-capped). Combined with the existing node-as-tool pattern (ctx.runNode(agent)), this covers multi-agent orchestration in workflows. Full Python chat/task delegation (FinishTaskTool, task-agent tools, isolation-scope content filtering, LlmAgent.mode) is out of scope here: it requires agents-package infrastructure that does not yet exist in adk-js. --- core/src/workflow/nodes/llm_agent_wrapper.ts | 43 +++++- core/test/workflow/multi_agent_test.ts | 150 +++++++++++++++++++ 2 files changed, 191 insertions(+), 2 deletions(-) create mode 100644 core/test/workflow/multi_agent_test.ts diff --git a/core/src/workflow/nodes/llm_agent_wrapper.ts b/core/src/workflow/nodes/llm_agent_wrapper.ts index 4af1aba00..d76088701 100644 --- a/core/src/workflow/nodes/llm_agent_wrapper.ts +++ b/core/src/workflow/nodes/llm_agent_wrapper.ts @@ -10,6 +10,9 @@ import {createEvent, Event} from '../../events/event.js'; import {BaseNode, BaseNodeConfig, isContent} from '../base_node.js'; import {NodeContext} from '../node_context.js'; +/** Safety cap on chained `transfer_to_agent` hand-offs. */ +const MAX_TRANSFER_DEPTH = 10; + /** Options for an {@link LLMAgentWrapper}. */ export interface LLMAgentWrapperConfig extends Partial< Omit @@ -56,10 +59,46 @@ export class LLMAgentWrapper extends BaseNode { ctx.session.events.push(userEvent); } - // Run the agent under the node's invocation context (it sets agent=itself). - for await (const event of this.agent.runAsync(ctx.invocationContext)) { + // Run the agent, following any transfer_to_agent hand-offs to peers. + yield* this.runWithTransfers(ctx, this.agent, 0); + } + + /** + * Runs `agent`; if it emits a `transfer_to_agent` action, resolves the target + * in the agent tree and continues with it (multi-agent hand-off). This is the + * portable slice of Python's chat mode; autonomous task delegation + * (FinishTaskTool / task tools / isolation scopes) is not yet supported. + */ + private async *runWithTransfers( + ctx: NodeContext, + agent: BaseAgent, + depth: number, + ): AsyncGenerator { + if (depth > MAX_TRANSFER_DEPTH) { + throw new Error( + `LLMAgentWrapper: transfer_to_agent depth exceeded ${MAX_TRANSFER_DEPTH} ` + + `(possible transfer loop starting at '${this.agent.name}').`, + ); + } + + let transferTarget: string | undefined; + for await (const event of agent.runAsync(ctx.invocationContext)) { this.maybeSetOutput(event); yield event; + if (event.actions?.transferToAgent) { + transferTarget = event.actions.transferToAgent; + break; + } + } + + if (transferTarget) { + const target = agent.rootAgent.findAgent(transferTarget); + if (!target) { + throw new Error( + `LLMAgentWrapper: transfer target agent '${transferTarget}' not found.`, + ); + } + yield* this.runWithTransfers(ctx, target, depth + 1); } } diff --git a/core/test/workflow/multi_agent_test.ts b/core/test/workflow/multi_agent_test.ts new file mode 100644 index 000000000..6d18869b0 --- /dev/null +++ b/core/test/workflow/multi_agent_test.ts @@ -0,0 +1,150 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {describe, expect, it} from 'vitest'; +import {BaseAgent} from '../../src/agents/base_agent.js'; +import {InvocationContext} from '../../src/agents/invocation_context.js'; +import {createEvent, Event} from '../../src/events/event.js'; +import {PluginManager} from '../../src/plugins/plugin_manager.js'; +import {Session} from '../../src/sessions/session.js'; +import {node} from '../../src/workflow/node.js'; +import {NodeContext} from '../../src/workflow/node_context.js'; +import {EventChannel} from '../../src/workflow/utils/event_channel.js'; +import {Workflow} from '../../src/workflow/workflow.js'; + +function createIc(): InvocationContext { + const session = { + id: 's1', + appName: 'app', + userId: 'u', + events: [], + state: {}, + lastUpdateTime: Date.now(), + } as unknown as Session; + return new InvocationContext({ + invocationId: 'inv-1', + session, + agent: { + name: 'wf', + runAsync: async function* () {}, + } as unknown as BaseAgent, + pluginManager: new PluginManager(), + }); +} + +async function driveWorkflow( + wf: Workflow, + input?: unknown, +): Promise<{output: unknown; events: Event[]}> { + const channel = new EventChannel(); + const root = new NodeContext({ + invocationContext: createIc(), + channel, + nodePath: '', + runId: 'root', + }); + const events: Event[] = []; + const settle = root.runNode(wf, input, {useAsOutput: true}).then( + () => channel.close(), + (err) => channel.fail(err), + ); + for await (const e of channel) { + events.push(e); + } + await settle; + return {output: root.output, events}; +} + +/** A fake agent that emits a fixed model text and optionally transfers. */ +class ScriptedAgent extends BaseAgent { + constructor( + name: string, + private readonly text: string, + private readonly transferTo?: string, + subAgents: BaseAgent[] = [], + ) { + super({name, subAgents}); + } + protected async *runAsyncImpl( + ctx: InvocationContext, + ): AsyncGenerator { + if (this.transferTo) { + yield createEvent({ + author: this.name, + invocationId: ctx.invocationId, + branch: ctx.branch, + actions: {transferToAgent: this.transferTo}, + }); + return; + } + yield createEvent({ + author: this.name, + invocationId: ctx.invocationId, + branch: ctx.branch, + content: {role: 'model', parts: [{text: this.text}]}, + }); + } + // eslint-disable-next-line require-yield + protected async *runLiveImpl(): AsyncGenerator { + return; + } +} + +describe('Phase 7b — multi-agent hand-off (transfer_to_agent)', () => { + it('follows a transfer to a peer agent and uses its output', async () => { + const specialist = new ScriptedAgent('specialist', 'specialist-answer'); + const coordinator = new ScriptedAgent( + 'coordinator', + '(unused)', + 'specialist', + [specialist], + ); + + const wf = new Workflow({ + name: 'transfer_wf', + edges: [['START', coordinator]], + }); + const {output, events} = await driveWorkflow(wf, 'question'); + + expect(output).toBe('specialist-answer'); + // Both the coordinator's transfer event and the specialist's answer stream. + expect( + events.some((e) => e.actions?.transferToAgent === 'specialist'), + ).toBe(true); + expect(events.some((e) => e.author === 'specialist')).toBe(true); + }); + + it('follows a chain of transfers', async () => { + const c = new ScriptedAgent('c_agent', 'final'); + const b = new ScriptedAgent('b_agent', '(unused)', 'c_agent', [c]); + const a = new ScriptedAgent('a_agent', '(unused)', 'b_agent', [b]); + + const wf = new Workflow({name: 'chain_wf', edges: [['START', a]]}); + expect((await driveWorkflow(wf, 'x')).output).toBe('final'); + }); +}); + +describe('Phase 7b — multi-agent orchestration via ctx.runNode', () => { + it('coordinates specialist agents imperatively (node-as-tool)', async () => { + const researcher = new ScriptedAgent('researcher', 'facts'); + const writer = new ScriptedAgent('writer', 'report'); + + // Idiomatic TS multi-agent: a coordinator drives sub-agents via runNode. + const wf = new Workflow({ + name: 'coordinator_wf', + dynamicEntry: async (ctx, input) => { + const research = await ctx.runNode(node(researcher), input); + const draft = await ctx.runNode(node(writer), research.output); + return {research: research.output, draft: draft.output}; + }, + }); + + expect(await driveWorkflow(wf, 'topic').then((r) => r.output)).toEqual({ + research: 'facts', + draft: 'report', + }); + }); +}); From 020729c353fc3ab24348599a00d5d37b623e158b Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Thu, 23 Jul 2026 15:57:32 -0700 Subject: [PATCH 18/41] test(workflow): add integration tests for major workflow use cases Port the major google/adk-python contributing/samples/workflows scenarios as end-to-end integration tests running through the real Runner: sequence, route (LLM classifier + output_schema), fan-out/fan-in, dynamic loop, parallel worker, state sharing, nested workflow, node-as-tool, retry, and request_input (HITL). LLM-backed samples mock model responses from JSON fixtures via GeminiWithMockResponses (the existing integration-test pattern). Also promote an LlmAgent node's structured output to a parsed object when it declares an outputSchema, matching Python. --- core/src/workflow/nodes/llm_agent_wrapper.ts | 15 +- .../workflows/core_workflows_test.ts | 249 ++++++++++++++++++ .../workflows/route.model_responses.json | 28 ++ tests/integration/workflows/route_llm_test.ts | 106 ++++++++ .../workflows/sequence.model_responses.json | 25 ++ .../workflows/sequence_llm_test.ts | 60 +++++ .../workflows/workflow_test_utils.ts | 107 ++++++++ 7 files changed, 589 insertions(+), 1 deletion(-) create mode 100644 tests/integration/workflows/core_workflows_test.ts create mode 100644 tests/integration/workflows/route.model_responses.json create mode 100644 tests/integration/workflows/route_llm_test.ts create mode 100644 tests/integration/workflows/sequence.model_responses.json create mode 100644 tests/integration/workflows/sequence_llm_test.ts create mode 100644 tests/integration/workflows/workflow_test_utils.ts diff --git a/core/src/workflow/nodes/llm_agent_wrapper.ts b/core/src/workflow/nodes/llm_agent_wrapper.ts index d76088701..fec88b06d 100644 --- a/core/src/workflow/nodes/llm_agent_wrapper.ts +++ b/core/src/workflow/nodes/llm_agent_wrapper.ts @@ -122,7 +122,20 @@ export class LLMAgentWrapper extends BaseNode { .map((p) => p.text) .join(''); - event.output = text; + // If the agent declares an output schema, its text is structured JSON; + // surface the parsed object as the node output (matching Python). + let output: unknown = text; + const hasOutputSchema = !!(this.agent as {outputSchema?: unknown}) + .outputSchema; + if (hasOutputSchema && text.trim()) { + try { + output = JSON.parse(text); + } catch { + output = text; + } + } + + event.output = output; event.nodeInfo = {...(event.nodeInfo ?? {}), messageAsOutput: true}; } } diff --git a/tests/integration/workflows/core_workflows_test.ts b/tests/integration/workflows/core_workflows_test.ts new file mode 100644 index 000000000..2d56100c1 --- /dev/null +++ b/tests/integration/workflows/core_workflows_test.ts @@ -0,0 +1,249 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Integration tests for the major (non-LLM) workflow use cases, mirroring the + * Python `contributing/samples/workflows` samples, run end-to-end through the + * real Runner. Workflows are started with a text prompt (as a user would), so + * structured inputs are produced inside nodes rather than passed as the prompt. + */ + +import { + createEvent, + DEFAULT_ROUTE, + FunctionNode, + JoinNode, + node, + NodeContext, + ParallelWorker, + RequestInput, + Workflow, +} from '@google/adk'; +import {describe, expect, it} from 'vitest'; +import { + collect, + createWorkflowRunner, + finalOutput, + runWorkflowOnce, +} from './workflow_test_utils.js'; + +describe('workflow integration — sequence', () => { + it('threads input through a linear chain', async () => { + const a = node((_c: NodeContext, i: string) => `${i}->A`, {name: 'a'}); + const b = node((_c: NodeContext, i: string) => `${i}->B`, {name: 'b'}); + const c = node((_c: NodeContext, i: string) => `${i}->C`, {name: 'c'}); + const wf = new Workflow({name: 'sequence', edges: [['START', a, b, c]]}); + + const events = await runWorkflowOnce(wf, 'INIT'); + expect(finalOutput(events)).toBe('INIT->A->B->C'); + }); +}); + +describe('workflow integration — route', () => { + it('routes to a branch and falls back to DEFAULT_ROUTE', async () => { + const routeNode = node( + (_c: NodeContext, input: string) => + createEvent( + input === 'jane' ? {output: input} : {route: 'retry', output: input}, + ), + {name: 'route_node'}, + ); + const retry = node((_c: NodeContext, i: string) => `RETRY:${i}`, { + name: 'retry_branch', + }); + const gen = node((_c: NodeContext, i: string) => `GEN:${i}`, { + name: 'generate', + }); + const wf = new Workflow({ + name: 'route', + edges: [ + ['START', routeNode], + [routeNode, {retry, [DEFAULT_ROUTE]: gen}], + ], + }); + + expect(finalOutput(await runWorkflowOnce(wf, 'john'))).toBe('RETRY:john'); + expect(finalOutput(await runWorkflowOnce(wf, 'jane'))).toBe('GEN:jane'); + }); +}); + +describe('workflow integration — fan-out / fan-in', () => { + it('runs branches in parallel and joins them', async () => { + const a = node((_c: NodeContext, i: string) => `A(${i})`, {name: 'a'}); + const b = node((_c: NodeContext, i: string) => `B(${i})`, {name: 'b'}); + const join = new JoinNode({name: 'join'}); + const wf = new Workflow({ + name: 'fan_out_fan_in', + edges: [['START', [a, b], join]], + }); + + expect(finalOutput(await runWorkflowOnce(wf, 'x'))).toEqual({ + a: 'A(x)', + b: 'B(x)', + }); + }); +}); + +describe('workflow integration — dynamic nodes & loop', () => { + it('runs an imperative loop that terminates', async () => { + const inc = new FunctionNode('inc', (_c, n: number) => (n as number) + 1); + const wf = new Workflow({ + name: 'dynamic_loop', + dynamicEntry: async (ctx) => { + let value = 0; + while (value < 3) { + value = (await ctx.runNode(inc, value)).output as number; + } + return value; + }, + }); + expect(finalOutput(await runWorkflowOnce(wf, 'go'))).toBe(3); + }); +}); + +describe('workflow integration — parallel worker', () => { + it('maps a list across the wrapped node with bounded concurrency', async () => { + // The list is produced inside the workflow, then fanned out. + const produce = node((): number[] => [1, 2, 3, 4], {name: 'produce'}); + const worker = new ParallelWorker( + new FunctionNode('double', (_c, n: number) => (n as number) * 2), + {maxParallelWorkers: 2}, + ); + const wf = new Workflow({ + name: 'parallel_worker', + edges: [['START', produce, worker]], + }); + expect(finalOutput(await runWorkflowOnce(wf, 'go'))).toEqual([2, 4, 6, 8]); + }); +}); + +describe('workflow integration — state', () => { + it('shares state across nodes', async () => { + const write = node( + (ctx: NodeContext, i: string) => { + ctx.state.set('greeting', `hi ${i}`); + return i; + }, + {name: 'write'}, + ); + const read = node((ctx: NodeContext) => ctx.state.get('greeting'), { + name: 'read', + }); + const wf = new Workflow({name: 'state', edges: [['START', write, read]]}); + expect(finalOutput(await runWorkflowOnce(wf, 'bob'))).toBe('hi bob'); + }); +}); + +describe('workflow integration — nested workflow', () => { + it('runs a workflow as a node inside another workflow', async () => { + const inner = new Workflow({ + name: 'inner', + edges: [ + [ + 'START', + node((_c: NodeContext, i: string) => `inner(${i})`, {name: 'in'}), + ], + ], + }); + const outer = new Workflow({ + name: 'outer', + edges: [ + [ + 'START', + inner, + node((_c: NodeContext, i: string) => `outer[${i}]`, {name: 'out'}), + ], + ], + }); + expect(finalOutput(await runWorkflowOnce(outer, 'x'))).toBe( + 'outer[inner(x)]', + ); + }); +}); + +describe('workflow integration — node as tool', () => { + it('lets a node call sub-nodes imperatively', async () => { + const add = new FunctionNode( + 'add', + (_c, args: {a: number; b: number}) => args.a + args.b, + ); + const orchestrator = node( + async (ctx: NodeContext) => { + const r1 = await ctx.runNode(add, {a: 2, b: 3}); + const r2 = await ctx.runNode(add, {a: 10, b: r1.output as number}); + return r2.output; + }, + {name: 'orchestrator'}, + ); + const wf = new Workflow({ + name: 'node_as_tool', + edges: [['START', orchestrator]], + }); + expect(finalOutput(await runWorkflowOnce(wf, 'go'))).toBe(15); + }); +}); + +describe('workflow integration — retry', () => { + it('retries a flaky node until it succeeds', async () => { + let attempts = 0; + const flaky = new FunctionNode( + 'flaky', + () => { + attempts++; + if (attempts < 3) { + throw new Error('transient'); + } + return 'ok'; + }, + {retryConfig: {maxAttempts: 3, initialDelay: 0.001, jitter: 0}}, + ); + const wf = new Workflow({name: 'retry', edges: [['START', flaky]]}); + expect(finalOutput(await runWorkflowOnce(wf, 'x'))).toBe('ok'); + expect(attempts).toBe(3); + }); +}); + +describe('workflow integration — request_input (HITL)', () => { + it('pauses for input and resumes on a function response', async () => { + const gate = node( + (ctx: NodeContext) => { + const answer = ctx.resumeInputs['confirm']; + if (answer === undefined) { + return new RequestInput({interruptId: 'confirm', message: 'ok?'}); + } + return `decided:${answer}`; + }, + {name: 'gate'}, + ); + const wf = new Workflow({name: 'request_input', edges: [['START', gate]]}); + const {run} = await createWorkflowRunner(wf); + + const turn1 = await collect(run('start')); + expect( + turn1.some((e) => + (e.content?.parts ?? []).some( + (p) => p.functionCall?.name === 'adk_request_input', + ), + ), + ).toBe(true); + + const turn2 = await collect( + run({ + role: 'user', + parts: [ + { + functionResponse: { + id: 'confirm', + name: 'adk_request_input', + response: {result: 'yes'}, + }, + }, + ], + }), + ); + expect(finalOutput(turn2)).toBe('decided:yes'); + }); +}); diff --git a/tests/integration/workflows/route.model_responses.json b/tests/integration/workflows/route.model_responses.json new file mode 100644 index 000000000..7b721d661 --- /dev/null +++ b/tests/integration/workflows/route.model_responses.json @@ -0,0 +1,28 @@ +{ + "classify_input": [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "{\"category\": \"question\"}"}] + }, + "finishReason": "STOP" + } + ] + } + ], + "answer_question": [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "The answer is 42."}] + }, + "finishReason": "STOP" + } + ] + } + ] +} diff --git a/tests/integration/workflows/route_llm_test.ts b/tests/integration/workflows/route_llm_test.ts new file mode 100644 index 000000000..d2c7a9215 --- /dev/null +++ b/tests/integration/workflows/route_llm_test.ts @@ -0,0 +1,106 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Integration test mirroring the Python `workflows/route` sample: an LLM + * classifier with an output schema drives conditional routing to branch agents. + */ + +import {createEvent, node, NodeContext, Workflow} from '@google/adk'; +import {describe, expect, it} from 'vitest'; +import {z} from 'zod'; +import {RawGenerateContentResponse} from '../test_case_utils.js'; +import modelResponses from './route.model_responses.json' with {type: 'json'}; +import { + finalOutput, + mockLlmAgent, + runWorkflowOnce, +} from './workflow_test_utils.js'; + +const responses = modelResponses as Record< + string, + RawGenerateContentResponse[] +>; + +describe('workflow integration — route via LLM classifier', () => { + it('classifies the input and routes to the matching branch agent', async () => { + const processInput = node( + (ctx: NodeContext, input: string) => { + ctx.state.set('input', input); + return input; + }, + {name: 'process_input'}, + ); + + const classifyInput = mockLlmAgent( + { + name: 'classify_input', + instruction: + 'Based on this input, decide which category it belongs to: {input}', + outputSchema: z.object({ + category: z.enum(['question', 'statement', 'other']), + }), + outputKey: 'category', + }, + responses['classify_input'], + ); + + const routeOnCategory = node( + (_c: NodeContext, input: unknown) => { + const category = + typeof input === 'string' + ? (JSON.parse(input) as {category: string}).category + : (input as {category: string}).category; + return createEvent({route: category}); + }, + {name: 'route_on_category'}, + ); + + const answerQuestion = mockLlmAgent( + {name: 'answer_question', instruction: 'Answer the question: {input}'}, + responses['answer_question'], + ); + const commentOnStatement = mockLlmAgent( + { + name: 'comment_on_statement', + instruction: 'Comment on the statement: {input}', + }, + [], + ); + const handleOther = node( + () => + createEvent({ + content: { + role: 'model', + parts: [{text: 'I can only answer questions or comment.'}], + }, + }), + {name: 'handle_other'}, + ); + + const wf = new Workflow({ + name: 'route_llm', + edges: [ + ['START', processInput, classifyInput, routeOnCategory], + [ + routeOnCategory, + { + question: answerQuestion, + statement: commentOnStatement, + other: handleOther, + }, + ], + ], + }); + + const events = await runWorkflowOnce(wf, 'What is the meaning of life?'); + + // Classified as "question" -> routed to answer_question. + expect(finalOutput(events)).toBe('The answer is 42.'); + expect(events.some((e) => e.author === 'answer_question')).toBe(true); + expect(events.some((e) => e.author === 'comment_on_statement')).toBe(false); + }); +}); diff --git a/tests/integration/workflows/sequence.model_responses.json b/tests/integration/workflows/sequence.model_responses.json new file mode 100644 index 000000000..4ef1af622 --- /dev/null +++ b/tests/integration/workflows/sequence.model_responses.json @@ -0,0 +1,25 @@ +{ + "generate_fruit_agent": [ + { + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": "apple"}]}, + "finishReason": "STOP" + } + ] + } + ], + "generate_benefit_agent": [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "Apples are rich in fiber."}] + }, + "finishReason": "STOP" + } + ] + } + ] +} diff --git a/tests/integration/workflows/sequence_llm_test.ts b/tests/integration/workflows/sequence_llm_test.ts new file mode 100644 index 000000000..2482c61db --- /dev/null +++ b/tests/integration/workflows/sequence_llm_test.ts @@ -0,0 +1,60 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Integration test mirroring the Python `workflows/sequence` sample: two LLM + * agents chained in a workflow, with model responses loaded from a JSON fixture. + */ + +import {Workflow} from '@google/adk'; +import {describe, expect, it} from 'vitest'; +import {RawGenerateContentResponse} from '../test_case_utils.js'; +import modelResponses from './sequence.model_responses.json' with {type: 'json'}; +import { + finalOutput, + mockLlmAgent, + runWorkflowOnce, +} from './workflow_test_utils.js'; + +const responses = modelResponses as Record< + string, + RawGenerateContentResponse[] +>; + +describe('workflow integration — sequence of LLM agents', () => { + it('chains two LLM agents, feeding the first output into the second', async () => { + const generateFruit = mockLlmAgent( + { + name: 'generate_fruit_agent', + instruction: + 'Return the name of a random fruit. Return only the name, nothing else.', + }, + responses['generate_fruit_agent'], + ); + const generateBenefit = mockLlmAgent( + { + name: 'generate_benefit_agent', + instruction: 'Tell me a health benefit about the specified fruit.', + }, + responses['generate_benefit_agent'], + ); + + const wf = new Workflow({ + name: 'sequence_llm', + edges: [['START', generateFruit, generateBenefit]], + }); + + const events = await runWorkflowOnce(wf, 'Give me a fruit fact'); + + // The final workflow output is the second agent's response. + expect(finalOutput(events)).toBe('Apples are rich in fiber.'); + // Both agents contributed events. + expect(events.some((e) => e.author === 'generate_fruit_agent')).toBe(true); + expect(events.some((e) => e.author === 'generate_benefit_agent')).toBe( + true, + ); + }); +}); diff --git a/tests/integration/workflows/workflow_test_utils.ts b/tests/integration/workflows/workflow_test_utils.ts new file mode 100644 index 000000000..701e2ff73 --- /dev/null +++ b/tests/integration/workflows/workflow_test_utils.ts @@ -0,0 +1,107 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { + Event, + InMemoryRunner, + LlmAgent, + LlmAgentConfig, + Workflow, + WorkflowAgent, +} from '@google/adk'; +import {Content, FinishReason} from '@google/genai'; +import { + GeminiWithMockResponses, + RawGenerateContentResponse, +} from '../test_case_utils.js'; + +/** + * Builds a raw generate-content response that returns plain model text. + */ +export function textResponse(text: string): RawGenerateContentResponse { + return { + candidates: [ + { + content: {role: 'model', parts: [{text}]}, + finishReason: FinishReason.STOP, + }, + ], + }; +} + +/** + * Constructs an {@link LlmAgent} whose model returns the given canned responses + * (loaded from a JSON fixture), so workflow integration tests are deterministic + * without a live model. + */ +export function mockLlmAgent( + config: Omit, + responses: RawGenerateContentResponse[], +): LlmAgent { + return new LlmAgent({ + ...config, + model: new GeminiWithMockResponses(responses), + }); +} + +/** + * Creates a runner for a {@link Workflow} (wrapped as a {@link WorkflowAgent}), + * bound to a single session so successive `run(...)` calls are additional turns + * (needed for HITL resume). Accepts a text prompt or a full `Content` (e.g. a + * function-response resume message). + */ +export async function createWorkflowRunner( + workflow: Workflow, +): Promise<{run: (message: string | Content) => AsyncGenerator}> { + const agent = new WorkflowAgent(workflow); + const runner = new InMemoryRunner({agent, appName: agent.name}); + const session = await runner.sessionService.createSession({ + appName: agent.name, + userId: 'u1', + }); + return { + run(message: string | Content): AsyncGenerator { + const newMessage: Content = + typeof message === 'string' + ? {role: 'user', parts: [{text: message}]} + : message; + return runner.runAsync({ + userId: 'u1', + sessionId: session.id, + newMessage, + }); + }, + }; +} + +/** Drains an event generator into an array. */ +export async function collect(gen: AsyncGenerator): Promise { + const events: Event[] = []; + for await (const event of gen) { + events.push(event); + } + return events; +} + +/** Runs a workflow for a single prompt and returns the emitted events. */ +export async function runWorkflowOnce( + workflow: Workflow, + prompt: string, +): Promise { + const {run} = await createWorkflowRunner(workflow); + return collect(run(prompt)); +} + +/** Returns the last non-undefined `output` across a list of events. */ +export function finalOutput(events: Event[]): unknown { + let output: unknown; + for (const event of events) { + if (event.output !== undefined) { + output = event.output; + } + } + return output; +} From aabcc0c8b1a5c87b697fa04d70bcf38ab0874fac Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Thu, 23 Jul 2026 16:06:37 -0700 Subject: [PATCH 19/41] fix(workflow): resume a waiting node with its original input On HITL resume, a waiting static-graph node was re-seeded from START with the new turn's message (e.g. the function response) instead of the input it had when it interrupted. Now the node's input is persisted on its interrupt event (actions.agentState.input), rehydrated by reconstructNodeStates, and used when the resolved node re-runs. The request_input integration test now asserts the correct 'start:yes' (original input preserved). The dynamic scheduler path is unaffected: a deterministic dynamicEntry re-passes the same input to ctx.runNode. --- core/src/workflow/node_runner.ts | 7 +++++++ core/src/workflow/utils/rehydration_utils.ts | 10 ++++++++++ core/src/workflow/workflow.ts | 10 +++++++++- tests/integration/workflows/core_workflows_test.ts | 8 +++++--- 4 files changed, 31 insertions(+), 4 deletions(-) diff --git a/core/src/workflow/node_runner.ts b/core/src/workflow/node_runner.ts index 528091eb2..ec7915642 100644 --- a/core/src/workflow/node_runner.ts +++ b/core/src/workflow/node_runner.ts @@ -153,6 +153,13 @@ async function runOnce( child.interruptIds.push(id); } } + // Persist the node's input on the interrupt event so a resumed + // (waiting) node re-runs with its ORIGINAL input, not the resume + // message. Rehydrated by reconstructNodeStates on the next turn. + event.actions.agentState = { + ...(event.actions.agentState ?? {}), + input, + }; } child.channel.push(event); } diff --git a/core/src/workflow/utils/rehydration_utils.ts b/core/src/workflow/utils/rehydration_utils.ts index 0ee0b307d..d35d49ec9 100644 --- a/core/src/workflow/utils/rehydration_utils.ts +++ b/core/src/workflow/utils/rehydration_utils.ts @@ -27,6 +27,8 @@ export interface RehydratedNode { route?: 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. */ @@ -107,6 +109,14 @@ function reconstruct( 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). + const agentState = event.actions?.agentState as + | {input?: unknown} + | undefined; + if (agentState && 'input' in agentState) { + node.input = agentState.input; + } } return nodes; diff --git a/core/src/workflow/workflow.ts b/core/src/workflow/workflow.ts index b7e4c9aa6..417a4d939 100644 --- a/core/src/workflow/workflow.ts +++ b/core/src/workflow/workflow.ts @@ -310,12 +310,20 @@ export class Workflow extends BaseNode { 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). const task: Promise = executeChildNode( ctx, node, - trigger.input, + nodeInput, { runId, useSubBranch: trigger.useSubBranch, diff --git a/tests/integration/workflows/core_workflows_test.ts b/tests/integration/workflows/core_workflows_test.ts index 2d56100c1..a41bd3e2d 100644 --- a/tests/integration/workflows/core_workflows_test.ts +++ b/tests/integration/workflows/core_workflows_test.ts @@ -209,12 +209,14 @@ describe('workflow integration — retry', () => { describe('workflow integration — request_input (HITL)', () => { it('pauses for input and resumes on a function response', async () => { const gate = node( - (ctx: NodeContext) => { + (ctx: NodeContext, input: string) => { const answer = ctx.resumeInputs['confirm']; if (answer === undefined) { return new RequestInput({interruptId: 'confirm', message: 'ok?'}); } - return `decided:${answer}`; + // On resume, `input` must still be the original 'start', not the + // function-response message. + return `${input}:${answer}`; }, {name: 'gate'}, ); @@ -244,6 +246,6 @@ describe('workflow integration — request_input (HITL)', () => { ], }), ); - expect(finalOutput(turn2)).toBe('decided:yes'); + expect(finalOutput(turn2)).toBe('start:yes'); }); }); From 9724e3400c8656324f720dcd8b5e6bcb73d7a9ed Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Thu, 23 Jul 2026 16:12:21 -0700 Subject: [PATCH 20/41] test(workflow): add essential coverage for parsing, validation, routing, schema, and workflow behaviors Add focused unit tests closing key coverage gaps: graph parsing (chains, fan-out, routing maps, numeric/DEFAULT_ROUTE keys, explicit Edges, node dedup), graph validation (duplicate names, missing/routed START, unreachable nodes, duplicate edges, DEFAULT_ROUTE constraints, conditional vs unconditional cycles), routing value types (numeric, list routes, fan-out-in-map, DEFAULT_ROUTE fallback), node input/output zod schema validation, and workflow-level concerns (maxConcurrency bound, error propagation halting downstream, three-way join, retry exception allow-list, and rerunOnResume vs fast-forward). Adds a shared test_helpers module. --- core/test/workflow/graph_parser_test.ts | 104 +++++++++++ core/test/workflow/graph_validation_test.ts | 105 +++++++++++ core/test/workflow/routing_test.ts | 93 ++++++++++ core/test/workflow/schema_validation_test.ts | 47 +++++ core/test/workflow/test_helpers.ts | 76 ++++++++ core/test/workflow/workflow_advanced_test.ts | 185 +++++++++++++++++++ 6 files changed, 610 insertions(+) create mode 100644 core/test/workflow/graph_parser_test.ts create mode 100644 core/test/workflow/graph_validation_test.ts create mode 100644 core/test/workflow/routing_test.ts create mode 100644 core/test/workflow/schema_validation_test.ts create mode 100644 core/test/workflow/test_helpers.ts create mode 100644 core/test/workflow/workflow_advanced_test.ts diff --git a/core/test/workflow/graph_parser_test.ts b/core/test/workflow/graph_parser_test.ts new file mode 100644 index 000000000..35edd1c02 --- /dev/null +++ b/core/test/workflow/graph_parser_test.ts @@ -0,0 +1,104 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {describe, expect, it} from 'vitest'; +import {DEFAULT_ROUTE, Edge, Graph} from '../../src/workflow/graph.js'; +import {parseEdgeItems} from '../../src/workflow/utils/graph_parser.js'; +import {FnNode} from './test_helpers.js'; + +const n = (name: string) => new FnNode(name, (_c, i) => i); + +/** Serializes edges to `from->to[:route]` strings for compact assertions. */ +function edgeStrings(edges: Edge[]): string[] { + return edges.map((e) => { + const route = + e.route === null || e.route === undefined + ? '' + : `:${JSON.stringify(e.route)}`; + return `${e.fromNode.name}->${e.toNode.name}${route}`; + }); +} + +describe('graph parser', () => { + it('parses a linear chain', () => { + const [a, b, c] = [n('a'), n('b'), n('c')]; + const edges = parseEdgeItems([['START', a, b, c]]); + expect(edgeStrings(edges)).toEqual(['__START__->a', 'a->b', 'b->c']); + }); + + it('parses a fan-out array to multiple edges', () => { + const [a, b, c] = [n('a'), n('b'), n('c')]; + const edges = parseEdgeItems([['START', [a, b], c]]); + expect(edgeStrings(edges)).toEqual([ + '__START__->a', + '__START__->b', + 'a->c', + 'b->c', + ]); + }); + + it('parses a routing map into conditional edges', () => { + const [a, b, c] = [n('a'), n('b'), n('c')]; + const edges = parseEdgeItems([[a, {question: b, statement: c}]]); + expect(edgeStrings(edges)).toEqual(['a->b:"question"', 'a->c:"statement"']); + }); + + it('parses numeric route keys', () => { + const [a, b, c] = [n('a'), n('b'), n('c')]; + const edges = parseEdgeItems([[a, {1: b, 2: c}]]); + expect(edgeStrings(edges)).toEqual(['a->b:1', 'a->c:2']); + }); + + it('parses fan-out inside a routing map', () => { + const [a, b, c] = [n('a'), n('b'), n('c')]; + const edges = parseEdgeItems([[a, {retry: [b, c]}]]); + expect(edgeStrings(edges)).toEqual(['a->b:"retry"', 'a->c:"retry"']); + }); + + it('parses DEFAULT_ROUTE keys', () => { + const [a, b, c] = [n('a'), n('b'), n('c')]; + const edges = parseEdgeItems([[a, {ok: b, [DEFAULT_ROUTE]: c}]]); + expect(edgeStrings(edges)).toEqual([ + 'a->b:"ok"', + `a->c:${JSON.stringify(DEFAULT_ROUTE)}`, + ]); + }); + + it('passes explicit Edge instances through', () => { + const [a, b] = [n('a'), n('b')]; + const edges = parseEdgeItems([new Edge(a, b, 'go')]); + expect(edgeStrings(edges)).toEqual(['a->b:"go"']); + }); + + it('rejects consecutive routing maps in a chain', () => { + const [a] = [n('a')]; + expect(() => parseEdgeItems([[a, {x: n('b')}, {y: n('c')}]])).toThrow( + /consecutive routing maps/i, + ); + }); + + it('rejects an empty routing map', () => { + const a = n('a'); + expect(() => parseEdgeItems([[a, {}]])).toThrow(/empty/i); + }); + + it('dedupes nodes by identity in the Graph', () => { + const [a, b, c] = [n('a'), n('b'), n('c')]; + // `a` referenced in two edge items -> one node in the graph. + const graph = Graph.fromEdgeItems([ + ['START', a, b], + [a, c], + ]); + expect(graph.nodes.map((node) => node.name).sort()).toEqual([ + '__START__', + 'a', + 'b', + 'c', + ]); + // Exactly one `a` instance. + expect(graph.nodes.filter((node) => node.name === 'a')).toHaveLength(1); + }); +}); diff --git a/core/test/workflow/graph_validation_test.ts b/core/test/workflow/graph_validation_test.ts new file mode 100644 index 000000000..73e5db000 --- /dev/null +++ b/core/test/workflow/graph_validation_test.ts @@ -0,0 +1,105 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {describe, expect, it} from 'vitest'; +import {START} from '../../src/workflow/base_node.js'; +import {DEFAULT_ROUTE, Edge, Graph} from '../../src/workflow/graph.js'; +import {FnNode} from './test_helpers.js'; + +const n = (name: string) => new FnNode(name, (_c, i) => i); + +/** Builds and validates a graph from edge items, returning the graph. */ +function build(edges: Parameters[0]): Graph { + const graph = Graph.fromEdgeItems(edges); + graph.validate(); + return graph; +} + +describe('graph validation', () => { + it('accepts a valid graph and computes terminal nodes', () => { + const [a, b] = [n('a'), n('b')]; + const graph = build([['START', a, b]]); + expect([...graph.terminalNodeNames]).toEqual(['b']); + }); + + it('rejects duplicate node names', () => { + // Two distinct instances sharing a name. + const a1 = new FnNode('dup', (_c, i) => i); + const a2 = new FnNode('dup', (_c, i) => i); + expect(() => + build([ + ['START', a1], + [a1, a2], + ]), + ).toThrow(/duplicate node names/i); + }); + + it('rejects a missing START node', () => { + const [a, b] = [n('a'), n('b')]; + expect(() => build([[a, b]])).toThrow(/START node.*not found/i); + }); + + it('rejects a routed edge from START', () => { + const a = n('a'); + expect(() => build([new Edge(START, a, 'go')])).toThrow( + /edges from START must not have routes/i, + ); + }); + + it('rejects unreachable nodes', () => { + const [a, orphan, b] = [n('a'), n('orphan'), n('b')]; + expect(() => + build([ + ['START', a], + [orphan, b], + ]), + ).toThrow(/unreachable/i); + }); + + it('rejects duplicate edges', () => { + const [a, b] = [n('a'), n('b')]; + expect(() => + build([ + ['START', a], + [a, b], + [a, b], + ]), + ).toThrow(/duplicate edge/i); + }); + + it('rejects multiple DEFAULT_ROUTE edges from one node', () => { + const [a, b, c] = [n('a'), n('b'), n('c')]; + expect(() => + build([ + ['START', a], + [a, {[DEFAULT_ROUTE]: b}], + new Edge(a, c, DEFAULT_ROUTE), + ]), + ).toThrow(/DEFAULT_ROUTE/i); + }); + + it('rejects an unconditional cycle', () => { + const [a, b] = [n('a'), n('b')]; + expect(() => + build([ + ['START', a], + [a, b], + [b, a], + ]), + ).toThrow(/cycle/i); + }); + + it('allows a routed (conditional) cycle', () => { + const [a, b] = [n('a'), n('b')]; + // a -> b unconditionally, b -> a only on route 'again' (conditional) => ok. + const graph = build([ + ['START', a], + [a, b], + [b, {again: a, [DEFAULT_ROUTE]: n('done')}], + ]); + expect(graph).toBeInstanceOf(Graph); + }); +}); diff --git a/core/test/workflow/routing_test.ts b/core/test/workflow/routing_test.ts new file mode 100644 index 000000000..c9398b46b --- /dev/null +++ b/core/test/workflow/routing_test.ts @@ -0,0 +1,93 @@ +/** + * @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)'); + }); +}); diff --git a/core/test/workflow/schema_validation_test.ts b/core/test/workflow/schema_validation_test.ts new file mode 100644 index 000000000..63dac9977 --- /dev/null +++ b/core/test/workflow/schema_validation_test.ts @@ -0,0 +1,47 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {describe, expect, it} from 'vitest'; +import {z} from 'zod'; +import {FunctionNode} from '../../src/workflow/nodes/function_node.js'; +import {driveNode} from './test_helpers.js'; + +describe('node schema validation', () => { + it('validates input against inputSchema', async () => { + const node = new FunctionNode('squared', (_c, n: number) => n * n, { + inputSchema: z.number(), + }); + expect((await driveNode(node, 5)).output).toBe(25); + await expect(driveNode(node, 'not-a-number')).rejects.toThrow(); + }); + + it('validates output against outputSchema', async () => { + const schema = z.object({total: z.number()}); + const good = new FunctionNode('g', () => ({total: 10}), { + outputSchema: schema, + }); + expect((await driveNode(good)).output).toEqual({total: 10}); + + const bad = new FunctionNode('b', () => ({total: 'oops'}), { + outputSchema: schema, + }); + await expect(driveNode(bad)).rejects.toThrow(); + }); + + it('coerces and validates a valid input, passing it to the handler', async () => { + let received: unknown; + const node = new FunctionNode( + 'capture', + (_c, value: {name: string}) => { + received = value; + return value.name; + }, + {inputSchema: z.object({name: z.string()})}, + ); + expect((await driveNode(node, {name: 'ada'})).output).toBe('ada'); + expect(received).toEqual({name: 'ada'}); + }); +}); diff --git a/core/test/workflow/test_helpers.ts b/core/test/workflow/test_helpers.ts new file mode 100644 index 000000000..e66ba7192 --- /dev/null +++ b/core/test/workflow/test_helpers.ts @@ -0,0 +1,76 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {BaseAgent} from '../../src/agents/base_agent.js'; +import {InvocationContext} from '../../src/agents/invocation_context.js'; +import {Event} from '../../src/events/event.js'; +import {PluginManager} from '../../src/plugins/plugin_manager.js'; +import {Session} from '../../src/sessions/session.js'; +import {BaseNode} from '../../src/workflow/base_node.js'; +import {NodeContext} from '../../src/workflow/node_context.js'; +import {EventChannel} from '../../src/workflow/utils/event_channel.js'; + +/** Builds a throwaway InvocationContext for driving nodes directly in tests. */ +export function createIc( + state: Record = {}, +): InvocationContext { + const session = { + id: 's1', + appName: 'app', + userId: 'u', + events: [], + state, + lastUpdateTime: Date.now(), + } as unknown as Session; + return new InvocationContext({ + invocationId: 'inv-1', + session, + agent: { + name: 'wf', + runAsync: async function* () {}, + } as unknown as BaseAgent, + pluginManager: new PluginManager(), + }); +} + +/** Runs a node (or workflow) to completion, returning its events and output. */ +export async function driveNode( + node: BaseNode, + input?: unknown, + ic: InvocationContext = createIc(), +): Promise<{events: Event[]; output: unknown; ctx: NodeContext}> { + const channel = new EventChannel(); + const root = new NodeContext({ + invocationContext: ic, + channel, + nodePath: '', + runId: 'root', + }); + const events: Event[] = []; + const settle = root.runNode(node, input, {useAsOutput: true}).then( + () => channel.close(), + (err) => channel.fail(err), + ); + for await (const ev of channel) { + events.push(ev); + } + await settle; + return {events, output: root.output, ctx: root}; +} + +/** A node whose behavior is a plain function returning a value or Event. */ +export class FnNode extends BaseNode { + constructor( + name: string, + private readonly fn: (ctx: NodeContext, input: unknown) => unknown, + config: {rerunOnResume?: boolean} = {}, + ) { + super({name, ...config}); + } + protected async *runImpl(ctx: NodeContext, input: unknown) { + yield await this.fn(ctx, input); + } +} diff --git a/core/test/workflow/workflow_advanced_test.ts b/core/test/workflow/workflow_advanced_test.ts new file mode 100644 index 000000000..8e7db53be --- /dev/null +++ b/core/test/workflow/workflow_advanced_test.ts @@ -0,0 +1,185 @@ +/** + * @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); + }); +}); + +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'); + }); +}); From a5a1a289998124f8f40c6322b7dfec4b576b5f1a Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Thu, 23 Jul 2026 16:17:53 -0700 Subject: [PATCH 21/41] test(workflow): add integration tests for dynamic, HITL-chain, and LLM pipelines New end-to-end (Runner) integration tests: dynamic fan-out/fan-in, conditional dynamic loop, mid-graph HITL resume (verifies upstream is not re-run and the gate resumes with its original input), three-way JoinNode fan-in, parallel-branch event streaming, a mixed function/LLM/function pipeline, and multi-agent orchestration via ctx.runNode (LLM responses from JSON fixtures). WorkflowAgent now emits a final result event carrying the workflow output so dynamicEntry workflows' return values are observable via the Runner (edges-based terminal output is unchanged in value). --- core/src/workflow/workflow_agent.ts | 31 +++- .../workflows/advanced_workflows_test.ts | 169 ++++++++++++++++++ .../agent_pipeline.model_responses.json | 41 +++++ .../workflows/agent_pipeline_test.ts | 81 +++++++++ 4 files changed, 317 insertions(+), 5 deletions(-) create mode 100644 tests/integration/workflows/advanced_workflows_test.ts create mode 100644 tests/integration/workflows/agent_pipeline.model_responses.json create mode 100644 tests/integration/workflows/agent_pipeline_test.ts diff --git a/core/src/workflow/workflow_agent.ts b/core/src/workflow/workflow_agent.ts index c5b5ea035..3a55d8c09 100644 --- a/core/src/workflow/workflow_agent.ts +++ b/core/src/workflow/workflow_agent.ts @@ -7,7 +7,8 @@ import {Content} from '@google/genai'; import {BaseAgent} from '../agents/base_agent.js'; import {InvocationContext} from '../agents/invocation_context.js'; -import {Event} from '../events/event.js'; +import {createEvent, Event} from '../events/event.js'; +import {toContent} from './base_node.js'; import {NodeContext} from './node_context.js'; import {EventChannel} from './utils/event_channel.js'; import {Workflow} from './workflow.js'; @@ -54,10 +55,30 @@ export class WorkflowAgent extends BaseAgent { const input = extractWorkflowInput(ic.userContent); - const settle = root.runNode(this.workflow, input, {useAsOutput: true}).then( - () => channel.close(), - (err) => channel.fail(err), - ); + 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); + } + })(); for await (const event of channel) { yield event; diff --git a/tests/integration/workflows/advanced_workflows_test.ts b/tests/integration/workflows/advanced_workflows_test.ts new file mode 100644 index 000000000..5fcc1bf7e --- /dev/null +++ b/tests/integration/workflows/advanced_workflows_test.ts @@ -0,0 +1,169 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Additional end-to-end (Runner) integration tests for advanced workflow + * scenarios: dynamic fan-out/fan-in, mid-graph HITL resume, multi-trigger + * re-execution, and a conditional dynamic loop. + */ + +import { + Event, + FunctionNode, + JoinNode, + node, + NodeContext, + RequestInput, + Workflow, +} from '@google/adk'; +import {describe, expect, it} from 'vitest'; +import { + collect, + createWorkflowRunner, + finalOutput, + runWorkflowOnce, +} from './workflow_test_utils.js'; + +describe('workflow integration — dynamic fan-out / fan-in', () => { + it('fans out concurrent ctx.runNode calls and aggregates results', async () => { + const worker = new FunctionNode('work', (_c, item: number) => item * 10); + const wf = new Workflow({ + name: 'dynamic_fan_out_fan_in', + dynamicEntry: async (ctx) => { + const items = [1, 2, 3]; + const results = await Promise.all( + items.map((i) => ctx.runNode(worker, i, {runId: `w${i}`})), + ); + return results.map((r) => r.output); + }, + }); + expect(finalOutput(await runWorkflowOnce(wf, 'go'))).toEqual([10, 20, 30]); + }); +}); + +describe('workflow integration — conditional dynamic loop', () => { + it('loops an LLM-free refiner until a condition is met', async () => { + const refine = new FunctionNode('refine', (_c, n: number) => n + 1); + const wf = new Workflow({ + name: 'conditional_loop', + dynamicEntry: async (ctx) => { + let value = 0; + let iterations = 0; + while (value < 5) { + value = (await ctx.runNode(refine, value, {runId: `r${iterations}`})) + .output as number; + iterations++; + } + return {value, iterations}; + }, + }); + expect(finalOutput(await runWorkflowOnce(wf, 'go'))).toEqual({ + value: 5, + iterations: 5, + }); + }); +}); + +describe('workflow integration — mid-graph HITL resume', () => { + it('pauses in the middle of a chain and resumes without re-running upstream', async () => { + let aRuns = 0; + let cRuns = 0; + const a = node( + (_c: NodeContext, i: string) => { + aRuns++; + return `A(${i})`; + }, + {name: 'a'}, + ); + const gate = node( + (ctx: NodeContext, input: string) => { + const answer = ctx.resumeInputs['approve']; + if (answer === undefined) { + return new RequestInput({interruptId: 'approve', message: 'ok?'}); + } + return `${input}|${answer}`; + }, + {name: 'gate'}, + ); + const c = node( + (_c: NodeContext, i: string) => { + cRuns++; + return `C(${i})`; + }, + {name: 'c'}, + ); + const wf = new Workflow({ + name: 'mid_graph_hitl', + edges: [['START', a, gate, c]], + }); + const {run} = await createWorkflowRunner(wf); + + // Turn 1: a runs, gate interrupts, c must not run. + const turn1 = await collect(run('start')); + expect(aRuns).toBe(1); + expect(cRuns).toBe(0); + expect( + turn1.some((e) => + (e.content?.parts ?? []).some( + (p) => p.functionCall?.name === 'adk_request_input', + ), + ), + ).toBe(true); + + // Turn 2: resume; a is fast-forwarded (not re-run), gate resolves, c runs. + const turn2 = await collect( + run({ + role: 'user', + parts: [ + { + functionResponse: { + id: 'approve', + name: 'adk_request_input', + response: {result: 'yes'}, + }, + }, + ], + }), + ); + expect(aRuns).toBe(1); + expect(cRuns).toBe(1); + // gate re-ran with its ORIGINAL input 'A(start)', resolved with 'yes'. + expect(finalOutput(turn2)).toBe('C(A(start)|yes)'); + }); +}); + +describe('workflow integration — multi-trigger fan-in with JoinNode', () => { + it('joins three parallel branches produced from START', async () => { + const mk = (name: string): FunctionNode => + new FunctionNode(name, (_c, i: string) => `${name}:${i}`); + const join = new JoinNode({name: 'join'}); + const wf = new Workflow({ + name: 'triple_fan_in', + edges: [['START', [mk('x'), mk('y'), mk('z')], join]], + }); + const output = finalOutput(await runWorkflowOnce(wf, 'v')) as Record< + string, + string + >; + expect(output).toEqual({x: 'x:v', y: 'y:v', z: 'z:v'}); + }); +}); + +describe('workflow integration — parallel branches emit independent events', () => { + it('streams events from all parallel branches', async () => { + const mk = (name: string): FunctionNode => + new FunctionNode(name, (_c, i: string) => `${name}(${i})`); + const join = new JoinNode({name: 'join'}); + const wf = new Workflow({ + name: 'parallel_events', + edges: [['START', [mk('p'), mk('q')], join]], + }); + const events: Event[] = await runWorkflowOnce(wf, 'x'); + expect(events.some((e) => e.author === 'p')).toBe(true); + expect(events.some((e) => e.author === 'q')).toBe(true); + expect(events.some((e) => e.author === 'join')).toBe(true); + }); +}); diff --git a/tests/integration/workflows/agent_pipeline.model_responses.json b/tests/integration/workflows/agent_pipeline.model_responses.json new file mode 100644 index 000000000..1dbe81ab5 --- /dev/null +++ b/tests/integration/workflows/agent_pipeline.model_responses.json @@ -0,0 +1,41 @@ +{ + "summarizer": [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "Summary: the report is positive."}] + }, + "finishReason": "STOP" + } + ] + } + ], + "researcher": [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "Findings: A, B, C."}] + }, + "finishReason": "STOP" + } + ] + } + ], + "writer": [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "Report drafted from the findings."}] + }, + "finishReason": "STOP" + } + ] + } + ] +} diff --git a/tests/integration/workflows/agent_pipeline_test.ts b/tests/integration/workflows/agent_pipeline_test.ts new file mode 100644 index 000000000..3c0f36fe4 --- /dev/null +++ b/tests/integration/workflows/agent_pipeline_test.ts @@ -0,0 +1,81 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Integration tests for LLM agents embedded in workflows (model responses from + * a JSON fixture): a mixed function/agent pipeline, and multi-agent + * orchestration driven imperatively via ctx.runNode. + */ + +import {FunctionNode, node, Workflow} from '@google/adk'; +import {describe, expect, it} from 'vitest'; +import {RawGenerateContentResponse} from '../test_case_utils.js'; +import modelResponses from './agent_pipeline.model_responses.json' with {type: 'json'}; +import { + finalOutput, + mockLlmAgent, + runWorkflowOnce, +} from './workflow_test_utils.js'; + +const responses = modelResponses as Record< + string, + RawGenerateContentResponse[] +>; + +describe('workflow integration — mixed function/agent pipeline', () => { + it('runs function -> LLM agent -> function, threading outputs', async () => { + const preprocess = new FunctionNode('preprocess', (_c, input: string) => + input.toUpperCase(), + ); + const summarizer = mockLlmAgent( + {name: 'summarizer', instruction: 'Summarize the provided text.'}, + responses['summarizer'], + ); + const postprocess = new FunctionNode( + 'postprocess', + (_c, input: string) => `[${input}]`, + ); + + const wf = new Workflow({ + name: 'agent_pipeline', + edges: [['START', preprocess, summarizer, postprocess]], + }); + + const events = await runWorkflowOnce(wf, 'the quarterly report'); + expect(finalOutput(events)).toBe('[Summary: the report is positive.]'); + expect(events.some((e) => e.author === 'summarizer')).toBe(true); + }); +}); + +describe('workflow integration — multi-agent orchestration (LLM)', () => { + it('coordinates two LLM agents via ctx.runNode', async () => { + const researcher = mockLlmAgent( + {name: 'researcher', instruction: 'Research the given topic.'}, + responses['researcher'], + ); + const writer = mockLlmAgent( + {name: 'writer', instruction: 'Write a report from the research.'}, + responses['writer'], + ); + + const wf = new Workflow({ + name: 'coordinator', + dynamicEntry: async (ctx, input) => { + const research = await ctx.runNode(node(researcher), input); + const report = await ctx.runNode(node(writer), research.output); + return {research: research.output, report: report.output}; + }, + }); + + const events = await runWorkflowOnce(wf, 'ADK workflows'); + expect(finalOutput(events)).toEqual({ + research: 'Findings: A, B, C.', + report: 'Report drafted from the findings.', + }); + expect(events.some((e) => e.author === 'researcher')).toBe(true); + expect(events.some((e) => e.author === 'writer')).toBe(true); + }); +}); From 4f5869b86b6701f600e8f235f0e99c7c77028b73 Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Thu, 23 Jul 2026 16:31:35 -0700 Subject: [PATCH 22/41] test(workflow): add integration tests for tools, retry exhaustion, nested HITL, and parallel LLM More end-to-end (Runner) integration coverage: a BaseTool run as a ToolNode with upstream-produced args; retry exhaustion failing the workflow; resuming a HITL interrupt raised inside a nested workflow; a tool-calling LLM agent as a workflow node (function-call -> tool -> final answer); and a ParallelWorker mapping an LLM agent across a list. LLM responses come from JSON fixtures. --- .../llm_tool_agent.model_responses.json | 34 +++++ .../workflows/llm_tool_agent_test.ts | 74 +++++++++++ .../parallel_llm.model_responses.json | 28 +++++ .../workflows/parallel_llm_test.ts | 51 ++++++++ .../workflows/tool_and_resilience_test.ts | 118 ++++++++++++++++++ 5 files changed, 305 insertions(+) create mode 100644 tests/integration/workflows/llm_tool_agent.model_responses.json create mode 100644 tests/integration/workflows/llm_tool_agent_test.ts create mode 100644 tests/integration/workflows/parallel_llm.model_responses.json create mode 100644 tests/integration/workflows/parallel_llm_test.ts create mode 100644 tests/integration/workflows/tool_and_resilience_test.ts diff --git a/tests/integration/workflows/llm_tool_agent.model_responses.json b/tests/integration/workflows/llm_tool_agent.model_responses.json new file mode 100644 index 000000000..54a00e68a --- /dev/null +++ b/tests/integration/workflows/llm_tool_agent.model_responses.json @@ -0,0 +1,34 @@ +{ + "assistant": [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "functionCall": { + "name": "lookup", + "args": {"key": "answer"}, + "id": "call-1" + } + } + ] + }, + "finishReason": "STOP" + } + ] + }, + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "The looked-up value is 42."}] + }, + "finishReason": "STOP" + } + ] + } + ] +} diff --git a/tests/integration/workflows/llm_tool_agent_test.ts b/tests/integration/workflows/llm_tool_agent_test.ts new file mode 100644 index 000000000..f38527740 --- /dev/null +++ b/tests/integration/workflows/llm_tool_agent_test.ts @@ -0,0 +1,74 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Integration test: an LLM agent that calls a tool, embedded as a workflow + * node. The mocked model first returns a function call, then a final answer + * after the tool runs. + */ + +import {FunctionTool, Workflow} from '@google/adk'; +import {describe, expect, it} from 'vitest'; +import {z} from 'zod'; +import {RawGenerateContentResponse} from '../test_case_utils.js'; +import modelResponses from './llm_tool_agent.model_responses.json' with {type: 'json'}; +import { + finalOutput, + mockLlmAgent, + runWorkflowOnce, +} from './workflow_test_utils.js'; + +const responses = modelResponses as Record< + string, + RawGenerateContentResponse[] +>; + +describe('workflow integration — LLM agent with tool calling', () => { + it('runs a tool-calling agent as a workflow node', async () => { + let toolCalled = false; + const lookup = new FunctionTool({ + name: 'lookup', + description: 'Looks up a value by key.', + parameters: z.object({key: z.string()}), + execute: async ({key}: {key: string}) => { + toolCalled = true; + return {key, value: 42}; + }, + }); + + const assistant = mockLlmAgent( + { + name: 'assistant', + instruction: 'Use the lookup tool to answer.', + tools: [lookup], + }, + responses['assistant'], + ); + + const wf = new Workflow({ + name: 'llm_tool_agent', + edges: [['START', assistant]], + }); + + const events = await runWorkflowOnce(wf, 'What is the answer?'); + + expect(toolCalled).toBe(true); + expect(finalOutput(events)).toBe('The looked-up value is 42.'); + // The tool call and its response both appear in the event stream. + expect( + events.some((e) => + (e.content?.parts ?? []).some((p) => p.functionCall?.name === 'lookup'), + ), + ).toBe(true); + expect( + events.some((e) => + (e.content?.parts ?? []).some( + (p) => p.functionResponse?.name === 'lookup', + ), + ), + ).toBe(true); + }); +}); diff --git a/tests/integration/workflows/parallel_llm.model_responses.json b/tests/integration/workflows/parallel_llm.model_responses.json new file mode 100644 index 000000000..fa051a50e --- /dev/null +++ b/tests/integration/workflows/parallel_llm.model_responses.json @@ -0,0 +1,28 @@ +{ + "classifier": [ + { + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": "processed"}]}, + "finishReason": "STOP" + } + ] + }, + { + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": "processed"}]}, + "finishReason": "STOP" + } + ] + }, + { + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": "processed"}]}, + "finishReason": "STOP" + } + ] + } + ] +} diff --git a/tests/integration/workflows/parallel_llm_test.ts b/tests/integration/workflows/parallel_llm_test.ts new file mode 100644 index 000000000..83e65f0d2 --- /dev/null +++ b/tests/integration/workflows/parallel_llm_test.ts @@ -0,0 +1,51 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Integration test: a ParallelWorker mapping an LLM agent over a list of items, + * with model responses from a JSON fixture. + */ + +import {node, ParallelWorker, Workflow} from '@google/adk'; +import {describe, expect, it} from 'vitest'; +import {RawGenerateContentResponse} from '../test_case_utils.js'; +import modelResponses from './parallel_llm.model_responses.json' with {type: 'json'}; +import { + finalOutput, + mockLlmAgent, + runWorkflowOnce, +} from './workflow_test_utils.js'; + +const responses = modelResponses as Record< + string, + RawGenerateContentResponse[] +>; + +describe('workflow integration — ParallelWorker over an LLM agent', () => { + it('maps an LLM agent across a list of items', async () => { + const classifier = mockLlmAgent( + {name: 'classifier', instruction: 'Classify the item.'}, + responses['classifier'], + ); + + // Produce the list inside the workflow, then map the agent across it. + const produce = node((): string[] => ['alpha', 'beta', 'gamma'], { + name: 'produce', + }); + const worker = new ParallelWorker(node(classifier) as never, { + maxParallelWorkers: 1, + }); + + const wf = new Workflow({ + name: 'parallel_llm', + edges: [['START', produce, worker]], + }); + + const output = finalOutput(await runWorkflowOnce(wf, 'go')) as string[]; + expect(output).toHaveLength(3); + expect(output).toEqual(['processed', 'processed', 'processed']); + }); +}); diff --git a/tests/integration/workflows/tool_and_resilience_test.ts b/tests/integration/workflows/tool_and_resilience_test.ts new file mode 100644 index 000000000..7805b890c --- /dev/null +++ b/tests/integration/workflows/tool_and_resilience_test.ts @@ -0,0 +1,118 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Integration tests for tools in workflows, retry exhaustion, and resuming a + * HITL interrupt raised inside a nested workflow. + */ + +import { + FunctionNode, + FunctionTool, + node, + NodeContext, + RequestInput, + ToolNode, + Workflow, +} from '@google/adk'; +import {describe, expect, it} from 'vitest'; +import {z} from 'zod'; +import { + collect, + createWorkflowRunner, + finalOutput, + runWorkflowOnce, +} from './workflow_test_utils.js'; + +describe('workflow integration — ToolNode', () => { + it('runs a BaseTool as a node with args from an upstream node', async () => { + const addTool = new FunctionTool({ + name: 'add', + description: 'Adds two numbers.', + parameters: z.object({a: z.number(), b: z.number()}), + execute: async ({a, b}: {a: number; b: number}) => ({sum: a + b}), + }); + const produceArgs = new FunctionNode('produce_args', () => ({a: 2, b: 3})); + const wf = new Workflow({ + name: 'tool_workflow', + edges: [['START', produceArgs, new ToolNode(addTool)]], + }); + expect(finalOutput(await runWorkflowOnce(wf, 'go'))).toEqual({sum: 5}); + }); +}); + +describe('workflow integration — retry exhaustion', () => { + it('fails the workflow when a node exhausts its retries', async () => { + let attempts = 0; + const flaky = new FunctionNode( + 'always_fails', + () => { + attempts++; + throw new Error('permanent failure'); + }, + {retryConfig: {maxAttempts: 3, initialDelay: 0.001, jitter: 0}}, + ); + const wf = new Workflow({name: 'retry_exhaust', edges: [['START', flaky]]}); + await expect(runWorkflowOnce(wf, 'go')).rejects.toThrow( + 'permanent failure', + ); + expect(attempts).toBe(3); + }); +}); + +describe('workflow integration — nested workflow HITL resume', () => { + it('resumes an interrupt raised inside a nested workflow', async () => { + const gate = node( + (ctx: NodeContext) => { + const answer = ctx.resumeInputs['approve']; + if (answer === undefined) { + return new RequestInput({interruptId: 'approve', message: 'ok?'}); + } + return `approved:${answer}`; + }, + {name: 'gate'}, + ); + const inner = new Workflow({name: 'inner', edges: [['START', gate]]}); + const outer = new Workflow({ + name: 'outer', + edges: [ + [ + 'START', + inner, + node((_c: NodeContext, i: string) => `wrapped(${i})`, {name: 'wrap'}), + ], + ], + }); + const {run} = await createWorkflowRunner(outer); + + // Turn 1: the nested gate interrupts; the interrupt bubbles up. + const turn1 = await collect(run('start')); + expect( + turn1.some((e) => + (e.content?.parts ?? []).some( + (p) => p.functionCall?.name === 'adk_request_input', + ), + ), + ).toBe(true); + + // Turn 2: resume; the nested gate resolves and the outer workflow finishes. + const turn2 = await collect( + run({ + role: 'user', + parts: [ + { + functionResponse: { + id: 'approve', + name: 'adk_request_input', + response: {result: 'yes'}, + }, + }, + ], + }), + ); + expect(finalOutput(turn2)).toBe('wrapped(approved:yes)'); + }); +}); From ccca0763a30aba2e7acad8480921be4c919e877a Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Thu, 23 Jul 2026 16:34:27 -0700 Subject: [PATCH 23/41] test(workflow): add integration tests for auth, routed loops, multi-trigger, and LLM loop More end-to-end (Runner) integration coverage: an API-key auth gate in a workflow (interrupt on first turn, run after credential supplied on resume); a routed self-loop and a DEFAULT_ROUTE loop exit; failure propagation from a parallel branch; multi-trigger re-execution of a non-join node (runs once per predecessor); and an imperative loop driven by an LLM agent's continue/done decision (JSON-mocked responses). --- .../workflows/auth_workflow_test.ts | 102 ++++++++++++++++ .../workflows/llm_loop.model_responses.json | 28 +++++ tests/integration/workflows/llm_loop_test.ts | 62 ++++++++++ .../workflows/loop_and_trigger_test.ts | 113 ++++++++++++++++++ 4 files changed, 305 insertions(+) create mode 100644 tests/integration/workflows/auth_workflow_test.ts create mode 100644 tests/integration/workflows/llm_loop.model_responses.json create mode 100644 tests/integration/workflows/llm_loop_test.ts create mode 100644 tests/integration/workflows/loop_and_trigger_test.ts diff --git a/tests/integration/workflows/auth_workflow_test.ts b/tests/integration/workflows/auth_workflow_test.ts new file mode 100644 index 000000000..cdadc0683 --- /dev/null +++ b/tests/integration/workflows/auth_workflow_test.ts @@ -0,0 +1,102 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Integration test: a workflow node that requires an API-key credential + * interrupts on the first turn and runs once the credential is supplied on + * resume (mirrors the Python `workflows/auth_api_key` sample). + */ + +import { + AuthConfig, + AuthCredential, + AuthCredentialTypes, + AuthScheme, + FunctionNode, + NodeContext, + Workflow, +} from '@google/adk'; +import {describe, expect, it} from 'vitest'; +import { + collect, + createWorkflowRunner, + finalOutput, +} from './workflow_test_utils.js'; + +const CREDENTIAL_KEY = 'weather_api'; + +function apiKeyAuthConfig(): AuthConfig { + return { + authScheme: {type: 'apiKey', in: 'header', name: 'X-API-Key'} as AuthScheme, + rawAuthCredential: {authType: AuthCredentialTypes.API_KEY}, + credentialKey: CREDENTIAL_KEY, + }; +} + +describe('workflow integration — auth gate (API key)', () => { + it('requests credentials, then runs after they are supplied on resume', async () => { + let runs = 0; + const secured = new FunctionNode( + 'secured', + (ctx: NodeContext) => { + runs++; + const cred = ctx.state.get('temp:' + CREDENTIAL_KEY); + return `weather(key=${cred?.apiKey})`; + }, + {authConfig: apiKeyAuthConfig()}, + ); + const wf = new Workflow({ + name: 'auth_api_key', + edges: [['START', secured]], + }); + const {run} = await createWorkflowRunner(wf); + + // Turn 1: no credential -> auth request interrupt; handler NOT run. + const turn1 = await collect(run('what is the weather?')); + expect(runs).toBe(0); + expect( + turn1.some((e) => + (e.content?.parts ?? []).some( + (p) => p.functionCall?.name === 'adk_request_credential', + ), + ), + ).toBe(true); + + // Turn 2: supply the credential -> node runs with it. + const credentialResponse: AuthConfig = { + authScheme: { + type: 'apiKey', + in: 'header', + name: 'X-API-Key', + } as AuthScheme, + credentialKey: CREDENTIAL_KEY, + exchangedAuthCredential: { + authType: AuthCredentialTypes.API_KEY, + apiKey: 'sk-test-123', + }, + }; + const turn2 = await collect( + run({ + role: 'user', + parts: [ + { + functionResponse: { + id: CREDENTIAL_KEY, + name: 'adk_request_credential', + response: credentialResponse as unknown as Record< + string, + unknown + >, + }, + }, + ], + }), + ); + + expect(runs).toBe(1); + expect(finalOutput(turn2)).toBe('weather(key=sk-test-123)'); + }); +}); diff --git a/tests/integration/workflows/llm_loop.model_responses.json b/tests/integration/workflows/llm_loop.model_responses.json new file mode 100644 index 000000000..2d40bf7a3 --- /dev/null +++ b/tests/integration/workflows/llm_loop.model_responses.json @@ -0,0 +1,28 @@ +{ + "decider": [ + { + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": "continue"}]}, + "finishReason": "STOP" + } + ] + }, + { + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": "continue"}]}, + "finishReason": "STOP" + } + ] + }, + { + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": "done"}]}, + "finishReason": "STOP" + } + ] + } + ] +} diff --git a/tests/integration/workflows/llm_loop_test.ts b/tests/integration/workflows/llm_loop_test.ts new file mode 100644 index 000000000..b3773693f --- /dev/null +++ b/tests/integration/workflows/llm_loop_test.ts @@ -0,0 +1,62 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Integration test: an imperative loop driven by an LLM agent's decision + * ("continue"/"done"), with model responses from a JSON fixture. + */ + +import {node, Workflow} from '@google/adk'; +import {describe, expect, it} from 'vitest'; +import {RawGenerateContentResponse} from '../test_case_utils.js'; +import modelResponses from './llm_loop.model_responses.json' with {type: 'json'}; +import { + finalOutput, + mockLlmAgent, + runWorkflowOnce, +} from './workflow_test_utils.js'; + +const responses = modelResponses as Record< + string, + RawGenerateContentResponse[] +>; + +describe('workflow integration — LLM-driven loop', () => { + it('loops until the LLM agent decides to stop', async () => { + const decider = mockLlmAgent( + { + name: 'decider', + instruction: 'Reply "continue" to keep going or "done" to stop.', + }, + responses['decider'], + ); + + const wf = new Workflow({ + name: 'llm_loop', + dynamicEntry: async (ctx) => { + let rounds = 0; + for (;;) { + const decision = await ctx.runNode(node(decider), `round ${rounds}`, { + runId: `d${rounds}`, + }); + rounds++; + if (String(decision.output).includes('done')) { + break; + } + if (rounds > 5) { + break; // safety valve + } + } + return {rounds}; + }, + }); + + // "continue", "continue", "done" -> 3 rounds. + expect(finalOutput(await runWorkflowOnce(wf, 'start'))).toEqual({ + rounds: 3, + }); + }); +}); diff --git a/tests/integration/workflows/loop_and_trigger_test.ts b/tests/integration/workflows/loop_and_trigger_test.ts new file mode 100644 index 000000000..cad0d3ab5 --- /dev/null +++ b/tests/integration/workflows/loop_and_trigger_test.ts @@ -0,0 +1,113 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Integration tests for routed (conditional) self-loops, error propagation from + * a parallel branch, and multi-trigger re-execution. + */ + +import { + createEvent, + DEFAULT_ROUTE, + FunctionNode, + JoinNode, + node, + NodeContext, + Workflow, +} from '@google/adk'; +import {describe, expect, it} from 'vitest'; +import {finalOutput, runWorkflowOnce} from './workflow_test_utils.js'; + +describe('workflow integration — routed self-loop', () => { + it('loops a node back to itself until a route condition ends it', async () => { + const init = new FunctionNode('init', () => 0); + // Counter threads via the node output; routes back to itself until >= 3. + const worker = node( + (_c: NodeContext, n: number) => { + const next = (n as number) + 1; + return createEvent({route: next < 3 ? 'again' : 'done', output: next}); + }, + {name: 'worker'}, + ); + const report = node((_c: NodeContext, n: number) => `final:${n}`, { + name: 'report', + }); + + const wf = new Workflow({ + name: 'loop_self', + edges: [ + ['START', init, worker], + [worker, {again: worker, done: report}], + ], + }); + + expect(finalOutput(await runWorkflowOnce(wf, 'go'))).toBe('final:3'); + }); + + it('supports a routed loop with a DEFAULT_ROUTE exit', async () => { + const init = new FunctionNode('init', () => 0); + const worker = node( + (_c: NodeContext, n: number) => { + const next = (n as number) + 1; + // Emit 'again' while looping; no route (=> DEFAULT) when done. + return next < 2 + ? createEvent({route: 'again', output: next}) + : createEvent({output: next}); + }, + {name: 'worker'}, + ); + const done = node((_c: NodeContext, n: number) => `done:${n}`, { + name: 'done', + }); + const wf = new Workflow({ + name: 'loop_default_exit', + edges: [ + ['START', init, worker], + [worker, {again: worker, [DEFAULT_ROUTE]: done}], + ], + }); + expect(finalOutput(await runWorkflowOnce(wf, 'go'))).toBe('done:2'); + }); +}); + +describe('workflow integration — parallel branch failure', () => { + it('fails the workflow when a parallel branch throws', async () => { + const good = new FunctionNode('good', (_c, i: string) => `good(${i})`); + const bad = new FunctionNode('bad', () => { + throw new Error('branch exploded'); + }); + const join = new JoinNode({name: 'join'}); + const wf = new Workflow({ + name: 'parallel_error', + edges: [['START', [good, bad], join]], + }); + await expect(runWorkflowOnce(wf, 'x')).rejects.toThrow('branch exploded'); + }); +}); + +describe('workflow integration — multi-trigger', () => { + it('re-executes a non-join node once per predecessor trigger', async () => { + let cRuns = 0; + const a = new FunctionNode('a', (_c, i: string) => `a(${i})`); + const b = new FunctionNode('b', (_c, i: string) => `b(${i})`); + const c = new FunctionNode('c', (_c, input: string) => { + cRuns++; + return `c(${input})`; + }); + // c has two predecessors and is NOT a JoinNode -> triggered twice. + const wf = new Workflow({ + name: 'multi_triggers', + edges: [ + ['START', [a, b]], + [a, c], + [b, c], + ], + }); + + await runWorkflowOnce(wf, 'x'); + expect(cRuns).toBe(2); + }); +}); From 53a757848845c1ebe32d76df157836eb8d81c52f Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Thu, 23 Jul 2026 16:41:25 -0700 Subject: [PATCH 24/41] docs(workflow): document public workflow types and suppress internal ones Fix typedoc --treatWarningsAsErrors failures introduced by the workflow module. Export the genuinely public types referenced by documented API (CreateEventParams, NodeInfo from events; BuildNodeOptions, FunctionNodeResult, RunNodeOptions from workflow), and mark internal machinery (EventChannel, ScheduleDynamicNode, NodeContextOptions) as intentionallyNotExported so they don't appear in the public API reference. --- core/src/common.ts | 2 +- core/src/workflow/index.ts | 3 +++ typedoc.json | 7 ++++++- 3 files changed, 10 insertions(+), 2 deletions(-) diff --git a/core/src/common.ts b/core/src/common.ts index 8b32cf0a6..f9f1e1ac6 100644 --- a/core/src/common.ts +++ b/core/src/common.ts @@ -139,7 +139,7 @@ export { isFinalResponse, stringifyContent, } from './events/event.js'; -export type {Event} from './events/event.js'; +export type {CreateEventParams, Event, NodeInfo} from './events/event.js'; export {createEventActions} from './events/event_actions.js'; export type {EventActions} from './events/event_actions.js'; export {EventType, toStructuredEvents} from './events/structured_events.js'; diff --git a/core/src/workflow/index.ts b/core/src/workflow/index.ts index 680753936..c3de73ade 100644 --- a/core/src/workflow/index.ts +++ b/core/src/workflow/index.ts @@ -26,6 +26,7 @@ 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 {LLMAgentWrapper} from './nodes/llm_agent_wrapper.js'; @@ -34,6 +35,7 @@ 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'; // --- Graph model --- export {DEFAULT_ROUTE, Edge, Graph} from './graph.js'; @@ -48,6 +50,7 @@ export type { // --- Execution context & state --- export {BranchPath} from './branch_path.js'; export {NodeContext} 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'; diff --git a/typedoc.json b/typedoc.json index 2ec80baa8..3d35f9b28 100644 --- a/typedoc.json +++ b/typedoc.json @@ -9,5 +9,10 @@ "tsconfig": "./core/tsconfig.json", "plugin": ["typedoc-theme-fresh"], "theme": "fresh", - "excludeExternals": true + "excludeExternals": true, + "intentionallyNotExported": [ + "EventChannel", + "ScheduleDynamicNode", + "NodeContextOptions" + ] } From 6f93eddf9c327b284aed24354f880446fa1f5c16 Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Thu, 23 Jul 2026 17:07:52 -0700 Subject: [PATCH 25/41] samples(workflow): add runnable ports of all Python workflow samples Recreate the google/adk-python contributing/samples/workflows samples as runnable TypeScript agents under samples/workflows//agent.ts, each exporting a WorkflowAgent so it runs via the ADK CLI (npm run sample -- ). Covers sequence, route, fan_out_fan_in, parallel_worker, dynamic_nodes, dynamic_fan_out_fan_in, loop, loop_self, multi_triggers, nested_workflow, node_as_tool, state, node_output, use_as_output, message, retry, request_input(+rerun/advanced), auth_api_key, auth_oauth, and agent_in_workflow. Most run offline; agent_in_workflow uses a live LlmAgent. Adds an 'npm run sample' script and a README.\n\nTo make HITL/auth samples resumable from an interactive text CLI, WorkflowAgent now feeds a plain-text reply to pending interrupts when the user does not send a structured function response. --- core/src/workflow/workflow_agent.ts | 42 ++++++++- package.json | 1 + samples/workflows/README.md | 87 +++++++++++++++++++ samples/workflows/agent_in_workflow/agent.ts | 45 ++++++++++ samples/workflows/auth_api_key/agent.ts | 52 +++++++++++ samples/workflows/auth_oauth/agent.ts | 59 +++++++++++++ .../workflows/dynamic_fan_out_fan_in/agent.ts | 37 ++++++++ samples/workflows/dynamic_nodes/agent.ts | 34 ++++++++ samples/workflows/fan_out_fan_in/agent.ts | 54 ++++++++++++ samples/workflows/loop/agent.ts | 68 +++++++++++++++ samples/workflows/loop_self/agent.ts | 67 ++++++++++++++ samples/workflows/message/agent.ts | 54 ++++++++++++ samples/workflows/multi_triggers/agent.ts | 38 ++++++++ samples/workflows/nested_workflow/agent.ts | 68 +++++++++++++++ samples/workflows/node_as_tool/agent.ts | 35 ++++++++ samples/workflows/node_output/agent.ts | 61 +++++++++++++ samples/workflows/parallel_worker/agent.ts | 37 ++++++++ samples/workflows/request_input/agent.ts | 85 ++++++++++++++++++ .../workflows/request_input_advanced/agent.ts | 69 +++++++++++++++ .../workflows/request_input_rerun/agent.ts | 61 +++++++++++++ samples/workflows/retry/agent.ts | 43 +++++++++ samples/workflows/route/agent.ts | 61 +++++++++++++ samples/workflows/sequence/agent.ts | 30 +++++++ samples/workflows/state/agent.ts | 63 ++++++++++++++ samples/workflows/use_as_output/agent.ts | 42 +++++++++ .../workflows/plain_text_resume_test.ts | 59 +++++++++++++ 26 files changed, 1349 insertions(+), 3 deletions(-) create mode 100644 samples/workflows/README.md create mode 100644 samples/workflows/agent_in_workflow/agent.ts create mode 100644 samples/workflows/auth_api_key/agent.ts create mode 100644 samples/workflows/auth_oauth/agent.ts create mode 100644 samples/workflows/dynamic_fan_out_fan_in/agent.ts create mode 100644 samples/workflows/dynamic_nodes/agent.ts create mode 100644 samples/workflows/fan_out_fan_in/agent.ts create mode 100644 samples/workflows/loop/agent.ts create mode 100644 samples/workflows/loop_self/agent.ts create mode 100644 samples/workflows/message/agent.ts create mode 100644 samples/workflows/multi_triggers/agent.ts create mode 100644 samples/workflows/nested_workflow/agent.ts create mode 100644 samples/workflows/node_as_tool/agent.ts create mode 100644 samples/workflows/node_output/agent.ts create mode 100644 samples/workflows/parallel_worker/agent.ts create mode 100644 samples/workflows/request_input/agent.ts create mode 100644 samples/workflows/request_input_advanced/agent.ts create mode 100644 samples/workflows/request_input_rerun/agent.ts create mode 100644 samples/workflows/retry/agent.ts create mode 100644 samples/workflows/route/agent.ts create mode 100644 samples/workflows/sequence/agent.ts create mode 100644 samples/workflows/state/agent.ts create mode 100644 samples/workflows/use_as_output/agent.ts create mode 100644 tests/integration/workflows/plain_text_resume_test.ts diff --git a/core/src/workflow/workflow_agent.ts b/core/src/workflow/workflow_agent.ts index 3a55d8c09..41825d78c 100644 --- a/core/src/workflow/workflow_agent.ts +++ b/core/src/workflow/workflow_agent.ts @@ -11,6 +11,7 @@ import {createEvent, Event} from '../events/event.js'; import {toContent} from './base_node.js'; import {NodeContext} from './node_context.js'; import {EventChannel} from './utils/event_channel.js'; +import {reconstructNodeStates} from './utils/rehydration_utils.js'; import {Workflow} from './workflow.js'; /** Options for a {@link WorkflowAgent}. */ @@ -48,9 +49,11 @@ export class WorkflowAgent extends BaseAgent { channel, nodePath: '', runId: this.name, - // TODO(phase-5b): reconstruct resumeInputs from session function - // responses so an interrupted workflow can resume via the Runner. - resumeInputs: {}, + // 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); @@ -92,6 +95,39 @@ export class WorkflowAgent extends BaseAgent { } } +/** + * When the workflow is paused on unresolved interrupt(s) and the incoming + * message is plain text (not a structured function response), maps that text to + * every pending interrupt id so an interactive client (e.g. `adk run`) can + * resume a HITL/auth pause by simply typing a reply. + */ +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); + } + } + } + + const resumeInputs: Record = {}; + for (const id of pending) { + resumeInputs[id] = text; + } + return resumeInputs; +} + /** * Derives the workflow input from the user message: plain text when the content * is text-only, otherwise the raw `Content` (nodes coerce as needed). diff --git a/package.json b/package.json index 86cee8241..960b33c9e 100644 --- a/package.json +++ b/package.json @@ -11,6 +11,7 @@ }, "scripts": { "build": "npm run build --workspaces", + "sample": "node dev/dist/esm/cli_entrypoint.js run", "clean": "npm run clean --workspaces", "clean:all": "rm package-lock.json && rm -rf ./node_modules && npm run clean:all --workspaces", "rebuild": "npm run clean:all && npm install && npm run build", diff --git a/samples/workflows/README.md b/samples/workflows/README.md new file mode 100644 index 000000000..7240bf3c4 --- /dev/null +++ b/samples/workflows/README.md @@ -0,0 +1,87 @@ +# Workflow samples + +Runnable TypeScript ports of the Python +[`contributing/samples/workflows`](https://github.com/google/adk-python/tree/main/contributing/samples/workflows) +samples, one per directory. Each exports a `rootAgent` (a `WorkflowAgent` +wrapping a `Workflow`) so it runs with the ADK CLI. + +## Running + +Build once, then run any sample by its `agent.ts` path: + +```bash +npm run build # builds @google/adk (and the CLI); needed once / after changes +npm run sample -- samples/workflows/sequence/agent.ts +``` + +`npm run sample -- ` is shorthand for +`node dev/dist/esm/cli_entrypoint.js run `. + +The CLI is interactive: type a message and press Enter to send it to the +workflow; type `exit` to quit. Node events are printed as +`[]: ` and the final line `[]: ...` is the +workflow's output. + +You can also pipe a single message: + +```bash +echo "hello world" | npm run sample -- samples/workflows/sequence/agent.ts +``` + +## API keys + +Most samples are **function-based and run offline** (no key needed). Samples +that call a live model are marked **(needs API key)** below — set +`GEMINI_API_KEY` (a `.env` file in the working directory is loaded +automatically) before running them. + +## Human-in-the-loop / auth samples + +For HITL and auth samples, the workflow **pauses** on the first turn (you'll see +an `adk_request_input` / `adk_request_credential` request). Simply **type your +reply on the next turn** — the plain-text reply is fed to the pending interrupt, +so you can approve/reject, give feedback, or supply an API key interactively. + +To script a multi-turn run non-interactively, use `--replay` with a JSON file of +queries: + +```bash +echo '{"state":{},"queries":["The product broke","approve"]}' > replay.json +npm run sample -- samples/workflows/request_input/agent.ts --replay replay.json +``` + +## Samples + +| Sample | What it shows | Offline? | +| ------------------------ | --------------------------------------------------- | ----------------- | +| `sequence` | Linear chain; each output feeds the next | ✅ | +| `route` | Classify input, route to a branch (+ DEFAULT_ROUTE) | ✅ | +| `fan_out_fan_in` | Parallel branches joined by a `JoinNode` | ✅ | +| `parallel_worker` | Map a node across a list with bounded concurrency | ✅ | +| `dynamic_nodes` | Imperative `dynamicEntry` driving `ctx.runNode()` | ✅ | +| `dynamic_fan_out_fan_in` | Concurrent `ctx.runNode()` + aggregate | ✅ | +| `loop` | Generate → evaluate → route back until it passes | ✅ | +| `loop_self` | A node routes back to itself (conditional cycle) | ✅ | +| `multi_triggers` | A non-join node runs once per predecessor trigger | ✅ | +| `nested_workflow` | A `Workflow` used as a node (+ parallel + join) | ✅ | +| `node_as_tool` | A node calls sub-nodes via `ctx.runNode()` | ✅ | +| `state` | Share data across nodes via `ctx.state` | ✅ | +| `node_output` | Raw value / `Event({output})` / structured output | ✅ | +| `use_as_output` | Promote a sub-node result via `useAsOutput` | ✅ | +| `message` | Emit a display message distinct from output | ✅ | +| `retry` | Retry a flaky node per `retryConfig` | ✅ | +| `request_input` | HITL: draft → review → approve/reject/revise | ✅ (interactive) | +| `request_input_rerun` | HITL single node with `rerunOnResume` | ✅ (interactive) | +| `request_input_advanced` | Auto-approve small / pause for large requests | ✅ (interactive) | +| `auth_api_key` | Pause to request an API-key credential | ✅ (interactive) | +| `auth_oauth` | Pause and emit an OAuth authorization URL | ✅ (request only) | +| `agent_in_workflow` | A real `LlmAgent` as a workflow node | ❗ needs API key | + +### Notes on faithfulness + +Some Python samples use `LlmAgent`s for steps like classification or +generation. To keep the ports runnable offline, those steps are implemented with +function nodes here (the workflow _structure_ is identical); swap a function +node for an `LlmAgent` to use a real model, as shown in `agent_in_workflow`. +`auth_oauth` emits a real authorization request, but completing the OAuth token +exchange requires a live provider, so its resume step is illustrative only. diff --git a/samples/workflows/agent_in_workflow/agent.ts b/samples/workflows/agent_in_workflow/agent.ts new file mode 100644 index 000000000..518d76dec --- /dev/null +++ b/samples/workflows/agent_in_workflow/agent.ts @@ -0,0 +1,45 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Agent in a workflow: a real LlmAgent as a workflow node, between two function + * nodes. Mirrors Python `workflows/agent_in_workflow`. + * + * REQUIRES an API key (this one calls a live model). Set GEMINI_API_KEY (a + * `.env` in the working directory is loaded automatically), then: + * node dev/dist/esm/cli_entrypoint.js run samples/workflows/agent_in_workflow/agent.ts + */ + +import { + LlmAgent, + node, + NodeContext, + Workflow, + WorkflowAgent, +} from '@google/adk'; + +const preprocess = node( + (_c: NodeContext, input: string) => `Please answer this concisely: ${input}`, + {name: 'preprocess'}, +); + +const assistant = new LlmAgent({ + name: 'assistant', + model: 'gemini-2.5-flash', + instruction: 'You are a helpful assistant. Answer the user concisely.', +}); + +const postprocess = node( + (_c: NodeContext, answer: string) => `Assistant replied:\n${answer}`, + {name: 'postprocess'}, +); + +export const rootAgent = new WorkflowAgent( + new Workflow({ + name: 'agent_in_workflow', + edges: [['START', preprocess, assistant, postprocess]], + }), +); diff --git a/samples/workflows/auth_api_key/agent.ts b/samples/workflows/auth_api_key/agent.ts new file mode 100644 index 000000000..af286bdbb --- /dev/null +++ b/samples/workflows/auth_api_key/agent.ts @@ -0,0 +1,52 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Auth (API key): a node requires a credential; it pauses to request one, then + * runs once supplied. Mirrors Python `workflows/auth_api_key`. + * + * Run: node dev/dist/esm/cli_entrypoint.js run samples/workflows/auth_api_key/agent.ts + * Turn 1: any prompt -> asks for an API key. Turn 2: type any API key value. + */ + +import { + AuthConfig, + AuthCredential, + AuthCredentialTypes, + AuthScheme, + node, + NodeContext, + Workflow, + WorkflowAgent, +} from '@google/adk'; + +const CREDENTIAL_KEY = 'weather_api'; + +const authConfig: AuthConfig = { + authScheme: {type: 'apiKey', in: 'header', name: 'X-Api-Key'} as AuthScheme, + rawAuthCredential: {authType: AuthCredentialTypes.API_KEY}, + credentialKey: CREDENTIAL_KEY, +}; + +const fetchWeather = node( + (ctx: NodeContext) => { + const cred = ctx.state.get('temp:' + CREDENTIAL_KEY); + return `Fetched weather using API key "${cred?.apiKey}": sunny, 25C.`; + }, + {name: 'fetch_weather', authConfig}, +); + +const summarize = node( + (_c: NodeContext, weather: string) => `Report: ${weather}`, + {name: 'summarize'}, +); + +export const rootAgent = new WorkflowAgent( + new Workflow({ + name: 'auth_api_key', + edges: [['START', fetchWeather, summarize]], + }), +); diff --git a/samples/workflows/auth_oauth/agent.ts b/samples/workflows/auth_oauth/agent.ts new file mode 100644 index 000000000..d22c7ffaf --- /dev/null +++ b/samples/workflows/auth_oauth/agent.ts @@ -0,0 +1,59 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Auth (OAuth2): a node requiring OAuth pauses and emits an authorization URL + * for the user to complete the flow. Mirrors Python `workflows/auth_oauth`. + * + * Run: node dev/dist/esm/cli_entrypoint.js run samples/workflows/auth_oauth/agent.ts + * Turn 1 emits an `adk_request_credential` interrupt with an auth URL. NOTE: + * completing the real OAuth token exchange requires a live provider, so the + * resume step is illustrative only. + */ + +import { + AuthConfig, + AuthCredentialTypes, + AuthScheme, + node, + NodeContext, + Workflow, + WorkflowAgent, +} from '@google/adk'; + +const authConfig: AuthConfig = { + authScheme: { + type: 'oauth2', + flows: { + authorizationCode: { + authorizationUrl: 'https://accounts.example.com/o/oauth2/v2/auth', + tokenUrl: 'https://oauth2.example.com/token', + scopes: {'https://example.com/auth/calendar.readonly': 'Read calendar'}, + }, + }, + } as AuthScheme, + rawAuthCredential: { + authType: AuthCredentialTypes.OAUTH2, + oauth2: { + clientId: 'demo-client-id', + clientSecret: 'demo-client-secret', + redirectUri: 'http://localhost:8080/callback', + }, + }, + credentialKey: 'example_calendar_oauth', +}; + +const fetchCalendar = node( + (_c: NodeContext) => 'Fetched 3 calendar events using OAuth credentials.', + {name: 'fetch_calendar', authConfig}, +); + +export const rootAgent = new WorkflowAgent( + new Workflow({ + name: 'auth_oauth', + edges: [['START', fetchCalendar]], + }), +); diff --git a/samples/workflows/dynamic_fan_out_fan_in/agent.ts b/samples/workflows/dynamic_fan_out_fan_in/agent.ts new file mode 100644 index 000000000..d9a6afc0a --- /dev/null +++ b/samples/workflows/dynamic_fan_out_fan_in/agent.ts @@ -0,0 +1,37 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Dynamic fan-out/fan-in: an imperative entry runs many nodes concurrently via + * `Promise.all(ctx.runNode(...))` and aggregates the results. Mirrors Python + * `workflows/dynamic_fan_out_fan_in`. + * + * Run: node dev/dist/esm/cli_entrypoint.js run samples/workflows/dynamic_fan_out_fan_in/agent.ts + */ + +import {node, NodeContext, Workflow, WorkflowAgent} from '@google/adk'; + +const square = node( + (_c: NodeContext, n: number) => (n as number) * (n as number), + { + name: 'square', + }, +); + +export const rootAgent = new WorkflowAgent( + new Workflow({ + name: 'dynamic_fan_out_fan_in', + dynamicEntry: async (ctx) => { + const items = [1, 2, 3, 4, 5]; + const results = await Promise.all( + items.map((n, i) => ctx.runNode(square, n, {runId: `sq-${i}`})), + ); + const squares = results.map((r) => r.output as number); + const total = squares.reduce((a, b) => a + b, 0); + return `Squares: ${squares.join(', ')} (sum = ${total})`; + }, + }), +); diff --git a/samples/workflows/dynamic_nodes/agent.ts b/samples/workflows/dynamic_nodes/agent.ts new file mode 100644 index 000000000..999451f37 --- /dev/null +++ b/samples/workflows/dynamic_nodes/agent.ts @@ -0,0 +1,34 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Dynamic nodes: an imperative entry drives execution with plain control flow + * and `ctx.runNode()`. Mirrors Python `workflows/dynamic_nodes`. + * + * Run: node dev/dist/esm/cli_entrypoint.js run samples/workflows/dynamic_nodes/agent.ts + */ + +import {node, NodeContext, Workflow, WorkflowAgent} from '@google/adk'; + +const step = node((_c: NodeContext, n: number) => (n as number) + 1, { + name: 'step', +}); + +export const rootAgent = new WorkflowAgent( + new Workflow({ + name: 'dynamic_nodes', + dynamicEntry: async (ctx) => { + let value = 0; + const trace: number[] = []; + for (let i = 0; i < 3; i++) { + const result = await ctx.runNode(step, value, {runId: `step-${i}`}); + value = result.output as number; + trace.push(value); + } + return `Ran ${trace.length} dynamic steps: ${trace.join(' -> ')}`; + }, + }), +); diff --git a/samples/workflows/fan_out_fan_in/agent.ts b/samples/workflows/fan_out_fan_in/agent.ts new file mode 100644 index 000000000..ae20d2da4 --- /dev/null +++ b/samples/workflows/fan_out_fan_in/agent.ts @@ -0,0 +1,54 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Fan-out / fan-in: run three nodes in parallel on the same input, then join + * their outputs and aggregate. Mirrors Python `workflows/fan_out_fan_in`. + * + * Run (offline, no API key): + * node dev/dist/esm/cli_entrypoint.js run samples/workflows/fan_out_fan_in/agent.ts + */ + +import { + JoinNode, + node, + NodeContext, + Workflow, + WorkflowAgent, +} from '@google/adk'; + +const makeUppercase = node((_c: NodeContext, s: string) => s.toUpperCase(), { + name: 'make_uppercase', +}); +const countCharacters = node((_c: NodeContext, s: string) => s.length, { + name: 'count_characters', +}); +const reverseString = node( + (_c: NodeContext, s: string) => s.split('').reverse().join(''), + {name: 'reverse_string'}, +); + +const aggregate = node( + (_c: NodeContext, results: Record) => + `Uppercase: ${results['make_uppercase']}\n` + + `Character Count: ${results['count_characters']}\n` + + `Reversed: ${results['reverse_string']}`, + {name: 'aggregate'}, +); + +export const rootAgent = new WorkflowAgent( + new Workflow({ + name: 'fan_out_fan_in', + edges: [ + [ + 'START', + [makeUppercase, countCharacters, reverseString], + new JoinNode({name: 'join_for_results'}), + aggregate, + ], + ], + }), +); diff --git a/samples/workflows/loop/agent.ts b/samples/workflows/loop/agent.ts new file mode 100644 index 000000000..bf7f950f8 --- /dev/null +++ b/samples/workflows/loop/agent.ts @@ -0,0 +1,68 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Loop: generate → evaluate → route back until the result passes. Mirrors + * Python `workflows/loop` (generate/evaluate kept function-based to run + * offline; swap for LlmAgents to use a model). + * + * Run: node dev/dist/esm/cli_entrypoint.js run samples/workflows/loop/agent.ts + */ + +import { + createEvent, + node, + NodeContext, + Workflow, + WorkflowAgent, +} from '@google/adk'; + +const processInput = node( + (ctx: NodeContext, topic: string) => { + ctx.state.set('topic', topic); + ctx.state.set('attempt', 0); + return topic; + }, + {name: 'process_input'}, +); + +const generateHeadline = node( + (ctx: NodeContext) => { + const attempt = (ctx.state.get('attempt') ?? 0) + 1; + ctx.state.set('attempt', attempt); + const topic = ctx.state.get('topic'); + return `Headline draft #${attempt} about "${topic}"`; + }, + {name: 'generate_headline'}, +); + +const evaluateHeadline = node( + (ctx: NodeContext, headline: string) => { + // Accept on the 3rd attempt (simulates a grader improving over iterations). + const attempt = ctx.state.get('attempt') ?? 0; + const grade = attempt >= 3 ? 'tech-related' : 'unrelated'; + return createEvent({route: grade, output: headline}); + }, + {name: 'evaluate_headline'}, +); + +const finalize = node( + (_c: NodeContext, headline: string) => `Final headline: ${headline}`, + {name: 'finalize'}, +); + +export const rootAgent = new WorkflowAgent( + new Workflow({ + name: 'loop_sample', + edges: [ + ['START', processInput, generateHeadline, evaluateHeadline], + [ + evaluateHeadline, + {unrelated: generateHeadline, 'tech-related': finalize}, + ], + ], + }), +); diff --git a/samples/workflows/loop_self/agent.ts b/samples/workflows/loop_self/agent.ts new file mode 100644 index 000000000..7c7b3e9a0 --- /dev/null +++ b/samples/workflows/loop_self/agent.ts @@ -0,0 +1,67 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Loop-self: a node routes back to ITSELF until a condition is met (a routed, + * i.e. conditional, cycle). Mirrors Python `workflows/loop_self` — guesses a + * random number until it matches the target. + * + * Run: node dev/dist/esm/cli_entrypoint.js run samples/workflows/loop_self/agent.ts + * Then type a number between 0 and 10. + */ + +import { + createEvent, + node, + NodeContext, + Workflow, + WorkflowAgent, +} from '@google/adk'; + +const validateInput = node( + (ctx: NodeContext, input: string) => { + const n = parseInt(String(input).trim(), 10); + if (Number.isNaN(n) || n < 0 || n > 10) { + throw new Error('Please provide a number between 0 and 10.'); + } + ctx.state.set('target_number', n); + return n; + }, + {name: 'validate_input'}, +); + +const guessNumber = node( + (ctx: NodeContext) => { + const target = ctx.state.get('target_number')!; + const guess = Math.floor(Math.random() * 11); + if (guess === target) { + return createEvent({route: 'correct', output: target}); + } + return createEvent({ + route: 'guessed_wrong', + content: { + role: 'model', + parts: [{text: `Guessed ${guess}, trying again...`}], + }, + }); + }, + {name: 'guess_number'}, +); + +const report = node( + (_c: NodeContext, target: number) => `Correct! The number was ${target}.`, + {name: 'report'}, +); + +export const rootAgent = new WorkflowAgent( + new Workflow({ + name: 'loop_self', + edges: [ + ['START', validateInput, guessNumber], + [guessNumber, {guessed_wrong: guessNumber, correct: report}], + ], + }), +); diff --git a/samples/workflows/message/agent.ts b/samples/workflows/message/agent.ts new file mode 100644 index 000000000..33a0782bf --- /dev/null +++ b/samples/workflows/message/agent.ts @@ -0,0 +1,54 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Message: a node emits a display message (event content) distinct from its + * structured output. Mirrors Python `workflows/message`. + * + * Run: node dev/dist/esm/cli_entrypoint.js run samples/workflows/message/agent.ts + */ + +import { + createEvent, + node, + NodeContext, + Workflow, + WorkflowAgent, +} from '@google/adk'; + +const greet = node( + (_c: NodeContext, name: string) => + createEvent({ + content: { + role: 'model', + parts: [ + { + text: `Hello, ${name}! This event carries a message for display, but no structured output.`, + }, + ], + }, + }), + {name: 'greet'}, +); + +const withOutput = node( + (_c: NodeContext, name: string) => + createEvent({ + content: { + role: 'model', + parts: [{text: `(also produced an output value)`}], + }, + output: {greeted: name}, + }), + {name: 'greet_with_output'}, +); + +export const rootAgent = new WorkflowAgent( + new Workflow({ + name: 'message_sample', + edges: [['START', greet, withOutput]], + }), +); diff --git a/samples/workflows/multi_triggers/agent.ts b/samples/workflows/multi_triggers/agent.ts new file mode 100644 index 000000000..24b900b29 --- /dev/null +++ b/samples/workflows/multi_triggers/agent.ts @@ -0,0 +1,38 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Multi-triggers: a non-join node with several predecessors runs once per + * incoming trigger. Mirrors Python `workflows/multi_triggers`. + * + * Run: node dev/dist/esm/cli_entrypoint.js run samples/workflows/multi_triggers/agent.ts + */ + +import {node, NodeContext, Workflow, WorkflowAgent} from '@google/adk'; + +const producerA = node((_c: NodeContext, i: string) => `A(${i})`, { + name: 'producer_a', +}); +const producerB = node((_c: NodeContext, i: string) => `B(${i})`, { + name: 'producer_b', +}); + +// `collector` is NOT a JoinNode, so it runs once for each predecessor trigger. +const collector = node( + (_c: NodeContext, input: string) => `collected: ${input}`, + {name: 'collector'}, +); + +export const rootAgent = new WorkflowAgent( + new Workflow({ + name: 'multi_triggers', + edges: [ + ['START', [producerA, producerB]], + [producerA, collector], + [producerB, collector], + ], + }), +); diff --git a/samples/workflows/nested_workflow/agent.ts b/samples/workflows/nested_workflow/agent.ts new file mode 100644 index 000000000..631a0d9e9 --- /dev/null +++ b/samples/workflows/nested_workflow/agent.ts @@ -0,0 +1,68 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Nested workflow: a Workflow used as a node inside another workflow, alongside + * a parallel branch and a JoinNode. Mirrors Python `workflows/nested_workflow`. + * + * Run: node dev/dist/esm/cli_entrypoint.js run samples/workflows/nested_workflow/agent.ts + */ + +import { + JoinNode, + node, + NodeContext, + Workflow, + WorkflowAgent, +} from '@google/adk'; + +const processInput = node( + (ctx: NodeContext, year: string) => { + ctx.state.set('year', year.trim()); + return year.trim(); + }, + {name: 'process_input'}, +); + +// A nested workflow: find a name, then a bio. +const findName = node( + (_c: NodeContext, year: string) => `A famous person born in ${year}`, + {name: 'find_name'}, +); +const generateBio = node( + (_c: NodeContext, name: string) => `${name} — a short 3-sentence biography.`, + {name: 'generate_bio'}, +); +const findFamousPerson = new Workflow({ + name: 'find_famous_person', + edges: [['START', findName, generateBio]], +}); + +const findHistoricalEvent = node( + (ctx: NodeContext) => `A significant event in ${ctx.state.get('year')}.`, + {name: 'find_historical_event'}, +); + +const aggregate = node( + (_c: NodeContext, results: Record) => + `Person: ${results['find_famous_person']}\n\nEvent: ${results['find_historical_event']}`, + {name: 'aggregate_results'}, +); + +export const rootAgent = new WorkflowAgent( + new Workflow({ + name: 'nested_workflow', + edges: [ + [ + 'START', + processInput, + [findFamousPerson, findHistoricalEvent], + new JoinNode({name: 'join'}), + aggregate, + ], + ], + }), +); diff --git a/samples/workflows/node_as_tool/agent.ts b/samples/workflows/node_as_tool/agent.ts new file mode 100644 index 000000000..fee6cdcae --- /dev/null +++ b/samples/workflows/node_as_tool/agent.ts @@ -0,0 +1,35 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Node-as-tool: a node imperatively calls other nodes via `ctx.runNode()`. + * Mirrors Python `workflows/node_as_tool`. + * + * Run: node dev/dist/esm/cli_entrypoint.js run samples/workflows/node_as_tool/agent.ts + */ + +import {node, NodeContext, Workflow, WorkflowAgent} from '@google/adk'; + +const add = node( + (_c: NodeContext, args: {a: number; b: number}) => args.a + args.b, + {name: 'add'}, +); + +const orchestrator = node( + async (ctx: NodeContext) => { + const first = await ctx.runNode(add, {a: 2, b: 3}); + const second = await ctx.runNode(add, {a: 10, b: first.output as number}); + return `2 + 3 = ${first.output}, then + 10 = ${second.output}`; + }, + {name: 'orchestrator'}, +); + +export const rootAgent = new WorkflowAgent( + new Workflow({ + name: 'node_as_tool', + edges: [['START', orchestrator]], + }), +); diff --git a/samples/workflows/node_output/agent.ts b/samples/workflows/node_output/agent.ts new file mode 100644 index 000000000..7916ade66 --- /dev/null +++ b/samples/workflows/node_output/agent.ts @@ -0,0 +1,61 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Node output styles: a raw value, an explicit `Event({output})`, and a + * structured object consumed downstream. Mirrors Python `workflows/node_output`. + * + * Run: node dev/dist/esm/cli_entrypoint.js run samples/workflows/node_output/agent.ts + */ + +import { + createEvent, + node, + NodeContext, + Workflow, + WorkflowAgent, +} from '@google/adk'; + +interface TopicDetails { + title: string; + description: string; + category: string; +} + +const stringOutput = node( + (_c: NodeContext, input: string) => `Processed input: ${input}`, + {name: 'generate_string_output'}, +); + +const eventOutput = node( + (_c: NodeContext, input: string) => + createEvent({output: `Event-wrapped output: ${input}`}), + {name: 'generate_event_output'}, +); + +const structuredOutput = node( + (_c: NodeContext, input: string): TopicDetails => ({ + title: 'Generated Topic', + description: `A creative topic based on: ${input}`, + category: 'general', + }), + {name: 'generate_structured_output'}, +); + +const consumeStructured = node( + (_c: NodeContext, details: TopicDetails) => + `Received structured output!\nTitle: ${details.title}\nDescription: ${details.description}\nCategory: ${details.category}`, + {name: 'consume_structured_output'}, +); + +export const rootAgent = new WorkflowAgent( + new Workflow({ + name: 'node_output', + edges: [ + ['START', stringOutput, eventOutput, structuredOutput, consumeStructured], + ], + }), +); diff --git a/samples/workflows/parallel_worker/agent.ts b/samples/workflows/parallel_worker/agent.ts new file mode 100644 index 000000000..8c8ffe6e3 --- /dev/null +++ b/samples/workflows/parallel_worker/agent.ts @@ -0,0 +1,37 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Parallel worker: `node(fn, {parallelWorker: true})` maps a node across each + * item of a list input with bounded concurrency. Mirrors Python + * `workflows/parallel_worker`. + * + * Run: node dev/dist/esm/cli_entrypoint.js run samples/workflows/parallel_worker/agent.ts + */ + +import {node, NodeContext, Workflow, WorkflowAgent} from '@google/adk'; + +const findTopics = node(() => ['ai', 'databases', 'networking'], { + name: 'find_related_topics', +}); + +const explainTopic = node( + (_c: NodeContext, topic: string) => + `${topic.toUpperCase()}: a short explanation of ${topic}.`, + {name: 'explain_topic', parallelWorker: true, maxParallelWorkers: 3}, +); + +const aggregate = node( + (_c: NodeContext, explanations: string[]) => explanations.join('\n\n---\n\n'), + {name: 'aggregate'}, +); + +export const rootAgent = new WorkflowAgent( + new Workflow({ + name: 'parallel_worker', + edges: [['START', findTopics, explainTopic, aggregate]], + }), +); diff --git a/samples/workflows/request_input/agent.ts b/samples/workflows/request_input/agent.ts new file mode 100644 index 000000000..c8a9bab53 --- /dev/null +++ b/samples/workflows/request_input/agent.ts @@ -0,0 +1,85 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Human-in-the-loop: draft an email, pause for human review, then route on the + * reply (approve / reject / feedback-to-revise). Mirrors Python + * `workflows/request_input`. + * + * Run: node dev/dist/esm/cli_entrypoint.js run samples/workflows/request_input/agent.ts + * Turn 1: type a complaint. Turn 2: type "approve", "reject", or feedback text. + */ + +import { + createEvent, + node, + NodeContext, + RequestInput, + Workflow, + WorkflowAgent, +} from '@google/adk'; + +const processInput = node( + (ctx: NodeContext, complaint: string) => { + ctx.state.set('complaint', complaint); + ctx.state.set('feedback', ''); + return complaint; + }, + {name: 'process_input'}, +); + +const draftEmail = node( + (ctx: NodeContext) => { + const complaint = ctx.state.get('complaint'); + const feedback = ctx.state.get('feedback'); + const base = `Dear Customer,\n\nRegarding your complaint ("${complaint}"), we sincerely apologize and will make it right.`; + return feedback ? `${base}\n\n[Revised per feedback: ${feedback}]` : base; + }, + {name: 'draft_email'}, +); + +const humanReview = node( + (ctx: NodeContext, draft: string) => { + const decision = ctx.resumeInputs['review']; + if (decision === undefined) { + return new RequestInput({ + interruptId: 'review', + message: `Please review this draft. Reply "approve", "reject", or give feedback:\n\n---\n${draft}\n---`, + }); + } + const d = String(decision).trim().toLowerCase(); + if (d === 'approve') { + return createEvent({route: 'approved', output: draft}); + } + if (d === 'reject') { + return createEvent({route: 'rejected'}); + } + ctx.state.set('feedback', decision); + return createEvent({route: 'revise'}); + }, + {name: 'human_review'}, +); + +const sendEmail = node( + (_c: NodeContext, draft: string) => `Approved and sent:\n\n${draft}`, + {name: 'send_email'}, +); +const rejectEmail = node(() => 'Draft rejected. No email sent.', { + name: 'reject_email', +}); + +export const rootAgent = new WorkflowAgent( + new Workflow({ + name: 'request_input', + edges: [ + ['START', processInput, draftEmail, humanReview], + [ + humanReview, + {approved: sendEmail, rejected: rejectEmail, revise: draftEmail}, + ], + ], + }), +); diff --git a/samples/workflows/request_input_advanced/agent.ts b/samples/workflows/request_input_advanced/agent.ts new file mode 100644 index 000000000..348fa0701 --- /dev/null +++ b/samples/workflows/request_input_advanced/agent.ts @@ -0,0 +1,69 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Advanced request-input: small requests are auto-approved (no interrupt); + * larger ones pause for manager approval. Mirrors Python + * `workflows/request_input_advanced`. + * + * Run: node dev/dist/esm/cli_entrypoint.js run samples/workflows/request_input_advanced/agent.ts + * Type a number of days. <=1 auto-approves; >1 asks for approval ("yes"/"no"). + */ + +import { + node, + NodeContext, + RequestInput, + Workflow, + WorkflowAgent, +} from '@google/adk'; + +interface Decision { + approved: boolean; + approvedDays: number; +} + +const processRequest = node( + (_c: NodeContext, input: string) => { + const days = Math.max(0, parseInt(String(input).trim(), 10) || 1); + return {days, reason: 'time off'}; + }, + {name: 'process_request'}, +); + +const evaluateRequest = node( + ( + ctx: NodeContext, + req: {days: number; reason: string}, + ): Decision | RequestInput => { + if (req.days <= 1) { + return {approved: true, approvedDays: req.days}; // auto-approve + } + const decision = ctx.resumeInputs['manager_approval']; + if (decision === undefined) { + return new RequestInput({ + interruptId: 'manager_approval', + message: `Approve ${req.days} day(s) off for "${req.reason}"? Reply "yes" or "no".`, + }); + } + const approved = String(decision).trim().toLowerCase().startsWith('y'); + return {approved, approvedDays: approved ? req.days : 0}; + }, + {name: 'evaluate_request'}, +); + +const processDecision = node( + (_c: NodeContext, d: Decision) => + d.approved ? `Approved for ${d.approvedDays} day(s).` : 'Request denied.', + {name: 'process_decision'}, +); + +export const rootAgent = new WorkflowAgent( + new Workflow({ + name: 'request_input_advanced', + edges: [['START', processRequest, evaluateRequest, processDecision]], + }), +); diff --git a/samples/workflows/request_input_rerun/agent.ts b/samples/workflows/request_input_rerun/agent.ts new file mode 100644 index 000000000..f170d8352 --- /dev/null +++ b/samples/workflows/request_input_rerun/agent.ts @@ -0,0 +1,61 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Request-input with rerun-on-resume: a single node both requests input and, + * when resumed, re-runs to consume the reply and route. Mirrors Python + * `workflows/request_input_rerun`. + * + * Run: node dev/dist/esm/cli_entrypoint.js run samples/workflows/request_input_rerun/agent.ts + * Turn 1: describe a task. Turn 2: type "approve" or "reject". + */ + +import { + createEvent, + node, + NodeContext, + RequestInput, + Workflow, + WorkflowAgent, +} from '@google/adk'; + +const plan = node((_c: NodeContext, task: string) => `Plan for: ${task}`, { + name: 'plan', +}); + +const humanReview = node( + (ctx: NodeContext, planText: string) => { + const reply = ctx.resumeInputs['human_review']; + if (reply === undefined) { + return new RequestInput({ + interruptId: 'human_review', + message: `Approve this plan? Reply "approve" or "reject":\n\n${planText}`, + }); + } + return String(reply).toLowerCase().startsWith('a') + ? createEvent({route: 'approved', output: planText}) + : createEvent({route: 'rejected'}); + }, + {name: 'human_review', rerunOnResume: true}, +); + +const execute = node( + (_c: NodeContext, planText: string) => `Executed: ${planText}`, + { + name: 'execute', + }, +); +const cancel = node(() => 'Plan rejected; nothing executed.', {name: 'cancel'}); + +export const rootAgent = new WorkflowAgent( + new Workflow({ + name: 'request_input_rerun', + edges: [ + ['START', plan, humanReview], + [humanReview, {approved: execute, rejected: cancel}], + ], + }), +); diff --git a/samples/workflows/retry/agent.ts b/samples/workflows/retry/agent.ts new file mode 100644 index 000000000..365ac101d --- /dev/null +++ b/samples/workflows/retry/agent.ts @@ -0,0 +1,43 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Retry: a flaky node is retried per its RetryConfig until it succeeds. Mirrors + * Python `workflows/retry`. + * + * Run: node dev/dist/esm/cli_entrypoint.js run samples/workflows/retry/agent.ts + */ + +import {node, NodeContext, Workflow, WorkflowAgent} from '@google/adk'; + +let attempts = 0; + +const getWeather = node( + () => { + attempts++; + if (attempts < 3) { + throw new Error(`Transient upstream error (attempt ${attempts}).`); + } + return 'sunny'; + }, + { + name: 'get_weather', + retryConfig: {maxAttempts: 5, initialDelay: 0.2, jitter: 0}, + }, +); + +const reportWeather = node( + (_c: NodeContext, weather: string) => + `The weather is ${weather} (after ${attempts} attempts).`, + {name: 'report_weather'}, +); + +export const rootAgent = new WorkflowAgent( + new Workflow({ + name: 'retry_sample', + edges: [['START', getWeather, reportWeather]], + }), +); diff --git a/samples/workflows/route/agent.ts b/samples/workflows/route/agent.ts new file mode 100644 index 000000000..abc002394 --- /dev/null +++ b/samples/workflows/route/agent.ts @@ -0,0 +1,61 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Route: classify the input, then route to the matching branch. Mirrors Python + * `workflows/route` (classifier kept function-based to run offline; swap for an + * LlmAgent to classify with a model). + * + * Run: node dev/dist/esm/cli_entrypoint.js run samples/workflows/route/agent.ts + * Try inputs like "What is ADK?" (question) or "ADK is great." (statement). + */ + +import { + createEvent, + DEFAULT_ROUTE, + node, + NodeContext, + Workflow, + WorkflowAgent, +} from '@google/adk'; + +const classify = node( + (_c: NodeContext, input: string) => { + const category = input.trim().endsWith('?') ? 'question' : 'statement'; + return createEvent({route: category, output: input}); + }, + {name: 'classify_input'}, +); + +const answerQuestion = node( + (_c: NodeContext, q: string) => `Answer to "${q}": 42.`, + {name: 'answer_question'}, +); +const commentOnStatement = node( + (_c: NodeContext, s: string) => `Nice statement: "${s}".`, + {name: 'comment_on_statement'}, +); +const handleOther = node( + () => 'I can only answer questions or comment on statements.', + {name: 'handle_other'}, +); + +export const rootAgent = new WorkflowAgent( + new Workflow({ + name: 'route_sample', + edges: [ + ['START', classify], + [ + classify, + { + question: answerQuestion, + statement: commentOnStatement, + [DEFAULT_ROUTE]: handleOther, + }, + ], + ], + }), +); diff --git a/samples/workflows/sequence/agent.ts b/samples/workflows/sequence/agent.ts new file mode 100644 index 000000000..99157f66e --- /dev/null +++ b/samples/workflows/sequence/agent.ts @@ -0,0 +1,30 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Sequence workflow: a linear chain of nodes where each node's output feeds the + * next. Mirrors the Python `workflows/sequence` sample (using function nodes so + * it runs offline without an API key). + * + * Run: npm run sample -- samples/workflows/sequence/agent.ts + */ + +import {node, NodeContext, Workflow, WorkflowAgent} from '@google/adk'; + +const generateFruit = node(() => 'apple', {name: 'generate_fruit'}); + +const describeFruit = node( + (_ctx: NodeContext, fruit: string) => + `A ${fruit} a day keeps the doctor away.`, + {name: 'describe_fruit'}, +); + +export const rootAgent = new WorkflowAgent( + new Workflow({ + name: 'sequence_workflow', + edges: [['START', generateFruit, describeFruit]], + }), +); diff --git a/samples/workflows/state/agent.ts b/samples/workflows/state/agent.ts new file mode 100644 index 000000000..9e2bd6085 --- /dev/null +++ b/samples/workflows/state/agent.ts @@ -0,0 +1,63 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * State: share data across nodes via `ctx.state`. Mirrors Python + * `workflows/state`. (TypeScript reads state explicitly via `ctx.state.get` + * rather than Python's by-name parameter injection.) + * + * Run: node dev/dist/esm/cli_entrypoint.js run samples/workflows/state/agent.ts + */ + +import {node, NodeContext, Workflow, WorkflowAgent} from '@google/adk'; + +const processInitialInput = node( + (ctx: NodeContext, input: string) => { + ctx.state.set('original_text', input); + return input; + }, + {name: 'process_initial_input'}, +); + +const updateStateViaEvent = node( + (ctx: NodeContext, input: string) => { + const upper = input.toUpperCase(); + ctx.state.set('uppercased_text', upper); + return upper; + }, + {name: 'update_state_via_event'}, +); + +const readStateViaCtx = node( + (ctx: NodeContext) => { + const upper = ctx.state.get('uppercased_text'); + const original = ctx.state.get('original_text'); + const appended = `${upper} (Original was: ${original})`; + ctx.state.set('appended_text', appended); + return appended; + }, + {name: 'read_state_via_ctx'}, +); + +const readState = node( + (ctx: NodeContext) => `Final Result: ${ctx.state.get('appended_text')}!`, + {name: 'read_state'}, +); + +export const rootAgent = new WorkflowAgent( + new Workflow({ + name: 'state_sample', + edges: [ + [ + 'START', + processInitialInput, + updateStateViaEvent, + readStateViaCtx, + readState, + ], + ], + }), +); diff --git a/samples/workflows/use_as_output/agent.ts b/samples/workflows/use_as_output/agent.ts new file mode 100644 index 000000000..77499cb08 --- /dev/null +++ b/samples/workflows/use_as_output/agent.ts @@ -0,0 +1,42 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * use_as_output: a node runs a sub-node via `ctx.runNode(..., {useAsOutput})` + * so the sub-node's result becomes the caller's output. Mirrors Python + * `workflows/use_as_output`. + * + * Run: node dev/dist/esm/cli_entrypoint.js run samples/workflows/use_as_output/agent.ts + */ + +import {node, NodeContext, Workflow, WorkflowAgent} from '@google/adk'; + +// Stands in for an LlmAgent summarizer (kept function-based to run offline). +const summarizer = node( + (_c: NodeContext, text: string) => + `Summary: ${String(text).split(/\s+/).slice(0, 6).join(' ')}...`, + {name: 'summarizer'}, +); + +const orchestrate = node( + async (ctx: NodeContext, input: string) => { + const child = await ctx.runNode(summarizer, input, {useAsOutput: true}); + return child.output; + }, + {name: 'orchestrate'}, +); + +const finalize = node( + (_c: NodeContext, summary: string) => `final: ${summary}`, + {name: 'finalize'}, +); + +export const rootAgent = new WorkflowAgent( + new Workflow({ + name: 'use_as_output', + edges: [['START', orchestrate, finalize]], + }), +); diff --git a/tests/integration/workflows/plain_text_resume_test.ts b/tests/integration/workflows/plain_text_resume_test.ts new file mode 100644 index 000000000..cb810717e --- /dev/null +++ b/tests/integration/workflows/plain_text_resume_test.ts @@ -0,0 +1,59 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Interactive resume: a HITL/auth pause can be resumed by a plain-text reply + * (not just a structured function response), which is what enables `adk run` to + * drive HITL workflows by typing a message. + */ + +import { + createEvent, + node, + NodeContext, + RequestInput, + Workflow, +} from '@google/adk'; +import {describe, expect, it} from 'vitest'; +import { + collect, + createWorkflowRunner, + finalOutput, +} from './workflow_test_utils.js'; + +describe('workflow integration — plain-text interactive resume', () => { + it('resumes a HITL node from a plain-text reply (preserving original input)', async () => { + const gate = node( + (ctx: NodeContext, input: string) => { + const reply = ctx.resumeInputs['review']; + if (reply === undefined) { + return new RequestInput({interruptId: 'review', message: 'ok?'}); + } + return createEvent({output: `input=${input} reply=${reply}`}); + }, + {name: 'gate'}, + ); + const wf = new Workflow({ + name: 'plain_text_resume', + edges: [['START', gate]], + }); + const {run} = await createWorkflowRunner(wf); + + // Turn 1: interrupts. + const turn1 = await collect(run('hello')); + expect( + turn1.some((e) => + (e.content?.parts ?? []).some( + (p) => p.functionCall?.name === 'adk_request_input', + ), + ), + ).toBe(true); + + // Turn 2: a plain-text reply resumes the pending interrupt. + const turn2 = await collect(run('approve')); + expect(finalOutput(turn2)).toBe('input=hello reply=approve'); + }); +}); From 180df609a42cd68c3aa62c6aa5eedae7a125dd83 Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Thu, 23 Jul 2026 17:38:46 -0700 Subject: [PATCH 26/41] feat(workflow): expose current attempt count on NodeContext Surface the running node's 1-based attempt number as ctx.attemptCount so node handlers can observe which retry attempt they are on (used by the retry sample). Set per attempt in executeChildNode's retry loop. --- core/src/workflow/node_context.ts | 3 +++ core/src/workflow/node_runner.ts | 1 + 2 files changed, 4 insertions(+) diff --git a/core/src/workflow/node_context.ts b/core/src/workflow/node_context.ts index 26d4f95a0..c1b8a1dff 100644 --- a/core/src/workflow/node_context.ts +++ b/core/src/workflow/node_context.ts @@ -65,6 +65,9 @@ export class NodeContext { */ scheduler?: ScheduleDynamicNode; + /** The current attempt number (1-based) for the running node (see retry). */ + attemptCount = 1; + private readonly _state: State; private readonly dynamicRunCounters = new Map(); diff --git a/core/src/workflow/node_runner.ts b/core/src/workflow/node_runner.ts index ec7915642..a1105fbdc 100644 --- a/core/src/workflow/node_runner.ts +++ b/core/src/workflow/node_runner.ts @@ -101,6 +101,7 @@ export async function executeChildNode( child.output = undefined; child.route = undefined; child.interruptIds = []; + child.attemptCount = nodeState.attemptCount; try { await runOnce(node, child, input, nodeName, branch, isolationScope); break; From 28f67adca16ef9e386a135ec5db9a4ef6bb702cd Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Thu, 23 Jul 2026 17:38:55 -0700 Subject: [PATCH 27/41] samples(workflow): port workflow samples faithfully from adk-python Rework the runnable workflow samples to mirror the Python contributing/samples/workflows sources: use real LlmAgents with zod output schemas where the Python versions do, expand docs/run instructions, and exercise ctx.attemptCount in the retry sample. --- .../workflows/dynamic_fan_out_fan_in/agent.ts | 80 +++++++++++---- samples/workflows/dynamic_nodes/agent.ts | 75 ++++++++++---- samples/workflows/fan_out_fan_in/agent.ts | 34 +++++-- samples/workflows/loop/agent.ts | 82 +++++++++------ samples/workflows/loop_self/agent.ts | 55 +++++------ samples/workflows/message/agent.ts | 99 +++++++++++++------ samples/workflows/multi_triggers/agent.ts | 48 ++++++--- samples/workflows/nested_workflow/agent.ts | 87 +++++++++++----- samples/workflows/node_output/agent.ts | 67 ++++++++----- samples/workflows/parallel_worker/agent.ts | 85 +++++++++++++--- samples/workflows/retry/agent.ts | 47 +++++---- samples/workflows/state/agent.ts | 66 ++++++------- samples/workflows/use_as_output/agent.ts | 41 +++++--- 13 files changed, 581 insertions(+), 285 deletions(-) diff --git a/samples/workflows/dynamic_fan_out_fan_in/agent.ts b/samples/workflows/dynamic_fan_out_fan_in/agent.ts index d9a6afc0a..a7c1136de 100644 --- a/samples/workflows/dynamic_fan_out_fan_in/agent.ts +++ b/samples/workflows/dynamic_fan_out_fan_in/agent.ts @@ -5,33 +5,77 @@ */ /** - * Dynamic fan-out/fan-in: an imperative entry runs many nodes concurrently via - * `Promise.all(ctx.runNode(...))` and aggregates the results. Mirrors Python - * `workflows/dynamic_fan_out_fan_in`. + * Dynamic fan-out / fan-in: an orchestrator node splits a comma-separated input + * into topics, fans out a worker LlmAgent per topic via `ctx.runNode()`, waits + * for all of them, and aggregates the results into a table. Faithful port of + * Python `contributing/samples/workflows/dynamic_fan_out_fan_in`. * - * Run: node dev/dist/esm/cli_entrypoint.js run samples/workflows/dynamic_fan_out_fan_in/agent.ts + * Requires an API key (calls a live model). Set GEMINI_API_KEY, then: + * npm run sample -- samples/workflows/dynamic_fan_out_fan_in/agent.ts + * Enter a comma-separated list of topics, e.g. "space, oceans, volcanoes". */ -import {node, NodeContext, Workflow, WorkflowAgent} from '@google/adk'; +import { + createEvent, + LlmAgent, + node, + NodeContext, + Workflow, + WorkflowAgent, +} from '@google/adk'; -const square = node( - (_c: NodeContext, n: number) => (n as number) * (n as number), - { - name: 'square', +// Worker agent to generate a headline for a single topic. +const generator = new LlmAgent({ + name: 'generator', + model: 'gemini-2.5-flash', + instruction: + 'Write a catchy one-line headline about the topic provided in the user message.', +}); +const generatorNode = node(generator); + +const orchestrator = node( + async function* (ctx: NodeContext, nodeInput: string) { + // Split input comma-separated string into topics. + const topics = String(nodeInput) + .split(',') + .map((t) => t.trim()) + .filter((t) => t.length > 0); + yield createEvent({ + content: { + role: 'model', + parts: [{text: `Processing ${topics.length} topics in parallel.`}], + }, + }); + + // Fan-out: schedule a dynamic node for each topic. + const tasks = topics.map((topic, i) => + ctx.runNode(generatorNode, topic, { + useSubBranch: true, + runId: `gen-${i}`, + }), + ); + + // Wait for all tasks to complete. + const results = await Promise.all(tasks); + + // Fan-in: aggregate results. + let aggregated = '### Aggregated Headlines\n\n'; + aggregated += '| Topic | Headline |\n'; + aggregated += '| :--- | :--- |\n'; + topics.forEach((topic, i) => { + aggregated += `| ${topic} | ${results[i].output} |\n`; + }); + + yield createEvent({ + content: {role: 'model', parts: [{text: aggregated}]}, + }); }, + {name: 'orchestrator', rerunOnResume: true}, ); export const rootAgent = new WorkflowAgent( new Workflow({ name: 'dynamic_fan_out_fan_in', - dynamicEntry: async (ctx) => { - const items = [1, 2, 3, 4, 5]; - const results = await Promise.all( - items.map((n, i) => ctx.runNode(square, n, {runId: `sq-${i}`})), - ); - const squares = results.map((r) => r.output as number); - const total = squares.reduce((a, b) => a + b, 0); - return `Squares: ${squares.join(', ')} (sum = ${total})`; - }, + edges: [['START', orchestrator]], }), ); diff --git a/samples/workflows/dynamic_nodes/agent.ts b/samples/workflows/dynamic_nodes/agent.ts index 999451f37..829f35d28 100644 --- a/samples/workflows/dynamic_nodes/agent.ts +++ b/samples/workflows/dynamic_nodes/agent.ts @@ -5,30 +5,71 @@ */ /** - * Dynamic nodes: an imperative entry drives execution with plain control flow - * and `ctx.runNode()`. Mirrors Python `workflows/dynamic_nodes`. + * Dynamic nodes: an imperative orchestrator drives LlmAgents with `ctx.runNode` + * in a loop until the generated headline is tech-related. Faithful port of + * Python `contributing/samples/workflows/dynamic_nodes`. * - * Run: node dev/dist/esm/cli_entrypoint.js run samples/workflows/dynamic_nodes/agent.ts + * Requires an API key. Set GEMINI_API_KEY, then: + * npm run sample -- samples/workflows/dynamic_nodes/agent.ts */ -import {node, NodeContext, Workflow, WorkflowAgent} from '@google/adk'; +import { + LlmAgent, + node, + NodeContext, + Workflow, + WorkflowAgent, +} from '@google/adk'; +import {z} from 'zod'; -const step = node((_c: NodeContext, n: number) => (n as number) + 1, { - name: 'step', +const feedbackSchema = z.object({ + grade: z.enum(['tech-related', 'unrelated']), + feedback: z.string(), }); +const generateHeadline = node( + new LlmAgent({ + name: 'generate_headline', + model: 'gemini-2.5-flash', + instruction: ` + Write a headline about the topic "{topic}". + If feedback is provided, take it into account. + The feedback: {feedback?} + `, + }), +); + +const evaluateHeadline = node( + new LlmAgent({ + name: 'evaluate_headline', + model: 'gemini-2.5-flash', + instruction: + 'Grade whether the headline is related to technology or software engineering.', + outputSchema: feedbackSchema, + outputKey: 'feedback', + }), +); + +const orchestrate = node( + async function* (ctx: NodeContext, nodeInput: string) { + ctx.state.set('topic', nodeInput); + + for (;;) { + const headline = (await ctx.runNode(generateHeadline)).output as string; + const feedback = (await ctx.runNode(evaluateHeadline, headline)) + .output as {grade: string}; + if (feedback.grade === 'tech-related') { + yield headline; + break; + } + } + }, + {name: 'orchestrate', rerunOnResume: true}, +); + export const rootAgent = new WorkflowAgent( new Workflow({ - name: 'dynamic_nodes', - dynamicEntry: async (ctx) => { - let value = 0; - const trace: number[] = []; - for (let i = 0; i < 3; i++) { - const result = await ctx.runNode(step, value, {runId: `step-${i}`}); - value = result.output as number; - trace.push(value); - } - return `Ran ${trace.length} dynamic steps: ${trace.join(' -> ')}`; - }, + name: 'root_agent', + edges: [['START', orchestrate]], }), ); diff --git a/samples/workflows/fan_out_fan_in/agent.ts b/samples/workflows/fan_out_fan_in/agent.ts index ae20d2da4..37af665ae 100644 --- a/samples/workflows/fan_out_fan_in/agent.ts +++ b/samples/workflows/fan_out_fan_in/agent.ts @@ -5,14 +5,15 @@ */ /** - * Fan-out / fan-in: run three nodes in parallel on the same input, then join - * their outputs and aggregate. Mirrors Python `workflows/fan_out_fan_in`. + * Fan-out / fan-in: run three functions in parallel on the same input, join + * their outputs, and aggregate. Faithful port of Python + * `contributing/samples/workflows/fan_out_fan_in`. * - * Run (offline, no API key): - * node dev/dist/esm/cli_entrypoint.js run samples/workflows/fan_out_fan_in/agent.ts + * Run (offline): npm run sample -- samples/workflows/fan_out_fan_in/agent.ts */ import { + createEvent, JoinNode, node, NodeContext, @@ -31,22 +32,35 @@ const reverseString = node( {name: 'reverse_string'}, ); +const joinNode = new JoinNode({name: 'join_for_results'}); + const aggregate = node( - (_c: NodeContext, results: Record) => - `Uppercase: ${results['make_uppercase']}\n` + - `Character Count: ${results['count_characters']}\n` + - `Reversed: ${results['reverse_string']}`, + async function* (_c: NodeContext, results: Record) { + yield createEvent({ + content: { + role: 'model', + parts: [ + { + text: + `Uppercase: ${results['make_uppercase']}\n\n` + + `Character Count: ${results['count_characters']}\n\n` + + `Reversed: ${results['reverse_string']}\n\n`, + }, + ], + }, + }); + }, {name: 'aggregate'}, ); export const rootAgent = new WorkflowAgent( new Workflow({ - name: 'fan_out_fan_in', + name: 'root_agent', edges: [ [ 'START', [makeUppercase, countCharacters, reverseString], - new JoinNode({name: 'join_for_results'}), + joinNode, aggregate, ], ], diff --git a/samples/workflows/loop/agent.ts b/samples/workflows/loop/agent.ts index bf7f950f8..7f6d9a570 100644 --- a/samples/workflows/loop/agent.ts +++ b/samples/workflows/loop/agent.ts @@ -5,64 +5,82 @@ */ /** - * Loop: generate → evaluate → route back until the result passes. Mirrors - * Python `workflows/loop` (generate/evaluate kept function-based to run - * offline; swap for LlmAgents to use a model). + * Loop: generate a headline, grade it, and route back until it is tech-related. + * Faithful port of Python `contributing/samples/workflows/loop`. * - * Run: node dev/dist/esm/cli_entrypoint.js run samples/workflows/loop/agent.ts + * Requires an API key. Set GEMINI_API_KEY, then: + * npm run sample -- samples/workflows/loop/agent.ts + * Enter a topic, e.g. "the ocean" (loops until the headline is tech-related). */ import { createEvent, + LlmAgent, node, NodeContext, Workflow, WorkflowAgent, } from '@google/adk'; +import {z} from 'zod'; + +const feedbackSchema = z.object({ + grade: z + .enum(['tech-related', 'unrelated']) + .describe( + 'Decide if the headline is related to technology or software engineering.', + ), + feedback: z + .string() + .describe( + 'If the headline is unrelated to technology, provide feedback on how to make it more tech-focused.', + ), +}); const processInput = node( - (ctx: NodeContext, topic: string) => { - ctx.state.set('topic', topic); - ctx.state.set('attempt', 0); - return topic; + (ctx: NodeContext, nodeInput: string) => { + ctx.state.set('topic', nodeInput); }, {name: 'process_input'}, ); -const generateHeadline = node( - (ctx: NodeContext) => { - const attempt = (ctx.state.get('attempt') ?? 0) + 1; - ctx.state.set('attempt', attempt); - const topic = ctx.state.get('topic'); - return `Headline draft #${attempt} about "${topic}"`; - }, - {name: 'generate_headline'}, -); +const generateHeadline = new LlmAgent({ + name: 'generate_headline', + model: 'gemini-2.5-flash', + instruction: ` + Write a headline about the topic "{topic}". + If feedback is provided, take it into account. + The feedback: {feedback?} + `, +}); -const evaluateHeadline = node( - (ctx: NodeContext, headline: string) => { - // Accept on the 3rd attempt (simulates a grader improving over iterations). - const attempt = ctx.state.get('attempt') ?? 0; - const grade = attempt >= 3 ? 'tech-related' : 'unrelated'; - return createEvent({route: grade, output: headline}); - }, - {name: 'evaluate_headline'}, -); +const evaluateHeadline = new LlmAgent({ + name: 'evaluate_headline', + model: 'gemini-2.5-flash', + instruction: ` + Grade whether the headline is related to technology or software engineering. + `, + outputSchema: feedbackSchema, + outputKey: 'feedback', +}); -const finalize = node( - (_c: NodeContext, headline: string) => `Final headline: ${headline}`, - {name: 'finalize'}, +const routeHeadline = node( + (_c: NodeContext, feedback: {grade: string}) => + createEvent({route: feedback.grade}), + {name: 'route_headline'}, ); export const rootAgent = new WorkflowAgent( new Workflow({ - name: 'loop_sample', + name: 'root_agent', edges: [ - ['START', processInput, generateHeadline, evaluateHeadline], [ + 'START', + processInput, + generateHeadline, evaluateHeadline, - {unrelated: generateHeadline, 'tech-related': finalize}, + routeHeadline, ], + [routeHeadline, {unrelated: generateHeadline}], ], }), ); diff --git a/samples/workflows/loop_self/agent.ts b/samples/workflows/loop_self/agent.ts index 7c7b3e9a0..eccb3c18c 100644 --- a/samples/workflows/loop_self/agent.ts +++ b/samples/workflows/loop_self/agent.ts @@ -5,12 +5,11 @@ */ /** - * Loop-self: a node routes back to ITSELF until a condition is met (a routed, - * i.e. conditional, cycle). Mirrors Python `workflows/loop_self` — guesses a - * random number until it matches the target. + * Loop-self: a node routes back to itself until it guesses the target number. + * Faithful port of Python `contributing/samples/workflows/loop_self`. * - * Run: node dev/dist/esm/cli_entrypoint.js run samples/workflows/loop_self/agent.ts - * Then type a number between 0 and 10. + * Run (offline): npm run sample -- samples/workflows/loop_self/agent.ts + * Enter a number between 0 and 10. */ import { @@ -22,46 +21,46 @@ import { } from '@google/adk'; const validateInput = node( - (ctx: NodeContext, input: string) => { - const n = parseInt(String(input).trim(), 10); - if (Number.isNaN(n) || n < 0 || n > 10) { - throw new Error('Please provide a number between 0 and 10.'); + function* (ctx: NodeContext, nodeInput: string) { + const parsed = parseInt(String(nodeInput).trim(), 10); + if (Number.isNaN(parsed) || parsed > 10 || parsed < 0) { + yield createEvent({ + content: { + role: 'model', + parts: [{text: 'Please provide a number between 0 and 10.'}], + }, + }); + throw new Error('Invalid input.'); } - ctx.state.set('target_number', n); - return n; + ctx.state.set('target_number', parsed); }, {name: 'validate_input'}, ); const guessNumber = node( - (ctx: NodeContext) => { - const target = ctx.state.get('target_number')!; + function* (ctx: NodeContext) { + const target = ctx.state.get('target_number'); const guess = Math.floor(Math.random() * 11); + yield createEvent({ + content: {role: 'model', parts: [{text: `Guessing ${guess}...`}]}, + }); if (guess === target) { - return createEvent({route: 'correct', output: target}); + yield createEvent({ + content: {role: 'model', parts: [{text: 'Correct!'}]}, + }); + } else { + yield createEvent({route: 'guessed_wrong'}); } - return createEvent({ - route: 'guessed_wrong', - content: { - role: 'model', - parts: [{text: `Guessed ${guess}, trying again...`}], - }, - }); }, {name: 'guess_number'}, ); -const report = node( - (_c: NodeContext, target: number) => `Correct! The number was ${target}.`, - {name: 'report'}, -); - export const rootAgent = new WorkflowAgent( new Workflow({ - name: 'loop_self', + name: 'root_agent', edges: [ ['START', validateInput, guessNumber], - [guessNumber, {guessed_wrong: guessNumber, correct: report}], + [guessNumber, {guessed_wrong: guessNumber}], ], }), ); diff --git a/samples/workflows/message/agent.ts b/samples/workflows/message/agent.ts index 33a0782bf..5483b8334 100644 --- a/samples/workflows/message/agent.ts +++ b/samples/workflows/message/agent.ts @@ -5,50 +5,89 @@ */ /** - * Message: a node emits a display message (event content) distinct from its - * structured output. Mirrors Python `workflows/message`. + * Message: the many ways a node can emit display messages — plain string, + * multi-modal (text + inline image), multiple messages, and streamed partial + * chunks. Faithful port of Python `contributing/samples/workflows/message`. * - * Run: node dev/dist/esm/cli_entrypoint.js run samples/workflows/message/agent.ts + * Run (offline): npm run sample -- samples/workflows/message/agent.ts */ -import { - createEvent, - node, - NodeContext, - Workflow, - WorkflowAgent, -} from '@google/adk'; - -const greet = node( - (_c: NodeContext, name: string) => - createEvent({ +import {createEvent, node, Workflow, WorkflowAgent} from '@google/adk'; + +// A 16x16 solid red PNG, base64 encoded. +const RED_SQUARE_PNG = + 'iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAAXElEQVR4nO2TSQ7AIAwD' + + '7fz/z+ZQtapwmrJc8QklmjBIgZJgIZMiAIl9KYbhjx4fgwosbNxgMrF0+4uhgHnYDM6' + + 'AzQHJeg5HYtyHFfgy2AztN/5tZWfrBtVzkl4DzfQkEPd+cEkAAAAASUVORK5CYII='; + +const sleep = (ms: number) => new Promise((r) => setTimeout(r, ms)); + +const sendString = node( + async function* () { + yield createEvent({ content: { role: 'model', - parts: [ - { - text: `Hello, ${name}! This event carries a message for display, but no structured output.`, - }, - ], + parts: [{text: '#1 This is a simple string message.'}], }, - }), - {name: 'greet'}, + }); + }, + {name: 'send_string'}, ); -const withOutput = node( - (_c: NodeContext, name: string) => - createEvent({ +const sendMultimodal = node( + async function* () { + yield createEvent({ content: { role: 'model', - parts: [{text: `(also produced an output value)`}], + parts: [ + {text: '#2 Here is a multi-modal message with an inline image:'}, + {inlineData: {data: RED_SQUARE_PNG, mimeType: 'image/png'}}, + ], }, - output: {greeted: name}, - }), - {name: 'greet_with_output'}, + }); + }, + {name: 'send_multimodal'}, +); + +const multipleMessages = node( + async function* () { + const msg = (text: string) => + createEvent({content: {role: 'model', parts: [{text}]}}); + yield msg('#3 Multiple messages'); + await sleep(300); + yield msg('Processing step 1...'); + await sleep(300); + yield msg('Processing step 2...'); + await sleep(300); + yield msg('Done processing.'); + }, + {name: 'multiple_messages'}, +); + +const streamSentence = node( + async function* () { + yield createEvent({ + content: {role: 'model', parts: [{text: '#4 Starting to stream...'}]}, + }); + const sentence = + 'This is a streaming message sent in chunks. ' + + 'You can stream markdown too.'; + for (let i = 0; i < sentence.length; i += 5) { + yield createEvent({ + partial: true, + content: {role: 'model', parts: [{text: sentence.slice(i, i + 5)}]}, + }); + await sleep(100); + } + }, + {name: 'stream_sentence'}, ); export const rootAgent = new WorkflowAgent( new Workflow({ - name: 'message_sample', - edges: [['START', greet, withOutput]], + name: 'message', + edges: [ + ['START', sendString, sendMultimodal, multipleMessages, streamSentence], + ], }), ); diff --git a/samples/workflows/multi_triggers/agent.ts b/samples/workflows/multi_triggers/agent.ts index 24b900b29..e4ac26b58 100644 --- a/samples/workflows/multi_triggers/agent.ts +++ b/samples/workflows/multi_triggers/agent.ts @@ -5,34 +5,50 @@ */ /** - * Multi-triggers: a non-join node with several predecessors runs once per - * incoming trigger. Mirrors Python `workflows/multi_triggers`. + * Multi-triggers: a node with several predecessors runs once per incoming + * trigger. Faithful port of Python `contributing/samples/workflows/multi_triggers`. * - * Run: node dev/dist/esm/cli_entrypoint.js run samples/workflows/multi_triggers/agent.ts + * Run (offline): npm run sample -- samples/workflows/multi_triggers/agent.ts */ -import {node, NodeContext, Workflow, WorkflowAgent} from '@google/adk'; +import { + createEvent, + node, + NodeContext, + Workflow, + WorkflowAgent, +} from '@google/adk'; +import {z} from 'zod'; -const producerA = node((_c: NodeContext, i: string) => `A(${i})`, { - name: 'producer_a', +const makeUppercase = node((_c: NodeContext, s: string) => s.toUpperCase(), { + name: 'make_uppercase', }); -const producerB = node((_c: NodeContext, i: string) => `B(${i})`, { - name: 'producer_b', +const countCharacters = node((_c: NodeContext, s: string) => s.length, { + name: 'count_characters', }); +const reverseString = node( + (_c: NodeContext, s: string) => s.split('').reverse().join(''), + {name: 'reverse_string'}, +); -// `collector` is NOT a JoinNode, so it runs once for each predecessor trigger. -const collector = node( - (_c: NodeContext, input: string) => `collected: ${input}`, - {name: 'collector'}, +const sendMessage = node( + async function* (_c: NodeContext, nodeInput: unknown) { + yield createEvent({ + content: { + role: 'model', + parts: [{text: `Triggered for input: ${nodeInput}`}], + }, + }); + }, + {name: 'send_message'}, ); export const rootAgent = new WorkflowAgent( new Workflow({ - name: 'multi_triggers', + name: 'root_agent', + inputSchema: z.string(), edges: [ - ['START', [producerA, producerB]], - [producerA, collector], - [producerB, collector], + ['START', [makeUppercase, countCharacters, reverseString], sendMessage], ], }), ); diff --git a/samples/workflows/nested_workflow/agent.ts b/samples/workflows/nested_workflow/agent.ts index 631a0d9e9..1ac81a87f 100644 --- a/samples/workflows/nested_workflow/agent.ts +++ b/samples/workflows/nested_workflow/agent.ts @@ -5,14 +5,19 @@ */ /** - * Nested workflow: a Workflow used as a node inside another workflow, alongside - * a parallel branch and a JoinNode. Mirrors Python `workflows/nested_workflow`. + * Nested workflow: a sub-Workflow used as a node, running in parallel with an + * agent, joined and aggregated. Faithful port of Python + * `contributing/samples/workflows/nested_workflow`. * - * Run: node dev/dist/esm/cli_entrypoint.js run samples/workflows/nested_workflow/agent.ts + * Requires an API key. Set GEMINI_API_KEY, then: + * npm run sample -- samples/workflows/nested_workflow/agent.ts + * Enter a 4-digit year, e.g. "1955". */ import { + createEvent, JoinNode, + LlmAgent, node, NodeContext, Workflow, @@ -20,48 +25,80 @@ import { } from '@google/adk'; const processInput = node( - (ctx: NodeContext, year: string) => { - ctx.state.set('year', year.trim()); - return year.trim(); + function* (ctx: NodeContext, nodeInput: string) { + const match = String(nodeInput).match(/\b\d{4}\b/); + if (!match) { + yield createEvent({ + content: { + role: 'model', + parts: [{text: 'Please provide a valid 4-digit year (e.g., 1955).'}], + }, + }); + throw new Error('Invalid year format.'); + } + ctx.state.set('year', match[0]); }, {name: 'process_input'}, ); -// A nested workflow: find a name, then a bio. -const findName = node( - (_c: NodeContext, year: string) => `A famous person born in ${year}`, - {name: 'find_name'}, -); -const generateBio = node( - (_c: NodeContext, name: string) => `${name} — a short 3-sentence biography.`, - {name: 'generate_bio'}, -); +const findName = new LlmAgent({ + name: 'find_name', + model: 'gemini-2.5-flash', + instruction: ` + Find the name of one famous person who was born in this year: {year}. + Return ONLY their name, nothing else. + `, +}); + +const generateBio = new LlmAgent({ + name: 'generate_bio', + model: 'gemini-2.5-flash', + instruction: ` + Write a short, engaging 3-sentence biography for the specified person. + `, +}); + +// Sub-workflow that acts as a single node in the parent workflow. const findFamousPerson = new Workflow({ name: 'find_famous_person', edges: [['START', findName, generateBio]], }); -const findHistoricalEvent = node( - (ctx: NodeContext) => `A significant event in ${ctx.state.get('year')}.`, - {name: 'find_historical_event'}, -); +const findHistoricalEvent = new LlmAgent({ + name: 'find_historical_event', + model: 'gemini-2.5-flash', + instruction: ` + Describe one highly significant historical event that occurred in this year: {year}. + Keep the description to 2 sentences. + `, +}); -const aggregate = node( - (_c: NodeContext, results: Record) => - `Person: ${results['find_famous_person']}\n\nEvent: ${results['find_historical_event']}`, +const joinForAggregation = new JoinNode({name: 'join_for_aggregation'}); + +const aggregateResults = node( + function* (ctx: NodeContext, nodeInput: Record) { + const year = ctx.state.get('year'); + const combined = + `# Year: ${year}\n\n` + + '## Famous Person Bio:\n\n' + + `${nodeInput['find_famous_person']}\n\n` + + '## Historical Event:\n\n' + + `${nodeInput['find_historical_event']}`; + yield createEvent({content: {role: 'model', parts: [{text: combined}]}}); + }, {name: 'aggregate_results'}, ); export const rootAgent = new WorkflowAgent( new Workflow({ - name: 'nested_workflow', + name: 'root_agent', edges: [ [ 'START', processInput, [findFamousPerson, findHistoricalEvent], - new JoinNode({name: 'join'}), - aggregate, + joinForAggregation, + aggregateResults, ], ], }), diff --git a/samples/workflows/node_output/agent.ts b/samples/workflows/node_output/agent.ts index 7916ade66..7aca0462b 100644 --- a/samples/workflows/node_output/agent.ts +++ b/samples/workflows/node_output/agent.ts @@ -5,57 +5,70 @@ */ /** - * Node output styles: a raw value, an explicit `Event({output})`, and a - * structured object consumed downstream. Mirrors Python `workflows/node_output`. + * Node output styles: a raw string, an explicit `Event({output})`, a + * schema-typed LlmAgent output, and a downstream node consuming it. Faithful + * port of Python `contributing/samples/workflows/node_output`. * - * Run: node dev/dist/esm/cli_entrypoint.js run samples/workflows/node_output/agent.ts + * Requires an API key. Set GEMINI_API_KEY, then: + * npm run sample -- samples/workflows/node_output/agent.ts */ import { createEvent, + LlmAgent, node, NodeContext, Workflow, WorkflowAgent, } from '@google/adk'; +import {z} from 'zod'; -interface TopicDetails { - title: string; - description: string; - category: string; -} +const topicDetails = z.object({ + title: z.string().describe('The title of the generated topic.'), + description: z.string().describe('A short description of the topic.'), + category: z.string().describe('The broad category of the topic.'), +}); -const stringOutput = node( - (_c: NodeContext, input: string) => `Processed input: ${input}`, +const generateStringOutput = node( + // Returns a simple string; the framework wraps it in an Event. + (_c: NodeContext, nodeInput: string) => `Processed input: ${nodeInput}`, {name: 'generate_string_output'}, ); -const eventOutput = node( - (_c: NodeContext, input: string) => - createEvent({output: `Event-wrapped output: ${input}`}), +const generateEventOutput = node( + // Explicitly returns an Event for more control. + (_c: NodeContext, nodeInput: string) => + createEvent({output: `Event wrapped output: ${nodeInput}`}), {name: 'generate_event_output'}, ); -const structuredOutput = node( - (_c: NodeContext, input: string): TopicDetails => ({ - title: 'Generated Topic', - description: `A creative topic based on: ${input}`, - category: 'general', - }), - {name: 'generate_structured_output'}, -); +const generatePydanticOutput = new LlmAgent({ + name: 'generate_pydantic_output', + model: 'gemini-2.5-flash', + instruction: 'Generate a creative topic based on the following input.', + outputSchema: topicDetails, +}); -const consumeStructured = node( - (_c: NodeContext, details: TopicDetails) => - `Received structured output!\nTitle: ${details.title}\nDescription: ${details.description}\nCategory: ${details.category}`, - {name: 'consume_structured_output'}, +const consumePydanticOutput = node( + (_c: NodeContext, nodeInput: z.infer) => + 'Received Pydantic Model!\n' + + `Title: ${nodeInput.title}\n` + + `Description: ${nodeInput.description}\n` + + `Category: ${nodeInput.category}`, + {name: 'consume_pydantic_output'}, ); export const rootAgent = new WorkflowAgent( new Workflow({ - name: 'node_output', + name: 'root_agent', edges: [ - ['START', stringOutput, eventOutput, structuredOutput, consumeStructured], + [ + 'START', + generateStringOutput, + generateEventOutput, + generatePydanticOutput, + consumePydanticOutput, + ], ], }), ); diff --git a/samples/workflows/parallel_worker/agent.ts b/samples/workflows/parallel_worker/agent.ts index 8c8ffe6e3..3e5b817e8 100644 --- a/samples/workflows/parallel_worker/agent.ts +++ b/samples/workflows/parallel_worker/agent.ts @@ -5,33 +5,94 @@ */ /** - * Parallel worker: `node(fn, {parallelWorker: true})` maps a node across each - * item of a list input with bounded concurrency. Mirrors Python - * `workflows/parallel_worker`. + * Parallel worker: an LlmAgent generates related topics, each is uppercased and + * explained by a parallel worker (a function and an agent with + * `parallelWorker: true`), then results are aggregated. Faithful port of Python + * `contributing/samples/workflows/parallel_worker`. * - * Run: node dev/dist/esm/cli_entrypoint.js run samples/workflows/parallel_worker/agent.ts + * Requires an API key. Set GEMINI_API_KEY, then: + * npm run sample -- samples/workflows/parallel_worker/agent.ts + * Enter a topic, e.g. "databases". */ -import {node, NodeContext, Workflow, WorkflowAgent} from '@google/adk'; +import { + createEvent, + LlmAgent, + node, + NodeContext, + Workflow, + WorkflowAgent, +} from '@google/adk'; +import {Type} from '@google/genai'; +import {z} from 'zod'; -const findTopics = node(() => ['ai', 'databases', 'networking'], { +interface TopicExplanation { + topic: string; + explanation: string; +} + +const processInput = node( + (ctx: NodeContext, nodeInput: string) => { + ctx.state.set('topic', nodeInput); + }, + {name: 'process_input'}, +); + +const findRelatedTopics = new LlmAgent({ name: 'find_related_topics', + model: 'gemini-2.5-flash', + instruction: + 'Given the specific topic "{topic}", generate a list of 3 related topics.', + outputSchema: {type: Type.ARRAY, items: {type: Type.STRING}}, }); +const makeUpperCase = node( + function* (_c: NodeContext, nodeInput: string) { + yield nodeInput.toUpperCase(); + }, + {name: 'make_upper_case', parallelWorker: true}, +); + const explainTopic = node( - (_c: NodeContext, topic: string) => - `${topic.toUpperCase()}: a short explanation of ${topic}.`, - {name: 'explain_topic', parallelWorker: true, maxParallelWorkers: 3}, + new LlmAgent({ + name: 'explain_topic', + model: 'gemini-2.5-flash', + instruction: + 'Explain how the following topic relates the the original topic: "{topic}".', + outputSchema: z.object({topic: z.string(), explanation: z.string()}), + }), + {parallelWorker: true}, ); const aggregate = node( - (_c: NodeContext, explanations: string[]) => explanations.join('\n\n---\n\n'), + (_c: NodeContext, nodeInput: TopicExplanation[]) => + createEvent({ + content: { + role: 'model', + parts: [ + { + text: nodeInput + .map((e) => `${e.topic}: ${e.explanation}`) + .join('\n\n---\n\n'), + }, + ], + }, + }), {name: 'aggregate'}, ); export const rootAgent = new WorkflowAgent( new Workflow({ - name: 'parallel_worker', - edges: [['START', findTopics, explainTopic, aggregate]], + name: 'root_agent', + edges: [ + [ + 'START', + processInput, + findRelatedTopics, + makeUpperCase, + explainTopic, + aggregate, + ], + ], }), ); diff --git a/samples/workflows/retry/agent.ts b/samples/workflows/retry/agent.ts index 365ac101d..7cc2963c4 100644 --- a/samples/workflows/retry/agent.ts +++ b/samples/workflows/retry/agent.ts @@ -5,39 +5,50 @@ */ /** - * Retry: a flaky node is retried per its RetryConfig until it succeeds. Mirrors - * Python `workflows/retry`. + * Retry: a mock task fails randomly (~70%) and is retried per its RetryConfig, + * using `ctx.attemptCount`. Faithful port of Python + * `contributing/samples/workflows/retry`. * - * Run: node dev/dist/esm/cli_entrypoint.js run samples/workflows/retry/agent.ts + * Run (offline): npm run sample -- samples/workflows/retry/agent.ts */ -import {node, NodeContext, Workflow, WorkflowAgent} from '@google/adk'; - -let attempts = 0; +import { + createEvent, + node, + NodeContext, + Workflow, + WorkflowAgent, +} from '@google/adk'; const getWeather = node( - () => { - attempts++; - if (attempts < 3) { - throw new Error(`Transient upstream error (attempt ${attempts}).`); + async function* (ctx: NodeContext) { + yield createEvent({ + content: { + role: 'model', + parts: [{text: `Getting weather... attempt ${ctx.attemptCount}`}], + }, + }); + if (Math.random() < 0.7) { + // 70% chance of failure + throw new Error('HTTP 500: Internal Server Error'); } - return 'sunny'; - }, - { - name: 'get_weather', - retryConfig: {maxAttempts: 5, initialDelay: 0.2, jitter: 0}, + yield 'sunny'; }, + {name: 'get_weather', retryConfig: {maxAttempts: 5, initialDelay: 1}}, ); const reportWeather = node( - (_c: NodeContext, weather: string) => - `The weather is ${weather} (after ${attempts} attempts).`, + async function* (_c: NodeContext, weather: string) { + yield createEvent({ + content: {role: 'model', parts: [{text: `The weather is ${weather}`}]}, + }); + }, {name: 'report_weather'}, ); export const rootAgent = new WorkflowAgent( new Workflow({ - name: 'retry_sample', + name: 'root_agent', edges: [['START', getWeather, reportWeather]], }), ); diff --git a/samples/workflows/state/agent.ts b/samples/workflows/state/agent.ts index 9e2bd6085..98c586633 100644 --- a/samples/workflows/state/agent.ts +++ b/samples/workflows/state/agent.ts @@ -5,47 +5,39 @@ */ /** - * State: share data across nodes via `ctx.state`. Mirrors Python - * `workflows/state`. (TypeScript reads state explicitly via `ctx.state.get` - * rather than Python's by-name parameter injection.) + * State: several ways to read/write shared workflow state. Faithful port of + * Python `contributing/samples/workflows/state`. (Python's final node reads a + * state value by automatic parameter injection; TypeScript reads it explicitly + * via `ctx.state`.) * - * Run: node dev/dist/esm/cli_entrypoint.js run samples/workflows/state/agent.ts + * Run (offline): npm run sample -- samples/workflows/state/agent.ts */ import {node, NodeContext, Workflow, WorkflowAgent} from '@google/adk'; -const processInitialInput = node( - (ctx: NodeContext, input: string) => { - ctx.state.set('original_text', input); - return input; - }, - {name: 'process_initial_input'}, -); +function processInitialInput(ctx: NodeContext, nodeInput: string): string { + // Set initial input in state via direct dictionary modification. + ctx.state.set('original_text', nodeInput); + return nodeInput; +} -const updateStateViaEvent = node( - (ctx: NodeContext, input: string) => { - const upper = input.toUpperCase(); - ctx.state.set('uppercased_text', upper); - return upper; - }, - {name: 'update_state_via_event'}, -); +function updateStateViaEvent(ctx: NodeContext, nodeInput: string): void { + // Implicitly update the shared workflow state (Python yields Event(state=)). + ctx.state.set('uppercased_text', nodeInput.toUpperCase()); +} -const readStateViaCtx = node( - (ctx: NodeContext) => { - const upper = ctx.state.get('uppercased_text'); - const original = ctx.state.get('original_text'); - const appended = `${upper} (Original was: ${original})`; - ctx.state.set('appended_text', appended); - return appended; - }, - {name: 'read_state_via_ctx'}, -); +function readStateViaCtx(ctx: NodeContext): string { + const original = ctx.state.get('original_text'); + const uppercased = ctx.state.get('uppercased_text'); + const result = `${uppercased} (Original was: ${original})`; + ctx.state.set('appended_text', result); + return result; +} -const readState = node( - (ctx: NodeContext) => `Final Result: ${ctx.state.get('appended_text')}!`, - {name: 'read_state'}, -); +function readStateViaParam(ctx: NodeContext): string { + const appendedText = ctx.state.get('appended_text'); + return `Final Result: ${appendedText}!`; +} export const rootAgent = new WorkflowAgent( new Workflow({ @@ -53,10 +45,10 @@ export const rootAgent = new WorkflowAgent( edges: [ [ 'START', - processInitialInput, - updateStateViaEvent, - readStateViaCtx, - readState, + node(processInitialInput, {name: 'process_initial_input'}), + node(updateStateViaEvent, {name: 'update_state_via_event'}), + node(readStateViaCtx, {name: 'read_state_via_ctx'}), + node(readStateViaParam, {name: 'read_state_via_param'}), ], ], }), diff --git a/samples/workflows/use_as_output/agent.ts b/samples/workflows/use_as_output/agent.ts index 77499cb08..846b42e55 100644 --- a/samples/workflows/use_as_output/agent.ts +++ b/samples/workflows/use_as_output/agent.ts @@ -5,38 +5,49 @@ */ /** - * use_as_output: a node runs a sub-node via `ctx.runNode(..., {useAsOutput})` - * so the sub-node's result becomes the caller's output. Mirrors Python - * `workflows/use_as_output`. + * use_as_output: an orchestrator runs a sub-agent with `useAsOutput`, so the + * sub-agent's result becomes the node's output. Faithful port of Python + * `contributing/samples/workflows/use_as_output`. * - * Run: node dev/dist/esm/cli_entrypoint.js run samples/workflows/use_as_output/agent.ts + * Requires an API key. Set GEMINI_API_KEY, then: + * npm run sample -- samples/workflows/use_as_output/agent.ts + * Paste some text to summarize. */ -import {node, NodeContext, Workflow, WorkflowAgent} from '@google/adk'; +import { + LlmAgent, + node, + NodeContext, + Workflow, + WorkflowAgent, +} from '@google/adk'; -// Stands in for an LlmAgent summarizer (kept function-based to run offline). const summarizer = node( - (_c: NodeContext, text: string) => - `Summary: ${String(text).split(/\s+/).slice(0, 6).join(' ')}...`, - {name: 'summarizer'}, + new LlmAgent({ + name: 'summarizer', + model: 'gemini-2.5-flash', + instruction: 'Summarize the following text in one sentence.', + }), ); const orchestrate = node( - async (ctx: NodeContext, input: string) => { - const child = await ctx.runNode(summarizer, input, {useAsOutput: true}); + async (ctx: NodeContext, nodeInput: string) => { + const child = await ctx.runNode(summarizer, nodeInput, {useAsOutput: true}); return child.output; }, - {name: 'orchestrate'}, + {name: 'orchestrate', rerunOnResume: true}, ); const finalize = node( - (_c: NodeContext, summary: string) => `final: ${summary}`, - {name: 'finalize'}, + (_c: NodeContext, nodeInput: string) => `final: ${nodeInput}`, + { + name: 'finalize', + }, ); export const rootAgent = new WorkflowAgent( new Workflow({ - name: 'use_as_output', + name: 'root_agent', edges: [['START', orchestrate, finalize]], }), ); From 6b08239de95bcfdaa02cab0517cb95b6f8e48c22 Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Thu, 23 Jul 2026 18:11:40 -0700 Subject: [PATCH 28/41] feat(workflow): complete an interrupted node with its resume value on resume For a node with rerunOnResume=false that interrupted last turn (raised interrupts, produced no output), resume now completes it with the resolved resume value(s) as its output instead of re-running its body, feeding the successor node. Mirrors Python's two-node request-input pattern (one node yields RequestInput, its successor receives the reply). --- core/src/workflow/workflow.ts | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/core/src/workflow/workflow.ts b/core/src/workflow/workflow.ts index 417a4d939..89e76d468 100644 --- a/core/src/workflow/workflow.ts +++ b/core/src/workflow/workflow.ts @@ -303,6 +303,36 @@ export class Workflow extends BaseNode { 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; + loop.pending.set( + nodeName, + Promise.resolve({ + name: nodeName, + childCtx: { + output, + route: undefined, + branch: prior.branch ?? ctx.branch, + interruptIds: [], + } as unknown as NodeContext, + }), + ); + return; + } + } + let runId = nodeState.runId; if (!runId) { nodeState.runCounter += 1; From b7104d0aa40097ff34293df3fe9ec55253f0d458 Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Thu, 23 Jul 2026 18:11:47 -0700 Subject: [PATCH 29/41] test(workflow): cover resume-value completion and plain-text resume Add/extend unit and integration coverage for resuming interrupted workflows: nodes completing with their resume value, plain-text resume, and auth-gate resume paths. --- core/test/workflow/auth_gate_test.ts | 6 +- core/test/workflow/resume_test.ts | 6 +- .../workflows/advanced_workflows_test.ts | 3 +- .../workflows/auth_workflow_test.ts | 4 +- .../workflows/core_workflows_test.ts | 57 ++++++++++++++++++- .../workflows/plain_text_resume_test.ts | 3 +- .../workflows/tool_and_resilience_test.ts | 3 +- 7 files changed, 75 insertions(+), 7 deletions(-) diff --git a/core/test/workflow/auth_gate_test.ts b/core/test/workflow/auth_gate_test.ts index 32878d015..96630544d 100644 --- a/core/test/workflow/auth_gate_test.ts +++ b/core/test/workflow/auth_gate_test.ts @@ -48,6 +48,10 @@ describe('Phase 5b-cont — FunctionNode auth gate', () => { 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) => { @@ -56,7 +60,7 @@ describe('Phase 5b-cont — FunctionNode auth gate', () => { sawApiKey = cred?.apiKey; return `data(${cred?.apiKey})`; }, - {authConfig: apiKeyAuthConfig()}, + {authConfig: apiKeyAuthConfig(), rerunOnResume: true}, ); const wf = new Workflow({name: 'auth_wf', edges: [['START', secured]]}); diff --git a/core/test/workflow/resume_test.ts b/core/test/workflow/resume_test.ts index 00cf17710..e3b0d1cb2 100644 --- a/core/test/workflow/resume_test.ts +++ b/core/test/workflow/resume_test.ts @@ -82,6 +82,10 @@ describe('Phase 5b — HITL resume via the Runner', () => { }, {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']; @@ -90,7 +94,7 @@ describe('Phase 5b — HITL resume via the Runner', () => { } return `${input}|${answer}`; }, - {name: 'gate'}, + {name: 'gate', rerunOnResume: true}, ); const c = node((_c: NodeContext, input: unknown) => `C(${input})`, { name: 'c', diff --git a/tests/integration/workflows/advanced_workflows_test.ts b/tests/integration/workflows/advanced_workflows_test.ts index 5fcc1bf7e..64cb097f9 100644 --- a/tests/integration/workflows/advanced_workflows_test.ts +++ b/tests/integration/workflows/advanced_workflows_test.ts @@ -86,7 +86,8 @@ describe('workflow integration — mid-graph HITL resume', () => { } return `${input}|${answer}`; }, - {name: 'gate'}, + // Single-node HITL gate: re-runs on resume to read its answer. + {name: 'gate', rerunOnResume: true}, ); const c = node( (_c: NodeContext, i: string) => { diff --git a/tests/integration/workflows/auth_workflow_test.ts b/tests/integration/workflows/auth_workflow_test.ts index cdadc0683..b01cf86c2 100644 --- a/tests/integration/workflows/auth_workflow_test.ts +++ b/tests/integration/workflows/auth_workflow_test.ts @@ -46,7 +46,9 @@ describe('workflow integration — auth gate (API key)', () => { const cred = ctx.state.get('temp:' + CREDENTIAL_KEY); return `weather(key=${cred?.apiKey})`; }, - {authConfig: apiKeyAuthConfig()}, + // Auth-gated nodes re-run on resume to store the credential and run their + // body (Python's auth samples set rerun_on_resume=True). + {authConfig: apiKeyAuthConfig(), rerunOnResume: true}, ); const wf = new Workflow({ name: 'auth_api_key', diff --git a/tests/integration/workflows/core_workflows_test.ts b/tests/integration/workflows/core_workflows_test.ts index a41bd3e2d..1ac202da0 100644 --- a/tests/integration/workflows/core_workflows_test.ts +++ b/tests/integration/workflows/core_workflows_test.ts @@ -218,7 +218,10 @@ describe('workflow integration — request_input (HITL)', () => { // function-response message. return `${input}:${answer}`; }, - {name: 'gate'}, + // Single-node HITL gate: re-runs on resume to read its answer (Python's + // rerun_on_resume=True). The default (two-node) semantics are covered by + // the request_input two-node test below. + {name: 'gate', rerunOnResume: true}, ); const wf = new Workflow({name: 'request_input', edges: [['START', gate]]}); const {run} = await createWorkflowRunner(wf); @@ -248,4 +251,56 @@ describe('workflow integration — request_input (HITL)', () => { ); expect(finalOutput(turn2)).toBe('start:yes'); }); + + it('two-node pattern: a rerun_on_resume=false node feeds its reply to the next node', async () => { + // Faithful port of Python's `request_input` two-node pattern: one node + // raises the interrupt and, on resume (with the default rerun_on_resume= + // false), does NOT re-run — its output becomes the resume value, which is + // passed as input to its successor. + let askRuns = 0; + const ask = node( + (_c: NodeContext) => { + askRuns++; + return new RequestInput({interruptId: 'review', message: 'reply?'}); + }, + {name: 'ask'}, + ); + const handle = node( + (_c: NodeContext, reply: string) => `handled(${reply})`, + {name: 'handle'}, + ); + const wf = new Workflow({ + name: 'request_input_two_node', + edges: [['START', ask, handle]], + }); + const {run} = await createWorkflowRunner(wf); + + const turn1 = await collect(run('start')); + expect( + turn1.some((e) => + (e.content?.parts ?? []).some( + (p) => p.functionCall?.name === 'adk_request_input', + ), + ), + ).toBe(true); + expect(askRuns).toBe(1); + + const turn2 = await collect( + run({ + role: 'user', + parts: [ + { + functionResponse: { + id: 'review', + name: 'adk_request_input', + response: {result: 'approve'}, + }, + }, + ], + }), + ); + // `ask` did NOT re-run; its reply flowed to `handle` as input. + expect(askRuns).toBe(1); + expect(finalOutput(turn2)).toBe('handled(approve)'); + }); }); diff --git a/tests/integration/workflows/plain_text_resume_test.ts b/tests/integration/workflows/plain_text_resume_test.ts index cb810717e..34942e5be 100644 --- a/tests/integration/workflows/plain_text_resume_test.ts +++ b/tests/integration/workflows/plain_text_resume_test.ts @@ -34,7 +34,8 @@ describe('workflow integration — plain-text interactive resume', () => { } return createEvent({output: `input=${input} reply=${reply}`}); }, - {name: 'gate'}, + // Single-node HITL gate: re-runs on resume to read its reply. + {name: 'gate', rerunOnResume: true}, ); const wf = new Workflow({ name: 'plain_text_resume', diff --git a/tests/integration/workflows/tool_and_resilience_test.ts b/tests/integration/workflows/tool_and_resilience_test.ts index 7805b890c..983c0b32a 100644 --- a/tests/integration/workflows/tool_and_resilience_test.ts +++ b/tests/integration/workflows/tool_and_resilience_test.ts @@ -73,7 +73,8 @@ describe('workflow integration — nested workflow HITL resume', () => { } return `approved:${answer}`; }, - {name: 'gate'}, + // Single-node HITL gate: re-runs on resume to read its answer. + {name: 'gate', rerunOnResume: true}, ); const inner = new Workflow({name: 'inner', edges: [['START', gate]]}); const outer = new Workflow({ From 459914ac7630ef42161d7e7063ff1f00242d01a3 Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Thu, 23 Jul 2026 18:11:55 -0700 Subject: [PATCH 30/41] samples(workflow): port HITL, auth, and routing samples faithfully from adk-python Rework the request-input, auth (API key/OAuth), route, and sequence samples to mirror the Python contributing/samples/workflows sources. --- samples/workflows/auth_api_key/agent.ts | 67 ++++++++-- samples/workflows/auth_oauth/agent.ts | 101 ++++++++++++--- samples/workflows/request_input/agent.ts | 99 +++++++++------ .../workflows/request_input_advanced/agent.ts | 115 +++++++++++++----- .../workflows/request_input_rerun/agent.ts | 83 +++++++++---- samples/workflows/route/agent.ts | 82 +++++++++---- samples/workflows/sequence/agent.ts | 33 +++-- 7 files changed, 428 insertions(+), 152 deletions(-) diff --git a/samples/workflows/auth_api_key/agent.ts b/samples/workflows/auth_api_key/agent.ts index af286bdbb..7f3951cdd 100644 --- a/samples/workflows/auth_api_key/agent.ts +++ b/samples/workflows/auth_api_key/agent.ts @@ -5,11 +5,22 @@ */ /** - * Auth (API key): a node requires a credential; it pauses to request one, then - * runs once supplied. Mirrors Python `workflows/auth_api_key`. + * Auth API Key: a FunctionNode with API-key authentication. The `fetch_weather` + * node declares an `authConfig`, so the framework pauses the workflow and + * requests a credential before running it; once supplied, the node runs with the + * credential available in session state. Faithful port of Python + * `contributing/samples/workflows/auth_api_key`. + * + * The node is `rerunOnResume: true` so that, on resume, it re-runs its body + * (storing the provided credential and fetching), rather than short-circuiting. + * + * TypeScript note: Python reads the credential via `ctx.get_auth_response(cfg)`. + * There is no such helper here; the framework stores the credential at + * `temp:` in state (see AuthHandler), which the node reads. * * Run: node dev/dist/esm/cli_entrypoint.js run samples/workflows/auth_api_key/agent.ts - * Turn 1: any prompt -> asks for an API key. Turn 2: type any API key value. + * Turn 1: any message -> the workflow requests an API key. + * Turn 2: type any API key value; the node runs and echoes it back (masked). */ import { @@ -17,30 +28,68 @@ import { AuthCredential, AuthCredentialTypes, AuthScheme, + createEvent, node, NodeContext, Workflow, WorkflowAgent, } from '@google/adk'; -const CREDENTIAL_KEY = 'weather_api'; +const CREDENTIAL_KEY = 'weather_api_key'; +// Uses API key auth: the simplest credential type. The user is prompted to +// provide an API key via the auth UI (or by typing it in `adk run`). const authConfig: AuthConfig = { authScheme: {type: 'apiKey', in: 'header', name: 'X-Api-Key'} as AuthScheme, - rawAuthCredential: {authType: AuthCredentialTypes.API_KEY}, + rawAuthCredential: { + authType: AuthCredentialTypes.API_KEY, + apiKey: 'placeholder', + }, credentialKey: CREDENTIAL_KEY, }; +interface Weather { + city: string; + temperature: string; + condition: string; + apiKeyUsed: string; +} + +// Fetches weather data using the authenticated API key. const fetchWeather = node( - (ctx: NodeContext) => { + (ctx: NodeContext): Weather => { + // After auth completes, the credential is available in state. const cred = ctx.state.get('temp:' + CREDENTIAL_KEY); - return `Fetched weather using API key "${cred?.apiKey}": sunny, 25C.`; + const apiKey = cred?.apiKey ?? 'unknown'; + + // In a real agent you would use the api_key to call an external API. For + // this sample we just echo it back (masked). + const masked = apiKey.length > 4 ? apiKey.slice(0, 4) + '****' : '****'; + return { + city: 'San Francisco', + temperature: '18C', + condition: 'Sunny', + apiKeyUsed: masked, + }; }, - {name: 'fetch_weather', authConfig}, + {name: 'fetch_weather', authConfig, rerunOnResume: true}, ); +// Displays the weather result. const summarize = node( - (_c: NodeContext, weather: string) => `Report: ${weather}`, + (_ctx: NodeContext, weather: Weather) => + createEvent({ + content: { + role: 'model', + parts: [ + { + text: + `Weather for ${weather.city}: ${weather.temperature}, ` + + `${weather.condition}. (Authenticated with key: ${weather.apiKeyUsed})`, + }, + ], + }, + }), {name: 'summarize'}, ); diff --git a/samples/workflows/auth_oauth/agent.ts b/samples/workflows/auth_oauth/agent.ts index d22c7ffaf..63ddabcf6 100644 --- a/samples/workflows/auth_oauth/agent.ts +++ b/samples/workflows/auth_oauth/agent.ts @@ -5,55 +5,124 @@ */ /** - * Auth (OAuth2): a node requiring OAuth pauses and emits an authorization URL - * for the user to complete the flow. Mirrors Python `workflows/auth_oauth`. + * OAuth Authentication: a FunctionNode with GitHub OAuth2 token request. The + * `list_github_repos` node declares an OAuth2 `authConfig`, so the framework + * pauses the workflow to request a GitHub OAuth token; once the user completes + * the flow, the node calls the GitHub API to list the user's repositories. + * Faithful port of Python `contributing/samples/workflows/auth_oauth`. + * + * To use this sample, register an OAuth application on GitHub and set the + * GITHUB_CLIENT_ID and GITHUB_CLIENT_SECRET environment variables. + * + * TypeScript notes: Python reads the credential via `ctx.get_auth_response(cfg)`; + * here the framework stores it at `temp:` in state. Python uses + * the `requests` library; this port uses the built-in `fetch`. * * Run: node dev/dist/esm/cli_entrypoint.js run samples/workflows/auth_oauth/agent.ts - * Turn 1 emits an `adk_request_credential` interrupt with an auth URL. NOTE: - * completing the real OAuth token exchange requires a live provider, so the - * resume step is illustrative only. + * Turn 1: any message ("start") -> requests GitHub OAuth credentials. + * Turn 2: complete the auth flow to list your repositories. */ import { AuthConfig, + AuthCredential, AuthCredentialTypes, AuthScheme, + createEvent, node, NodeContext, Workflow, WorkflowAgent, } from '@google/adk'; +// Uses GitHub OAuth2 authorization code flow. const authConfig: AuthConfig = { authScheme: { type: 'oauth2', flows: { authorizationCode: { - authorizationUrl: 'https://accounts.example.com/o/oauth2/v2/auth', - tokenUrl: 'https://oauth2.example.com/token', - scopes: {'https://example.com/auth/calendar.readonly': 'Read calendar'}, + authorizationUrl: 'https://github.com/login/oauth/authorize', + tokenUrl: 'https://github.com/login/oauth/access_token', + scopes: { + user: 'Read user profile', + repo: 'Access public repositories', + }, }, }, } as AuthScheme, rawAuthCredential: { authType: AuthCredentialTypes.OAUTH2, oauth2: { - clientId: 'demo-client-id', - clientSecret: 'demo-client-secret', - redirectUri: 'http://localhost:8080/callback', + clientId: process.env.GITHUB_CLIENT_ID ?? 'YOUR_GITHUB_CLIENT_ID', + clientSecret: + process.env.GITHUB_CLIENT_SECRET ?? 'YOUR_GITHUB_CLIENT_SECRET', }, }, - credentialKey: 'example_calendar_oauth', + credentialKey: 'github_oauth_token', }; -const fetchCalendar = node( - (_c: NodeContext) => 'Fetched 3 calendar events using OAuth credentials.', - {name: 'fetch_calendar', authConfig}, +interface RepoResult { + status: 'Success' | 'Error'; + repos?: string[]; + message?: string; +} + +// Fetches GitHub repositories for the authenticated user. +const listGithubRepos = node( + async (ctx: NodeContext): Promise => { + // After auth completes, the credential is available in state. + const cred = ctx.state.get('temp:github_oauth_token'); + const accessToken = cred?.oauth2?.accessToken; + + if (!accessToken) { + return {status: 'Error', message: 'No access token found'}; + } + + // GitHub API requires a User-Agent header. + const headers = { + Authorization: `Bearer ${accessToken}`, + 'User-Agent': 'ADK-Sample-Agent', + Accept: 'application/json', + }; + + try { + const response = await fetch('https://api.github.com/user/repos', { + headers, + }); + if (!response.ok) { + throw new Error(`HTTP ${response.status}`); + } + const reposData = (await response.json()) as Array<{name: string}>; + return {status: 'Success', repos: reposData.map((repo) => repo.name)}; + } catch (e) { + return {status: 'Error', message: `Failed to fetch repos: ${e}`}; + } + }, + {name: 'list_github_repos', authConfig, rerunOnResume: true}, +); + +/** Emits a plain display message (Python `Event(message=...)`). */ +const message = (text: string) => + createEvent({content: {role: 'model', parts: [{text}]}}); + +// Displays the result of accessing the resource. +const displayResult = node( + (_ctx: NodeContext, nodeInput: RepoResult) => { + if (nodeInput.status === 'Success') { + return message( + `Successfully fetched repositories: ${(nodeInput.repos ?? []).join(', ')}`, + ); + } + return message( + `Failed to fetch repositories. Error: ${nodeInput.message ?? 'Unknown error'}`, + ); + }, + {name: 'display_result'}, ); export const rootAgent = new WorkflowAgent( new Workflow({ name: 'auth_oauth', - edges: [['START', fetchCalendar]], + edges: [['START', listGithubRepos, displayResult]], }), ); diff --git a/samples/workflows/request_input/agent.ts b/samples/workflows/request_input/agent.ts index c8a9bab53..4d617da32 100644 --- a/samples/workflows/request_input/agent.ts +++ b/samples/workflows/request_input/agent.ts @@ -5,16 +5,24 @@ */ /** - * Human-in-the-loop: draft an email, pause for human review, then route on the - * reply (approve / reject / feedback-to-revise). Mirrors Python - * `workflows/request_input`. + * Human-in-the-loop (two-node pattern). An LlmAgent drafts a reply to a customer + * complaint; `request_human_review` pauses the workflow (RequestInput) and, on + * resume, its successor `handle_human_review` receives the human's reply as its + * input and routes on it (approve / reject / feedback-to-revise). Faithful port + * of Python `contributing/samples/workflows/request_input`. * - * Run: node dev/dist/esm/cli_entrypoint.js run samples/workflows/request_input/agent.ts + * The two-node split relies on `rerun_on_resume=false` semantics: the node that + * raised the interrupt does NOT re-run on resume; instead it completes with the + * resume value as its output, which feeds the next node. + * + * REQUIRES an API key (draft_email calls a live model). Set GEMINI_API_KEY, then: + * node dev/dist/esm/cli_entrypoint.js run samples/workflows/request_input/agent.ts * Turn 1: type a complaint. Turn 2: type "approve", "reject", or feedback text. */ import { createEvent, + LlmAgent, node, NodeContext, RequestInput, @@ -22,63 +30,80 @@ import { WorkflowAgent, } from '@google/adk'; +/** Emits a plain display message (Python `Event(message=...)`). */ +const message = (text: string) => + createEvent({content: {role: 'model', parts: [{text}]}}); + +// Takes the initial customer complaint and seeds it into workflow state. const processInput = node( (ctx: NodeContext, complaint: string) => { ctx.state.set('complaint', complaint); ctx.state.set('feedback', ''); - return complaint; }, {name: 'process_input'}, ); -const draftEmail = node( - (ctx: NodeContext) => { - const complaint = ctx.state.get('complaint'); - const feedback = ctx.state.get('feedback'); - const base = `Dear Customer,\n\nRegarding your complaint ("${complaint}"), we sincerely apologize and will make it right.`; - return feedback ? `${base}\n\n[Revised per feedback: ${feedback}]` : base; - }, - {name: 'draft_email'}, +const draftEmail = new LlmAgent({ + name: 'draft_email', + model: 'gemini-2.5-flash', + instruction: ` + Please write a polite, helpful response email to the following customer complaint: "{complaint}" + + If there is any feedback from the manager to revise the draft, please incorporate it: "{feedback?}" + `, + outputKey: 'draft', +}); + +// Pauses the workflow to request human review of the draft. With the default +// rerun_on_resume=false, on resume this node does NOT re-run — it completes with +// the reviewer's reply as its output, which becomes handle_human_review's input. +const requestHumanReview = node( + (_ctx: NodeContext, draft: string) => + new RequestInput({ + message: + "Please review the following draft email and provide 'approve', " + + `'reject', or feedback to revise.\n\n---\n${draft}\n---`, + }), + {name: 'request_human_review'}, ); -const humanReview = node( - (ctx: NodeContext, draft: string) => { - const decision = ctx.resumeInputs['review']; - if (decision === undefined) { - return new RequestInput({ - interruptId: 'review', - message: `Please review this draft. Reply "approve", "reject", or give feedback:\n\n---\n${draft}\n---`, - }); - } - const d = String(decision).trim().toLowerCase(); - if (d === 'approve') { - return createEvent({route: 'approved', output: draft}); - } - if (d === 'reject') { +// Receives the human's reply (the resume value) as its input and routes on it. +const handleHumanReview = node( + (ctx: NodeContext, nodeInput: string) => { + if (nodeInput === 'reject') { return createEvent({route: 'rejected'}); } - ctx.state.set('feedback', decision); + if (nodeInput === 'approve') { + return createEvent({route: 'approved'}); + } + ctx.state.set('feedback', nodeInput); return createEvent({route: 'revise'}); }, - {name: 'human_review'}, + {name: 'handle_human_review'}, ); -const sendEmail = node( - (_c: NodeContext, draft: string) => `Approved and sent:\n\n${draft}`, - {name: 'send_email'}, -); -const rejectEmail = node(() => 'Draft rejected. No email sent.', { +const rejectEmail = node(() => message('Draft rejected.'), { name: 'reject_email', }); +const sendEmail = node(() => message('Draft approved and sent successfully.'), { + name: 'send_email', +}); + export const rootAgent = new WorkflowAgent( new Workflow({ name: 'request_input', edges: [ - ['START', processInput, draftEmail, humanReview], [ - humanReview, - {approved: sendEmail, rejected: rejectEmail, revise: draftEmail}, + 'START', + processInput, + draftEmail, + requestHumanReview, + handleHumanReview, + ], + [ + handleHumanReview, + {revise: draftEmail, approved: sendEmail, rejected: rejectEmail}, ], ], }), diff --git a/samples/workflows/request_input_advanced/agent.ts b/samples/workflows/request_input_advanced/agent.ts index 348fa0701..8f5b18bd6 100644 --- a/samples/workflows/request_input_advanced/agent.ts +++ b/samples/workflows/request_input_advanced/agent.ts @@ -5,62 +5,115 @@ */ /** - * Advanced request-input: small requests are auto-approved (no interrupt); - * larger ones pause for manager approval. Mirrors Python - * `workflows/request_input_advanced`. + * Advanced human-in-the-loop with structured schemas. An LlmAgent extracts a + * structured time-off request; `evaluate_request` auto-approves short requests + * (<= 1 day) and otherwise pauses for manager approval (RequestInput carrying a + * response schema). `process_decision` renders the outcome. Faithful port of + * Python `contributing/samples/workflows/request_input_advanced`. * - * Run: node dev/dist/esm/cli_entrypoint.js run samples/workflows/request_input_advanced/agent.ts - * Type a number of days. <=1 auto-approves; >1 asks for approval ("yes"/"no"). + * The manager-approval branch uses `rerun_on_resume=false`: evaluate_request + * does not re-run on resume; its output becomes the manager's decision, which + * feeds process_decision. + * + * REQUIRES an API key (process_request calls a live model). Set GEMINI_API_KEY: + * node dev/dist/esm/cli_entrypoint.js run samples/workflows/request_input_advanced/agent.ts + * Turn 1: e.g. "I need 3 days off next week for a family trip". + * Turn 2 (only if > 1 day): type "yes" or "no" to approve/deny. */ import { + createEvent, + LlmAgent, node, NodeContext, RequestInput, Workflow, WorkflowAgent, } from '@google/adk'; +import {z} from 'zod'; -interface Decision { - approved: boolean; - approvedDays: number; -} +const timeOffRequestSchema = z.object({ + days: z.number().describe('Number of days requested.'), + reason: z.string().describe('Reason for the time off.'), +}); +type TimeOffRequest = z.infer; -const processRequest = node( - (_c: NodeContext, input: string) => { - const days = Math.max(0, parseInt(String(input).trim(), 10) || 1); - return {days, reason: 'time off'}; - }, - {name: 'process_request'}, -); +const timeOffDecisionSchema = z.object({ + approved: z.boolean().describe('Whether the time off is approved.'), + approvedDays: z.number().nullish().describe('Number of days approved.'), +}); +type TimeOffDecision = z.infer; +/** Emits a plain display message (Python `Event(message=...)`). */ +const message = (text: string) => + createEvent({content: {role: 'model', parts: [{text}]}}); + +const processRequest = new LlmAgent({ + name: 'process_request', + model: 'gemini-2.5-flash', + instruction: + "Extract the number of days and the reason from the user's natural " + + 'language time off request.', + outputSchema: timeOffRequestSchema, + outputKey: 'request', +}); + +// If days <= 1, it's auto-approved. Otherwise, route to manager review by +// raising a RequestInput that carries the request as payload and declares the +// expected response schema. const evaluateRequest = node( ( ctx: NodeContext, - req: {days: number; reason: string}, - ): Decision | RequestInput => { - if (req.days <= 1) { - return {approved: true, approvedDays: req.days}; // auto-approve - } - const decision = ctx.resumeInputs['manager_approval']; - if (decision === undefined) { - return new RequestInput({ - interruptId: 'manager_approval', - message: `Approve ${req.days} day(s) off for "${req.reason}"? Reply "yes" or "no".`, - }); + request: TimeOffRequest, + ): TimeOffDecision | RequestInput => { + // Persist the request so process_decision can read it back (TypeScript has + // no signature-based state injection like Python's `request` parameter). + ctx.state.set('request', request); + + if (request.days <= 1) { + return {approved: true, approvedDays: request.days}; } - const approved = String(decision).trim().toLowerCase().startsWith('y'); - return {approved, approvedDays: approved ? req.days : 0}; + return new RequestInput({ + interruptId: 'manager_approval', + message: 'Please review this time off request.', + payload: request, + responseSchema: timeOffDecisionSchema, + }); }, {name: 'evaluate_request'}, ); const processDecision = node( - (_c: NodeContext, d: Decision) => - d.approved ? `Approved for ${d.approvedDays} day(s).` : 'Request denied.', + (ctx: NodeContext, nodeInput: TimeOffDecision | string) => { + const request = ctx.state.get('request'); + const decision = normalizeDecision(nodeInput); + + if (decision.approved) { + const approvedDays = decision.approvedDays ?? request?.days ?? 0; + return message( + `Time Off Approved! ${approvedDays} out of ${request?.days ?? approvedDays} days granted.`, + ); + } + return message('Time Off Denied.'); + }, {name: 'process_decision'}, ); +/** + * Accepts either a structured {@link TimeOffDecision} (auto-approve path or a + * structured resume) or a plain-string reply typed by an interactive user + * (e.g. "yes"/"no"), and normalizes it to a decision. + */ +function normalizeDecision(input: TimeOffDecision | string): TimeOffDecision { + if (typeof input === 'string') { + const yes = ['yes', 'y', 'true', 'approve', 'approved'].includes( + input.trim().toLowerCase(), + ); + return {approved: yes}; + } + return input; +} + export const rootAgent = new WorkflowAgent( new Workflow({ name: 'request_input_advanced', diff --git a/samples/workflows/request_input_rerun/agent.ts b/samples/workflows/request_input_rerun/agent.ts index f170d8352..f0f8f3587 100644 --- a/samples/workflows/request_input_rerun/agent.ts +++ b/samples/workflows/request_input_rerun/agent.ts @@ -5,16 +5,20 @@ */ /** - * Request-input with rerun-on-resume: a single node both requests input and, - * when resumed, re-runs to consume the reply and route. Mirrors Python - * `workflows/request_input_rerun`. + * Human-in-the-loop (single-node, rerun-on-resume). An LlmAgent drafts a reply; + * one `human_review` node both raises the RequestInput and, because it is marked + * `rerunOnResume: true`, RE-RUNS on resume to consume the reply (via + * `ctx.resumeInputs`) and route. Faithful port of Python + * `contributing/samples/workflows/request_input_rerun`. * - * Run: node dev/dist/esm/cli_entrypoint.js run samples/workflows/request_input_rerun/agent.ts - * Turn 1: describe a task. Turn 2: type "approve" or "reject". + * REQUIRES an API key (draft_email calls a live model). Set GEMINI_API_KEY, then: + * node dev/dist/esm/cli_entrypoint.js run samples/workflows/request_input_rerun/agent.ts + * Turn 1: type a complaint. Turn 2: type "approve", "reject", or feedback text. */ import { createEvent, + LlmAgent, node, NodeContext, RequestInput, @@ -22,40 +26,73 @@ import { WorkflowAgent, } from '@google/adk'; -const plan = node((_c: NodeContext, task: string) => `Plan for: ${task}`, { - name: 'plan', +/** Emits a plain display message (Python `Event(message=...)`). */ +const message = (text: string) => + createEvent({content: {role: 'model', parts: [{text}]}}); + +// Takes the initial customer complaint and seeds it into workflow state. +const processInput = node( + (ctx: NodeContext, complaint: string) => { + ctx.state.set('complaint', complaint); + ctx.state.set('feedback', ''); + }, + {name: 'process_input'}, +); + +const draftEmail = new LlmAgent({ + name: 'draft_email', + model: 'gemini-2.5-flash', + instruction: ` + Please write a polite, helpful response email to the following customer complaint: "{complaint}" + + If there is any feedback from the manager to revise the draft, please incorporate it: "{feedback?}" + `, + outputKey: 'draft', }); +// A single node that both requests input and, on resume (it re-runs because +// rerunOnResume is true), consumes the reply from ctx.resumeInputs and routes. const humanReview = node( - (ctx: NodeContext, planText: string) => { - const reply = ctx.resumeInputs['human_review']; - if (reply === undefined) { + (ctx: NodeContext, draft: string) => { + const resumeInput = ctx.resumeInputs['human_review']; + if (!resumeInput) { return new RequestInput({ interruptId: 'human_review', - message: `Approve this plan? Reply "approve" or "reject":\n\n${planText}`, + message: + "Please review the following draft email and provide 'approve', " + + `'reject', or feedback to revise.\n\n---\n${draft}\n---`, }); } - return String(reply).toLowerCase().startsWith('a') - ? createEvent({route: 'approved', output: planText}) - : createEvent({route: 'rejected'}); + + if (resumeInput === 'reject') { + return createEvent({route: 'rejected'}); + } + if (resumeInput === 'approve') { + return createEvent({route: 'approved'}); + } + ctx.state.set('feedback', resumeInput); + return createEvent({route: 'revise'}); }, {name: 'human_review', rerunOnResume: true}, ); -const execute = node( - (_c: NodeContext, planText: string) => `Executed: ${planText}`, - { - name: 'execute', - }, -); -const cancel = node(() => 'Plan rejected; nothing executed.', {name: 'cancel'}); +const rejectEmail = node(() => message('Draft rejected.'), { + name: 'reject_email', +}); + +const sendEmail = node(() => message('Draft approved and sent successfully.'), { + name: 'send_email', +}); export const rootAgent = new WorkflowAgent( new Workflow({ name: 'request_input_rerun', edges: [ - ['START', plan, humanReview], - [humanReview, {approved: execute, rejected: cancel}], + ['START', processInput, draftEmail, humanReview], + [ + humanReview, + {revise: draftEmail, approved: sendEmail, rejected: rejectEmail}, + ], ], }), ); diff --git a/samples/workflows/route/agent.ts b/samples/workflows/route/agent.ts index abc002394..1a20a6c26 100644 --- a/samples/workflows/route/agent.ts +++ b/samples/workflows/route/agent.ts @@ -5,55 +5,91 @@ */ /** - * Route: classify the input, then route to the matching branch. Mirrors Python - * `workflows/route` (classifier kept function-based to run offline; swap for an - * LlmAgent to classify with a model). + * Route: an LlmAgent classifies the input into a category, a routing node emits + * that category as the route, and the matching branch (an LlmAgent, or a + * function for the fallback) handles it. Faithful port of Python + * `contributing/samples/workflows/route`. * - * Run: node dev/dist/esm/cli_entrypoint.js run samples/workflows/route/agent.ts - * Try inputs like "What is ADK?" (question) or "ADK is great." (statement). + * REQUIRES an API key (classification and answers call a live model). Set + * GEMINI_API_KEY, then: + * node dev/dist/esm/cli_entrypoint.js run samples/workflows/route/agent.ts + * Try "What is ADK?" (question) or "ADK is great." (statement). */ import { createEvent, - DEFAULT_ROUTE, + LlmAgent, node, NodeContext, Workflow, WorkflowAgent, } from '@google/adk'; +import {z} from 'zod'; -const classify = node( - (_c: NodeContext, input: string) => { - const category = input.trim().endsWith('?') ? 'question' : 'statement'; - return createEvent({route: category, output: input}); +const inputCategorySchema = z.object({ + category: z.enum(['question', 'statement', 'other']), +}); +type InputCategory = z.infer; + +const processInput = node( + (ctx: NodeContext, nodeInput: string) => { + ctx.state.set('input', nodeInput); }, - {name: 'classify_input'}, + {name: 'process_input'}, ); -const answerQuestion = node( - (_c: NodeContext, q: string) => `Answer to "${q}": 42.`, - {name: 'answer_question'}, -); -const commentOnStatement = node( - (_c: NodeContext, s: string) => `Nice statement: "${s}".`, - {name: 'comment_on_statement'}, +const classifyInput = new LlmAgent({ + name: 'classify_input', + model: 'gemini-2.5-flash', + instruction: + 'Based on this input, decide which category it belongs to: {input}', + outputSchema: inputCategorySchema, + outputKey: 'category', +}); + +// Yields an Event with a specific route based on the classification. +const routeOnCategory = node( + (_ctx: NodeContext, category: InputCategory) => + createEvent({route: category.category}), + {name: 'route_on_category'}, ); + +const answerQuestion = new LlmAgent({ + name: 'answer_question', + model: 'gemini-2.5-flash', + instruction: 'Answer the question: {input}', +}); + +const commentOnStatement = new LlmAgent({ + name: 'comment_on_statement', + model: 'gemini-2.5-flash', + instruction: 'Comment on the statement: {input}', +}); + const handleOther = node( - () => 'I can only answer questions or comment on statements.', + () => + createEvent({ + content: { + role: 'model', + parts: [ + {text: 'Sorry I can only answer questions or comment on statements.'}, + ], + }, + }), {name: 'handle_other'}, ); export const rootAgent = new WorkflowAgent( new Workflow({ - name: 'route_sample', + name: 'root_agent', edges: [ - ['START', classify], + ['START', processInput, classifyInput, routeOnCategory], [ - classify, + routeOnCategory, { question: answerQuestion, statement: commentOnStatement, - [DEFAULT_ROUTE]: handleOther, + other: handleOther, }, ], ], diff --git a/samples/workflows/sequence/agent.ts b/samples/workflows/sequence/agent.ts index 99157f66e..4e272aabf 100644 --- a/samples/workflows/sequence/agent.ts +++ b/samples/workflows/sequence/agent.ts @@ -5,26 +5,33 @@ */ /** - * Sequence workflow: a linear chain of nodes where each node's output feeds the - * next. Mirrors the Python `workflows/sequence` sample (using function nodes so - * it runs offline without an API key). + * Simple sequential workflow with LLM agents: the first agent names a random + * fruit, and its output feeds the second agent, which describes a health benefit + * of that fruit. Faithful port of Python + * `contributing/samples/workflows/sequence`. * - * Run: npm run sample -- samples/workflows/sequence/agent.ts + * REQUIRES an API key (both nodes call a live model). Set GEMINI_API_KEY, then: + * node dev/dist/esm/cli_entrypoint.js run samples/workflows/sequence/agent.ts */ -import {node, NodeContext, Workflow, WorkflowAgent} from '@google/adk'; +import {LlmAgent, Workflow, WorkflowAgent} from '@google/adk'; -const generateFruit = node(() => 'apple', {name: 'generate_fruit'}); +const generateFruitAgent = new LlmAgent({ + name: 'generate_fruit_agent', + model: 'gemini-2.5-flash', + instruction: `Return the name of a random fruit. + Return only the name, nothing else.`, +}); -const describeFruit = node( - (_ctx: NodeContext, fruit: string) => - `A ${fruit} a day keeps the doctor away.`, - {name: 'describe_fruit'}, -); +const generateBenefitAgent = new LlmAgent({ + name: 'generate_benefit_agent', + model: 'gemini-2.5-flash', + instruction: 'Tell me a health benefit about the specified fruit.', +}); export const rootAgent = new WorkflowAgent( new Workflow({ - name: 'sequence_workflow', - edges: [['START', generateFruit, describeFruit]], + name: 'root_agent', + edges: [['START', generateFruitAgent, generateBenefitAgent]], }), ); From 2a01c46669fc7b0e79b3f7005ad9a8a83a2f7dd3 Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Mon, 27 Jul 2026 12:00:34 -0700 Subject: [PATCH 31/41] feat(tools): add FunctionTool require_confirmation for HITL tool approval A tool can now declare require_confirmation (a boolean or a predicate over the call args). When set, the tool pauses the run via the existing tool-confirmation interrupt and only executes once the user approves, mirroring Python's FunctionTool(require_confirmation=...). --- core/src/tools/function_tool.ts | 76 +++++++++ .../tools/function_tool_confirmation_test.ts | 151 ++++++++++++++++++ 2 files changed, 227 insertions(+) create mode 100644 core/test/tools/function_tool_confirmation_test.ts diff --git a/core/src/tools/function_tool.ts b/core/src/tools/function_tool.ts index c51a15c13..638de576e 100644 --- a/core/src/tools/function_tool.ts +++ b/core/src/tools/function_tool.ts @@ -57,6 +57,20 @@ export type ToolOptions = { parameters?: TParameters; execute: ToolExecuteFunction; isLongRunning?: boolean; + /** + * Whether this tool requires user confirmation before it runs. A boolean, or + * a predicate over the (validated) call arguments and tool context returning + * a boolean. When confirmation is required the tool pauses the run (HITL): + * the framework emits an `adk_request_confirmation` interrupt, and the tool + * only executes once the user approves. Mirrors Python's + * `FunctionTool(require_confirmation=...)`. + */ + requireConfirmation?: + | boolean + | (( + input: ToolExecuteArgument, + tool_context?: Context, + ) => boolean | Promise); }; function toSchema( @@ -111,6 +125,13 @@ export class FunctionTool< private readonly execute: ToolExecuteFunction; // Typed input parameters. private readonly parameters?: TParameters; + // Whether the tool requires user confirmation before running. + private readonly requireConfirmation: + | boolean + | (( + input: ToolExecuteArgument, + tool_context?: Context, + ) => boolean | Promise); /** * The constructor acts as the user-friendly factory. @@ -130,6 +151,7 @@ export class FunctionTool< }); this.execute = options.execute; this.parameters = options.parameters; + this.requireConfirmation = options.requireConfirmation ?? false; } /** @@ -157,6 +179,19 @@ export class FunctionTool< if (isZodObject(this.parameters)) { validatedArgs = this.parameters.parse(req.args); } + + // HITL confirmation gate (Python `require_confirmation`). On the first + // pass we record a confirmation request and pause; on resume the tool + // context carries the user's decision. + const confirmationResult = this.checkConfirmation( + validatedArgs as ToolExecuteArgument, + req.toolContext, + ); + const pending = await confirmationResult; + if (pending !== undefined) { + return pending; + } + return await this.execute( validatedArgs as ToolExecuteArgument, req.toolContext, @@ -167,4 +202,45 @@ export class FunctionTool< throw new Error(`Error in tool '${this.name}': ${errorMessage}`); } } + + /** + * Evaluates the confirmation gate. Returns `undefined` if the tool may + * proceed; otherwise returns the function response payload to surface instead + * of running (a request-for-confirmation on the first pass, or a rejection + * once the user declined). + */ + private async checkConfirmation( + input: ToolExecuteArgument, + toolContext?: Context, + ): Promise<{error: string} | undefined> { + const requireConfirmation = + typeof this.requireConfirmation === 'function' + ? await this.requireConfirmation(input, toolContext) + : this.requireConfirmation; + if (!requireConfirmation) { + return undefined; + } + if (!toolContext) { + throw new Error( + `Tool '${this.name}' requires confirmation but no tool context was provided.`, + ); + } + if (!toolContext.toolConfirmation) { + toolContext.requestConfirmation({ + hint: + `Please approve or reject the tool call ${this.name}() by ` + + 'responding with a FunctionResponse with an expected ' + + 'ToolConfirmation payload.', + }); + toolContext.actions.skipSummarization = true; + return { + error: + 'This tool call requires confirmation, please approve or reject.', + }; + } + if (!toolContext.toolConfirmation.confirmed) { + return {error: 'This tool call is rejected.'}; + } + return undefined; + } } diff --git a/core/test/tools/function_tool_confirmation_test.ts b/core/test/tools/function_tool_confirmation_test.ts new file mode 100644 index 000000000..de9cfa912 --- /dev/null +++ b/core/test/tools/function_tool_confirmation_test.ts @@ -0,0 +1,151 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { + Context, + createSession, + FunctionTool, + InvocationContext, + PluginManager, + ToolConfirmation, +} from '@google/adk'; +import {describe, expect, it} from 'vitest'; +import {z} from 'zod/v3'; + +function makeContext(options: { + functionCallId?: string; + toolConfirmation?: ToolConfirmation; +}): Context { + const session = createSession({ + id: 's1', + appName: 'app', + userId: 'u1', + }); + const invocationContext = new InvocationContext({ + invocationId: 'inv-1', + agent: {name: 'a', runAsync: async function* () {}} as never, + session, + pluginManager: new PluginManager([]), + }); + return new Context({invocationContext, ...options}); +} + +describe('FunctionTool require_confirmation', () => { + function makeTool() { + let ran = false; + const tool = new FunctionTool({ + name: 'delete_file', + description: 'Deletes a file.', + parameters: z.object({path: z.string()}), + execute: () => { + ran = true; + return 'deleted'; + }, + requireConfirmation: true, + }); + return {tool, didRun: () => ran}; + } + + it('pauses and requests confirmation on first call', async () => { + const {tool, didRun} = makeTool(); + const ctx = makeContext({functionCallId: 'fc-1'}); + + const result = await tool.runAsync({ + args: {path: '/tmp/x'}, + toolContext: ctx, + }); + + expect(result).toEqual({ + error: 'This tool call requires confirmation, please approve or reject.', + }); + expect(didRun()).toBe(false); + expect(ctx.actions.requestedToolConfirmations['fc-1']).toBeDefined(); + expect(ctx.actions.skipSummarization).toBe(true); + }); + + it('runs the tool once the call is confirmed', async () => { + const {tool, didRun} = makeTool(); + const ctx = makeContext({ + functionCallId: 'fc-1', + toolConfirmation: new ToolConfirmation({confirmed: true}), + }); + + const result = await tool.runAsync({ + args: {path: '/tmp/x'}, + toolContext: ctx, + }); + + expect(result).toBe('deleted'); + expect(didRun()).toBe(true); + }); + + it('rejects the tool call when confirmation is denied', async () => { + const {tool, didRun} = makeTool(); + const ctx = makeContext({ + functionCallId: 'fc-1', + toolConfirmation: new ToolConfirmation({confirmed: false}), + }); + + const result = await tool.runAsync({ + args: {path: '/tmp/x'}, + toolContext: ctx, + }); + + expect(result).toEqual({error: 'This tool call is rejected.'}); + expect(didRun()).toBe(false); + }); + + it('runs immediately when confirmation is not required', async () => { + let ran = false; + const tool = new FunctionTool({ + name: 'noop', + description: 'no-op', + execute: () => { + ran = true; + return 'ok'; + }, + }); + const ctx = makeContext({functionCallId: 'fc-1'}); + + const result = await tool.runAsync({args: {}, toolContext: ctx}); + + expect(result).toBe('ok'); + expect(ran).toBe(true); + }); + + it('supports a predicate to decide confirmation per-args', async () => { + let ran = false; + const tool = new FunctionTool({ + name: 'transfer', + description: 'Transfers money.', + parameters: z.object({amount: z.number()}), + execute: () => { + ran = true; + return 'sent'; + }, + requireConfirmation: (input) => input.amount > 100, + }); + + // Small amount: no confirmation required, runs directly. + const smallCtx = makeContext({functionCallId: 'fc-small'}); + expect( + await tool.runAsync({args: {amount: 10}, toolContext: smallCtx}), + ).toBe('sent'); + expect(ran).toBe(true); + + // Large amount: confirmation required, pauses. + ran = false; + const largeCtx = makeContext({functionCallId: 'fc-large'}); + const result = await tool.runAsync({ + args: {amount: 1000}, + toolContext: largeCtx, + }); + expect(result).toEqual({ + error: 'This tool call requires confirmation, please approve or reject.', + }); + expect(ran).toBe(false); + }); +}); From 9efb28c42bcffc58d6ad67e0e403758d09a3a756 Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Mon, 27 Jul 2026 12:00:48 -0700 Subject: [PATCH 32/41] feat(agents): run nodes/workflows as agent tools and add LlmAgent task mode Ports the remaining LlmAgent<->workflow integration from adk-python so the last two Python workflow samples can be recreated faithfully: - NodeTool: a BaseNode/Workflow passed in an LlmAgent's tools is auto-wrapped as a callable tool. An optional event queue on InvocationContext lets a node-tool stream its intermediate and interrupt events into the agent's output; a node that raises RequestInput pauses the invocation and resumes on the next turn (RequestInputLlmRequestProcessor re-runs it with the reply as resumeInputs). - LlmAgent task mode: a task-mode agent is given a finish_task tool (schema mirrors outputSchema) and loops until it calls finish_task, whose args become the node output. JSON response mode is suppressed in task mode so function calling works. - Tool confirmations can now be resolved by a plain-text reply (parity with interactive request-input resume). --- core/src/agents/invocation_context.ts | 10 + core/src/agents/llm_agent.ts | 84 ++++++- .../processors/basic_llm_request_processor.ts | 9 +- ...uest_confirmation_llm_request_processor.ts | 90 ++++++++ .../request_input_llm_request_processor.ts | 205 ++++++++++++++++++ core/src/common.ts | 5 + core/src/tools/finish_task_tool.ts | 140 ++++++++++++ core/src/workflow/index.ts | 1 + core/src/workflow/nodes/llm_agent_wrapper.ts | 75 ++++++- core/src/workflow/nodes/node_tool.ts | 143 ++++++++++++ .../workflow/utils/workflow_graph_utils.ts | 1 + .../workflows/node_as_tool_hitl_test.ts | 121 +++++++++++ .../workflows/node_as_tool_test.ts | 80 +++++++ tests/integration/workflows/task_mode_test.ts | 69 ++++++ .../workflows/workflow_test_utils.ts | 18 ++ 15 files changed, 1040 insertions(+), 11 deletions(-) create mode 100644 core/src/agents/processors/request_input_llm_request_processor.ts create mode 100644 core/src/tools/finish_task_tool.ts create mode 100644 core/src/workflow/nodes/node_tool.ts create mode 100644 tests/integration/workflows/node_as_tool_hitl_test.ts create mode 100644 tests/integration/workflows/node_as_tool_test.ts create mode 100644 tests/integration/workflows/task_mode_test.ts diff --git a/core/src/agents/invocation_context.ts b/core/src/agents/invocation_context.ts index 6b3312481..6f8df4665 100644 --- a/core/src/agents/invocation_context.ts +++ b/core/src/agents/invocation_context.ts @@ -8,11 +8,13 @@ import {Content} from '@google/genai'; import {SessionArtifactService} from '../artifacts/session_artifact_service.js'; import {BaseCredentialService} from '../auth/credential_service/base_credential_service.js'; +import {Event} from '../events/event.js'; import {BaseMemoryService} from '../memory/base_memory_service.js'; import {PluginManager} from '../plugins/plugin_manager.js'; import {BaseSessionService} from '../sessions/base_session_service.js'; import {Session} from '../sessions/session.js'; import {randomUUID} from '../utils/env_aware_utils.js'; +import {EventChannel} from '../workflow/utils/event_channel.js'; import {ActiveStreamingTool} from './active_streaming_tool.js'; import {BaseAgent} from './base_agent.js'; @@ -187,6 +189,14 @@ export class InvocationContext { readonly abortSignal?: AbortSignal; + /** + * An optional channel into which a running tool can push events to be + * interleaved into the agent's output stream. Set by the LLM flow around tool + * execution so a {@link NodeTool} (running a node/workflow) can surface the + * node's intermediate and interrupt events. Cleared once tools finish. + */ + eventQueue?: EventChannel; + /** * Checkpointed states for workflow nodes under this invocation. */ diff --git a/core/src/agents/llm_agent.ts b/core/src/agents/llm_agent.ts index d1c794de2..c79bbbab1 100644 --- a/core/src/agents/llm_agent.ts +++ b/core/src/agents/llm_agent.ts @@ -6,7 +6,11 @@ import {GenerateContentConfig, Schema} from '@google/genai'; import {context, trace} from '@opentelemetry/api'; +import {FinishTaskTool} from '../tools/finish_task_tool.js'; import {FunctionTool} from '../tools/function_tool.js'; +import {BaseNode} from '../workflow/base_node.js'; +import {NodeTool} from '../workflow/nodes/node_tool.js'; +import {EventChannel} from '../workflow/utils/event_channel.js'; import {z as z3} from 'zod/v3'; import {z as z4} from 'zod/v4'; @@ -67,6 +71,7 @@ import {IDENTITY_LLM_REQUEST_PROCESSOR} from './processors/identity_llm_request_ import {INSTRUCTIONS_LLM_REQUEST_PROCESSOR} from './processors/instructions_llm_request_processor.js'; import {INTERACTIONS_REQUEST_PROCESSOR} from './processors/interactions_request_processor.js'; import {REQUEST_CONFIRMATION_LLM_REQUEST_PROCESSOR} from './processors/request_confirmation_llm_request_processor.js'; +import {REQUEST_INPUT_LLM_REQUEST_PROCESSOR} from './processors/request_input_llm_request_processor.js'; import {TOOL_FILTER_REQUEST_PROCESSOR} from './processors/tool_filter_request_processor.js'; import {ReadonlyContext} from './readonly_context.js'; import {StreamingMode} from './run_config.js'; @@ -192,7 +197,7 @@ export type AfterToolCallback = export type ExamplesUnion = Example[] | BaseExampleProvider; /** A union of tool types that can be provided to an agent. */ -export type ToolUnion = BaseTool | BaseToolset; +export type ToolUnion = BaseTool | BaseToolset | BaseNode; const ADK_AGENT_NAME_LABEL_KEY = 'adk_agent_name'; @@ -258,6 +263,16 @@ export interface LlmAgentConfig extends BaseAgentConfig { */ includeContents?: 'default' | 'none'; + /** + * The agent's execution mode when run as a workflow node. + * + * - `single_turn` (default): the agent runs once against the node input. + * - `task`: the agent is given a `finish_task` tool and runs a multi-round + * loop until it calls `finish_task`, whose arguments (conforming to + * `outputSchema`) become the node output. Mirrors Python's `Agent(mode=...)`. + */ + mode?: 'single_turn' | 'task'; + /** The input schema when agent is used as a tool. */ inputSchema?: LlmAgentSchema; @@ -322,6 +337,11 @@ async function convertToolUnionToTools( if (isBaseTool(toolUnion)) { return [toolUnion]; } + if (toolUnion instanceof BaseNode) { + // A node/Workflow passed as a tool is auto-wrapped as a NodeTool so the + // model can call it (mirrors Python's Agent(tools=[node/workflow])). + return [new NodeTool(toolUnion)]; + } return await toolUnion.getTools(context); } @@ -361,9 +381,11 @@ export class LlmAgent extends BaseAgent { disallowTransferToParent: boolean; disallowTransferToPeers: boolean; includeContents: 'default' | 'none'; + mode?: 'single_turn' | 'task'; inputSchema?: Schema; outputSchema?: Schema; outputKey?: string; + private _finishTaskTool?: FinishTaskTool; beforeModelCallback?: BeforeModelCallback; afterModelCallback?: AfterModelCallback; beforeToolCallback?: BeforeToolCallback; @@ -388,6 +410,7 @@ export class LlmAgent extends BaseAgent { this.outputSchema = isZodObject(config.outputSchema) ? zodObjectToSchema(config.outputSchema) : config.outputSchema; + this.mode = config.mode; this.outputKey = config.outputKey; this.beforeModelCallback = config.beforeModelCallback; this.afterModelCallback = config.afterModelCallback; @@ -403,6 +426,7 @@ export class LlmAgent extends BaseAgent { IDENTITY_LLM_REQUEST_PROCESSOR, INSTRUCTIONS_LLM_REQUEST_PROCESSOR, REQUEST_CONFIRMATION_LLM_REQUEST_PROCESSOR, + REQUEST_INPUT_LLM_REQUEST_PROCESSOR, CONTENT_REQUEST_PROCESSOR, INTERACTIONS_REQUEST_PROCESSOR, CODE_EXECUTION_REQUEST_PROCESSOR, @@ -499,6 +523,17 @@ export class LlmAgent extends BaseAgent { throw new Error(`No model found for ${this.name}.`); } + /** + * The `finish_task` tool for this agent (task mode). Lazily created and cached + * so its declaration (derived from `outputSchema`) is stable across turns. + */ + get finishTaskTool(): FinishTaskTool { + if (!this._finishTaskTool) { + this._finishTaskTool = new FinishTaskTool(this.outputSchema); + } + return this._finishTaskTool; + } + /** * The resolved instruction field to construct instruction for this * agent. @@ -787,7 +822,11 @@ export class LlmAgent extends BaseAgent { // TODO - b/425992518: check if tool preprocessors can be simplified. // Run pre-processors for tools. const allTools = [...this.tools]; - if (this.outputSchema && allTools.length > 0) { + if (this.mode === 'task') { + // Task mode: the agent completes by calling `finish_task` (whose params + // mirror the output schema) rather than emitting structured output. + allTools.push(this.finishTaskTool); + } else if (this.outputSchema && allTools.length > 0) { const setModelResponseTool = new FunctionTool({ name: 'set_model_response', description: @@ -973,13 +1012,40 @@ export class LlmAgent extends BaseAgent { // Call functions // TODO - b/425992518: bloated funciton input, fix. // Tool callback passed to get rid of cyclic dependency. - const functionResponseEvent = await handleFunctionCallsAsync({ - invocationContext: invocationContext, - functionCallEvent: mergedEvent, - toolsDict: llmRequest.toolsDict, - beforeToolCallbacks: this.canonicalBeforeToolCallbacks, - afterToolCallbacks: this.canonicalAfterToolCallbacks, - }); + // A NodeTool (running a node/workflow) streams the node's intermediate and + // interrupt events into `invocationContext.eventQueue`; drain it concurrently + // so those events interleave into this agent's output stream. The tool runs + // in a self-contained task that captures its result/error and always closes + // the queue, so there is a single error path (no unhandled rejection). + const eventQueue = new EventChannel(); + invocationContext.eventQueue = eventQueue; + const toolTask = (async (): Promise<{ + event: Event | null; + error?: unknown; + }> => { + try { + const event = await handleFunctionCallsAsync({ + invocationContext: invocationContext, + functionCallEvent: mergedEvent, + toolsDict: llmRequest.toolsDict, + beforeToolCallbacks: this.canonicalBeforeToolCallbacks, + afterToolCallbacks: this.canonicalAfterToolCallbacks, + }); + return {event}; + } catch (error) { + return {event: null, error}; + } finally { + eventQueue.close(); + } + })(); + for await (const queuedEvent of eventQueue) { + yield queuedEvent; + } + const {event: functionResponseEvent, error: toolError} = await toolTask; + invocationContext.eventQueue = undefined; + if (toolError) { + throw toolError; + } if (!functionResponseEvent || invocationContext.abortSignal?.aborted) { return; diff --git a/core/src/agents/processors/basic_llm_request_processor.ts b/core/src/agents/processors/basic_llm_request_processor.ts index a2fa2594f..070ceba0a 100644 --- a/core/src/agents/processors/basic_llm_request_processor.ts +++ b/core/src/agents/processors/basic_llm_request_processor.ts @@ -37,7 +37,14 @@ export class BasicLlmRequestProcessor extends BaseLlmRequestProcessor { llmRequest.model = agent.canonicalModel.model; llmRequest.config = {...(agent.generateContentConfig ?? {})}; - if (agent.outputSchema && (!agent.tools || agent.tools.length === 0)) { + // Task-mode agents complete via the `finish_task` tool, so the JSON response + // mode must not be set (function calling is incompatible with a JSON + // response mime type). + if ( + agent.outputSchema && + agent.mode !== 'task' && + (!agent.tools || agent.tools.length === 0) + ) { setOutputSchema(llmRequest, agent.outputSchema); } diff --git a/core/src/agents/processors/request_confirmation_llm_request_processor.ts b/core/src/agents/processors/request_confirmation_llm_request_processor.ts index bff4d1a95..bcc74ab61 100644 --- a/core/src/agents/processors/request_confirmation_llm_request_processor.ts +++ b/core/src/agents/processors/request_confirmation_llm_request_processor.ts @@ -99,6 +99,17 @@ export class RequestConfirmationLlmRequestProcessor extends BaseLlmRequestProces } } + // Plain-text fallback: an interactive user (e.g. `adk run`) can approve or + // deny a pending confirmation by simply typing a reply (yes/no) instead of + // sending a structured confirmation response. + if (Object.keys(requestConfirmationFunctionResponses).length === 0) { + const fallback = mapPlainTextConfirmation(events); + Object.assign(requestConfirmationFunctionResponses, fallback.responses); + if (fallback.turnIndex >= 0) { + confirmationEventIndex = fallback.turnIndex; + } + } + if (Object.keys(requestConfirmationFunctionResponses).length === 0) { return; } @@ -190,5 +201,84 @@ export class RequestConfirmationLlmRequestProcessor extends BaseLlmRequestProces } } +/** Words interpreted as an approval when a user confirms by plain text. */ +const AFFIRMATIVE = new Set([ + 'yes', + 'y', + 'true', + 'approve', + 'approved', + 'ok', + 'okay', + 'confirm', + 'confirmed', +]); + +/** + * Maps a plain-text user reply to confirmations for any still-pending + * `adk_request_confirmation` calls, so a user can approve/deny by typing. + * Returns the synthesized confirmations keyed by the confirmation call id, and + * the index of the plain-text user turn (or -1 when not applicable). + */ +function mapPlainTextConfirmation(events: Event[]): { + responses: Record; + turnIndex: number; +} { + const answered = new Set(); + for (const event of events) { + if (event.author !== 'user') { + continue; + } + for (const fr of getFunctionResponses(event)) { + if (fr.id) { + answered.add(fr.id); + } + } + } + const pendingIds: string[] = []; + for (const event of events) { + for (const fc of getFunctionCalls(event)) { + if ( + fc.name === REQUEST_CONFIRMATION_FUNCTION_CALL_NAME && + fc.id && + !answered.has(fc.id) + ) { + pendingIds.push(fc.id); + } + } + } + if (pendingIds.length === 0) { + return {responses: {}, turnIndex: -1}; + } + + // Only the most recent user turn is considered, and only if it is plain text. + let turnIndex = -1; + let text = ''; + for (let i = events.length - 1; i >= 0; i--) { + const event = events[i]; + if (event.author !== 'user') { + continue; + } + const parts = event.content?.parts ?? []; + const isPlainText = + parts.length > 0 && parts.every((p) => typeof p.text === 'string'); + if (isPlainText) { + turnIndex = i; + text = parts.map((p) => p.text).join(''); + } + break; + } + if (turnIndex < 0) { + return {responses: {}, turnIndex: -1}; + } + + const confirmed = AFFIRMATIVE.has(text.trim().toLowerCase()); + const responses: Record = {}; + for (const id of pendingIds) { + responses[id] = new ToolConfirmation({confirmed}); + } + return {responses, turnIndex}; +} + export const REQUEST_CONFIRMATION_LLM_REQUEST_PROCESSOR = new RequestConfirmationLlmRequestProcessor(); diff --git a/core/src/agents/processors/request_input_llm_request_processor.ts b/core/src/agents/processors/request_input_llm_request_processor.ts new file mode 100644 index 000000000..6dbc42257 --- /dev/null +++ b/core/src/agents/processors/request_input_llm_request_processor.ts @@ -0,0 +1,205 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {FunctionCall} from '@google/genai'; + +import { + Event, + getFunctionCalls, + getFunctionResponses, +} from '../../events/event.js'; +import {ToolConfirmation} from '../../tools/tool_confirmation.js'; +import {NodeTool} from '../../workflow/nodes/node_tool.js'; +import {EventChannel} from '../../workflow/utils/event_channel.js'; +import {REQUEST_INPUT_FUNCTION_CALL_NAME} from '../../workflow/utils/hitl_utils.js'; +import {unwrapResponse} from '../../workflow/utils/rehydration_utils.js'; +import {handleFunctionCallList} from '../functions.js'; +import {InvocationContext} from '../invocation_context.js'; +import {isLlmAgent} from '../llm_agent.js'; +import {ReadonlyContext} from '../readonly_context.js'; +import {BaseLlmRequestProcessor} from './base_llm_processor.js'; + +/** + * Resumes a {@link NodeTool} call that paused for input. When a node/workflow + * run inside a node-tool raises a `RequestInput` interrupt, the node-tool's + * function call is left pending (no response) and the invocation pauses. On the + * next turn, once the user answers the `adk_request_input` interrupt, this + * processor re-runs the pending node-tool with the answer(s) threaded as + * `resumeInputs`, then emits the tool's function response so the agent can + * continue. Analogous to {@link RequestConfirmationLlmRequestProcessor}. + */ +export class RequestInputLlmRequestProcessor extends BaseLlmRequestProcessor { + override async *runAsync( + invocationContext: InvocationContext, + ): AsyncGenerator { + const agent = invocationContext.agent; + if (!isLlmAgent(agent)) { + return; + } + const events = invocationContext.session.events; + if (!events || events.length === 0) { + return; + } + + // 1. Collect resume inputs (interruptId -> value): prefer structured + // `adk_request_input` responses, else map a plain-text reply to any + // pending interrupt (so an interactive client can resume by typing). + const resumeInputs = collectResumeInputs(events); + if (Object.keys(resumeInputs).length === 0) { + return; + } + + // 2. Resolve the agent's node-tools. + const toolsList = await agent.canonicalTools( + new ReadonlyContext(invocationContext), + ); + const toolsDict = Object.fromEntries(toolsList.map((t) => [t.name, t])); + const nodeToolNames = new Set( + toolsList.filter((t) => t instanceof NodeTool).map((t) => t.name), + ); + if (nodeToolNames.size === 0) { + return; + } + + // 3. Find pending node-tool function calls (raised but not yet answered). + const answeredIds = new Set(); + for (const event of events) { + for (const fr of getFunctionResponses(event)) { + if (fr.id) { + answeredIds.add(fr.id); + } + } + } + const pending: Record = {}; + for (const event of events) { + for (const fc of getFunctionCalls(event)) { + if ( + fc.id && + fc.name && + nodeToolNames.has(fc.name) && + !answeredIds.has(fc.id) + ) { + pending[fc.id] = fc; + } + } + } + if (Object.keys(pending).length === 0) { + return; + } + + // 4. Re-run each pending node-tool, threading the resume inputs through the + // tool confirmation payload (read by NodeTool as the node's resumeInputs). + const toolConfirmationDict: Record = {}; + for (const id of Object.keys(pending)) { + toolConfirmationDict[id] = new ToolConfirmation({ + confirmed: true, + payload: resumeInputs, + }); + } + + const eventQueue = new EventChannel(); + invocationContext.eventQueue = eventQueue; + const task = (async (): Promise => { + try { + return await handleFunctionCallList({ + invocationContext, + functionCalls: Object.values(pending), + toolsDict, + beforeToolCallbacks: agent.canonicalBeforeToolCallbacks, + afterToolCallbacks: agent.canonicalAfterToolCallbacks, + filters: new Set(Object.keys(pending)), + toolConfirmationDict, + }); + } finally { + eventQueue.close(); + } + })(); + for await (const queuedEvent of eventQueue) { + yield queuedEvent; + } + const functionResponseEvent = await task; + invocationContext.eventQueue = undefined; + if (functionResponseEvent) { + yield functionResponseEvent; + } + } +} + +/** + * Collects resume inputs from the session: structured `adk_request_input` + * function responses take precedence; otherwise a plain-text reply is mapped to + * every still-pending interrupt id. + */ +function collectResumeInputs(events: Event[]): Record { + for (let i = events.length - 1; i >= 0; i--) { + const event = events[i]; + if (event.author !== 'user') { + continue; + } + const structured: Record = {}; + let found = false; + for (const fr of getFunctionResponses(event)) { + if (fr.name === REQUEST_INPUT_FUNCTION_CALL_NAME && fr.id) { + structured[fr.id] = unwrapResponse(fr.response); + found = true; + } + } + if (found) { + return structured; + } + } + + // Plain-text fallback: map the latest plain-text user turn to pending + // interrupts (mirrors WorkflowAgent's interactive resume). + const pending = pendingInterruptIds(events); + if (pending.size === 0) { + return {}; + } + const lastUser = [...events].reverse().find((e) => e.author === 'user'); + const parts = lastUser?.content?.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 inputs: Record = {}; + for (const id of pending) { + inputs[id] = text; + } + return inputs; +} + +/** Interrupt ids raised via `adk_request_input` that have no user response. */ +function pendingInterruptIds(events: Event[]): Set { + const answered = new Set(); + for (const event of events) { + if (event.author !== 'user') { + continue; + } + for (const fr of getFunctionResponses(event)) { + if (fr.id) { + answered.add(fr.id); + } + } + } + const pending = new Set(); + for (const event of events) { + for (const fc of getFunctionCalls(event)) { + if ( + fc.name === REQUEST_INPUT_FUNCTION_CALL_NAME && + fc.id && + !answered.has(fc.id) + ) { + pending.add(fc.id); + } + } + } + return pending; +} + +export const REQUEST_INPUT_LLM_REQUEST_PROCESSOR = + new RequestInputLlmRequestProcessor(); diff --git a/core/src/common.ts b/core/src/common.ts index f9f1e1ac6..286ecf697 100644 --- a/core/src/common.ts +++ b/core/src/common.ts @@ -233,6 +233,11 @@ export {BaseToolset, isBaseToolset} from './tools/base_toolset.js'; export type {ToolPredicate} from './tools/base_toolset.js'; export {ConsolidateContextTool} from './tools/consolidate_context_tool.js'; export {EXIT_LOOP, ExitLoopTool} from './tools/exit_loop_tool.js'; +export { + FINISH_TASK_SUCCESS_RESULT, + FINISH_TASK_TOOL_NAME, + FinishTaskTool, +} from './tools/finish_task_tool.js'; export {FunctionTool, isFunctionTool} from './tools/function_tool.js'; export type { ToolExecuteArgument, diff --git a/core/src/tools/finish_task_tool.ts b/core/src/tools/finish_task_tool.ts new file mode 100644 index 000000000..203992a0f --- /dev/null +++ b/core/src/tools/finish_task_tool.ts @@ -0,0 +1,140 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {FunctionDeclaration, Schema, Type} from '@google/genai'; + +import {appendInstructions} from '../models/llm_request.js'; +import { + BaseTool, + RunAsyncToolRequest, + ToolProcessLlmRequest, +} from './base_tool.js'; + +/** The name of the finish_task tool. */ +export const FINISH_TASK_TOOL_NAME = 'finish_task'; + +/** + * The result returned by {@link FinishTaskTool.runAsync} when validation passes. + * The task-mode wrapper uses this to distinguish a successful completion from a + * validation-error retry signal. + */ +export const FINISH_TASK_SUCCESS_RESULT = 'Task completed.'; + +/** The default output schema when the task agent declares none. */ +const DEFAULT_TASK_OUTPUT_SCHEMA: Schema = { + type: Type.OBJECT, + properties: { + result: { + type: Type.STRING, + description: 'A brief summary of what the agent accomplished.', + }, + }, + required: ['result'], +}; + +/** + * Tool for signaling that a task-mode {@link LlmAgent} has completed its task. + * + * The tool's parameters mirror the agent's `outputSchema` (or a default single + * `result` string). The task-mode wrapper sniffs the `finish_task` function call + * and, on a successful function response, promotes the call's arguments to the + * node's output. + * + * Ported from `google/adk-python` + * `agents/llm/task/_finish_task_tool.py::FinishTaskTool`. + */ +export class FinishTaskTool extends BaseTool { + /** The schema describing the expected task output. */ + private readonly outputSchema: Schema; + /** + * When the output schema is a non-object (primitive/array), the value is + * wrapped under this key (the GenAI API requires object-typed parameters). + * `undefined` for object schemas (the value lives at the top level of args). + */ + readonly wrapperKey?: string; + + constructor(outputSchema?: Schema) { + const schema = outputSchema ?? DEFAULT_TASK_OUTPUT_SCHEMA; + let description = + 'Signal that this agent has completed its delegated task. Call this' + + ' when you have finished your delegated task.'; + if (outputSchema) { + description += ' Pass the required output data in the parameters.'; + } + super({name: FINISH_TASK_TOOL_NAME, description}); + this.outputSchema = schema; + this.wrapperKey = schema.type === Type.OBJECT ? undefined : 'result'; + } + + override _getDeclaration(): FunctionDeclaration { + const parameters: Schema = this.wrapperKey + ? { + type: Type.OBJECT, + properties: {[this.wrapperKey]: this.outputSchema}, + required: [this.wrapperKey], + } + : this.outputSchema; + return {name: this.name, description: this.description, parameters}; + } + + override async processLlmRequest( + request: ToolProcessLlmRequest, + ): Promise { + await super.processLlmRequest(request); + // Tell the model when to call finish_task (mirrors Python's tool + // instruction), so it completes the task deliberately. + appendInstructions(request.llmRequest, [ + 'Do NOT call `finish_task` prematurely. Use your available tools to fully' + + ' complete every aspect of the task first. If the task is unclear, ask' + + ' the user for clarification before proceeding. Once the task is fully' + + ' complete, call `finish_task` by itself with no accompanying text' + + ' output.', + ]); + } + + /** + * Extracts the task output from a `finish_task` call's arguments, applying the + * wrapper-key unwrapping when the schema is a non-object. + */ + extractOutput(args: Record): unknown { + if (this.wrapperKey) { + return args[this.wrapperKey]; + } + return args; + } + + override async runAsync({args}: RunAsyncToolRequest): Promise { + const value = this.wrapperKey ? args[this.wrapperKey] : args; + const missing = this.missingRequiredKeys(value); + if (missing.length > 0) { + return { + error: + `Invoking \`${this.name}()\` failed due to missing required ` + + `parameters: ${missing.join(', ')}. You could retry calling this ` + + 'tool, but it is IMPORTANT for you to provide all the mandatory ' + + 'parameters with correct types.', + }; + } + return FINISH_TASK_SUCCESS_RESULT; + } + + /** Returns any `required` keys the schema declares that are absent. */ + private missingRequiredKeys(value: unknown): string[] { + const required = this.wrapperKey + ? value === undefined || value === null + ? [this.wrapperKey] + : [] + : (this.outputSchema.required ?? []); + if (this.wrapperKey) { + return required; + } + if (typeof value !== 'object' || value === null) { + return required; + } + const obj = value as Record; + return required.filter((key) => obj[key] === undefined); + } +} diff --git a/core/src/workflow/index.ts b/core/src/workflow/index.ts index c3de73ade..431f17cf4 100644 --- a/core/src/workflow/index.ts +++ b/core/src/workflow/index.ts @@ -31,6 +31,7 @@ export type { export {JoinNode} from './nodes/join_node.js'; export {LLMAgentWrapper} from './nodes/llm_agent_wrapper.js'; export type {LLMAgentWrapperConfig} from './nodes/llm_agent_wrapper.js'; +export {NodeTool} from './nodes/node_tool.js'; export {ParallelWorker} from './nodes/parallel_worker.js'; export type {ParallelWorkerConfig} from './nodes/parallel_worker.js'; export {ToolNode} from './nodes/tool_node.js'; diff --git a/core/src/workflow/nodes/llm_agent_wrapper.ts b/core/src/workflow/nodes/llm_agent_wrapper.ts index fec88b06d..e76cda844 100644 --- a/core/src/workflow/nodes/llm_agent_wrapper.ts +++ b/core/src/workflow/nodes/llm_agent_wrapper.ts @@ -6,7 +6,17 @@ import {Content} from '@google/genai'; import {BaseAgent} from '../../agents/base_agent.js'; -import {createEvent, Event} from '../../events/event.js'; +import {isLlmAgent, LlmAgent} from '../../agents/llm_agent.js'; +import { + createEvent, + Event, + getFunctionCalls, + getFunctionResponses, +} from '../../events/event.js'; +import { + FINISH_TASK_SUCCESS_RESULT, + FINISH_TASK_TOOL_NAME, +} from '../../tools/finish_task_tool.js'; import {BaseNode, BaseNodeConfig, isContent} from '../base_node.js'; import {NodeContext} from '../node_context.js'; @@ -59,10 +69,58 @@ export class LLMAgentWrapper extends BaseNode { ctx.session.events.push(userEvent); } + // Task mode: run a multi-round loop until the agent calls `finish_task`, + // whose arguments become the node output. + if (isLlmAgent(this.agent) && this.agent.mode === 'task') { + yield* this.runTaskMode(ctx, this.agent); + return; + } + // Run the agent, following any transfer_to_agent hand-offs to peers. yield* this.runWithTransfers(ctx, this.agent, 0); } + /** + * Runs a `task`-mode agent: the agent loops (LLM ↔ tools) until it calls the + * `finish_task` tool. The wrapper sniffs the `finish_task` function call and, + * on its successful function response, promotes the call's arguments to the + * node output (and to `outputKey` state, if set). Mirrors Python's + * `run_llm_agent_as_node` task branch. + */ + private async *runTaskMode( + ctx: NodeContext, + agent: LlmAgent, + ): AsyncGenerator { + const finishTool = agent.finishTaskTool; + let pendingArgs: Record | undefined; + + for await (const event of agent.runAsync(ctx.invocationContext)) { + const finishCall = getFunctionCalls(event).find( + (fc) => fc.name === FINISH_TASK_TOOL_NAME, + ); + if (finishCall) { + // Remember the latest finish_task args; wait for the success function + // response before terminating (a validation error lets the LLM retry). + pendingArgs = {...(finishCall.args ?? {})}; + yield event; + continue; + } + + if (pendingArgs !== undefined && isFinishTaskSuccessResponse(event)) { + const output = finishTool.extractOutput(pendingArgs); + event.output = output; + event.nodeInfo = {...(event.nodeInfo ?? {}), messageAsOutput: true}; + if (agent.outputKey && output !== undefined) { + ctx.actions.stateDelta[agent.outputKey] = output; + } + yield event; + return; + } + + yield event; + } + } + /** * Runs `agent`; if it emits a `transfer_to_agent` action, resolves the target * in the agent tree and continues with it (multi-agent hand-off). This is the @@ -144,6 +202,21 @@ function hasFunctionCalls(event: Event): boolean { return (event.content?.parts ?? []).some((p) => p.functionCall); } +/** + * Whether an event carries the success function response from `finish_task`. + * A non-success response (e.g. a validation error) returns false so the caller + * keeps iterating and the LLM gets a chance to retry. + */ +function isFinishTaskSuccessResponse(event: Event): boolean { + return getFunctionResponses(event).some((fr) => { + if (fr.name !== FINISH_TASK_TOOL_NAME) { + return false; + } + const response = (fr.response ?? {}) as {result?: unknown}; + return response.result === FINISH_TASK_SUCCESS_RESULT; + }); +} + /** Converts an arbitrary node input into a user-role `Content`. */ function toUserContent(input: unknown): Content { if (isContent(input)) { diff --git a/core/src/workflow/nodes/node_tool.ts b/core/src/workflow/nodes/node_tool.ts new file mode 100644 index 000000000..df02c4200 --- /dev/null +++ b/core/src/workflow/nodes/node_tool.ts @@ -0,0 +1,143 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {FunctionDeclaration, Schema, Type} from '@google/genai'; + +import {Context} from '../../agents/context.js'; +import {Event} from '../../events/event.js'; +import {BaseTool, RunAsyncToolRequest} from '../../tools/base_tool.js'; +import { + isZodObject, + zodObjectToSchema, +} from '../../utils/simple_zod_to_json.js'; +import {BaseNode} from '../base_node.js'; +import {NodeContext} from '../node_context.js'; +import {executeChildNode} from '../node_runner.js'; +import {EventChannel} from '../utils/event_channel.js'; + +/** + * A tool that executes a {@link BaseNode} (e.g. a `Workflow` or a function node) + * on behalf of an `LlmAgent`. This is the inverse of {@link ToolNode} (which + * exposes a tool as a workflow node): here a node/workflow is exposed to a model + * as a callable tool. + * + * The wrapped node MUST declare an `inputSchema` (the tool's parameter schema is + * derived from it). When the model calls the tool, the node runs with a + * {@link NodeContext} bridged from the tool's agent context (sharing the + * invocation, session, and state); the node's structured output becomes the + * tool result. + * + * Ported from `google/adk-python` `tools/_node_tool.py::NodeTool`. + * + * The tool is marked long-running so a node that pauses for input + * (`RequestInput`) does not force a synthetic empty response. + */ +export class NodeTool extends BaseTool { + readonly node: BaseNode; + + constructor(node: BaseNode, name?: string, description?: string) { + if (!node.inputSchema) { + throw new Error( + `Node '${node.name}' does not have an inputSchema defined. NodeTool ` + + 'requires an explicit input schema on the wrapped node.', + ); + } + super({ + name: name ?? node.name, + description: + description || node.description || `Executes the node: ${node.name}`, + isLongRunning: true, + }); + this.node = node; + } + + /** Whether the node's input schema is a (Zod) object rather than a scalar. */ + private get inputIsObject(): boolean { + return isZodObject(this.node.inputSchema); + } + + override _getDeclaration(): FunctionDeclaration { + let parameters: Schema; + if (this.inputIsObject) { + parameters = zodObjectToSchema(this.node.inputSchema as never); + } else { + // The GenAI API requires object-typed parameters; wrap a scalar schema + // under a single `request` property. + parameters = { + type: Type.OBJECT, + properties: {request: {type: Type.STRING}}, + required: ['request'], + }; + } + return {name: this.name, description: this.description, parameters}; + } + + override async runAsync({ + args, + toolContext, + }: RunAsyncToolRequest): Promise { + const nodeInput = this.inputIsObject ? args : args['request']; + + const child = await this.runNode(toolContext, nodeInput); + + if (child.interruptIds.length > 0) { + // The node paused for input. Returning undefined leaves the (long-running) + // tool call pending; the interrupt event has been surfaced separately so + // the invocation can pause and resume. (Resume wiring is layered on top.) + return undefined; + } + + return child.output === undefined ? {result: null} : child.output; + } + + /** + * Runs the wrapped node with a {@link NodeContext} bridged from the agent's + * tool context. Node events are streamed into the invocation's event queue + * when one is present (so intermediate/interrupt events surface to the agent); + * otherwise they are buffered and dropped (completion-only path). + */ + private async runNode( + toolContext: Context, + nodeInput: unknown, + ): Promise { + const ic = toolContext.invocationContext; + const runId = toolContext.functionCallId ?? this.node.name; + const channel = + (ic as {eventQueue?: EventChannel}).eventQueue ?? + new EventChannel(); + + const nodeCtx = new NodeContext({ + invocationContext: ic, + channel, + nodePath: this.node.name, + runId, + resumeInputs: collectResumeInputs(toolContext), + }); + + const base = ic.branch; + const segment = `${this.name}@${runId}`; + const overrideBranch = base ? `${base}.${segment}` : segment; + + return executeChildNode(nodeCtx, this.node, nodeInput, { + runId, + overrideBranch, + }); + } +} + +/** + * Collects resume inputs for the node from the tool context. When the tool call + * is being resumed after a `RequestInput`, the user's response is threaded + * through `toolConfirmation.payload` keyed by interrupt id (see the request-input + * resume processor). + */ +function collectResumeInputs(toolContext: Context): Record { + const payload = toolContext.toolConfirmation?.payload; + if (payload && typeof payload === 'object') { + return payload as Record; + } + return {}; +} diff --git a/core/src/workflow/utils/workflow_graph_utils.ts b/core/src/workflow/utils/workflow_graph_utils.ts index e9be061df..ea6821084 100644 --- a/core/src/workflow/utils/workflow_graph_utils.ts +++ b/core/src/workflow/utils/workflow_graph_utils.ts @@ -21,6 +21,7 @@ import {RetryConfig} from '../retry_config.js'; */ export interface BuildNodeOptions { name?: string; + description?: string; rerunOnResume?: boolean; retryConfig?: RetryConfig; timeout?: number; diff --git a/tests/integration/workflows/node_as_tool_hitl_test.ts b/tests/integration/workflows/node_as_tool_hitl_test.ts new file mode 100644 index 000000000..a17f641da --- /dev/null +++ b/tests/integration/workflows/node_as_tool_hitl_test.ts @@ -0,0 +1,121 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Integration test for HITL through a node-tool: an `LlmAgent` calls a node + * (passed as a tool) that raises a `RequestInput` while running. The invocation + * pauses; on the next turn the user answers the interrupt, the node-tool is + * re-run with the answer threaded as `resumeInputs`, and the tool result flows + * back to the model. Mirrors the `node_as_tool` `calculate_discount` pattern. + */ + +import { + getFunctionCalls, + getFunctionResponses, + InMemoryRunner, + node, + NodeContext, + RequestInput, +} from '@google/adk'; +import {Content} from '@google/genai'; +import {describe, expect, it} from 'vitest'; +import {z} from 'zod'; +import { + collect, + functionCallResponse, + mockLlmAgent, + textResponse, +} from './workflow_test_utils.js'; + +describe('workflow integration — HITL through a node-tool', () => { + it('pauses on RequestInput raised inside a node-tool and resumes', async () => { + const calculateDiscount = node( + (ctx: NodeContext, args: {tier: string}) => { + const resume = ctx.resumeInputs['confirm_vip_discount']; + if (!args.tier.includes('VIP')) { + return '5% off'; + } + if (resume === undefined) { + return new RequestInput({ + interruptId: 'confirm_vip_discount', + message: `Apply VIP discount for tier '${args.tier}'?`, + }); + } + const answer = + typeof resume === 'object' && resume !== null + ? (resume as {text?: string}).text + : resume; + return String(answer).toLowerCase() === 'yes' + ? '20% off' + : '5% off (VIP declined)'; + }, + { + name: 'calculate_discount', + inputSchema: z.object({tier: z.string()}), + rerunOnResume: true, + }, + ); + + const agent = mockLlmAgent( + { + name: 'discount_agent', + instruction: 'Compute the discount.', + tools: [calculateDiscount], + }, + [ + functionCallResponse('calculate_discount', { + tier: 'Verified VIP Member', + }), + textResponse('You get 20% off.'), + ], + ); + + const runner = new InMemoryRunner({agent, appName: agent.name}); + const session = await runner.sessionService.createSession({ + appName: agent.name, + userId: 'u1', + }); + + // Turn 1: the model calls the node-tool; the node interrupts for input. + const turn1 = await collect( + runner.runAsync({ + userId: 'u1', + sessionId: session.id, + newMessage: {role: 'user', parts: [{text: 'What discount do I get?'}]}, + }), + ); + const raisedInterrupt = turn1 + .flatMap((e) => getFunctionCalls(e)) + .some((fc) => fc.name === 'adk_request_input'); + expect(raisedInterrupt).toBe(true); + + // Turn 2: the user answers the interrupt; the node-tool re-runs and resolves. + const resume: Content = { + role: 'user', + parts: [ + { + functionResponse: { + id: 'confirm_vip_discount', + name: 'adk_request_input', + response: {result: 'yes'}, + }, + }, + ], + }; + const turn2 = await collect( + runner.runAsync({ + userId: 'u1', + sessionId: session.id, + newMessage: resume, + }), + ); + + const discountResult = turn2 + .flatMap((e) => getFunctionResponses(e)) + .find((fr) => fr.name === 'calculate_discount'); + expect(discountResult?.response).toMatchObject({result: '20% off'}); + }); +}); diff --git a/tests/integration/workflows/node_as_tool_test.ts b/tests/integration/workflows/node_as_tool_test.ts new file mode 100644 index 000000000..bcb0cc179 --- /dev/null +++ b/tests/integration/workflows/node_as_tool_test.ts @@ -0,0 +1,80 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Integration test for node/workflow-as-tool: an `LlmAgent` is given a + * `Workflow` (and a function node) in its `tools`; the framework auto-wraps them + * as `NodeTool`s so the model can call them, and the node's structured output + * becomes the tool result. Mirrors the `node_as_tool` sample. + */ + +import { + getFunctionResponses, + InMemoryRunner, + node, + NodeContext, + Workflow, +} from '@google/adk'; +import {describe, expect, it} from 'vitest'; +import {z} from 'zod'; +import { + collect, + functionCallResponse, + mockLlmAgent, + textResponse, +} from './workflow_test_utils.js'; + +describe('workflow integration — node/workflow as an agent tool', () => { + it('lets an LlmAgent call a Workflow passed as a tool', async () => { + const lookup = node( + (_c: NodeContext, args: {userId: string}) => ({ + userId: args.userId, + tier: 'Verified VIP Member', + }), + {name: 'lookup_customer', inputSchema: z.object({userId: z.string()})}, + ); + const lookupWorkflow = new Workflow({ + name: 'customer_lookup_workflow', + description: 'Looks up customer status and tier by user_id.', + inputSchema: z.object({userId: z.string()}), + edges: [['START', lookup]], + }); + + const agent = mockLlmAgent( + { + name: 'customer_service_agent', + instruction: 'Help the customer.', + tools: [lookupWorkflow], + }, + [ + functionCallResponse('customer_lookup_workflow', {userId: 'u123'}), + textResponse('The customer is a Verified VIP Member.'), + ], + ); + + const runner = new InMemoryRunner({agent, appName: agent.name}); + const session = await runner.sessionService.createSession({ + appName: agent.name, + userId: 'u1', + }); + const events = await collect( + runner.runAsync({ + userId: 'u1', + sessionId: session.id, + newMessage: {role: 'user', parts: [{text: 'look up user u123'}]}, + }), + ); + + // The workflow tool ran and returned the tier as the function response. + const toolResult = events + .flatMap((e) => getFunctionResponses(e)) + .find((fr) => fr.name === 'customer_lookup_workflow'); + expect(toolResult?.response).toMatchObject({ + tier: 'Verified VIP Member', + userId: 'u123', + }); + }); +}); diff --git a/tests/integration/workflows/task_mode_test.ts b/tests/integration/workflows/task_mode_test.ts new file mode 100644 index 000000000..0c1e02022 --- /dev/null +++ b/tests/integration/workflows/task_mode_test.ts @@ -0,0 +1,69 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Integration test for `LlmAgent` task mode as a workflow node: the agent is + * given a `finish_task` tool and runs until it calls it; the call's arguments + * (conforming to the agent's output schema) become the node output and feed the + * next node. Mirrors the `agent_in_workflow` intake pattern. + */ + +import {node, NodeContext, Workflow} from '@google/adk'; +import {describe, expect, it} from 'vitest'; +import {z} from 'zod'; +import { + collect, + createWorkflowRunner, + finalOutput, + functionCallResponse, + mockLlmAgent, +} from './workflow_test_utils.js'; + +describe('workflow integration — LlmAgent task mode', () => { + it('runs until finish_task and promotes its args to the node output', async () => { + const intake = mockLlmAgent( + { + name: 'intake_agent', + mode: 'task', + instruction: 'Collect the patient name and phone number.', + outputSchema: z.object({name: z.string(), phoneNumber: z.string()}), + outputKey: 'identity', + }, + [ + functionCallResponse('finish_task', { + name: 'Jane Doe', + phoneNumber: '555-1234', + }), + ], + ); + + const check = node( + (_ctx: NodeContext, identity: {name: string; phoneNumber: string}) => + `checked:${identity.name} (${identity.phoneNumber})`, + {name: 'check'}, + ); + + const wf = new Workflow({ + name: 'task_wf', + edges: [['START', intake, check]], + }); + + const {run} = await createWorkflowRunner(wf); + const events = await collect(run('Hi, I am Jane Doe, 555-1234.')); + + // finish_task args flowed to `check` as its input. + expect(finalOutput(events)).toBe('checked:Jane Doe (555-1234)'); + // The finish_task args were also promoted to the node output. + expect( + events.some( + (e) => + typeof e.output === 'object' && + e.output !== null && + (e.output as {name?: string}).name === 'Jane Doe', + ), + ).toBe(true); + }); +}); diff --git a/tests/integration/workflows/workflow_test_utils.ts b/tests/integration/workflows/workflow_test_utils.ts index 701e2ff73..5dc61365f 100644 --- a/tests/integration/workflows/workflow_test_utils.ts +++ b/tests/integration/workflows/workflow_test_utils.ts @@ -32,6 +32,24 @@ export function textResponse(text: string): RawGenerateContentResponse { }; } +/** + * Builds a raw generate-content response that returns a single function call. + */ +export function functionCallResponse( + name: string, + args: Record, + id?: string, +): RawGenerateContentResponse { + return { + candidates: [ + { + content: {role: 'model', parts: [{functionCall: {name, args, id}}]}, + finishReason: FinishReason.STOP, + }, + ], + }; +} + /** * Constructs an {@link LlmAgent} whose model returns the given canned responses * (loaded from a JSON fixture), so workflow integration tests are deterministic From 06501d3e21c8d36bb14b23d692aab9602ad66246 Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Mon, 27 Jul 2026 12:00:59 -0700 Subject: [PATCH 33/41] samples(workflow): faithfully port node_as_tool and agent_in_workflow Replaces the earlier simplified stand-ins with real ports of the adk-python samples, now that the framework supports them: - node_as_tool: an LlmAgent uses a Workflow and a node as tools; the node streams a status message and raises RequestInput to confirm a VIP discount (HITL through a tool), resolved by a plain-text reply on the next turn. - agent_in_workflow: a task-mode intake agent completes via finish_task, an identity check routes retry vs proceed, and generate_instruction calls a require_confirmation tool. Both verified end-to-end with a live model. Updates the samples README to reflect that every sample is now a faithful port (real LlmAgents where Python uses them) and documents the new feature coverage. --- samples/workflows/README.md | 123 ++++++++++--------- samples/workflows/agent_in_workflow/agent.ts | 97 ++++++++++++--- samples/workflows/node_as_tool/agent.ts | 108 +++++++++++++--- 3 files changed, 239 insertions(+), 89 deletions(-) diff --git a/samples/workflows/README.md b/samples/workflows/README.md index 7240bf3c4..f9f7177cf 100644 --- a/samples/workflows/README.md +++ b/samples/workflows/README.md @@ -2,8 +2,13 @@ Runnable TypeScript ports of the Python [`contributing/samples/workflows`](https://github.com/google/adk-python/tree/main/contributing/samples/workflows) -samples, one per directory. Each exports a `rootAgent` (a `WorkflowAgent` -wrapping a `Workflow`) so it runs with the ADK CLI. +samples, one per directory. These are **faithful** ports: where the Python +sample uses an `LlmAgent`, the TypeScript port uses a real `LlmAgent` (calling a +live model) with the same node/graph structure — not an offline stand-in. + +Each directory exports a `rootAgent` that runs with the ADK CLI. Most wrap a +`Workflow` in a `WorkflowAgent`; `node_as_tool` exports a plain `LlmAgent` that +uses a node and a workflow as tools. ## Running @@ -19,69 +24,77 @@ npm run sample -- samples/workflows/sequence/agent.ts The CLI is interactive: type a message and press Enter to send it to the workflow; type `exit` to quit. Node events are printed as -`[]: ` and the final line `[]: ...` is the -workflow's output. +`[]: ` and the final line is the workflow's output. -You can also pipe a single message: +You can also pipe a single message, or script a multi-turn run with `--replay` +(a JSON file of queries, resolved relative to the working directory): ```bash echo "hello world" | npm run sample -- samples/workflows/sequence/agent.ts + +echo '{"state":{},"queries":["The product broke","approve"]}' > replay.json +npm run sample -- samples/workflows/request_input/agent.ts --replay replay.json ``` ## API keys -Most samples are **function-based and run offline** (no key needed). Samples -that call a live model are marked **(needs API key)** below — set -`GEMINI_API_KEY` (a `.env` file in the working directory is loaded -automatically) before running them. - -## Human-in-the-loop / auth samples +Samples marked **needs API key** below call a live model. Set `GEMINI_API_KEY` +(a `.env` file in the working directory is loaded automatically) before running +them. The rest are function-based and run offline with no key. -For HITL and auth samples, the workflow **pauses** on the first turn (you'll see -an `adk_request_input` / `adk_request_credential` request). Simply **type your -reply on the next turn** — the plain-text reply is fed to the pending interrupt, -so you can approve/reject, give feedback, or supply an API key interactively. +## Human-in-the-loop / auth / confirmation -To script a multi-turn run non-interactively, use `--replay` with a JSON file of -queries: - -```bash -echo '{"state":{},"queries":["The product broke","approve"]}' > replay.json -npm run sample -- samples/workflows/request_input/agent.ts --replay replay.json -``` +HITL, auth, and tool-confirmation samples **pause** mid-run (you'll see an +`adk_request_input`, `adk_request_credential`, or `adk_request_confirmation` +request). Simply **type your reply on the next turn** — the plain-text reply is +routed to the pending interrupt, so you can approve/reject, give feedback, supply +an API key, or confirm a tool call interactively. ## Samples -| Sample | What it shows | Offline? | -| ------------------------ | --------------------------------------------------- | ----------------- | -| `sequence` | Linear chain; each output feeds the next | ✅ | -| `route` | Classify input, route to a branch (+ DEFAULT_ROUTE) | ✅ | -| `fan_out_fan_in` | Parallel branches joined by a `JoinNode` | ✅ | -| `parallel_worker` | Map a node across a list with bounded concurrency | ✅ | -| `dynamic_nodes` | Imperative `dynamicEntry` driving `ctx.runNode()` | ✅ | -| `dynamic_fan_out_fan_in` | Concurrent `ctx.runNode()` + aggregate | ✅ | -| `loop` | Generate → evaluate → route back until it passes | ✅ | -| `loop_self` | A node routes back to itself (conditional cycle) | ✅ | -| `multi_triggers` | A non-join node runs once per predecessor trigger | ✅ | -| `nested_workflow` | A `Workflow` used as a node (+ parallel + join) | ✅ | -| `node_as_tool` | A node calls sub-nodes via `ctx.runNode()` | ✅ | -| `state` | Share data across nodes via `ctx.state` | ✅ | -| `node_output` | Raw value / `Event({output})` / structured output | ✅ | -| `use_as_output` | Promote a sub-node result via `useAsOutput` | ✅ | -| `message` | Emit a display message distinct from output | ✅ | -| `retry` | Retry a flaky node per `retryConfig` | ✅ | -| `request_input` | HITL: draft → review → approve/reject/revise | ✅ (interactive) | -| `request_input_rerun` | HITL single node with `rerunOnResume` | ✅ (interactive) | -| `request_input_advanced` | Auto-approve small / pause for large requests | ✅ (interactive) | -| `auth_api_key` | Pause to request an API-key credential | ✅ (interactive) | -| `auth_oauth` | Pause and emit an OAuth authorization URL | ✅ (request only) | -| `agent_in_workflow` | A real `LlmAgent` as a workflow node | ❗ needs API key | - -### Notes on faithfulness - -Some Python samples use `LlmAgent`s for steps like classification or -generation. To keep the ports runnable offline, those steps are implemented with -function nodes here (the workflow _structure_ is identical); swap a function -node for an `LlmAgent` to use a real model, as shown in `agent_in_workflow`. -`auth_oauth` emits a real authorization request, but completing the OAuth token -exchange requires a live provider, so its resume step is illustrative only. +| Sample | What it shows | Needs API key | +| ------------------------ | -------------------------------------------------------- | ------------- | +| `sequence` | Linear chain of two `LlmAgent`s | ✅ | +| `route` | `LlmAgent` classifier (schema) routes to a branch | ✅ | +| `fan_out_fan_in` | Parallel branches joined by a `JoinNode` | — | +| `parallel_worker` | Map a node across a list with bounded concurrency | ✅ | +| `dynamic_nodes` | Imperative `dynamicEntry` driving `ctx.runNode()` | ✅ | +| `dynamic_fan_out_fan_in` | Concurrent `ctx.runNode()` + aggregate | ✅ | +| `loop` | Generate → evaluate → route back until it passes | ✅ | +| `loop_self` | A node routes back to itself (conditional cycle) | — | +| `multi_triggers` | A non-join node runs once per predecessor trigger | — | +| `nested_workflow` | A `Workflow` used as a node (+ parallel + join) | ✅ | +| `node_as_tool` | An `LlmAgent` uses a node + a `Workflow` as tools (HITL) | ✅ | +| `state` | Share data across nodes via `ctx.state` | — | +| `node_output` | Raw value / `Event({output})` / structured LLM output | ✅ | +| `use_as_output` | Promote a sub-node result via `useAsOutput` | ✅ | +| `message` | Emit a display message distinct from output | — | +| `retry` | Retry a flaky node per `retryConfig` | — | +| `request_input` | HITL two-node: draft → review → approve/reject/revise | ✅ | +| `request_input_rerun` | HITL single node with `rerunOnResume` | ✅ | +| `request_input_advanced` | Structured HITL: auto-approve small / pause for large | ✅ | +| `auth_api_key` | Pause to request an API-key credential | — | +| `auth_oauth` | Pause and request GitHub OAuth credentials | — | +| `agent_in_workflow` | `task`-mode agent + identity check + confirmation tool | ✅ | + +## Feature coverage + +These ports exercise the full workflow + agent-integration surface: + +- **`task` mode** (`agent_in_workflow`): an `LlmAgent` runs a multi-round loop and + completes via a `finish_task` tool whose arguments become the node output. +- **node / workflow as a tool** (`node_as_tool`): a `BaseNode`/`Workflow` passed + in an agent's `tools` is auto-wrapped as a `NodeTool`; a node may even pause for + input (`RequestInput`) mid-tool-call and resume on the next turn. +- **`require_confirmation`** (`agent_in_workflow`): a `FunctionTool` pauses for + user approval before it runs. +- **`rerun_on_resume`** semantics: the default two-node HITL pattern + (`request_input`) vs. the single-node re-run pattern (`request_input_rerun`). + +## Notes + +`auth_oauth` issues a real GitHub OAuth authorization request; completing the +token exchange requires registering an OAuth app and setting `GITHUB_CLIENT_ID` / +`GITHUB_CLIENT_SECRET`, so its resume step needs a live provider. A couple of +samples read state via `ctx.state` where Python injects it as a function +parameter (an intended API difference) — this is noted in those files. diff --git a/samples/workflows/agent_in_workflow/agent.ts b/samples/workflows/agent_in_workflow/agent.ts index 518d76dec..929e31b3c 100644 --- a/samples/workflows/agent_in_workflow/agent.ts +++ b/samples/workflows/agent_in_workflow/agent.ts @@ -5,41 +5,106 @@ */ /** - * Agent in a workflow: a real LlmAgent as a workflow node, between two function - * nodes. Mirrors Python `workflows/agent_in_workflow`. + * Agent in a workflow: a `task`-mode `LlmAgent` (`intake_agent`) chats to collect + * a structured identity and completes via `finish_task`; a function node routes + * on the result (retry the intake, or proceed); and a second `LlmAgent` + * (`generate_instruction`) uses a `require_confirmation` tool. Faithful port of + * Python `contributing/samples/workflows/agent_in_workflow`. * - * REQUIRES an API key (this one calls a live model). Set GEMINI_API_KEY (a - * `.env` in the working directory is loaded automatically), then: + * REQUIRES an API key. Set GEMINI_API_KEY, then: * node dev/dist/esm/cli_entrypoint.js run samples/workflows/agent_in_workflow/agent.ts + * Provide a name + phone (use "Jane Doe" to pass the identity check). The + * `find_orders` tool pauses for confirmation before running. */ import { + createEvent, + DEFAULT_ROUTE, + FunctionTool, LlmAgent, node, NodeContext, Workflow, WorkflowAgent, } from '@google/adk'; +import {z} from 'zod'; -const preprocess = node( - (_c: NodeContext, input: string) => `Please answer this concisely: ${input}`, - {name: 'preprocess'}, -); +const patientIdentity = z.object({ + name: z.string().describe("The patient's full name."), + phoneNumber: z.string().describe("The patient's phone number."), +}); +type PatientIdentity = z.infer; -const assistant = new LlmAgent({ - name: 'assistant', +/** Emits a plain display message (Python `Event(message=...)`). */ +const message = (text: string) => + createEvent({content: {role: 'model', parts: [{text}]}}); + +// A task-mode agent: it chats to gather the identity and calls finish_task with +// the structured result, which becomes this node's output. +const intakeAgent = new LlmAgent({ + name: 'intake_agent', model: 'gemini-2.5-flash', - instruction: 'You are a helpful assistant. Answer the user concisely.', + mode: 'task', + outputSchema: patientIdentity, + instruction: `You are a medical lab intake assistant. Your job is to chat with +the user to get their full name and phone number. Do not make up +information. Once you have both, finish your task. +If identity check failed, ask for another name.`, }); -const postprocess = node( - (_c: NodeContext, answer: string) => `Assistant replied:\n${answer}`, - {name: 'postprocess'}, +// Mocks checking the database for the patient. Routes back to intake_agent if +// the name is not Jane Doe. +const checkIdentity = node( + (_ctx: NodeContext, identity: PatientIdentity) => { + if (identity.name.toLowerCase() !== 'jane doe') { + return createEvent({ + route: 'retry', + content: { + role: 'model', + parts: [ + { + text: `Could not find matching records for ${identity.name}. Let's try again.`, + }, + ], + }, + }); + } + return message(`Hello ${identity.name}! Let me look up your orders.`); + }, + {name: 'check_identity'}, ); +// A tool that requires confirmation before it runs (HITL). +const findOrders = new FunctionTool({ + name: 'find_orders', + description: "Finds the patient's lab orders.", + execute: () => ['CBC (Complete Blood Count)', 'Lipid Panel'], + requireConfirmation: true, +}); + +const generateInstruction = new LlmAgent({ + name: 'generate_instruction', + model: 'gemini-2.5-flash', + instruction: `You MUST call the find_orders tool to get the patient's actual +orders. Do not invent orders. After the tool returns, list the orders found and +then generate a concise instruction about how to prepare based on those orders.`, + tools: [findOrders], +}); + export const rootAgent = new WorkflowAgent( new Workflow({ - name: 'agent_in_workflow', - edges: [['START', preprocess, assistant, postprocess]], + name: 'task_in_workflow', + edges: [ + ['START', intakeAgent, checkIdentity], + [ + checkIdentity, + { + retry: intakeAgent, + // generate_instruction re-runs on resume so its tool-confirmation can + // resolve after the user approves the find_orders call. + [DEFAULT_ROUTE]: node(generateInstruction, {rerunOnResume: true}), + }, + ], + ], }), ); diff --git a/samples/workflows/node_as_tool/agent.ts b/samples/workflows/node_as_tool/agent.ts index fee6cdcae..ff6cbef72 100644 --- a/samples/workflows/node_as_tool/agent.ts +++ b/samples/workflows/node_as_tool/agent.ts @@ -5,31 +5,103 @@ */ /** - * Node-as-tool: a node imperatively calls other nodes via `ctx.runNode()`. - * Mirrors Python `workflows/node_as_tool`. + * Node-as-tool: an `LlmAgent` uses a `Workflow` AND a function node as tools. + * The framework auto-wraps each as a `NodeTool`, so the model can call them like + * any other tool; a node may even pause for input (HITL) mid-tool-call. Faithful + * port of Python `contributing/samples/workflows/node_as_tool`. * - * Run: node dev/dist/esm/cli_entrypoint.js run samples/workflows/node_as_tool/agent.ts + * `customer_lookup_workflow` looks up a customer's tier; `calculate_discount` + * (a node) streams a status message and, for VIP tiers, raises a `RequestInput` + * to confirm the discount — pausing the agent until the user responds. + * + * REQUIRES an API key. Set GEMINI_API_KEY, then: + * node dev/dist/esm/cli_entrypoint.js run samples/workflows/node_as_tool/agent.ts + * Turn 1: "Look up user u123 and tell me my discount." + * Turn 2 (VIP): resume by answering the confirmation (a function response to the + * `confirm_vip_discount` interrupt, e.g. via the web UI). */ -import {node, NodeContext, Workflow, WorkflowAgent} from '@google/adk'; +import { + createEvent, + LlmAgent, + node, + NodeContext, + RequestInput, + Workflow, +} from '@google/adk'; +import {z} from 'zod'; -const add = node( - (_c: NodeContext, args: {a: number; b: number}) => args.a + args.b, - {name: 'add'}, -); +const customerLookupArgs = z.object({ + userId: z.string().describe("The customer's unique identifier."), +}); + +/** Emits a plain display message (Python `Event(message=...)`). */ +const message = (text: string) => + createEvent({content: {role: 'model', parts: [{text}]}}); + +// A node exposed as a tool. It streams an intermediate message and, for VIP +// tiers, pauses (RequestInput) to confirm the discount. rerunOnResume=true so it +// re-runs on resume and reads the reply from ctx.resumeInputs. +const calculateDiscount = node( + function* (ctx: NodeContext, args: {tier: string}) { + yield message(`Checking discount rules for tier '${args.tier}'...`); -const orchestrator = node( - async (ctx: NodeContext) => { - const first = await ctx.runNode(add, {a: 2, b: 3}); - const second = await ctx.runNode(add, {a: 10, b: first.output as number}); - return `2 + 3 = ${first.output}, then + 10 = ${second.output}`; + const resume = ctx.resumeInputs['confirm_vip_discount']; + if (args.tier.includes('VIP')) { + if (resume === undefined) { + yield new RequestInput({ + interruptId: 'confirm_vip_discount', + message: `Apply VIP discount for tier '${args.tier}'?`, + }); + return; + } + const answer = + typeof resume === 'object' && resume !== null + ? (resume as {text?: string}).text + : resume; + yield ['yes', 'y', 'true'].includes(String(answer).toLowerCase()) + ? '20% off' + : '5% off (VIP declined)'; + } else { + yield '5% off'; + } + }, + { + name: 'calculate_discount', + description: + 'Calculates the discount percentage based on the customer tier.', + inputSchema: z.object({ + tier: z.string().describe('The customer membership tier (e.g. VIP).'), + }), + rerunOnResume: true, }, - {name: 'orchestrator'}, ); -export const rootAgent = new WorkflowAgent( - new Workflow({ - name: 'node_as_tool', - edges: [['START', orchestrator]], +// A Workflow exposed as a tool: looks up customer status/tier by user_id. +const lookupCustomerData = node( + (_ctx: NodeContext, args: {userId: string}) => ({ + userId: args.userId, + tier: 'Verified VIP Member', }), + {name: 'lookup_customer_data'}, ); + +const customerLookupWorkflow = new Workflow({ + name: 'customer_lookup_workflow', + description: 'Looks up customer status and tier by user_id.', + inputSchema: customerLookupArgs, + edges: [['START', lookupCustomerData]], +}); + +// The agent uses both the Workflow and the node as tools. +export const rootAgent = new LlmAgent({ + name: 'customer_service_agent', + model: 'gemini-2.5-flash', + instruction: ` + You are a customer service assistant. + 1. First, call \`customer_lookup_workflow\` using the user_id to get their membership tier. + 2. Then, call \`calculate_discount\` with that tier to find out what discount they get. + Summarize these details for the customer. + `, + tools: [customerLookupWorkflow, calculateDiscount], +}); From eec18701e666c7af9e036c75648f9157376cd2db Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Mon, 27 Jul 2026 12:25:59 -0700 Subject: [PATCH 34/41] fix(workflow): boolean/multi-value routes and serialization fidelity Close three graph-workflow parity gaps found against the ADK graph docs: - Boolean routes in a routing map: normalize 'true'/'false' object keys back to booleans (as numeric keys already were), so a node emitting route: true matches {true: a, false: b} (mirrors Python dict key True). Previously such a node silently matched nothing and produced no output. - Multi-value routes: widen Event/EventActions/NodeContext route to accept an array so one node can fire several branches at once. The graph matcher already handled arrays; only the public types blocked it. - Serialization: preserve arbitrary output and actions.agentState payloads across snake/camel round-trips (DB/Vertex session backends) instead of mangling their keys, fixing node-output fidelity and HITL resume input. Adds unit coverage for each (routing_test, event_test). --- core/src/events/event.ts | 17 ++++++-- core/src/events/event_actions.ts | 7 +-- core/src/workflow/node_context.ts | 5 ++- core/src/workflow/utils/graph_parser.ts | 22 ++++++++-- core/src/workflow/utils/rehydration_utils.ts | 6 +-- core/src/workflow/workflow.ts | 2 +- core/test/events/event_test.ts | 45 ++++++++++++++++++++ core/test/workflow/routing_test.ts | 34 +++++++++++++++ 8 files changed, 122 insertions(+), 16 deletions(-) diff --git a/core/src/events/event.ts b/core/src/events/event.ts index c022f14a4..58316a8ad 100644 --- a/core/src/events/event.ts +++ b/core/src/events/event.ts @@ -92,10 +92,12 @@ export interface Event extends LlmResponse { output?: unknown; /** - * Workflow: the route key emitted by a routing node, used by the graph to - * select the matching outgoing edge. Mirrors Python `Event.route`. + * Workflow: the route key(s) emitted by a routing node, used by the graph to + * select the matching outgoing edge(s). A single value fires one branch; an + * array fires every branch whose route matches any listed value (multi-route + * dispatch). Mirrors Python `Event.route`. */ - route?: string | number | boolean; + route?: string | number | boolean | Array; /** * Workflow: provenance of the emitting node. Mirrors Python `Event.node_info`. @@ -311,6 +313,11 @@ 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). + 'output', + 'actions.agentState', ]; /** @@ -330,6 +337,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). + 'output', + 'actions.agent_state', ]; /** diff --git a/core/src/events/event_actions.ts b/core/src/events/event_actions.ts index 47711d9c8..d1fdf8db6 100644 --- a/core/src/events/event_actions.ts +++ b/core/src/events/event_actions.ts @@ -77,10 +77,11 @@ export interface EventActions { endOfAgent?: boolean; /** - * Workflow: route key selected by a routing node (alternative carrier to the - * top-level `Event.route`, used by callbacks/tools). + * Workflow: route key(s) selected by a routing node (alternative carrier to + * the top-level `Event.route`, used by callbacks/tools). A single value or an + * array for multi-route dispatch. */ - route?: string | number | boolean; + route?: string | number | boolean | Array; [key: string]: unknown; } diff --git a/core/src/workflow/node_context.ts b/core/src/workflow/node_context.ts index c1b8a1dff..e6785176b 100644 --- a/core/src/workflow/node_context.ts +++ b/core/src/workflow/node_context.ts @@ -9,6 +9,7 @@ import {Event} from '../events/event.js'; import {createEventActions, EventActions} from '../events/event_actions.js'; import {State} from '../sessions/state.js'; import type {BaseNode} from './base_node.js'; +import type {RouteValue} from './graph.js'; import {executeChildNode, RunNodeOptions} from './node_runner.js'; import type {ScheduleDynamicNode} from './schedule_dynamic_node.js'; import {EventChannel} from './utils/event_channel.js'; @@ -51,8 +52,8 @@ export class NodeContext { /** The structured output produced by the node during its run. */ output: unknown = undefined; - /** The route key emitted by the node, if any. */ - route?: string | number | boolean; + /** The route key(s) emitted by the node, if any (array = multi-route). */ + route?: RouteValue | RouteValue[]; /** Interrupt ids the node is currently blocked on (HITL). */ interruptIds: string[] = []; diff --git a/core/src/workflow/utils/graph_parser.ts b/core/src/workflow/utils/graph_parser.ts index cae1ebaeb..5cace2977 100644 --- a/core/src/workflow/utils/graph_parser.ts +++ b/core/src/workflow/utils/graph_parser.ts @@ -26,6 +26,23 @@ function isRouteValue(value: unknown): value is RouteValue { return t === 'string' || t === 'number' || t === 'boolean'; } +/** + * Normalizes a routing-map key back to its typed {@link RouteValue}. JS object + * keys are always strings, so integer route keys arrive as numeric strings and + * boolean route keys as `'true'`/`'false'`. Reconstructing the typed value lets + * a node emitting `2` or `true` match `{2: ...}` / `{true: ...}` (mirroring + * Python dict keys `2` / `True`). + */ +function normalizeRouteKey(routeKey: string): RouteValue { + if (/^-?\d+$/.test(routeKey)) { + return Number(routeKey); + } + if (routeKey === 'true' || routeKey === 'false') { + return routeKey === 'true'; + } + return routeKey; +} + /** Expands a routing map into individual (from, to, route) triples. */ function expandRoutingMap( fromElement: ChainElement, @@ -42,10 +59,7 @@ function expandRoutingMap( [ChainElement, NodeLike | readonly NodeLike[], RouteValue] > = []; for (const routeKey of keys) { - // Object keys are strings; numeric route keys arrive as numeric strings. - const normalizedKey: RouteValue = /^-?\d+$/.test(routeKey) - ? Number(routeKey) - : routeKey; + const normalizedKey: RouteValue = normalizeRouteKey(routeKey); const target = routingMap[routeKey]; if (Array.isArray(target)) { for (const node of target) { diff --git a/core/src/workflow/utils/rehydration_utils.ts b/core/src/workflow/utils/rehydration_utils.ts index d35d49ec9..3fb2a1f75 100644 --- a/core/src/workflow/utils/rehydration_utils.ts +++ b/core/src/workflow/utils/rehydration_utils.ts @@ -23,8 +23,8 @@ const RESULT_KEY = 'result'; export interface RehydratedNode { /** The node's cached output from a prior run, if it produced one. */ output?: unknown; - /** The route the node emitted, if any. */ - route?: RouteValue; + /** 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). */ @@ -103,7 +103,7 @@ function reconstruct( node.branch = event.branch; } if (event.route !== undefined) { - node.route = event.route as RouteValue; + node.route = event.route as RouteValue | RouteValue[]; } for (const id of event.longRunningToolIds ?? []) { node.interruptIds.add(id); diff --git a/core/src/workflow/workflow.ts b/core/src/workflow/workflow.ts index 89e76d468..b93a022e8 100644 --- a/core/src/workflow/workflow.ts +++ b/core/src/workflow/workflow.ts @@ -412,7 +412,7 @@ export class Workflow extends BaseNode { loop: LoopState, nodeName: string, output: unknown, - route: RouteValue | undefined, + route: RouteValue | RouteValue[] | undefined, branch: string | undefined, ): void { const nextNodes = this.graph!.getNextPendingNodes(nodeName, route ?? null); diff --git a/core/test/events/event_test.ts b/core/test/events/event_test.ts index 7a7e547f0..4fac62afe 100644 --- a/core/test/events/event_test.ts +++ b/core/test/events/event_test.ts @@ -322,5 +322,50 @@ describe('Event Utils', () => { NestedKey: 'value2', }); }); + + it('preserves workflow output and agentState keys verbatim', () => { + const camelEvent = createEvent({ + id: '123', + invocationId: 'inv1', + output: {cityName: 'Paris', timeInfo: '10:10 AM'}, + actions: createEventActions({ + agentState: {input: {userId: 42, requestedItems: ['a']}}, + }), + }); + const snakeEvent = transformToSnakeCaseEvent(camelEvent); + // Arbitrary payloads must NOT be snake_cased. + expect(snakeEvent.output).toEqual({ + cityName: 'Paris', + timeInfo: '10:10 AM', + }); + expect( + (snakeEvent.actions as Record).agent_state, + ).toEqual({input: {userId: 42, requestedItems: ['a']}}); + }); + }); + + describe('event round-trip serialization', () => { + it('round-trips workflow output and agentState without mangling keys', () => { + const original = createEvent({ + id: '123', + invocationId: 'inv1', + output: {cityName: 'Paris', nested: {timeInfo: '10:10 AM'}}, + route: ['BUG', 'LOGISTICS'], + actions: createEventActions({ + agentState: {input: {userId: 42, camelKey: 'v'}}, + }), + }); + const restored = transformToCamelCaseEvent( + transformToSnakeCaseEvent(original), + ); + expect(restored.output).toEqual({ + cityName: 'Paris', + nested: {timeInfo: '10:10 AM'}, + }); + expect(restored.route).toEqual(['BUG', 'LOGISTICS']); + expect(restored.actions?.agentState).toEqual({ + input: {userId: 42, camelKey: 'v'}, + }); + }); }); }); diff --git a/core/test/workflow/routing_test.ts b/core/test/workflow/routing_test.ts index c9398b46b..0fa5db644 100644 --- a/core/test/workflow/routing_test.ts +++ b/core/test/workflow/routing_test.ts @@ -90,4 +90,38 @@ describe('workflow routing values', () => { }); 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)', + }); + }); }); From 95348c1ce99e9f69791467151e92a8ce5f2bb384 Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Mon, 27 Jul 2026 12:32:45 -0700 Subject: [PATCH 35/41] fix(workflow): actually cancel a node when its timeout elapses MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Previously a per-node `timeout` only raced the node body against a timer: on timeout the runner rejected with NodeTimeoutError but the node's generator kept executing, continuing to push events into the channel (leaking into a retry or the next node) and consuming resources — with `retryConfig` it could even run overlapping attempts. runOnce now drives the generator step-by-step raced against the deadline. On timeout it stops consuming events, closes the generator so its finally/cleanup runs, and aborts a new `ctx.abortSignal` (linked to the invocation's abort signal so cooperative node bodies can cancel their own in-flight work. Mirrors asyncio.wait_for cancellation semantics. Adds a test asserting a timed-out node aborts its signal and its post-deadline events are dropped. EOF ) --- core/src/workflow/node_context.ts | 10 ++ core/src/workflow/node_runner.ts | 131 +++++++++++++--------- core/test/workflow/node_execution_test.ts | 59 ++++++++++ 3 files changed, 146 insertions(+), 54 deletions(-) diff --git a/core/src/workflow/node_context.ts b/core/src/workflow/node_context.ts index e6785176b..1f7175cb0 100644 --- a/core/src/workflow/node_context.ts +++ b/core/src/workflow/node_context.ts @@ -58,6 +58,16 @@ export class NodeContext { /** Interrupt ids the node is currently blocked on (HITL). */ 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. + */ + abortSignal?: AbortSignal; + /** * The dynamic-node scheduler for this subtree. When set, `ctx.runNode()` * routes through it (dedup/resume/fresh); otherwise it runs the child diff --git a/core/src/workflow/node_runner.ts b/core/src/workflow/node_runner.ts index a1105fbdc..d109414fc 100644 --- a/core/src/workflow/node_runner.ts +++ b/core/src/workflow/node_runner.ts @@ -128,7 +128,15 @@ export async function executeChildNode( /** * Drives one attempt of `node.run()`, enriching and pushing each event and - * tracking the child's output/route. Wrapped in a timeout when configured. + * 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 + * `asyncio.wait_for`. */ async function runOnce( node: BaseNode, @@ -138,38 +146,79 @@ async function runOnce( branch: string | undefined, isolationScope: string | undefined, ): Promise { - const body = (async () => { - for await (const event of node.run(child, input)) { - enrichEvent(event, child, nodeName, branch, isolationScope); - if (event.output !== undefined) { - child.output = event.output; - } - if (event.route !== undefined) { - child.route = event.route; - } - // HITL: an interrupt event marks its ids as long-running tool ids. - if (event.longRunningToolIds && event.longRunningToolIds.length > 0) { - for (const id of event.longRunningToolIds) { - if (!child.interruptIds.includes(id)) { - child.interruptIds.push(id); - } + const consume = (event: Event): void => { + enrichEvent(event, child, nodeName, branch, isolationScope); + if (event.output !== undefined) { + child.output = event.output; + } + if (event.route !== undefined) { + child.route = event.route; + } + // HITL: an interrupt event marks its ids as long-running tool ids. + if (event.longRunningToolIds && event.longRunningToolIds.length > 0) { + for (const id of event.longRunningToolIds) { + if (!child.interruptIds.includes(id)) { + child.interruptIds.push(id); } - // Persist the node's input on the interrupt event so a resumed - // (waiting) node re-runs with its ORIGINAL input, not the resume - // message. Rehydrated by reconstructNodeStates on the next turn. - event.actions.agentState = { - ...(event.actions.agentState ?? {}), - input, - }; } - child.channel.push(event); + // Persist the node's input on the interrupt event so a resumed + // (waiting) node re-runs with its ORIGINAL input, not the resume + // message. Rehydrated by reconstructNodeStates on the next turn. + event.actions.agentState = { + ...(event.actions.agentState ?? {}), + input, + }; } - })(); + child.channel.push(event); + }; - if (node.timeout && node.timeout > 0) { - await withTimeout(body, node.timeout, nodeName); + if (!(typeof node.timeout === 'number' && node.timeout > 0)) { + for await (const event of node.run(child, input)) { + consume(event); + } + return; + } + + const timeoutSeconds = node.timeout; + const controller = new AbortController(); + const parentSignal = child.invocationContext.abortSignal; + const onParentAbort = () => controller.abort(); + if (parentSignal?.aborted) { + controller.abort(); } else { - await body; + parentSignal?.addEventListener('abort', onParentAbort, {once: true}); + } + const timer = setTimeout(() => controller.abort(), timeoutSeconds * 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})); + if (controller.signal.aborted) { + fail(); + } else { + controller.signal.addEventListener('abort', fail, {once: true}); + } + }); + + const iterator = node.run(child, input)[Symbol.asyncIterator](); + try { + for (;;) { + const result = await Promise.race([iterator.next(), aborted]); + if (result.done) { + break; + } + consume(result.value); + } + } finally { + clearTimeout(timer); + parentSignal?.removeEventListener('abort', onParentAbort); + child.abortSignal = undefined; + // Best-effort: close the generator so its `finally`/cleanup runs. This is + // queued behind any in-flight `next()`; its result is discarded. + void Promise.resolve(iterator.return?.(undefined)).catch(() => {}); } } @@ -210,32 +259,6 @@ function withBranch( }); } -/** - * Rejects with {@link NodeTimeoutError} if `promise` does not settle within - * `timeoutSeconds`. - */ -function withTimeout( - promise: Promise, - timeoutSeconds: number, - nodeName: string, -): Promise { - return new Promise((resolve, reject) => { - const timer = setTimeout(() => { - reject(new NodeTimeoutError({nodeName, timeout: timeoutSeconds})); - }, timeoutSeconds * 1000); - promise.then( - (value) => { - clearTimeout(timer); - resolve(value); - }, - (err) => { - clearTimeout(timer); - reject(err); - }, - ); - }); -} - /** * Promise-based delay that rejects early if the abort signal fires. */ diff --git a/core/test/workflow/node_execution_test.ts b/core/test/workflow/node_execution_test.ts index ec588ff63..285c73225 100644 --- a/core/test/workflow/node_execution_test.ts +++ b/core/test/workflow/node_execution_test.ts @@ -226,4 +226,63 @@ describe('Phase 1 — node execution & the push/pull bridge', () => { NodeTimeoutError, ); }); + + it('cancels a timed-out node: aborts the signal and drops post-deadline events', async () => { + let captured: AbortSignal | undefined; + class SlowStream extends BaseNode { + protected async *runImpl(ctx: NodeContext) { + captured = ctx.abortSignal; + yield createEvent({ + author: 'slow', + content: {role: 'model', parts: [{text: 'early'}]}, + }); + // Cooperative wait that ends on abort; well past the 20ms timeout. + await new Promise((resolve) => { + const t = setTimeout(resolve, 500); + ctx.abortSignal?.addEventListener( + 'abort', + () => { + clearTimeout(t); + resolve(); + }, + {once: true}, + ); + }); + yield createEvent({ + author: 'slow', + content: {role: 'model', parts: [{text: 'late'}]}, + output: 'late', + }); + } + } + const slow = new SlowStream({name: 'slow', timeout: 0.02}); + + // Custom harness: capture events even though the run rejects. + const channel = new EventChannel(); + const root = new NodeContext({ + invocationContext: createIc(), + channel, + nodePath: '', + runId: 'root', + }); + const seen: string[] = []; + const orchestration = root.runNode(slow, 'x').then( + () => channel.close(), + (err) => channel.fail(err), + ); + let thrown: unknown; + try { + for await (const ev of channel) { + seen.push(ev.content?.parts?.[0]?.text ?? ''); + } + } catch (err) { + thrown = err; + } + await orchestration.catch(() => {}); + + expect(thrown).toBeInstanceOf(NodeTimeoutError); + expect(seen).toContain('early'); + expect(seen).not.toContain('late'); // produced after the deadline -> dropped + expect(captured?.aborted).toBe(true); // signal fired for cooperative cancel + }); }); From dfeb412a89caf855bacf8f64e60a491e7d955d14 Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Mon, 27 Jul 2026 12:45:09 -0700 Subject: [PATCH 36/41] feat(workflow): resolve {Class.field} and instruction placeholders An LlmAgent run as a workflow node can now reference its structured input and predecessor outputs directly in its instruction, matching Python's data-selection syntax: - {ClassName.field} -> field on the node's input - -> field on a named predecessor output LLMAgentWrapper attaches a WorkflowInstructionScope (the node input plus predecessor outputs by node name, gathered from session events) to a child InvocationContext, and injectSessionState resolves the two placeholder forms from it. Ordinary (non-workflow) agents have no scope, so their instructions are unchanged. Closes the last substantive graph-workflow parity gap (#3). Adds unit tests for the resolver and integration tests through the wrapper. --- core/src/agents/instructions.ts | 68 ++++++++++++++--- core/src/agents/invocation_context.ts | 22 ++++++ core/src/workflow/nodes/llm_agent_wrapper.ts | 63 +++++++++++++-- core/test/agents/instructions_test.ts | 61 +++++++++++++++ core/test/workflow/llm_agent_test.ts | 80 ++++++++++++++++++++ 5 files changed, 279 insertions(+), 15 deletions(-) diff --git a/core/src/agents/instructions.ts b/core/src/agents/instructions.ts index 093586d3a..1f907a381 100644 --- a/core/src/agents/instructions.ts +++ b/core/src/agents/instructions.ts @@ -5,10 +5,36 @@ */ import {State} from '../sessions/state.js'; +import type {WorkflowInstructionScope} from './invocation_context.js'; import {ReadonlyContext} from './readonly_context.js'; const ARTIFACT_PREFIX = 'artifact.'; +/** Matches a `{Class.field}` workflow placeholder key (dotted identifier pair). */ +const WORKFLOW_FIELD_KEY = /^[A-Za-z_]\w*\.[A-Za-z_]\w*$/; + +/** Matches a `` workflow placeholder. */ +const SOURCE_NODE_PLACEHOLDER = + /<\s*[A-Za-z_]\w*\.([A-Za-z_]\w*)\s+from\s+([A-Za-z_]\w*)\s*>/g; + +/** + * Resolves `` placeholders against a workflow + * scope (predecessor outputs by node name). Synchronous; unresolved placeholders + * are left untouched. Mirrors Python's source-node-qualified data selection. + */ +function resolveSourceNodePlaceholders( + template: string, + scope: WorkflowInstructionScope, +): string { + return template.replace(SOURCE_NODE_PLACEHOLDER, (raw, field, nodeName) => { + const out = scope.outputsByNode?.[nodeName]; + if (out && typeof out === 'object' && field in (out as object)) { + return formatValue((out as Record)[field], false); + } + return raw; + }); +} + /** * Resolves a single key from the context (state or artifact). */ @@ -39,19 +65,30 @@ async function resolveKey( } // Step 3: Handle state variable injection. - if (!isValidStateName(key)) { - return rawMatch; - } - - if (key in invocationContext.session.state) { - return formatValue(invocationContext.session.state[key], false); + if (isValidStateName(key)) { + if (key in invocationContext.session.state) { + return formatValue(invocationContext.session.state[key], false); + } + if (isOptional) { + return ''; + } + throw new Error(`Context variable not found: \`${key}\`.`); } - if (isOptional) { - return ''; + // Step 4: Workflow — resolve `{Class.field}` from the current node input. + const scope = invocationContext.workflowInstructionScope; + if (scope && WORKFLOW_FIELD_KEY.test(key)) { + const field = key.slice(key.indexOf('.') + 1); + const input = scope.input; + if (input && typeof input === 'object' && field in (input as object)) { + return formatValue((input as Record)[field], false); + } + if (isOptional) { + return ''; + } } - throw new Error(`Context variable not found: \`${key}\`.`); + return rawMatch; } /** @@ -115,6 +152,14 @@ export async function injectSessionState( template: string, readonlyContext: ReadonlyContext, ): Promise { + // Workflow: first resolve `` placeholders, and + // enable `{Class.field}` resolution below. Both are no-ops (placeholders left + // untouched) for ordinary agents, which have no workflow scope. + const scope = readonlyContext.invocationContext.workflowInstructionScope; + if (scope) { + template = resolveSourceNodePlaceholders(template, scope); + } + const pattern = /\{+[^{}]*}+/g; const matches = Array.from(template.matchAll(pattern)); @@ -130,7 +175,10 @@ export async function injectSessionState( if (isOptional) { key = key.slice(0, -1); } - const isValid = key.startsWith(ARTIFACT_PREFIX) || isValidStateName(key); + const isValid = + key.startsWith(ARTIFACT_PREFIX) || + isValidStateName(key) || + (!!scope && WORKFLOW_FIELD_KEY.test(key)); return { raw, key, diff --git a/core/src/agents/invocation_context.ts b/core/src/agents/invocation_context.ts index 6f8df4665..89a08a9be 100644 --- a/core/src/agents/invocation_context.ts +++ b/core/src/agents/invocation_context.ts @@ -21,6 +21,19 @@ import {BaseAgent} from './base_agent.js'; import {RunConfig} from './run_config.js'; import {TranscriptionEntry} from './transcription_entry.js'; +/** + * Workflow: data exposed to `{Class.field}` and `` + * instruction placeholders when an LlmAgent runs as a workflow node. Populated by + * `LLMAgentWrapper`; absent for ordinary (non-workflow) agent runs, in which case + * those placeholders are left untouched. + */ +export interface WorkflowInstructionScope { + /** The current node's input, exposing fields for `{Class.field}`. */ + input?: unknown; + /** Predecessor node outputs keyed by node name, for ``. */ + outputsByNode?: Record; +} + /** * The parameters for creating an invocation context. */ @@ -42,6 +55,7 @@ export interface InvocationContextParams { abortSignal?: AbortSignal; agentStates?: Record; endOfAgents?: Record; + workflowInstructionScope?: WorkflowInstructionScope; } /** @@ -208,6 +222,13 @@ export class InvocationContext { endOfAgents: Record; + /** + * Workflow: field-resolution scope for `{Class.field}` / + * `` instruction placeholders (set by + * `LLMAgentWrapper`). + */ + workflowInstructionScope?: WorkflowInstructionScope; + /** * @param params The parameters for creating an invocation context. */ @@ -228,6 +249,7 @@ export class InvocationContext { this.abortSignal = params.abortSignal; this.agentStates = params.agentStates ?? {}; this.endOfAgents = params.endOfAgents ?? {}; + this.workflowInstructionScope = params.workflowInstructionScope; // Inherit the parent invocation's cost manager when one is available. // Child contexts created for sub-agents, agent transfers and loop diff --git a/core/src/workflow/nodes/llm_agent_wrapper.ts b/core/src/workflow/nodes/llm_agent_wrapper.ts index e76cda844..b95bbc20d 100644 --- a/core/src/workflow/nodes/llm_agent_wrapper.ts +++ b/core/src/workflow/nodes/llm_agent_wrapper.ts @@ -6,6 +6,11 @@ import {Content} from '@google/genai'; import {BaseAgent} from '../../agents/base_agent.js'; +import { + InvocationContext, + InvocationContextParams, + WorkflowInstructionScope, +} from '../../agents/invocation_context.js'; import {isLlmAgent, LlmAgent} from '../../agents/llm_agent.js'; import { createEvent, @@ -69,15 +74,24 @@ export class LLMAgentWrapper extends BaseNode { ctx.session.events.push(userEvent); } + // Expose the node input and predecessor outputs to `{Class.field}` and + // `` instruction placeholders (Python's + // data-selection syntax). Carried on a child context so it never leaks to + // sibling/ordinary agent runs. + const agentIc = withWorkflowInstructionScope(ctx.invocationContext, { + input, + outputsByNode: collectPredecessorOutputs(ctx), + }); + // Task mode: run a multi-round loop until the agent calls `finish_task`, // whose arguments become the node output. if (isLlmAgent(this.agent) && this.agent.mode === 'task') { - yield* this.runTaskMode(ctx, this.agent); + yield* this.runTaskMode(ctx, agentIc, this.agent); return; } // Run the agent, following any transfer_to_agent hand-offs to peers. - yield* this.runWithTransfers(ctx, this.agent, 0); + yield* this.runWithTransfers(ctx, agentIc, this.agent, 0); } /** @@ -89,12 +103,13 @@ export class LLMAgentWrapper extends BaseNode { */ private async *runTaskMode( ctx: NodeContext, + agentIc: InvocationContext, agent: LlmAgent, ): AsyncGenerator { const finishTool = agent.finishTaskTool; let pendingArgs: Record | undefined; - for await (const event of agent.runAsync(ctx.invocationContext)) { + for await (const event of agent.runAsync(agentIc)) { const finishCall = getFunctionCalls(event).find( (fc) => fc.name === FINISH_TASK_TOOL_NAME, ); @@ -129,6 +144,7 @@ export class LLMAgentWrapper extends BaseNode { */ private async *runWithTransfers( ctx: NodeContext, + agentIc: InvocationContext, agent: BaseAgent, depth: number, ): AsyncGenerator { @@ -140,7 +156,7 @@ export class LLMAgentWrapper extends BaseNode { } let transferTarget: string | undefined; - for await (const event of agent.runAsync(ctx.invocationContext)) { + for await (const event of agent.runAsync(agentIc)) { this.maybeSetOutput(event); yield event; if (event.actions?.transferToAgent) { @@ -156,7 +172,7 @@ export class LLMAgentWrapper extends BaseNode { `LLMAgentWrapper: transfer target agent '${transferTarget}' not found.`, ); } - yield* this.runWithTransfers(ctx, target, depth + 1); + yield* this.runWithTransfers(ctx, agentIc, target, depth + 1); } } @@ -198,6 +214,43 @@ export class LLMAgentWrapper extends BaseNode { } } +/** + * Creates a child InvocationContext carrying a workflow instruction scope, + * preserving the shared session/services/cost manager (like `withBranch`). + */ +function withWorkflowInstructionScope( + ic: InvocationContext, + scope: WorkflowInstructionScope, +): InvocationContext { + return new InvocationContext({ + ...(ic as unknown as InvocationContextParams), + workflowInstructionScope: scope, + }); +} + +/** + * Collects predecessor node outputs (keyed by node name) for the current + * invocation from the session events, for `` + * resolution. Node names are the leaf of each event's `nodeInfo.path` (with any + * `@runId` suffix stripped). + */ +function collectPredecessorOutputs(ctx: NodeContext): Record { + const outputs: Record = {}; + for (const event of ctx.session.events) { + if (event.invocationId !== ctx.invocationId || event.output === undefined) { + continue; + } + const path = event.nodeInfo?.path; + if (!path) { + continue; + } + const leaf = path.slice(path.lastIndexOf('.') + 1); + const name = leaf.includes('@') ? leaf.slice(0, leaf.indexOf('@')) : leaf; + outputs[name] = event.output; + } + return outputs; +} + function hasFunctionCalls(event: Event): boolean { return (event.content?.parts ?? []).some((p) => p.functionCall); } diff --git a/core/test/agents/instructions_test.ts b/core/test/agents/instructions_test.ts index 6b7ed862e..50c01ecf9 100644 --- a/core/test/agents/instructions_test.ts +++ b/core/test/agents/instructions_test.ts @@ -15,6 +15,7 @@ import {injectSessionState} from '../../src/agents/instructions.js'; function makeContext( state: Record = {}, artifactService?: unknown, + workflowInstructionScope?: unknown, ): ReadonlyContext { const fakeInvocationContext = { session: { @@ -24,6 +25,7 @@ function makeContext( state, }, artifactService, + workflowInstructionScope, } as unknown as InvocationContext; return new ReadonlyContext(fakeInvocationContext); @@ -304,4 +306,63 @@ describe('injectSessionState', () => { 'Data: {"inlineData":{"mimeType":"text/plain","data":"abc"}}', ); }); + + describe('workflow field placeholders', () => { + it('resolves {Class.field} from the node input', async () => { + const ctx = makeContext({}, undefined, { + input: {time_info: '10:10 AM', city: 'Paris'}, + }); + expect( + await injectSessionState( + 'It is {CityTime.time_info} in {CityTime.city} right now.', + ctx, + ), + ).toBe('It is 10:10 AM in Paris right now.'); + }); + + it('resolves from predecessor outputs', async () => { + const ctx = makeContext({}, undefined, { + outputsByNode: { + lookup_time_function: {time_info: '9:00 AM', city: 'Rome'}, + }, + }); + expect( + await injectSessionState( + 'It is in ' + + '.', + ctx, + ), + ).toBe('It is 9:00 AM in Rome.'); + }); + + it('resolves workflow fields alongside normal state keys', async () => { + const ctx = makeContext({tone: 'formal'}, undefined, { + input: {city: 'Paris'}, + }); + expect(await injectSessionState('{tone}: {City.city}', ctx)).toBe( + 'formal: Paris', + ); + }); + + it('leaves {Class.field} untouched when there is no workflow scope', async () => { + const ctx = makeContext(); + expect(await injectSessionState('It is {CityTime.time_info}.', ctx)).toBe( + 'It is {CityTime.time_info}.', + ); + }); + + it('leaves untouched when there is no workflow scope', async () => { + const ctx = makeContext(); + expect(await injectSessionState('X Y', ctx)).toBe( + 'X Y', + ); + }); + + it('leaves an unknown {Class.field} untouched when the field is absent', async () => { + const ctx = makeContext({}, undefined, {input: {city: 'Paris'}}); + expect(await injectSessionState('{CityTime.missing}', ctx)).toBe( + '{CityTime.missing}', + ); + }); + }); }); diff --git a/core/test/workflow/llm_agent_test.ts b/core/test/workflow/llm_agent_test.ts index a53796093..bdc8432c8 100644 --- a/core/test/workflow/llm_agent_test.ts +++ b/core/test/workflow/llm_agent_test.ts @@ -6,7 +6,9 @@ import {describe, expect, it} from 'vitest'; import {BaseAgent} from '../../src/agents/base_agent.js'; +import {injectSessionState} from '../../src/agents/instructions.js'; import {InvocationContext} from '../../src/agents/invocation_context.js'; +import {ReadonlyContext} from '../../src/agents/readonly_context.js'; import {createEvent, Event} from '../../src/events/event.js'; import {PluginManager} from '../../src/plugins/plugin_manager.js'; import {Session} from '../../src/sessions/session.js'; @@ -89,6 +91,39 @@ class EchoAgent extends BaseAgent { } } +/** + * A fake agent that resolves a given instruction template against its context + * (the way the real instruction request-processor does) and yields the result — + * so we can assert workflow `{Class.field}` / `` placeholders + * resolve from the scope the wrapper attaches to the invocation context. + */ +class TemplateProbeAgent extends BaseAgent { + constructor( + private readonly template: string, + name = 'probe', + ) { + super({name}); + } + protected async *runAsyncImpl( + ctx: InvocationContext, + ): AsyncGenerator { + const resolved = await injectSessionState( + this.template, + new ReadonlyContext(ctx), + ); + yield createEvent({ + author: this.name, + invocationId: ctx.invocationId, + branch: ctx.branch, + content: {role: 'model', parts: [{text: resolved}]}, + }); + } + // eslint-disable-next-line require-yield + protected async *runLiveImpl(): AsyncGenerator { + return; + } +} + describe('Phase 7 — LlmAgent as a node (single_turn)', () => { it('runs an agent as a node and extracts its text output', async () => { const wf = new Workflow({ @@ -115,6 +150,51 @@ describe('Phase 7 — LlmAgent as a node (single_turn)', () => { expect((await driveWorkflow(wf, 'hi')).output).toBe('ECHO:HI'); }); + it('resolves {Class.field} instruction placeholders from the node input', async () => { + const probe = new TemplateProbeAgent( + 'It is {CityTime.time_info} in {CityTime.city} right now.', + ); + const wf = new Workflow({name: 'tmpl_input', edges: [['START', probe]]}); + const {output} = await driveWorkflow(wf, { + time_info: '10:10 AM', + city: 'Paris', + }); + expect(output).toBe('It is 10:10 AM in Paris right now.'); + }); + + it('resolves from a predecessor output event', async () => { + const ic = createIc(); + // Seed a predecessor output the way the Runner persists node events. + ic.session.events.push( + createEvent({ + author: 'lookup_time_function', + invocationId: ic.invocationId, + nodeInfo: {path: 'wf.lookup_time_function'}, + output: {time_info: '9:00 AM', city: 'Rome'}, + }), + ); + const probe = new TemplateProbeAgent( + 'It is in ' + + '.', + ); + const channel = new EventChannel(); + const root = new NodeContext({ + invocationContext: ic, + channel, + nodePath: '', + runId: 'root', + }); + const run = root.runNode(node(probe), undefined, {useAsOutput: true}).then( + () => channel.close(), + (err) => channel.fail(err), + ); + for await (const _ev of channel) { + // drain + } + await run; + expect(root.output).toBe('It is 9:00 AM in Rome.'); + }); + it('node(agent) produces an LLMAgentWrapper carrying the agent name', () => { const wrapped = node(new EchoAgent('assistant')); expect(wrapped).toBeInstanceOf(LLMAgentWrapper); From e8ccf982cb8e446fb96d169cceaa2a79d3f24cbe Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Mon, 27 Jul 2026 12:48:05 -0700 Subject: [PATCH 37/41] fix(workflow): assign ParallelWorker child run ids by item index ParallelWorker scheduled each item via ctx.runNode without an explicit run id, so the run id was assigned by call order. Under bounded concurrency the run-id to item mapping was nondeterministic, which could mismatch cached outputs to items on resume (contradicting the "only failed/interrupted workers re-run" guarantee). Pass runId=String(index) so each item's run is keyed by its index and fast-forwards from its own cached run. --- core/src/workflow/nodes/parallel_worker.ts | 5 +++++ core/test/workflow/parallel_test.ts | 19 +++++++++++++++++++ 2 files changed, 24 insertions(+) diff --git a/core/src/workflow/nodes/parallel_worker.ts b/core/src/workflow/nodes/parallel_worker.ts index 7d8e0e57b..675396ab1 100644 --- a/core/src/workflow/nodes/parallel_worker.ts +++ b/core/src/workflow/nodes/parallel_worker.ts @@ -75,8 +75,13 @@ export class ParallelWorker extends BaseNode { return; } try { + // Key each child run by its item index (not call order) so the + // run id -> item mapping is deterministic. On resume this lets each + // item fast-forward from its own cached run rather than being matched + // to a differently-ordered run id. const child = await ctx.runNode(this.inner, items[i], { useSubBranch: true, + runId: String(i), }); results[i] = child.output; } catch (err) { diff --git a/core/test/workflow/parallel_test.ts b/core/test/workflow/parallel_test.ts index 5cfdebbe2..9c85347e8 100644 --- a/core/test/workflow/parallel_test.ts +++ b/core/test/workflow/parallel_test.ts @@ -152,4 +152,23 @@ describe('Phase 6 — ParallelWorker', () => { }), ).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'])).toEqual([ + 'a#0', + 'b#1', + 'c#2', + 'd#3', + ]); + }); }); From fee379edc16a841f866b82845d2e1220a88ef5c6 Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Mon, 27 Jul 2026 12:49:38 -0700 Subject: [PATCH 38/41] test(workflow): guard resume against DB serialization key-mangling Adds a rehydration test that round-trips resume events through the snake/camel transforms a persistent (DB/Vertex) session store applies, then asserts reconstructNodeStates still recovers a completed node's structured output and a waiting node's stashed input verbatim. Regression guard for the output / actions.agentState preserve-list fix. --- core/test/workflow/resume_test.ts | 45 ++++++++++++++++++++++++++++++- 1 file changed, 44 insertions(+), 1 deletion(-) diff --git a/core/test/workflow/resume_test.ts b/core/test/workflow/resume_test.ts index e3b0d1cb2..cd315efc1 100644 --- a/core/test/workflow/resume_test.ts +++ b/core/test/workflow/resume_test.ts @@ -5,7 +5,13 @@ */ import {describe, expect, it} from 'vitest'; -import {createEvent, Event} from '../../src/events/event.js'; +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'; @@ -70,6 +76,43 @@ describe('Phase 5b — rehydration utility', () => { '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'], + }); + }); }); describe('Phase 5b — HITL resume via the Runner', () => { From da541a7844fd8613dcea12dff7d7b22c3d7dbc4a Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Mon, 27 Jul 2026 14:36:01 -0700 Subject: [PATCH 39/41] fix(workflow): persist LLMAgentWrapper's injected user turn MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The wrapper appended the node input as a user turn by pushing directly onto session.events, which is in-memory only — on DB/Vertex session backends the turn was never persisted and was lost on resume. Route it through sessionService.appendEvent instead (which also adds it to session.events, deduped by id, so the agent still reads it synchronously), falling back to a direct push when no session service is wired. --- core/src/workflow/nodes/llm_agent_wrapper.ts | 15 ++++++- core/test/workflow/llm_agent_test.ts | 41 ++++++++++++++++++++ 2 files changed, 55 insertions(+), 1 deletion(-) diff --git a/core/src/workflow/nodes/llm_agent_wrapper.ts b/core/src/workflow/nodes/llm_agent_wrapper.ts index b95bbc20d..37b1fb9f2 100644 --- a/core/src/workflow/nodes/llm_agent_wrapper.ts +++ b/core/src/workflow/nodes/llm_agent_wrapper.ts @@ -71,7 +71,20 @@ export class LLMAgentWrapper extends BaseNode { if (ctx.isolationScope) { userEvent.isolationScope = ctx.isolationScope; } - ctx.session.events.push(userEvent); + // Persist the injected user turn (not just push it in-memory) so it + // survives on DB/Vertex session backends and is present on resume. + // appendEvent also adds it to `session.events` (deduped by id), which the + // agent reads synchronously to build its request. Falls back to a direct + // push when no session service is wired (e.g. in unit tests). + const sessionService = ctx.invocationContext.sessionService; + if (sessionService) { + await sessionService.appendEvent({ + session: ctx.session, + event: userEvent, + }); + } else { + ctx.session.events.push(userEvent); + } } // Expose the node input and predecessor outputs to `{Class.field}` and diff --git a/core/test/workflow/llm_agent_test.ts b/core/test/workflow/llm_agent_test.ts index bdc8432c8..ce2a6278c 100644 --- a/core/test/workflow/llm_agent_test.ts +++ b/core/test/workflow/llm_agent_test.ts @@ -195,6 +195,47 @@ describe('Phase 7 — LlmAgent as a node (single_turn)', () => { expect(root.output).toBe('It is 9:00 AM in Rome.'); }); + it('persists the injected user turn through the session service', async () => { + const ic = createIc(); + const appended: Event[] = []; + (ic as unknown as {sessionService: unknown}).sessionService = { + appendEvent: async ({ + session, + event, + }: { + session: {events: Event[]}; + event: Event; + }) => { + appended.push(event); + session.events.push(event); // mimic the base service adding to the list + return event; + }, + }; + const channel = new EventChannel(); + const root = new NodeContext({ + invocationContext: ic, + channel, + nodePath: '', + runId: 'root', + }); + const run = root + .runNode(node(new EchoAgent()), 'hi', {useAsOutput: true}) + .then( + () => channel.close(), + (err) => channel.fail(err), + ); + for await (const _ev of channel) { + // drain + } + await run; + + // The user turn went through the persistence path (appendEvent), not just a + // silent in-memory push, and the agent still saw it. + const userTurn = appended.find((e) => e.author === 'user'); + expect(userTurn?.content?.parts?.[0]?.text).toBe('hi'); + expect(root.output).toBe('echo:hi'); + }); + it('node(agent) produces an LLMAgentWrapper carrying the agent name', () => { const wrapped = node(new EchoAgent('assistant')); expect(wrapped).toBeInstanceOf(LLMAgentWrapper); From 1cfdb633bebed9f43d557662a8c997bad7142522 Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Mon, 27 Jul 2026 14:39:02 -0700 Subject: [PATCH 40/41] fix(workflow): scope resume rehydration to a workflow's own child nodes reconstructNodeStates keyed static nodes by their path leaf, so a nested workflow whose node shared a name with an outer node collided on resume (last one won, mis-attributing cached output/interrupts). reconstructNodeStates now accepts a parentPath and, when given, only considers that path's direct children (keyed by child name); Workflow.runImpl passes its own nodePath. The no-parentPath leaf-name mode is kept for the utility/tests. --- core/src/workflow/utils/rehydration_utils.ts | 32 +++++++++++++++++++- core/src/workflow/workflow.ts | 9 ++++-- core/test/workflow/resume_test.ts | 20 ++++++++++++ 3 files changed, 58 insertions(+), 3 deletions(-) diff --git a/core/src/workflow/utils/rehydration_utils.ts b/core/src/workflow/utils/rehydration_utils.ts index 3fb2a1f75..2a5298178 100644 --- a/core/src/workflow/utils/rehydration_utils.ts +++ b/core/src/workflow/utils/rehydration_utils.ts @@ -38,11 +38,23 @@ export interface RehydratedNode { /** * 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 { - // Key static-graph nodes by their name (path leaf), robust to author rewrite. + 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, ); @@ -161,6 +173,24 @@ export function nodeNameFromPath(path: string): string { 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]; +} + /** Unwraps a `{result: value}` FunctionResponse envelope to the bare value. */ export function unwrapResponse(response: unknown): unknown { if ( diff --git a/core/src/workflow/workflow.ts b/core/src/workflow/workflow.ts index b93a022e8..09508a9a3 100644 --- a/core/src/workflow/workflow.ts +++ b/core/src/workflow/workflow.ts @@ -120,8 +120,13 @@ export class Workflow extends BaseNode { // --- REHYDRATE (resume) --- // Reconstruct node state from prior session events and surface resolved - // interrupt responses so waiting nodes can resume. - const rehydrated = reconstructNodeStates(ctx.session?.events ?? []); + // 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) { diff --git a/core/test/workflow/resume_test.ts b/core/test/workflow/resume_test.ts index cd315efc1..d4212460f 100644 --- a/core/test/workflow/resume_test.ts +++ b/core/test/workflow/resume_test.ts @@ -113,6 +113,26 @@ describe('Phase 5b — rehydration utility', () => { 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', () => { From 6fab87329ff4cc6dd19c5ab6890bfbaf02420fbc Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Mon, 27 Jul 2026 14:42:59 -0700 Subject: [PATCH 41/41] refactor(events)!: drop the EventActions catch-all index signature MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit EventActions declared `[key: string]: unknown`, which turned off excess-property and typo checking for the type across the whole codebase (every property access compiled, and Object.values(actions) was unknown[]). All real fields — including the workflow ones (output, agentState, route, ...) — are declared explicitly, and nothing accesses arbitrary keys, so the index signature is removed. Workspace typecheck stays at zero errors. --- core/src/events/event_actions.ts | 2 -- 1 file changed, 2 deletions(-) diff --git a/core/src/events/event_actions.ts b/core/src/events/event_actions.ts index d1fdf8db6..54c89fa71 100644 --- a/core/src/events/event_actions.ts +++ b/core/src/events/event_actions.ts @@ -82,8 +82,6 @@ export interface EventActions { * array for multi-route dispatch. */ route?: string | number | boolean | Array; - - [key: string]: unknown; } /**