diff --git a/core/src/workflow/node_context.ts b/core/src/workflow/node_context.ts index 00670b20e..d2d56166d 100644 --- a/core/src/workflow/node_context.ts +++ b/core/src/workflow/node_context.ts @@ -117,9 +117,10 @@ export class NodeContext { 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._state = new NodeStateView( + invocationOverlay(opts.invocationContext), this.actions.stateDelta, + opts.invocationContext.session.state, ); } @@ -182,3 +183,112 @@ export class NodeContext { return executeChildNode({parent: this, node, input, options}); } } + +/** + * Per-invocation record of every `ctx.state` write made by the nodes of one + * invocation, keyed by the session's live state object. + * + * `session.state` object identity is stable for the duration of a turn (the + * session services mutate it in place, and hand out a fresh object per + * `getSession`), so the WeakMap entry naturally scopes to one turn. + * + * The `invocationId` guard covers *sequential* reuse of one state object: a + * later invocation replaces the entry rather than inheriting the earlier one's + * writes. It does not let two *live* invocations share a state object — they + * would evict each other's entry and fall back to reading committed state. + * Nothing does that today: `getSession` hands out a fresh state object per + * turn, and a sub-invocation (`AgentTool`) re-fetches the session, so it keys a + * different entry. + */ +const invocationOverlays = new WeakMap< + Record, + {invocationId: string; values: Record} +>(); + +/** Returns the write overlay for `ic`'s invocation, creating it on first use. */ +function invocationOverlay(ic: InvocationContext): Record { + const sessionState = ic.session.state; + const existing = invocationOverlays.get(sessionState); + if (existing && existing.invocationId === ic.invocationId) { + return existing.values; + } + const created = {invocationId: ic.invocationId, values: {}}; + invocationOverlays.set(sessionState, created); + return created.values; +} + +/** + * The state view a workflow node sees: its own pending delta, then the + * invocation's write overlay, then the session's committed state. + * + * The overlay exists to keep node-to-node reads honest. `session.state` is + * mutated from two directions during a run: nodes write through it + * immediately, while the runner separately re-applies each event's + * `actions.stateDelta` as it commits that event — and that commit lags node + * execution. Re-applying an earlier node's delta therefore rolls back a later + * node's write, and any node reading in that window observes the stale value: + * + * a: set('attempts', 0) + * b: get -> 0, set('attempts', 1) + * commit(a) re-applies attempts=0 <- rolls back b's write + * c: get -> 0 <- wrong; b already set 1 + * commit(b) re-applies attempts=1 <- rolls forward, too late + * + * Reads are served from the overlay, which only ever moves forward, so `c` + * sees `1`. Writes still land in `session.state` as well, so consumers that + * read it directly — notably `{key}` instruction templating in an agent node — + * behave exactly as before. + * + * Two gaps remain. They are mirror images of each other, share the root cause + * above, and are both only really fixable at the source (version-stamping keys + * as the session services apply a delta), so neither is closed here: + * + * - Workflow writes still reach an outside reader out of order. `session.state` + * is written through *and* re-applied by the event commit, so an agent + * resolving `{key}` for a key two nodes wrote can observe the stale value + * inside the same window. + * - Outside writes no longer reach a workflow reader. Once a node writes `k`, + * every later read of `k` in this invocation is served from the overlay, so a + * tool or callback that writes `k` straight to `session.state` mid-run is + * invisible to those reads for the rest of the invocation. Before the overlay + * they saw it, subject to the rollback race being fixed. + */ +class NodeStateView extends State { + constructor( + overlay: Record, + delta: Record, + private readonly committed: Record, + ) { + super(overlay, delta); + } + + override get(key: string, defaultValue?: T): T | undefined { + if (super.has(key)) { + return super.get(key, defaultValue); + } + return key in this.committed ? (this.committed[key] as T) : defaultValue; + } + + override has(key: string): boolean { + return super.has(key) || key in this.committed; + } + + // Both of `State`'s write paths (`set` and `update`) keep writing through to + // the session, so readers of `session.state` (agent instruction templates, + // callbacks, tools) see the write immediately, as they did before the overlay + // existed. + + override set(key: string, value: unknown): void { + super.set(key, value); + this.committed[key] = value; + } + + override update(delta: Record): void { + super.update(delta); + Object.assign(this.committed, delta); + } + + override toRecord(): Record { + return {...this.committed, ...super.toRecord()}; + } +} diff --git a/core/src/workflow/nodes/llm_agent_wrapper.ts b/core/src/workflow/nodes/llm_agent_wrapper.ts index fadf5c7b4..eb51fe991 100644 --- a/core/src/workflow/nodes/llm_agent_wrapper.ts +++ b/core/src/workflow/nodes/llm_agent_wrapper.ts @@ -138,7 +138,14 @@ export class LLMAgentWrapper extends BaseNode { event.output = output; event.nodeInfo = {...(event.nodeInfo ?? {}), messageAsOutput: true}; if (agent.outputKey && output !== undefined) { - ctx.actions.stateDelta[agent.outputKey] = output; + // Route the write through `ctx.state` so it lands in the invocation + // overlay and on `session.state` right away — a downstream node + // reading `outputKey` runs before this event is committed, and used + // to see `undefined`. Also stamp it on the event being emitted: + // unlike `FunctionNode`, this node never drains `ctx.actions` onto an + // event, so the event delta is what actually gets committed. + ctx.state.set(agent.outputKey, output); + event.actions.stateDelta[agent.outputKey] = output; } yield event; return; diff --git a/core/test/workflow/state_consistency_test.ts b/core/test/workflow/state_consistency_test.ts new file mode 100644 index 000000000..7f3accf24 --- /dev/null +++ b/core/test/workflow/state_consistency_test.ts @@ -0,0 +1,292 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {describe, expect, it} from 'vitest'; +import {LlmAgent} from '../../src/agents/llm_agent.js'; +import {Event} from '../../src/events/event.js'; +import {BaseLlm} from '../../src/models/base_llm.js'; +import {BaseLlmConnection} from '../../src/models/base_llm_connection.js'; +import {LlmRequest} from '../../src/models/llm_request.js'; +import {LlmResponse} from '../../src/models/llm_response.js'; +import {Runner} from '../../src/runner/runner.js'; +import {InMemorySessionService} from '../../src/sessions/in_memory_session_service.js'; +import {AsyncQueue} from '../../src/utils/async_queue.js'; +import {NodeContext} from '../../src/workflow/node_context.js'; +import {FunctionNode} from '../../src/workflow/nodes/function_node.js'; +import {WorkflowAgent} from '../../src/workflow/workflow_agent.js'; +import {createIc} from './test_helpers.js'; + +/** A model that replays canned responses, one per call. */ +class MockLlm extends BaseLlm { + private callCount = 0; + + constructor(private readonly responses: LlmResponse[]) { + super({model: 'mock-llm'}); + } + + async *generateContentAsync( + _request: LlmRequest, + ): AsyncGenerator { + const response = this.responses[this.callCount++]; + if (response) { + yield response; + } + } + + async connect(_request: LlmRequest): Promise { + throw new Error('not implemented'); + } +} + +async function drain(gen: AsyncGenerator): Promise { + const out: Event[] = []; + for await (const event of gen) { + out.push(event); + } + return out; +} + +async function runOnce(agent: WorkflowAgent, text = 'x') { + 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 = await drain( + runner.runAsync({ + userId: 'u1', + sessionId: session.id, + newMessage: {role: 'user', parts: [{text}]}, + }), + ); + const finalSession = await sessionService.getSession({ + appName: 'test_app', + userId: 'u1', + sessionId: session.id, + }); + return {events, state: finalSession?.state ?? {}}; +} + +describe('workflow state consistency across nodes', () => { + it('a node observes the most recent write to a key, not an earlier one', async () => { + // Regression: the runner's event commit lags node execution and re-applies + // each event's stateDelta to the live session state. Re-applying `a`'s + // delta used to roll back `b`'s write, so `c` read 0 instead of 1. + const reads: Array = []; + + const a = new FunctionNode('a', (ctx: NodeContext) => { + ctx.state.set('attempts', 0); + return 'a'; + }); + const b = new FunctionNode('b', (ctx: NodeContext) => { + const seen = ctx.state.get('attempts') ?? -1; + ctx.state.set('attempts', seen + 1); + return 'b'; + }); + const c = new FunctionNode('c', (ctx: NodeContext) => { + reads.push(ctx.state.get('attempts')); + return 'c'; + }); + + const {state} = await runOnce( + new WorkflowAgent({name: 'state_wf', edges: [['START', a, b, c]]}), + ); + + expect(reads).toEqual([1]); + // The committed session state still ends up correct. + expect(state['attempts']).toBe(1); + }); + + it('survives a long read-modify-write chain', async () => { + const reads: number[] = []; + const bump = (name: string) => + new FunctionNode(name, (ctx: NodeContext) => { + const next = (ctx.state.get('n') ?? 0) + 1; + ctx.state.set('n', next); + reads.push(next); + return name; + }); + + const nodes = ['n1', 'n2', 'n3', 'n4', 'n5'].map(bump); + const {state} = await runOnce( + new WorkflowAgent({name: 'chain_wf', edges: [['START', ...nodes]]}), + ); + + expect(reads).toEqual([1, 2, 3, 4, 5]); + expect(state['n']).toBe(5); + }); + + it('still reads state that was seeded on the session before the run', async () => { + const seen: Array = []; + const read = new FunctionNode('read', (ctx: NodeContext) => { + seen.push(ctx.state.get('seeded')); + return 'read'; + }); + + const sessionService = new InMemorySessionService(); + const session = await sessionService.createSession({ + appName: 'test_app', + userId: 'u1', + state: {seeded: 'from-before'}, + }); + const agent = new WorkflowAgent({ + name: 'seeded_wf', + edges: [['START', read]], + }); + const runner = new Runner({appName: 'test_app', agent, sessionService}); + await drain( + runner.runAsync({ + userId: 'u1', + sessionId: session.id, + newMessage: {role: 'user', parts: [{text: 'x'}]}, + }), + ); + + expect(seen).toEqual(['from-before']); + }); + + it('writes remain visible on session.state for instruction templating', async () => { + // Agent instruction templates resolve `{key}` against + // `invocationContext.session.state`, so node writes must still land there + // synchronously — the overlay must not divert them. + let observed: unknown; + const write = new FunctionNode('write', (ctx: NodeContext) => { + ctx.state.set('topic', 'oceans'); + return 'write'; + }); + const peek = new FunctionNode('peek', (ctx: NodeContext) => { + observed = ctx.invocationContext.session.state['topic']; + return 'peek'; + }); + + await runOnce( + new WorkflowAgent({name: 'tmpl_wf', edges: [['START', write, peek]]}), + ); + + expect(observed).toBe('oceans'); + }); + + it('writes made through update() are visible the same way as set()', async () => { + // `update` is State's other write path and is public on `ctx.state`, so it + // has to honour both halves of the contract: later nodes read it back, and + // it lands on `session.state` for direct readers. + let readBack: unknown; + let onSession: unknown; + const write = new FunctionNode('write', (ctx: NodeContext) => { + ctx.state.update({topic: 'oceans'}); + return 'write'; + }); + const peek = new FunctionNode('peek', (ctx: NodeContext) => { + readBack = ctx.state.get('topic'); + onSession = ctx.invocationContext.session.state['topic']; + return 'peek'; + }); + + const {state} = await runOnce( + new WorkflowAgent({name: 'update_wf', edges: [['START', write, peek]]}), + ); + + expect(readBack).toBe('oceans'); + expect(onSession).toBe('oceans'); + expect(state['topic']).toBe('oceans'); + }); + + it('an agent node\u2019s outputKey is visible to the next node', async () => { + // `outputKey` is the wrapper's own write path into state, so it has to + // honour the same contract as `ctx.state`: readable by a later node, on + // `session.state` for direct readers, and committed. + let readBack: unknown; + let onSession: unknown; + const writer = new LlmAgent({ + name: 'writer', + model: new MockLlm([ + { + content: { + role: 'model', + parts: [ + { + functionCall: { + id: 'fc_1', + name: 'finish_task', + args: {result: 'done'}, + }, + }, + ], + }, + }, + ]), + mode: 'task', + outputKey: 'summary', + }); + const peek = new FunctionNode('peek', (ctx: NodeContext) => { + readBack = ctx.state.get('summary'); + onSession = ctx.invocationContext.session.state['summary']; + return 'peek'; + }); + + const {state} = await runOnce( + new WorkflowAgent({ + name: 'output_key_wf', + edges: [['START', writer, peek]], + }), + ); + + expect(readBack).toEqual({result: 'done'}); + expect(onSession).toEqual({result: 'done'}); + expect(state['summary']).toEqual({result: 'done'}); + }); + + it('serves a second invocation over the same session from committed state', () => { + // The overlay is keyed by the session's live state object, so sequential + // invocations that reuse one session object must not inherit each other's + // writes — that is what the invocation-id guard is for. + const ic1 = createIc(); + const channel = new AsyncQueue(); + const mkCtx = (ic = ic1) => + new NodeContext({ + invocationContext: ic, + channel, + nodePath: '', + runId: 'root', + }); + + mkCtx().state.set('k', 'from-inv-1'); + // Stand in for the runner re-applying a stale delta to the live session + // state while the invocation is still running. + ic1.session.state['k'] = 'stale'; + + // Same invocation: the overlay wins, which is the whole point of the fix. + expect(mkCtx().state.get('k')).toBe('from-inv-1'); + + // Next invocation on the same session object: fresh overlay, so the read + // falls through to committed state rather than the previous run's write. + const ic2 = ic1.clone({invocationId: 'inv-2'}); + expect(ic2.session.state).toBe(ic1.session.state); + expect(mkCtx(ic2).state.get('k')).toBe('stale'); + }); + + it('does not leak one invocation\u2019s overlay into another session', async () => { + // The overlay is keyed by the session's live state object and guarded by + // invocation id, so a second, unrelated run starts from a clean slate. + const seen: Array = []; + const bump = new FunctionNode('bump', (ctx: NodeContext) => { + const next = (ctx.state.get('count') ?? 0) + 1; + ctx.state.set('count', next); + seen.push(next); + return next; + }); + const agent = new WorkflowAgent({ + name: 'isolated_wf', + edges: [['START', bump]], + }); + + await runOnce(agent, 'one'); + await runOnce(agent, 'two'); + + expect(seen).toEqual([1, 1]); + }); +});