fix(workflow): stop node state writes from being rolled back mid-run

[kalenkevich]

Summary

A workflow node that reads a session-state key an earlier node also wrote can observe the earlier, already-superseded value:

a: ctx.state.set('attempts', 0);
b: ctx.state.get('attempts'); // 0
ctx.state.set('attempts', 1);
c: ctx.state.get('attempts'); // -> 0   WRONG, expected 1

Root cause

session.state is written from two directions during a run:

1. NodeContext builds its State directly over the live session.state (node_context.ts:120), and State.set writes through immediately (state.ts:49-52).
2. The runner separately re-applies each event's actions.stateDelta onto that same object as it commits the event (base_session_service.ts:195-209).

The event drain lags node execution, so an earlier node's delta gets re-applied after a later node has already written a newer value. Traced with an appendEvent probe:

a: set 0
b: read 0, set 1
  >>> appendEvent(author=a) re-applies attempts=0   <- rolls back b's write
c: read 0                                            <- reads the rolled-back value
  >>> appendEvent(author=b) re-applies attempts=1   <- rolls forward, too late

The value converges once the backlog drains — the same read returns 1 after a 300 ms delay — so committed state was already correct. Only in-run reads were wrong, which made this silent, timing-dependent, and invisible unless a key has more than one writer. That is exactly the counter pattern the data-handling docs demonstrate (attempts + 1), so the documented example does not work as written.

Fix

Node reads are served from a per-invocation write overlay that only ever moves forward; reads fall through to committed session state for anything this run has not written. The overlay is keyed by the session's live state object (stable within a turn) and guarded by invocation id.

Writes still land in session.state as well, so consumers that read it directly are unaffected — notably {key} instruction templating in an agent node, which resolves against invocationContext.session.state. There is a test pinning that.

Contained to the workflow package; State and the session services are untouched, and committed state is byte-for-byte what it was before.

A second, older bug found in review

An agent node's outputKey bypassed all of this: task mode wrote it straight to ctx.actions.stateDelta (llm_agent_wrapper.ts:141). Only FunctionNode ever drains that object onto an event (function_node.ts:182); BaseNode.toEvent does not, and nothing else merges it — so for this node it is write-only. The value reached neither the overlay, nor session.state, nor committed state. A downstream node read undefined, and the key never persisted at all. Pre-dates this PR, and outputKey had no workflow test.

It now writes through ctx.state (overlay + session, for in-run readers) and stamps the delta on the event being emitted (for the commit). Both halves are load-bearing: with only the first, committed state stays undefined; with only the second, the next node still reads undefined.

Known residual

Two gaps remain, and they are mirror images of each other. Both have the same root cause and the same real fix — version-stamping keys in updateSessionState, which touches every session service — so both are deliberately left out of this PR.

• Workflow writes still reach an outside reader out of order. Agent instruction templating reads session.state directly, so an agent resolving {key} for a key written by two different nodes is still exposed to the same window. Same for LlmAgent's own outputKey (llm_agent.ts:708), which writes the event delta and so is committed rather than dropped — I probed a single_turn agent node followed by a reader node and the reader does see the value, because the runner commits the agent's event before the next node starts. Exposed to the window, not broken.
• Outside writes no longer reach a workflow reader. Once a node writes k, later reads of k in that invocation are served from the overlay, so a tool or callback writing k straight to session.state mid-run is invisible to them for the rest of the invocation. This one is a behaviour change, not just an untouched gap: before the overlay those reads did see the write, subject to the race being fixed here. The trade looks right — a mid-run outside writer to a key the workflow also owns is already in undefined-ordering territory, whereas node-to-node read-modify-write is the documented pattern — but it is a trade. Both are now written down in the NodeStateView doc comment.

Happy to follow up with the version-stamping fix if you want it.

Tests

New core/test/workflow/state_consistency_test.ts:

• the read-modify-write regression above
• a 5-node chain reading [1,2,3,4,5]
• state seeded on the session before the run is still readable
• node writes still visible on session.state (instruction-templating contract)
• ctx.state.update() honours both halves of that contract (added in review)
• an agent node's outputKey is readable by the next node, on session.state, and committed (added in review; first workflow coverage of outputKey)
• a second invocation over the same session object is served from committed state, not the previous invocation's overlay — the invocation-id guard, which nothing exercised before (added in review)
• the overlay does not leak across runs

unit:core green: 206 files, 2833 tests. unit:dev has 1 failure (cli_create_test) that reproduces identically on unmodified main.

Unrelated issue found while testing

reconstructNodeStates scans all session events with no invocation filter (workflow.ts:192). A workflow that completed in turn 1 therefore has every node fast-forwarded from cache on turn 2 — re-invoking it in the same session is a no-op. Reproduces on clean main; not touched here. Already covered by #637, which fixes it (delimiting runs by pausing rather than by invocation id) — so no separate issue needed unless you want one for tracking.

[AmaadMartin]

The diagnosis is convincing and the mechanism works. I checked the load-bearing part: State.set writes to both value and delta (state.ts:49-52), so super.set populates the shared overlay and a later node's super.has finds an earlier node's write. toRecord composes in the right order.

I also checked the case your comment calls out — two invocations sharing a session object, where the guard evicts rather than coexists. It is not reachable: AgentTool builds a new Runner and re-fetches the session, and getSession hands back a fresh state object, so a sub-invocation keys a different WeakMap entry.

Two things inline, one of them one line. On the unrelated issue you found — the missing invocation filter in reconstructNodeStates — that is #637, already open.

[AmaadMartin]

Not a nit. State has two write paths and this overrides one of them.

override set(key: string, value: unknown): void {
  super.set(key, value);
  this.committed[key] = value;
}

update() is the other one (state.ts:73), and it is public on ctx.state:

update(delta: Record<string, unknown>) {
  Object.assign(this.delta, delta);
  Object.assign(this.value, delta);   // overlay only
}

So ctx.state.update({...}) lands in the overlay and the node delta but never in session.state — the write-through your set comment protects, and which "writes remain visible on session.state for instruction templating" tests. That test passes because it uses set.

Nothing calls update on a NodeContext today, so this is latent. One line closes it:

override update(delta: Record<string, unknown>): void {
  super.update(delta);
  Object.assign(this.committed, delta);
}

[kalenkevich]

Good catch — fixed in a00c183. Added the override exactly as you wrote it:

override update(delta: Record<string, unknown>): void {
  super.update(delta);
  Object.assign(this.committed, delta);
}

Also moved the write-through comment above both overrides so it reads as a property of NodeStateView, not of set specifically — that framing is what let the gap through in the first place.

Pinned with a test (writes made through update() are visible the same way as set()) asserting both halves of the contract: a later node reads it back, and it lands on session.state. Verified it fails without the override (expected undefined to be 'oceans').

[AmaadMartin]

Nit. The overlay also shadows writes coming the other way, which is worth naming next to the known residual.

Once a node writes k, every later read of k in this invocation is served from the overlay:

override get<T>(key: string, defaultValue?: T): T | undefined {
  if (super.has(key)) {
    return super.get<T>(key, defaultValue);
  }
  return key in this.committed ? (this.committed[key] as T) : defaultValue;
}

A tool or callback that writes the same key straight to session.state mid-run is therefore invisible to later nodes, for the rest of the invocation. Before the overlay they saw it, subject to the race you are fixing.

This is the mirror of the residual you already document: that one is workflow writes not reaching an outside reader, this one is outside writes not reaching a workflow reader. Same root cause, same real fix (version-stamping). Worth a sentence in the same place, so the next person does not have to derive it.

[kalenkevich]

Agreed, and you framed it better than I would have. Documented both directions together in the NodeStateView doc comment (a00c183) rather than only the one I'd noticed:

 * 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. ...
 * - 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.

Worth stating plainly that this second one is a behaviour change, not just an untouched gap: before the overlay those reads did see the outside write (racily). I think it's the right trade — a mid-run outside writer to a key the workflow also owns is already in undefined-ordering territory, whereas node-to-node read-modify-write is the documented pattern — but it is a trade, and the doc now says so. I'll also carry both into the PR description's Known residual section.

[AmaadMartin]

The core is sound: I checked the write-through (state.ts:49), the delta re-apply (base_session_service.ts:204), and the new state object each getSession returns. Keying the overlay on that object is the right scope, and it matches FunctionNode.attachedStateByCtx (function_node.ts:76). Two comments below: one write path still bypasses the overlay, and the invocation-id guard is untested. Windows CI fails on unsafe_local_code_executor_test.ts, which also fails on main. Please open a separate issue for the reconstructNodeStates gap.

[AmaadMartin]

Not a nit. A third path misses the overlay: llm_agent_wrapper.ts:141 writes the delta directly.

if (agent.outputKey && output !== undefined) {
  ctx.actions.stateDelta[agent.outputKey] = output;
}

That write skips the overlay and skips session.state, so the next node reads undefined until the event commits. Use ctx.state.set(agent.outputKey, output), which writes the same delta key, plus the overlay and the session. The fault is older than this PR, and no workflow test covers outputKey. llm_agent.ts:708 repeats the pattern outside this package.

[kalenkevich]

Fixed in 5cc73fc — and you undersold it: that write does not just miss the overlay, it lands nowhere at all.

ctx.actions.stateDelta is only ever drained onto an event by FunctionNode.toEvent (function_node.ts:182, via pendingStateDelta). BaseNode.toEvent doesn't, enrichEvent doesn't, and nothing constructs a NodeContext with a shared actions, so for this node the object is write-only. Task-mode outputKey therefore reached neither the overlay, nor session.state, nor committed state — it was dropped, not merely late.

Measured with the new test, which asserts all three (later node reads it, it is on session.state, it is in the committed session):

wrapper writes	later node	session.state	committed
ctx.actions.stateDelta[k] = v (before)	✗ undefined	✗	✗
ctx.state.set(k, v) (as suggested)	✓	✓	✗ undefined
both (landed)	✓	✓	✓

So it needs the write-through and the event stamp:

ctx.state.set(agent.outputKey, output);
event.actions.stateDelta[agent.outputKey] = output;

Test: an agent node's outputKey is visible to the next node — a task-mode LlmAgent over a mock model that calls finish_task, feeding a node that reads the key. First workflow coverage of outputKey, as you said.

On llm_agent.ts:708: left alone deliberately, and it is not the same severity. I probed a single_turn agent node with an outputKey followed by a reader node — the reader does see the value today, because the runner commits the agent's event before the next node starts. That path is exposed to the commit-lag window (the first residual in the NodeStateView doc comment), not to a dropped write, and closing it properly is the version-stamping change rather than a one-liner here.

[AmaadMartin]

Not a nit. I retract two parts of my comment above; your framing is correct.

The value never committed, so "reads undefined until the event commits" was wrong. Nothing drains ctx.actions for this node:

• BaseNode.toEvent never reads it (base_node.ts:191).
• enrichEvent only writes agentState (node_runner.ts:250).
• No caller passes actions to a NodeContext.

ctx.state.set alone was also incomplete — committed state would stay empty. The landed pair is right, and the new test pins all three effects.

[AmaadMartin]

Nit. The guard covers sequential reuse only, not two live invocations.

const created = {invocationId: ic.invocationId, values: {}};
invocationOverlays.set(sessionState, created);

Two invocations on one session object overwrite each other's entry, so both fall back to the old behaviour. Please narrow this sentence to "sequential invocations". Nothing tests the guard: runOnce makes a new session service per call, so the two runs never share a state object. Two NodeContexts over one session, with different invocation ids, would cover it.

[kalenkevich]

Right, the sentence claimed more than the code does. Narrowed in 5cc73fc:

 * 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.

Also tested it, the way you described — two NodeContexts over one session object with different invocation ids (serves a second invocation over the same session from committed state). The one wrinkle is that a fresh overlay is invisible on its own, since the read just falls through to session.state and finds the same value the previous invocation wrote through. So the test first makes committed state disagree with the overlay (standing in for the runner re-applying a stale delta) and then pins both sides:

mkCtx().state.set('k', 'from-inv-1');
ic1.session.state['k'] = 'stale';

expect(mkCtx().state.get('k')).toBe('from-inv-1');        // same invocation: overlay wins

const ic2 = ic1.clone({invocationId: 'inv-2'});
expect(ic2.session.state).toBe(ic1.session.state);        // same state object
expect(mkCtx(ic2).state.get('k')).toBe('stale');          // fresh overlay: committed state

Fails without the guard (the second read returns from-inv-1).

[kalenkevich]

Both comments addressed in 5cc73fc; replies inline.

• outputKey bypass — fixed, and it was worse than late: ctx.actions.stateDelta is drained onto an event only by FunctionNode, so for an agent node that write landed nowhere — not the overlay, not session.state, not committed state. Needed ctx.state.set(...) and a delta stamp on the emitted event; the table in the inline reply shows what each half buys. First workflow test for outputKey comes with it. llm_agent.ts:708 left alone — probed it, a single_turn agent node's outputKey does reach the next node today (the runner commits the event first), so it is in the commit-lag window rather than dropped; noted under Known residual.
• Guard doc — narrowed to sequential reuse, and tested with two NodeContexts over one session object with different invocation ids. Fails without the guard.

On reconstructNodeStates: it is #637, which fixes it rather than just tracking it — you pointed me there in the first review. Happy to open an issue too if you would rather have one for tracking.

Windows unsafe_local_code_executor_test.ts — agreed, fails on main.

unit:core green: 206 files, 2833 tests.

[AmaadMartin]

Both findings are closed at 5cc73fc, verified against source, and the fix commit is clean. The narrowed doc matches the code, and the new guard test uses two invocation ids over one shared session.state. The outputKey fix writes the overlay, session.state, and the event delta; no test covered task mode before, so nothing depended on the old path. Leaving llm_agent.ts:708 alone is right — that write reaches the event, so it is late, not lost. I am not approving only because run-tests (windows-latest) is red; zero tests failed there and the A2A suite's server exited at startup, so please re-run that job.

[kalenkevich]

Windows is green on the re-run — every check on the PR now passes.

Confirming the flake was what you read it as: the failed run had Tests 3277 passed | 47 skipped, zero failures, and one suite failure at setup —

FAIL  integration  tests/integration/a2a/input_required/input_required_test.ts > A2A: RemoteAgent InputRequired
Error: CLI exited prematurely with code 1
 ❯ ChildProcess.<anonymous> tests/integration/test_case_utils.ts:341:13

i.e. the spawned server never came up; ubuntu and macOS ran the same suite fine on the same commit, and the previous commit on this branch (8e050d7) passed on Windows too. Nothing in the diff goes near that path.

Re-ran job: https://github.com/google/adk-js/actions/runs/31448921496/job/93981457475 — pass, 10m22s.

Ready for another look whenever you have a moment.

[AmaadMartin]

Approved at 5cc73fc.

Both findings are closed. llm_agent_wrapper.ts:147-148 now writes through ctx.state.set and stamps event.actions.stateDelta, and the doc matches the guard.

One correction to my earlier review: I wrote that the value read undefined until the event commits. That was wrong. BaseNode.toEvent never reads ctx.actions, so the value was dropped, not late. My one-line fix would have left committed state empty. Your framing was correct.

run-tests (windows-latest) failed three different ways across four runs on this branch and main. It passes on the re-run, so the job is unstable and the failures are not from this diff.