Skip to content

fix(graph): rollback from incomplete checkpoint messages#4791

Open
Jupiter363 wants to merge 1 commit into
alibaba:mainfrom
Jupiter363:fix/4789-checkpoint-rollback
Open

fix(graph): rollback from incomplete checkpoint messages#4791
Jupiter363 wants to merge 1 commit into
alibaba:mainfrom
Jupiter363:fix/4789-checkpoint-rollback

Conversation

@Jupiter363

Copy link
Copy Markdown
Contributor

Describe what this PR does / why we need it

Fixes #4789.

When a graph stream is cancelled or times out after a tool node has already persisted ToolResponseMessage, the following streaming model node may never persist the final AssistantMessage. The next user turn can then load an incomplete checkpoint and build an invalid message sequence such as:

assistant(tool_calls) -> tool -> user

That sequence is not accepted by OpenAI/DashScope-compatible chat APIs and can make the next model request fail.

This PR replaces #4790 with a smaller rollback-oriented fix. The previous background-continuation approach tried to keep the cancelled run alive internally, but that introduced extra concurrency risks between an old background run and a newer turn on the same thread. This PR keeps normal Reactor cancellation semantics and instead makes the next turn recover from the latest safe checkpoint.

Does this pull request fix one issue?

Fixes #4789

Describe how you did it

This PR adds a checkpoint message-sequence validator and uses it when CompiledGraph.getInitialState(...) loads checkpoint state for a new user input.

If the latest checkpoint contains an incomplete tool-call sequence, the graph skips it and rolls back to the nearest earlier checkpoint whose messages history can safely accept another user message.

The validator rejects persisted histories such as:

  • assistant tool calls without matching tool responses;
  • partial tool responses for multiple tool calls;
  • tool responses not followed by a final assistant message;
  • tool -> user continuation.

It accepts histories where the tool-call chain has been closed by a final assistant message.

This PR does not rewrite checkpoint records, delete checkpoint records, fabricate a placeholder assistant message, or keep cancelled graph executions running in the background. It only changes which checkpoint is selected as the base state for the next user input.

Why rollback instead of background continuation

A background-continuation approach can preserve the cancelled turn's final assistant output when it eventually completes, but it also creates a race: the old cancelled run may keep appending checkpoints while the user immediately sends the next message on the same thread. That requires additional lineage guards around checkpoint writes and releaseThread(true), and still leaves backend-level CAS/distributed-locking concerns for multi-instance deployments.

The rollback approach is more conservative for this issue:

  • downstream timeout/cancellation still aborts the current run;
  • incomplete checkpoints can remain in storage for history/debugging;
  • the next user turn never resumes from an unsafe tool-response checkpoint;
  • no detached subscription or background writer is introduced.

Describe how to verify it

I verified this PR with:

mvn -pl spring-ai-alibaba-graph-core -Dtest=CheckpointRollbackConsistencyTest,MessageSequenceValidatorTest,GraphRunnerContextResumeMessageOrderTest,StreamFunctionCallTest test
mvn -pl spring-ai-alibaba-graph-core checkstyle:check
mvn -pl spring-ai-alibaba-graph-core spotless:check
git diff --check

Latest local result:

Tests run: 9, Failures: 0, Errors: 0, Skipped: 0
You have 0 Checkstyle violations.
spotless:check: 0 needs changes
git diff --check: no whitespace errors

Regression coverage includes:

  • latest checkpoint ending with ToolResponseMessage rolls back to the previous safe checkpoint;
  • cancellation after a tool node writes its checkpoint, followed by the next user turn;
  • timeout after the streaming model node starts, followed by the next user turn;
  • next model call receives OpenAI-compatible message ordering;
  • validator rejects unclosed assistant tool calls, partial tool responses, and tool histories without a final assistant;
  • validator accepts a fully closed assistant(tool_calls) -> tool -> assistant chain;
  • existing resume message ordering still passes;
  • existing streaming tool-call merge behavior still passes.

Signed-off-by: Jupiter363 <1404785625@qq.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area/graph SAA Grpah module

Projects

None yet

Development

Successfully merging this pull request may close these issues.

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

1 participant