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;