Skip to content

Commit 5334b45

Browse files
authored
feat(workflow): LLM-agent-as-node, task mode, and node-as-tool (Part 6) (#593)
* feat(workflow): add LLM-agent-as-node, task mode, and node-as-tool Part 6/9 of the feature/workflows split. Lets agents participate in workflows. - nodes/llm_agent_wrapper: runs a BaseAgent as a workflow node — streaming, transfer_to_agent hand-offs, workflow instruction scope, and task mode (loops until the agent calls finish_task, whose args become the node output). It registers the agent node-builder, explicitly excluding BaseTool (which also exposes runAsync) so tool-before-agent precedence holds regardless of registration order. - nodes/node_tool: exposes a node/workflow as a tool an agent can call. - tools/finish_task_tool: the finish_task tool backing task mode. - agents/llm_agent: task mode (mode/finishTaskTool) and registers the request-input + request-confirmation LLM request processors. - agents/processors/request_input_llm_request_processor: agent-side HITL (request user input mid-run). - agents/{invocation_context,instructions}, basic_llm_request_processor: workflow instruction scope and {Class.field}/<field from node> placeholder resolution. Wired into the public barrel (LLMAgentWrapper, NodeTool) and register_builtin_nodes; llm_agent 3-way merged onto current main. Tests: node_api (16), llm_agent (7), multi_agent (3), instructions (37), plus the workflow integration suite (15 files / 35 tests, recorded model responses). Full core suite green (2476), integration workflows green, docs:check + tsc clean. * fix(workflow): reconcile Part 6 with the updated engine Rebased Part 6 onto the current Part 5 and reconciled with the Parts 2-5 conventions: - executeChildNode now takes a single params object (node_tool.ts). - Move the LLM-agent-wrapper builder into the static node_builders.ts const list; drop its registerNodeBuilder self-registration and the now-obsolete register_builtin_nodes.ts (the const list replaces side-effect registration). * fix(workflow): address Part 6 review comments node-as-tool (node_tool.ts): - Bound node -> tool -> node recursion with a MAX_NODE_TOOL_DEPTH cap, tracked via a new immutable InvocationContext.nodeToolDepth carried through a depth+1 clone (so it survives agent-run ic clones). - Require the invocation event queue and a function-call id (throw otherwise) instead of falling back to a dead queue / a collapsing runId; drop the structural cast on eventQueue. - Pass an empty parent nodePath so the child path is a single segment, not the node name doubled. - Derive the tool parameter schema by narrowing isZodObject inline (drops the `as never`). instanceof -> brand guards: - llm_agent.ts uses isBaseNode; the request-input processor uses isNodeTool (new brand + guard on NodeTool). Both imports become type-only, removing the agents -> workflow value cycle. task mode (llm_agent_wrapper.ts): throw when a task-mode agent ends without a successful finish_task, instead of reporting the node COMPLETE with no output (its turn loop stays bounded by the invocation maxLlmCalls). InvocationContext: add clone(overrides) and use it for both the workflow-instruction-scope and node-tool-depth children (removes the `as unknown as InvocationContextParams` cast); remove the unused agentStates / endOfAgents fields; export WorkflowInstructionScope from common.ts / index.ts (public field type) instead of suppressing the TypeDoc warning. * test(workflow): use the shared createIc fixture (drop hand-rolled casts) Replace the duplicated hand-rolled createIc (with `as unknown as Session` / `as unknown as BaseAgent`) in the Part 6 workflow tests with the shared test_helpers.createIc (createSession + a real BaseAgent), removing the repeated double-casts the review flagged. * docs(workflow): de-link isNodeTool in NodeTool brand comment * fix(workflow): re-publish Part 6 nodes after the Part 5 export change Part 5 replaced the `export * from './workflow/index.js'` star in index.ts with explicit named re-exports in common.ts, so Part 6's public additions must be listed there too: - Add LLMAgentWrapper / NodeTool and the LLMAgentWrapperConfig type to the common.ts workflow block (they reach the web entry point this way, and the block now mirrors the barrel exactly). - Update node_api_test to subclass the renamed `WorkflowNode` base class.
1 parent 1b90f3b commit 5334b45

38 files changed

Lines changed: 3614 additions & 21 deletions

core/src/agents/instructions.ts

Lines changed: 58 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -5,10 +5,36 @@
55
*/
66

77
import {State} from '../sessions/state.js';
8+
import type {WorkflowInstructionScope} from './invocation_context.js';
89
import {ReadonlyContext} from './readonly_context.js';
910

1011
const ARTIFACT_PREFIX = 'artifact.';
1112

13+
/** Matches a `{Class.field}` workflow placeholder key (dotted identifier pair). */
14+
const WORKFLOW_FIELD_KEY = /^[A-Za-z_]\w*\.[A-Za-z_]\w*$/;
15+
16+
/** Matches a `<Class.field from source_node>` workflow placeholder. */
17+
const SOURCE_NODE_PLACEHOLDER =
18+
/<\s*[A-Za-z_]\w*\.([A-Za-z_]\w*)\s+from\s+([A-Za-z_]\w*)\s*>/g;
19+
20+
/**
21+
* Resolves `<Class.field from source_node>` placeholders against a workflow
22+
* scope (predecessor outputs by node name). Synchronous; unresolved placeholders
23+
* are left untouched. Mirrors Python's source-node-qualified data selection.
24+
*/
25+
function resolveSourceNodePlaceholders(
26+
template: string,
27+
scope: WorkflowInstructionScope,
28+
): string {
29+
return template.replace(SOURCE_NODE_PLACEHOLDER, (raw, field, nodeName) => {
30+
const out = scope.outputsByNode?.[nodeName];
31+
if (out && typeof out === 'object' && field in (out as object)) {
32+
return formatValue((out as Record<string, unknown>)[field], false);
33+
}
34+
return raw;
35+
});
36+
}
37+
1238
/**
1339
* Resolves a single key from the context (state or artifact).
1440
*/
@@ -39,19 +65,30 @@ async function resolveKey(
3965
}
4066

4167
// Step 3: Handle state variable injection.
42-
if (!isValidStateName(key)) {
43-
return rawMatch;
44-
}
45-
46-
if (key in invocationContext.session.state) {
47-
return formatValue(invocationContext.session.state[key], false);
68+
if (isValidStateName(key)) {
69+
if (key in invocationContext.session.state) {
70+
return formatValue(invocationContext.session.state[key], false);
71+
}
72+
if (isOptional) {
73+
return '';
74+
}
75+
throw new Error(`Context variable not found: \`${key}\`.`);
4876
}
4977

50-
if (isOptional) {
51-
return '';
78+
// Step 4: Workflow — resolve `{Class.field}` from the current node input.
79+
const scope = invocationContext.workflowInstructionScope;
80+
if (scope && WORKFLOW_FIELD_KEY.test(key)) {
81+
const field = key.slice(key.indexOf('.') + 1);
82+
const input = scope.input;
83+
if (input && typeof input === 'object' && field in (input as object)) {
84+
return formatValue((input as Record<string, unknown>)[field], false);
85+
}
86+
if (isOptional) {
87+
return '';
88+
}
5289
}
5390

54-
throw new Error(`Context variable not found: \`${key}\`.`);
91+
return rawMatch;
5592
}
5693

5794
/**
@@ -115,6 +152,14 @@ export async function injectSessionState(
115152
template: string,
116153
readonlyContext: ReadonlyContext,
117154
): Promise<string> {
155+
// Workflow: first resolve `<Class.field from source_node>` placeholders, and
156+
// enable `{Class.field}` resolution below. Both are no-ops (placeholders left
157+
// untouched) for ordinary agents, which have no workflow scope.
158+
const scope = readonlyContext.invocationContext.workflowInstructionScope;
159+
if (scope) {
160+
template = resolveSourceNodePlaceholders(template, scope);
161+
}
162+
118163
const pattern = /\{+[^{}]*}+/g;
119164
const matches = Array.from(template.matchAll(pattern));
120165

@@ -130,7 +175,10 @@ export async function injectSessionState(
130175
if (isOptional) {
131176
key = key.slice(0, -1);
132177
}
133-
const isValid = key.startsWith(ARTIFACT_PREFIX) || isValidStateName(key);
178+
const isValid =
179+
key.startsWith(ARTIFACT_PREFIX) ||
180+
isValidStateName(key) ||
181+
(!!scope && WORKFLOW_FIELD_KEY.test(key));
134182
return {
135183
raw,
136184
key,

core/src/agents/invocation_context.ts

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,17 +8,32 @@ import {Content} from '@google/genai';
88

99
import {SessionArtifactService} from '../artifacts/session_artifact_service.js';
1010
import {BaseCredentialService} from '../auth/credential_service/base_credential_service.js';
11+
import {Event} from '../events/event.js';
1112
import {BaseMemoryService} from '../memory/base_memory_service.js';
1213
import {PluginManager} from '../plugins/plugin_manager.js';
1314
import {BaseSessionService} from '../sessions/base_session_service.js';
1415
import {Session} from '../sessions/session.js';
16+
import {AsyncQueue} from '../utils/async_queue.js';
1517
import {randomUUID} from '../utils/env_aware_utils.js';
1618

1719
import {ActiveStreamingTool} from './active_streaming_tool.js';
1820
import {BaseAgent} from './base_agent.js';
1921
import {RunConfig} from './run_config.js';
2022
import {TranscriptionEntry} from './transcription_entry.js';
2123

24+
/**
25+
* Workflow: data exposed to `{Class.field}` and `<Class.field from source_node>`
26+
* instruction placeholders when an LlmAgent runs as a workflow node. Populated by
27+
* `LLMAgentWrapper`; absent for ordinary (non-workflow) agent runs, in which case
28+
* those placeholders are left untouched.
29+
*/
30+
export interface WorkflowInstructionScope {
31+
/** The current node's input, exposing fields for `{Class.field}`. */
32+
input?: unknown;
33+
/** Predecessor node outputs keyed by node name, for `<Class.field from node>`. */
34+
outputsByNode?: Record<string, unknown>;
35+
}
36+
2237
/**
2338
* The parameters for creating an invocation context.
2439
*/
@@ -38,6 +53,9 @@ export interface InvocationContextParams {
3853
activeStreamingTools?: Record<string, ActiveStreamingTool>;
3954
pluginManager: PluginManager;
4055
abortSignal?: AbortSignal;
56+
workflowInstructionScope?: WorkflowInstructionScope;
57+
/** Nesting depth of node-as-tool executions; used to bound recursion. */
58+
nodeToolDepth?: number;
4159
}
4260

4361
/**
@@ -185,6 +203,28 @@ export class InvocationContext {
185203

186204
readonly abortSignal?: AbortSignal;
187205

206+
/**
207+
* An optional channel into which a running tool can push events to be
208+
* interleaved into the agent's output stream. Set by the LLM flow around tool
209+
* execution so a {@link NodeTool} (running a node/workflow) can surface the
210+
* node's intermediate and interrupt events. Cleared once tools finish.
211+
*/
212+
eventQueue?: AsyncQueue<Event>;
213+
214+
/**
215+
* Workflow: field-resolution scope for `{Class.field}` /
216+
* `<Class.field from node>` instruction placeholders (set by
217+
* `LLMAgentWrapper`).
218+
*/
219+
workflowInstructionScope?: WorkflowInstructionScope;
220+
221+
/**
222+
* Nesting depth of node-as-tool ({@link NodeTool}) executions in this
223+
* invocation. Incremented each time a node runs as a tool (via a depth+1
224+
* clone), so `NodeTool` can bound `node -> tool -> node` recursion.
225+
*/
226+
readonly nodeToolDepth: number;
227+
188228
/**
189229
* @param params The parameters for creating an invocation context.
190230
*/
@@ -203,7 +243,10 @@ export class InvocationContext {
203243
this.activeStreamingTools = params.activeStreamingTools;
204244
this.pluginManager = params.pluginManager;
205245
this.abortSignal = params.abortSignal;
246+
this.workflowInstructionScope = params.workflowInstructionScope;
247+
this.nodeToolDepth = params.nodeToolDepth ?? 0;
206248
// Inherit the parent invocation's cost manager when one is available.
249+
207250
// Child contexts created for sub-agents, agent transfers and loop
208251
// iterations (via createInvocationContext / createBranchCtxForSubAgent)
209252
// carry the parent context's fields over, so reusing its cost manager
@@ -236,6 +279,19 @@ export class InvocationContext {
236279
incrementLlmCallCount() {
237280
this.invocationCostManager.incrementAndEnforceLlmCallsLimit(this.runConfig);
238281
}
282+
283+
/**
284+
* Returns a copy of this context with `overrides` applied. The spread carries
285+
* every own field over (including the shared cost manager), so the copy keeps
286+
* a single LLM-call counter for the invocation.
287+
*
288+
* Note: this copies own enumerable fields by value — scalar mutable fields
289+
* (e.g. `endInvocation`) are decoupled from the original, while object-valued
290+
* fields (`session`, …) stay shared by reference.
291+
*/
292+
clone(overrides: Partial<InvocationContextParams> = {}): InvocationContext {
293+
return new InvocationContext({...this, ...overrides});
294+
}
239295
}
240296

241297
export function newInvocationContextId(): string {

core/src/agents/llm_agent.ts

Lines changed: 75 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,11 @@
66

77
import {GenerateContentConfig, Schema} from '@google/genai';
88
import {context, trace} from '@opentelemetry/api';
9+
import {FinishTaskTool} from '../tools/finish_task_tool.js';
910
import {FunctionTool} from '../tools/function_tool.js';
11+
import {AsyncQueue} from '../utils/async_queue.js';
12+
import {isBaseNode, type BaseNode} from '../workflow/base_node.js';
13+
import {NodeTool} from '../workflow/nodes/node_tool.js';
1014

1115
import {z as z3} from 'zod/v3';
1216
import {z as z4} from 'zod/v4';
@@ -68,6 +72,7 @@ import {IDENTITY_LLM_REQUEST_PROCESSOR} from './processors/identity_llm_request_
6872
import {INSTRUCTIONS_LLM_REQUEST_PROCESSOR} from './processors/instructions_llm_request_processor.js';
6973
import {INTERACTIONS_REQUEST_PROCESSOR} from './processors/interactions_request_processor.js';
7074
import {REQUEST_CONFIRMATION_LLM_REQUEST_PROCESSOR} from './processors/request_confirmation_llm_request_processor.js';
75+
import {REQUEST_INPUT_LLM_REQUEST_PROCESSOR} from './processors/request_input_llm_request_processor.js';
7176
import {TOOL_FILTER_REQUEST_PROCESSOR} from './processors/tool_filter_request_processor.js';
7277
import {ReadonlyContext} from './readonly_context.js';
7378
import {StreamingMode} from './run_config.js';
@@ -193,7 +198,7 @@ export type AfterToolCallback =
193198
export type ExamplesUnion = Example[] | BaseExampleProvider;
194199

195200
/** A union of tool types that can be provided to an agent. */
196-
export type ToolUnion = BaseTool | BaseToolset;
201+
export type ToolUnion = BaseTool | BaseToolset | BaseNode;
197202

198203
const ADK_AGENT_NAME_LABEL_KEY = 'adk_agent_name';
199204

@@ -259,6 +264,16 @@ export interface LlmAgentConfig extends BaseAgentConfig {
259264
*/
260265
includeContents?: 'default' | 'none';
261266

267+
/**
268+
* The agent's execution mode when run as a workflow node.
269+
*
270+
* - `single_turn` (default): the agent runs once against the node input.
271+
* - `task`: the agent is given a `finish_task` tool and runs a multi-round
272+
* loop until it calls `finish_task`, whose arguments (conforming to
273+
* `outputSchema`) become the node output. Mirrors Python's `Agent(mode=...)`.
274+
*/
275+
mode?: 'single_turn' | 'task';
276+
262277
/** The input schema when agent is used as a tool. */
263278
inputSchema?: LlmAgentSchema;
264279

@@ -323,6 +338,11 @@ async function convertToolUnionToTools(
323338
if (isBaseTool(toolUnion)) {
324339
return [toolUnion];
325340
}
341+
if (isBaseNode(toolUnion)) {
342+
// A node/Workflow passed as a tool is auto-wrapped as a NodeTool so the
343+
// model can call it (mirrors Python's Agent(tools=[node/workflow])).
344+
return [new NodeTool(toolUnion)];
345+
}
326346
return await toolUnion.getTools(context);
327347
}
328348

@@ -362,9 +382,11 @@ export class LlmAgent extends BaseAgent<LlmAgentConfig> {
362382
disallowTransferToParent: boolean;
363383
disallowTransferToPeers: boolean;
364384
includeContents: 'default' | 'none';
385+
mode?: 'single_turn' | 'task';
365386
inputSchema?: Schema;
366387
outputSchema?: Schema;
367388
outputKey?: string;
389+
private _finishTaskTool?: FinishTaskTool;
368390
beforeModelCallback?: BeforeModelCallback;
369391
afterModelCallback?: AfterModelCallback;
370392
beforeToolCallback?: BeforeToolCallback;
@@ -389,6 +411,7 @@ export class LlmAgent extends BaseAgent<LlmAgentConfig> {
389411
this.outputSchema = isZodObject(config.outputSchema)
390412
? zodObjectToSchema(config.outputSchema)
391413
: config.outputSchema;
414+
this.mode = config.mode;
392415
this.outputKey = config.outputKey;
393416
this.beforeModelCallback = config.beforeModelCallback;
394417
this.afterModelCallback = config.afterModelCallback;
@@ -404,6 +427,7 @@ export class LlmAgent extends BaseAgent<LlmAgentConfig> {
404427
IDENTITY_LLM_REQUEST_PROCESSOR,
405428
INSTRUCTIONS_LLM_REQUEST_PROCESSOR,
406429
REQUEST_CONFIRMATION_LLM_REQUEST_PROCESSOR,
430+
REQUEST_INPUT_LLM_REQUEST_PROCESSOR,
407431
CONTENT_REQUEST_PROCESSOR,
408432
INTERACTIONS_REQUEST_PROCESSOR,
409433
CODE_EXECUTION_REQUEST_PROCESSOR,
@@ -500,6 +524,17 @@ export class LlmAgent extends BaseAgent<LlmAgentConfig> {
500524
throw new Error(`No model found for ${this.name}.`);
501525
}
502526

527+
/**
528+
* The `finish_task` tool for this agent (task mode). Lazily created and cached
529+
* so its declaration (derived from `outputSchema`) is stable across turns.
530+
*/
531+
get finishTaskTool(): FinishTaskTool {
532+
if (!this._finishTaskTool) {
533+
this._finishTaskTool = new FinishTaskTool(this.outputSchema);
534+
}
535+
return this._finishTaskTool;
536+
}
537+
503538
/**
504539
* The resolved instruction field to construct instruction for this
505540
* agent.
@@ -788,7 +823,11 @@ export class LlmAgent extends BaseAgent<LlmAgentConfig> {
788823
// TODO - b/425992518: check if tool preprocessors can be simplified.
789824
// Run pre-processors for tools.
790825
const allTools = [...this.tools];
791-
if (
826+
if (this.mode === 'task') {
827+
// Task mode: the agent completes by calling `finish_task` (whose params
828+
// mirror the output schema) rather than emitting structured output.
829+
allTools.push(this.finishTaskTool);
830+
} else if (
792831
this.outputSchema &&
793832
allTools.length > 0 &&
794833
!canUseOutputSchemaWithTools(this.canonicalModel.model)
@@ -978,13 +1017,40 @@ export class LlmAgent extends BaseAgent<LlmAgentConfig> {
9781017
// Call functions
9791018
// TODO - b/425992518: bloated funciton input, fix.
9801019
// Tool callback passed to get rid of cyclic dependency.
981-
const functionResponseEvent = await handleFunctionCallsAsync({
982-
invocationContext: invocationContext,
983-
functionCallEvent: mergedEvent,
984-
toolsDict: llmRequest.toolsDict,
985-
beforeToolCallbacks: this.canonicalBeforeToolCallbacks,
986-
afterToolCallbacks: this.canonicalAfterToolCallbacks,
987-
});
1020+
// A NodeTool (running a node/workflow) streams the node's intermediate and
1021+
// interrupt events into `invocationContext.eventQueue`; drain it concurrently
1022+
// so those events interleave into this agent's output stream. The tool runs
1023+
// in a self-contained task that captures its result/error and always closes
1024+
// the queue, so there is a single error path (no unhandled rejection).
1025+
const eventQueue = new AsyncQueue<Event>();
1026+
invocationContext.eventQueue = eventQueue;
1027+
const toolTask = (async (): Promise<{
1028+
event: Event | null;
1029+
error?: unknown;
1030+
}> => {
1031+
try {
1032+
const event = await handleFunctionCallsAsync({
1033+
invocationContext: invocationContext,
1034+
functionCallEvent: mergedEvent,
1035+
toolsDict: llmRequest.toolsDict,
1036+
beforeToolCallbacks: this.canonicalBeforeToolCallbacks,
1037+
afterToolCallbacks: this.canonicalAfterToolCallbacks,
1038+
});
1039+
return {event};
1040+
} catch (error) {
1041+
return {event: null, error};
1042+
} finally {
1043+
eventQueue.close();
1044+
}
1045+
})();
1046+
for await (const queuedEvent of eventQueue) {
1047+
yield queuedEvent;
1048+
}
1049+
const {event: functionResponseEvent, error: toolError} = await toolTask;
1050+
invocationContext.eventQueue = undefined;
1051+
if (toolError) {
1052+
throw toolError;
1053+
}
9881054

9891055
if (!functionResponseEvent || invocationContext.abortSignal?.aborted) {
9901056
return;

core/src/agents/processors/basic_llm_request_processor.ts

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,8 +41,12 @@ export class BasicLlmRequestProcessor extends BaseLlmRequestProcessor {
4141
// Models that cannot take an output schema alongside tools get the
4242
// prompt-based `set_model_response` workaround instead, injected by
4343
// `LlmAgent.runOneStepAsync` and the instructions processor.
44+
// Task-mode agents complete via the `finish_task` tool, so the JSON response
45+
// mode must not be set (function calling is incompatible with a JSON
46+
// response mime type).
4447
if (
4548
agent.outputSchema &&
49+
agent.mode !== 'task' &&
4650
(!agent.tools?.length ||
4751
canUseOutputSchemaWithTools(agent.canonicalModel.model))
4852
) {

0 commit comments

Comments
 (0)