Skip to content

fix(agent): rewind checkpoint to turn boundary on stream cancellation - #79

Merged
yuluo-yx merged 2 commits into
agentic-ai-java:mainfrom
zeng-bohan:fix/stream-cancel-checkpoint-rewind
Sep 9, 2026
Merged

yuluo-yx merged 2 commits into
agentic-ai-java:mainfrom
zeng-bohan:fix/stream-cancel-checkpoint-rewind

Conversation

@zeng-bohan

Copy link
Copy Markdown
Contributor

Describe what this PR does / why we need it

When a subscriber cancels a streaming run — SSE disconnect, flux.timeout, debugger pause — the in-flight turn could end up half-persisted in the checkpoint. Two observed variants:

  • a dangling ASSISTANT(tool_call) with no following tool response (reproduced offline in this repository, see test), or
  • a ToolResponseMessage whose preceding assistant tool_call never got persisted (reported in the linked issue).

Either way, the next turn loads an invalid message sequence and mainstream LLM APIs reject the request with 400 ("roles must alternate…").

Root cause: state merge + checkpoint persistence sit in processedFlux.concatWith(updateContextMono…) inside NodeExecutor.processGraphResponseFlux, so they only run when downstream demand pulls the tail — cancellation skips them, while nested async continuations may still land their own checkpoints.

This PR makes a cancelled turn behave as if it never happened:

  1. GraphRunner.run marks the GraphRunnerContext cancelled on the stream's doOnCancel; addCheckpoint skips persistence when cancelled — with a re-check after the (potentially expensive) state clone, so in-flight completions racing the cancel cannot record a half turn.
  2. ReactAgent.doStream snapshots the thread's pre-turn checkpoint at assembly time and, on cancellation, rewinds the saver to that snapshot via saver.put(config, Checkpoint.copyOf(...)) (an empty-state checkpoint for a first-ever turn).

Normal completion and error paths are untouched (doOnCancel does not fire on them).

Does this pull request fix one issue?

Fixes #25

Describe how you did it

  • GraphRunnerContext: volatile cancelled flag + markCancelled(); addCheckpoint early-returns when cancelled (checked on entry and again right before the saver put).
  • GraphRunner.run: attaches doOnCancel(context::markCancelled).
  • ReactAgent.doStream override: pre-turn checkpoint snapshot + doOnCancel rewind; rewind failures are logged-and-swallowed so they never mask the cancellation.
  • Issue25StreamCancelCheckpointTest (offline, stub streaming model + tool + MemorySaver):
    • full run persists a valid paired sequence (baseline);
    • cancelling mid-turn leaves no dangling fragment — after the fix the checkpoint is rewound to the pre-turn state (against main this test fails with USER -> ASSISTANT(tool_call));
    • disposing after completion keeps the full 4-message turn (guards against over-rewinding).

Describe how to verify it

mvn -pl agentic-spring-ai-agent-framework -am test -Dtest=Issue25StreamCancelCheckpointTest

Full module suite (mvn -pl agentic-spring-ai-agent-framework -am test, 649 tests) shows the same failure set as unmodified main (the pre-existing Windows symlink-privilege tests in FilesystemInterceptorTest / FileSystemToolsTest); no new failures.

Special notes for reviewers

Design note: rewinding to the turn boundary was chosen over salvage-on-cancel (cannot fix the dangling tool_call variant) and load-time sequence sanitization (masks the problem, silently rewrites history). A design summary was posted on issue #25 before this PR. Scope note: the rewind is wired at the Agent streaming layer; generalizing the cancelled-flag handling to other entry points can follow if the approach looks right.

When a subscriber cancels a streaming run (SSE disconnect, flux.timeout),
in-flight node completions could still persist part of the current turn,
leaving a dangling assistant tool_call or an orphan tool response in the
checkpoint. The next turn then loads an invalid message sequence that
LLM APIs reject with 400.

- GraphRunnerContext marks the run cancelled on the stream's doOnCancel;
  addCheckpoint skips persistence (with a re-check after the state clone)
  so in-flight completions cannot record a cancelled turn.
- ReactAgent snapshots the thread's pre-turn checkpoint when assembling
  the stream and rewinds the saver to that snapshot on cancellation, so a
  cancelled turn is treated as if it never happened.

Fixes #25
@github-actions github-actions Bot added area/core 核心:智能体运行时基础与公共接口 area/graph 图引擎:图工作流、路由、节点与嵌套图 labels Sep 4, 2026

@yuluo-yx yuluo-yx left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

本地审查发现可复现的行为回归,建议修复下述问题后再合并。测试基线为 b726cbd021f79e5ef4f9dd84fdfd22935abeee74,验证使用 JDK 17;未触发远端 CI。

CompiledGraph compiledGraph = getAndCompileGraph();
RunnableConfig config = buildStreamConfig(runnableConfig);
Optional<BaseCheckpointSaver> saver = compiledGraph.compileConfig.checkpointSaver();
Optional<Checkpoint> preTurnCheckpoint = saver.flatMap(s -> {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P1] 在订阅开始时捕获回退快照,避免丢失已提交状态

这里在 agent.stream(...) 创建 Flux 时立即读取 checkpoint,但图执行在 Flux.defer 中直到订阅才开始。若创建流之后、订阅之前同一 thread 完成了另一次状态更新,本次取消会把已经提交的新状态覆盖成创建流时的旧快照。重复订阅/重试同一个冷流也会复用旧快照。

已用公共 API 本地复现:保存 marker=at-assembly → 创建 Flux<NodeOutput> deferred = agent.stream(..., config) → 保存 marker=at-subscription → deferred.take(1).blockLast()。取消后期望仍为 at-subscription,本 PR 实际变回 at-assembly;相同测试在 main 上通过。PR 原有 3 个测试通过,但没有覆盖创建与订阅之间的状态变化。

建议把 checkpoint 读取及对应执行一起放进每次订阅的 Flux.defer 中,并让取消只回退该次执行拥有的状态,补充延迟订阅和重复订阅测试后再合并。

…cution

doStream() read the pre-turn checkpoint while assembling the Flux, but the
graph only executes on subscribe. A state commit landing between assembly
and subscription was therefore rolled back by a cancellation that restored
the stale assembly-time snapshot, and re-subscribing a cold Flux reused
that same stale read (review finding on #79, reproduced through the public
API: marker=at-assembly, assemble, marker=at-subscription, subscribe+cancel
ended with at-assembly state).

Wrap the pipeline in Flux.defer so the snapshot is read once per
subscription in the same scope as the execution, and move the rewind from
doOnCancel to doFinally(CANCEL): cancellation is cooperative and the
cancelled run can still persist one more checkpoint while unwinding, which
could overwrite a doOnCancel rewind; doFinally runs after the upstream has
fully terminated, making the rewind the last writer of that execution.

Adds two regression tests: an end-to-end scenario (complete turn, assemble,
complete another turn, subscribe+cancel mid-flight — the second turn's
committed state must survive) and a saver read-timing test (no snapshot
read at assembly). Both fail on the previous code (assembly-time read
count 1; rewind restored the 2-message stale state) and pass now, and the
existing Issue25 cancellation tests stay green.

Signed-off-by: zengbohan1 <[email protected]>
@zeng-bohan

Copy link
Copy Markdown
Contributor Author

确认,已修复(350db5933),并比你指出的多修了一层。

  1. 快照读取移入 Flux.defer,与执行处于同一次订阅作用域——组装期不再读 checkpoint,重复订阅冷流时每次都会取新的回退基线。
  2. 验证时发现另一个问题:cancel 之后协同取消的图线程仍可能再落一个 checkpoint,把 doOnCancel 里的 rewind 覆盖掉(原有 Issue25StreamCancelCheckpointTest 在此场景转红)。因此 rewind 从 doOnCancel 移到 doFinally(CANCEL)——上游完全终止后,rewind 必然是本次执行的最后一个写入者,与"取消只回退该次执行拥有的状态"一致。
  3. 新增两个回归测试:端到端场景(完成 turn1 → 组装 Flux → 完成 turn2 → 订阅并中途取消,turn2 已提交状态必须存活)和 saver 读时序(组装期读次数必须为 0)。旧代码上分别以"回滚到 2 条消息的陈旧状态""组装期读次数=1"失败,与你本地复现一致;修复后新测试 + Issue25 原有 3 个测试全部通过(复跑两轮稳定)。

@yuluo-yx
yuluo-yx merged commit 1b5feee into agentic-ai-java:main Sep 9, 2026
1 check passed
yuluo-yx pushed a commit that referenced this pull request Oct 17, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area/core 核心:智能体运行时基础与公共接口 area/graph 图引擎:图工作流、路由、节点与嵌套图

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] 流式节点未正常完成时,下一轮对话消息序列结构错乱(导致LLM 拒绝请求)

2 participants