From 4c7360460255981a727139ce11180e08e2725d2c Mon Sep 17 00:00:00 2001 From: Mohammad Mseet Date: Sun, 9 Aug 2026 11:28:41 +0200 Subject: [PATCH] fix(workflow): scope node rehydration to the current invocation Workflow rehydration matched prior session events by node path alone, so a completed node from one invocation was fast-forwarded into a later, unrelated invocation on the same session. A workflow that gates protected work behind a node (an authorization or credential check, as in auth_gate_test) would skip that node on the next message and continue with the earlier decision. The Python implementation this module is ported from does not have this: it threads the invocation id into reconstruction and skips events belonging to another invocation (workflow/utils/_rehydration_utils.py), passing ic.invocation_id from the scheduler, node runner and replay manager. The port kept the path matching but dropped the invocation scoping, and three pieces were missing to make it work: 1. Interrupt events were not stamped with an invocation id. createRequestInputEvent and createAuthRequestEvent were the only workflow event constructors not setting invocationId, so the events that carry longRunningToolIds could not be attributed to an invocation at all. Every other workflow event site already sets it. 2. A resume did not continue the invocation it resumed. Runner always minted a new invocation id, so rehydration could not have been scoped without breaking resume. Python resolves the id from the resume message before building the context; resolveResumedInvocationId does the same, matching a function response to the interrupt that raised it and otherwise continuing an invocation that still has an unanswered interrupt so plain-text replies keep working. 3. Reconstruction did not filter. reconstructNodeStates and reconstructNodeStatesByPath now take an optional invocation id and skip events from other invocations, mirroring the Python filter including its "no id means no filtering" utility behaviour. The three production call sites pass the current invocation. Two unit tests seeded a prior event without an invocation id and relied on it being visible to a different invocation. They now seed the same invocation, so they still cover fast-forward and plain-text resume, which is what they are about. Adds invocation_scoping_test.ts: a node completed in one invocation is not reused by another, its own invocation still sees it, the unscoped utility form is unchanged, a second message re-runs a gate rather than inheriting its result, and the resume-id resolver adopts, continues or declines an invocation. --- core/src/runner/runner.ts | 63 +++++- core/src/workflow/base_node.ts | 2 +- core/src/workflow/dynamic_node_scheduler.ts | 7 +- core/src/workflow/nodes/function_node.ts | 6 +- core/src/workflow/utils/hitl_utils.ts | 8 +- core/src/workflow/utils/rehydration_utils.ts | 33 +++- core/src/workflow/workflow.ts | 1 + core/src/workflow/workflow_agent.ts | 6 +- core/test/workflow/invocation_scoping_test.ts | 185 ++++++++++++++++++ core/test/workflow/workflow_advanced_test.ts | 3 + core/test/workflow/workflow_agent_test.ts | 2 + 11 files changed, 301 insertions(+), 15 deletions(-) create mode 100644 core/test/workflow/invocation_scoping_test.ts diff --git a/core/src/runner/runner.ts b/core/src/runner/runner.ts index ba43f2a93..4966bb317 100644 --- a/core/src/runner/runner.ts +++ b/core/src/runner/runner.ts @@ -296,7 +296,9 @@ export class Runner { sessionService: this.sessionService, memoryService: this.memoryService, credentialService: this.credentialService, - invocationId: newInvocationContextId(), + invocationId: + resolveResumedInvocationId(session.events, newMessage) ?? + newInvocationContextId(), agent: this.agent, session, userContent: newMessage, @@ -663,6 +665,65 @@ export function isRoutableLlmAgent(agentToRun: BaseAgent): boolean { return true; } +/** + * Resolves the invocation that `newMessage` resumes, if any. + * + * Mirrors `google/adk-python` `runners.py`, which resolves the invocation id + * from a resume message before building the invocation context instead of + * always minting a new one. Workflow node rehydration is scoped by invocation + * id, so this is what lets a genuine resume see the nodes it already ran while + * an unrelated new message in the same session cannot. + * + * Returns `undefined` when the message does not continue anything, in which + * case the caller mints a fresh invocation id as before. + */ +export function resolveResumedInvocationId( + events: Event[], + newMessage?: Content, +): string | undefined { + // 1. The message explicitly answers a pending interrupt. + const responseIds = new Set(); + for (const part of newMessage?.parts ?? []) { + const id = part.functionResponse?.id; + if (id) { + responseIds.add(id); + } + } + if (responseIds.size > 0) { + for (let i = events.length - 1; i >= 0; i--) { + const event = events[i]; + for (const id of event.longRunningToolIds ?? []) { + if (responseIds.has(id) && event.invocationId) { + return event.invocationId; + } + } + } + } + + // 2. An invocation is still waiting on an unresolved interrupt, so a reply + // that carries no function response (a plain-text answer to a single + // pending request) still continues it. + const answered = new Set(); + for (const event of events) { + for (const part of event.content?.parts ?? []) { + const id = part.functionResponse?.id; + if (id) { + answered.add(id); + } + } + } + for (let i = events.length - 1; i >= 0; i--) { + const event = events[i]; + for (const id of event.longRunningToolIds ?? []) { + if (!answered.has(id) && event.invocationId) { + return event.invocationId; + } + } + } + + return undefined; +} + /** * It iterates through the events in reverse order, and returns the event * containing a function call with a functionCall.id matching the diff --git a/core/src/workflow/base_node.ts b/core/src/workflow/base_node.ts index 119e40014..d98eb5c35 100644 --- a/core/src/workflow/base_node.ts +++ b/core/src/workflow/base_node.ts @@ -150,7 +150,7 @@ export abstract class BaseNode { 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); + yield createRequestInputEvent(item, ctx.invocationContext.invocationId); continue; } const event = this.toEvent(ctx, item); diff --git a/core/src/workflow/dynamic_node_scheduler.ts b/core/src/workflow/dynamic_node_scheduler.ts index 7cfbe1ec0..d6ae0018a 100644 --- a/core/src/workflow/dynamic_node_scheduler.ts +++ b/core/src/workflow/dynamic_node_scheduler.ts @@ -62,9 +62,10 @@ export class DynamicNodeScheduler implements ScheduleDynamicNode { // 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, - ); + const prior = reconstructNodeStatesByPath( + ctx.session?.events ?? [], + ctx.invocationId, + ).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, { diff --git a/core/src/workflow/nodes/function_node.ts b/core/src/workflow/nodes/function_node.ts index d4df8957e..41c72830c 100644 --- a/core/src/workflow/nodes/function_node.ts +++ b/core/src/workflow/nodes/function_node.ts @@ -148,7 +148,11 @@ export class FunctionNode extends BaseNode< } // The credential key doubles as a deterministic interrupt id so the resume // response matches across turns. - return createAuthRequestEvent(authConfig, authConfig.credentialKey); + return createAuthRequestEvent( + authConfig, + authConfig.credentialKey, + ctx.invocationId, + ); } /** diff --git a/core/src/workflow/utils/hitl_utils.ts b/core/src/workflow/utils/hitl_utils.ts index f87d5a7d6..2af9a481a 100644 --- a/core/src/workflow/utils/hitl_utils.ts +++ b/core/src/workflow/utils/hitl_utils.ts @@ -34,7 +34,10 @@ export const REQUEST_CREDENTIAL_FUNCTION_CALL_NAME = 'adk_request_credential'; * carries an `adk_request_input` function call and marks the interrupt id as a * long-running tool id. */ -export function createRequestInputEvent(requestInput: RequestInput): Event { +export function createRequestInputEvent( + requestInput: RequestInput, + invocationId?: string, +): Event { const args: Record = { interruptId: requestInput.interruptId, payload: requestInput.payload ?? null, @@ -58,6 +61,7 @@ export function createRequestInputEvent(requestInput: RequestInput): Event { ], }, longRunningToolIds: [requestInput.interruptId], + invocationId, }); } @@ -125,6 +129,7 @@ export function hasAuthCredential( export function createAuthRequestEvent( authConfig: AuthConfig, interruptId: string, + invocationId?: string, ): Event { const authRequest = new AuthHandler(authConfig).generateAuthRequest(); const args: Record = { @@ -146,6 +151,7 @@ export function createAuthRequestEvent( ], }, longRunningToolIds: [interruptId], + invocationId, }); } diff --git a/core/src/workflow/utils/rehydration_utils.ts b/core/src/workflow/utils/rehydration_utils.ts index 3b3bf87c0..b2167624a 100644 --- a/core/src/workflow/utils/rehydration_utils.ts +++ b/core/src/workflow/utils/rehydration_utils.ts @@ -47,16 +47,25 @@ export interface RehydratedNode { export function reconstructNodeStates( events: Event[], parentPath?: string, + invocationId?: string, ): Map { if (parentPath) { - return reconstruct(events, (event) => - event.nodeInfo?.path - ? directChildName(event.nodeInfo.path, parentPath) - : undefined, + return reconstruct( + events, + (event) => + event.nodeInfo?.path + ? directChildName(event.nodeInfo.path, parentPath) + : undefined, + invocationId, ); } - return reconstruct(events, (event) => - event.nodeInfo?.path ? nodeNameFromPath(event.nodeInfo.path) : event.author, + return reconstruct( + events, + (event) => + event.nodeInfo?.path + ? nodeNameFromPath(event.nodeInfo.path) + : event.author, + invocationId, ); } @@ -67,14 +76,20 @@ export function reconstructNodeStates( */ export function reconstructNodeStatesByPath( events: Event[], + invocationId?: string, ): Map { - return reconstruct(events, (event) => event.nodeInfo?.path ?? event.author); + return reconstruct( + events, + (event) => event.nodeInfo?.path ?? event.author, + invocationId, + ); } /** Shared scan that groups node events by the key returned by `keyFor`. */ function reconstruct( events: Event[], keyFor: (event: Event) => string | undefined, + invocationId?: string, ): Map { const nodes = new Map(); const interruptOwner = new Map(); @@ -89,6 +104,10 @@ function reconstruct( }; for (const event of events) { + if (invocationId && event.invocationId !== invocationId) { + continue; + } + // 1. User function responses resolving prior interrupts. if (event.author === 'user' && event.content?.parts) { for (const part of event.content.parts) { diff --git a/core/src/workflow/workflow.ts b/core/src/workflow/workflow.ts index e3152916f..2bc7aadd0 100644 --- a/core/src/workflow/workflow.ts +++ b/core/src/workflow/workflow.ts @@ -192,6 +192,7 @@ export class Workflow extends BaseNode { const rehydrated = reconstructNodeStates( ctx.session?.events ?? [], ctx.nodePath || undefined, + ctx.invocationId, ); this.applyResumeInputs(ctx, rehydrated); diff --git a/core/src/workflow/workflow_agent.ts b/core/src/workflow/workflow_agent.ts index 9c75c4ac0..cb6951ff0 100644 --- a/core/src/workflow/workflow_agent.ts +++ b/core/src/workflow/workflow_agent.ts @@ -153,7 +153,11 @@ function resumeInputsFromPlainText( const text = parts.map((p) => p.text).join(''); const pending = new Set(); - for (const node of reconstructNodeStates(ic.session?.events ?? []).values()) { + for (const node of reconstructNodeStates( + ic.session?.events ?? [], + undefined, + ic.invocationId, + ).values()) { for (const id of node.interruptIds) { if (!node.resolvedResponses.has(id)) { pending.add(id); diff --git a/core/test/workflow/invocation_scoping_test.ts b/core/test/workflow/invocation_scoping_test.ts new file mode 100644 index 000000000..b15a1755d --- /dev/null +++ b/core/test/workflow/invocation_scoping_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 {resolveResumedInvocationId, 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 {createRequestInputEvent} from '../../src/workflow/utils/hitl_utils.js'; +import {reconstructNodeStatesByPath} from '../../src/workflow/utils/rehydration_utils.js'; +import {Workflow} from '../../src/workflow/workflow.js'; +import {WorkflowAgent} from '../../src/workflow/workflow_agent.js'; + +describe('rehydration — invocation scoping', () => { + it('does not reuse a completed node from a different invocation', () => { + const events = [ + createEvent({ + author: 'authorize', + invocationId: 'inv-a', + nodeInfo: {path: 'wf.authorize@1'}, + output: {authorized: true}, + }), + ]; + + // Same session, a later invocation: the earlier node must not be visible. + expect( + reconstructNodeStatesByPath(events, 'inv-b').get('wf.authorize@1'), + ).toBeUndefined(); + + // Its own invocation still sees it, which is what resume relies on. + expect( + reconstructNodeStatesByPath(events, 'inv-a').get('wf.authorize@1') + ?.output, + ).toEqual({authorized: true}); + + // Omitting the invocation id keeps the unscoped utility behaviour. + expect( + reconstructNodeStatesByPath(events).get('wf.authorize@1')?.output, + ).toEqual({authorized: true}); + }); + + it('re-runs a node in a new invocation on the same session', async () => { + let authorized = true; + let authChecks = 0; + let sinkRuns = 0; + + const authorize = node( + (_ctx: NodeContext, _input: unknown) => { + authChecks++; + if (!authorized) { + throw new Error('NOT_AUTHORIZED'); + } + return {authorized: true}; + }, + {name: 'authorize'}, + ); + + const workflow = new Workflow({ + name: 'gated_workflow', + dynamicEntry: async (ctx: NodeContext, input: unknown) => { + await ctx.runNode(authorize, input); + sinkRuns++; + return 'done'; + }, + }); + + const sessions = new InMemorySessionService(); + const runner = new Runner({ + appName: 'app', + agent: new WorkflowAgent(workflow), + sessionService: sessions, + }); + const session = await sessions.createSession({appName: 'app', userId: 'u'}); + + const runTurn = async (text: string) => { + const events: Event[] = []; + for await (const event of runner.runAsync({ + userId: 'u', + sessionId: session.id, + newMessage: {role: 'user', parts: [{text}]}, + })) { + events.push(event); + } + return events; + }; + + await runTurn('first'); + expect(authChecks).toBe(1); + expect(sinkRuns).toBe(1); + + // The caller keeps the session but loses the entitlement. A second message + // is a new invocation, so the gate must run again rather than being + // fast-forwarded from the first invocation's cached output. + authorized = false; + await expect(runTurn('second')).rejects.toThrow('NOT_AUTHORIZED'); + expect(authChecks).toBe(2); + expect(sinkRuns).toBe(1); + }); +}); + +describe('runner — resumed invocation id', () => { + it('adopts the invocation that raised the interrupt being answered', () => { + const events = [ + createRequestInputEvent( + new RequestInput({interruptId: 'gate-1', message: '?'}), + 'inv-a', + ), + ]; + const resumed = resolveResumedInvocationId(events, { + role: 'user', + parts: [ + { + functionResponse: { + id: 'gate-1', + name: 'adk_request_input', + response: {}, + }, + }, + ], + }); + expect(resumed).toBe('inv-a'); + }); + + it('adopts the invocation still waiting on an unanswered interrupt', () => { + const events = [ + createRequestInputEvent( + new RequestInput({interruptId: 'gate-1', message: '?'}), + 'inv-a', + ), + ]; + // A plain-text reply carries no function response, but there is exactly one + // pending interrupt to continue. + expect( + resolveResumedInvocationId(events, { + role: 'user', + parts: [{text: 'yes'}], + }), + ).toBe('inv-a'); + }); + + it('starts a new invocation when nothing is pending', () => { + const events = [ + createEvent({author: 'a', invocationId: 'inv-a', output: 'done'}), + ]; + expect( + resolveResumedInvocationId(events, {role: 'user', parts: [{text: 'hi'}]}), + ).toBeUndefined(); + }); + + it('starts a new invocation once every interrupt has been answered', () => { + const events = [ + createRequestInputEvent( + new RequestInput({interruptId: 'gate-1', message: '?'}), + 'inv-a', + ), + createEvent({ + author: 'user', + invocationId: 'inv-a', + content: { + role: 'user', + parts: [ + { + functionResponse: { + id: 'gate-1', + name: 'adk_request_input', + response: {result: 'ok'}, + }, + }, + ], + }, + }), + ]; + expect( + resolveResumedInvocationId(events, { + role: 'user', + parts: [{text: 'next'}], + }), + ).toBeUndefined(); + }); +}); diff --git a/core/test/workflow/workflow_advanced_test.ts b/core/test/workflow/workflow_advanced_test.ts index 1ceb1f7ce..fb3fa0b97 100644 --- a/core/test/workflow/workflow_advanced_test.ts +++ b/core/test/workflow/workflow_advanced_test.ts @@ -228,6 +228,9 @@ describe('workflow — rerunOnResume', () => { author: 'once', nodeInfo: {path: 'ff.once'}, output: 'cached', + // Same invocation as createIc(): a resume of this invocation, which is + // what fast-forward is for. A different invocation must not be reused. + invocationId: 'inv-1', }); const ic = createIc(); ic.session.events.push(priorEvent); diff --git a/core/test/workflow/workflow_agent_test.ts b/core/test/workflow/workflow_agent_test.ts index 83fed4ad8..d136638c5 100644 --- a/core/test/workflow/workflow_agent_test.ts +++ b/core/test/workflow/workflow_agent_test.ts @@ -21,6 +21,8 @@ import {WorkflowAgent} from '../../src/workflow/workflow_agent.js'; function pendingInterruptEvent(id: string): Event { const event = createRequestInputEvent( new RequestInput({interruptId: id, message: '?'}), + // Same invocation as createIc(), i.e. an interrupt this turn resumes. + 'inv-1', ); event.author = id; return event;