Skip to content

fix(tui): persist session before stall/cancel recovery so --continue keeps history (#2739) - #3285

Merged
Hmbown merged 2 commits into
codewhale-hq:mainfrom
LeoLin990405:fix/issue-2739-persist-session-on-recovery
Jun 18, 2026
Merged

Hmbown merged 2 commits into
codewhale-hq:mainfrom
LeoLin990405:fix/issue-2739-persist-session-on-recovery

Conversation

@LeoLin990405

Copy link
Copy Markdown
Contributor

Problem

Fixes #2739 (partial — the data-loss part) — after a long turn stalls and is recovered (or is cancelled with Esc), running --continue loads the previous session and the entire in-progress turn is lost.

Root cause

The stall watchdog and cancel paths clear turn bookkeeping without first writing the in-progress turn to disk. Three recovery paths in crates/tui/src/tui/ui.rs were affected:

  • recover_stalled_runtime_turn (stall watchdog)
  • mark_active_turn_cancelled_locally (Esc cancel)
  • recover_engine_event_disconnect (engine disconnect)

Each finalizes the in-memory cells and resets turn state, but the partial turn lives only in app.api_messages and is never persisted. So --continue reads the last completed-turn save and the interrupted turn vanishes.

Fix

Add persist_recovery_snapshot() and call it on all three recovery paths before turn state is cleared. It uses the exact same mechanism as a normally-completed turn's auto-save:

if let Ok(manager) = SessionManager::default_location() {
    let session = build_session_snapshot(app, &manager);
    if app.current_session_id.is_none() {
        app.current_session_id = Some(session.metadata.id.clone());
    }
    persistence_actor::persist(PersistRequest::SessionSnapshot(session));
}

build_session_snapshot loads and updates the session keyed by current_session_id, so it writes back to the same session file --continue reads. This mirrors the existing completed-turn auto-save (see the SessionSnapshot persist at the top of the turn-complete handler), so it introduces no new persistence path or risk.

Testing

  • New colocated tests issue_2739_stalled_turn_snapshot_preserves_api_messages and issue_2739_esc_cancel_preserves_session_messages_before_clear: verify the snapshot captures the in-progress api_messages and that cancel clears turn bookkeeping only after the messages are captured.
  • cargo clippy -p codewhale-tui --all-features clean.
  • cargo test -p codewhale-tui --all-features green (4857 passed / 0 failed).

Note

This addresses the session-loss symptom in #2739. The freeze/stall itself is already mitigated by the existing liveness watchdogs (reconcile_turn_liveness); this PR ensures those recoveries no longer cost the user their in-progress turn.

…2739)

Stall recovery, Esc cancel, and engine-disconnect recovery cleared
turn state without first writing the in-progress turn to disk, so the
partial turn lived only in api_messages and --continue loaded the
previous save -- losing the entire interrupted turn. Snapshot the
session (same build_session_snapshot + persist path used by normal
turn completion) before clearing state on all three recovery paths.
@LeoLin990405
LeoLin990405 requested a review from Hmbown as a code owner June 17, 2026 07:32
@github-actions

Copy link
Copy Markdown
Contributor

Thanks @LeoLin990405 for taking the time to contribute.

This repository is observing a maintainer-managed PR intake gate in dry-run mode, so this pull request is staying open. This note helps maintainers prepare the allowlist before any enforcement is considered.

Please read CONTRIBUTING.md for the expected contribution shape. A maintainer can grant recurring PR access by commenting /lgtm on a pull request.

@greptile-apps greptile-apps Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Your free trial has ended. If you'd like to continue receiving code reviews, you can add a payment method here.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Code Review

This pull request addresses issue #2739 by persisting the current in-memory session state before clearing turn bookkeeping during cancellation, stalled turns, or engine disconnects, ensuring that partial or cancelled turns are saved to disk so that --continue can resume from the correct state. Unit tests were also added to verify this behavior. Feedback highlights a potential issue where persisting an interrupted turn with incomplete tool calls can result in mismatched tool messages, leading to API 400 Bad Request errors on resume. It also suggests calling persist_recovery_snapshot in the turn dispatch timeout branch of reconcile_turn_liveness.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

Comment thread crates/tui/src/tui/ui.rs
Comment on lines +4962 to +4970
fn persist_recovery_snapshot(app: &mut App) {
if let Ok(manager) = SessionManager::default_location() {
let session = build_session_snapshot(app, &manager);
if app.current_session_id.is_none() {
app.current_session_id = Some(session.metadata.id.clone());
}
persistence_actor::persist(PersistRequest::SessionSnapshot(session));
}
}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

high

Potential API 400 Bad Request on Resume due to Mismatched Tool Messages\n\nWhen a turn is cancelled (via Esc) or stalls while a tool call is in-flight or after some tool calls have completed but before the assistant message is fully completed, app.api_messages can end up in an invalid state.\n\nSpecifically, if a tool call completed, its ToolResult is appended to app.api_messages (via EngineEvent::ToolCallComplete). However, because the turn was interrupted before MessageComplete was received, the preceding assistant message containing the matching ToolUse block is never appended to app.api_messages.\n\nThis results in a ToolResult message without a preceding ToolUse block in the persisted session history. When the user resumes the session using --continue, the next API request to the OpenAI/DeepSeek API will fail immediately with a 400 Bad Request error due to the strict message ordering constraints of these APIs (which require every tool role message to have a matching tool_call_id in the immediately preceding assistant message).\n\nTo prevent this, we should ensure that persist_recovery_snapshot or build_session_snapshot cleans up any trailing mismatched ToolResult messages or incomplete assistant messages before persisting the snapshot. For example, we can truncate app.api_messages back to the last clean state (e.g., the last user prompt before the interrupted turn).\n\nAdditionally, note that the turn dispatch timeout branch in reconcile_turn_liveness (which clears turn state when dispatch times out) does not call persist_recovery_snapshot(app). This means the user's prompt (which was already appended to app.api_messages during dispatch) will be lost on --continue. We should call persist_recovery_snapshot(app) in that branch as well.

Address review feedback: reconcile_turn_liveness's dispatch-timeout
branch cleared turn state without persisting, so a prompt that was
appended to api_messages but never dispatched was lost on --continue.
Snapshot before clearing, matching the other recovery paths. Add a
regression test covering the dispatch-timeout branch.

Note: persisting a partially-interrupted tool round is safe because the
send path (build_chat_messages) already strips orphaned tool_use /
tool_result blocks before every request, so --continue cannot produce a
400 from a dangling tool exchange.

@greptile-apps greptile-apps Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Your free trial has ended. If you'd like to continue receiving code reviews, you can add a payment method here.

@LeoLin990405

Copy link
Copy Markdown
Contributor Author

Thanks for the review!

Dispatch-timeout branch — good catch. reconcile_turn_liveness's dispatch-timeout branch cleared turn state without persisting, so a prompt appended to api_messages but never dispatched was lost on --continue. Fixed in 34ea946 (persist before clearing, matching the other recovery paths) + a regression test (issue_2739_dispatch_timeout_preserves_user_prompt).

400-on-resume from a partial tool round — this is already covered by the send path. Every request goes through build_chat_messages_for_request_and_provider (crates/tui/src/client/chat.rs), which strips orphaned tool_use and tool_result blocks in both directions before the request is built. So a partially-interrupted tool round persisted by this PR and reloaded via --continue is sanitized before the first request — no 400. Covered by existing tests chat_messages_strips_orphaned_tool_calls_after_compaction, chat_messages_drop_orphan_tool_results, and chat_messages_strips_partial_tool_results. I kept the snapshot logic minimal rather than duplicating that orphan-stripping at persist time.

@Hmbown

Hmbown commented Jun 18, 2026

Copy link
Copy Markdown
Collaborator

this is awesome - this is something that's bugged me too in my own usage but I was always working on something and wasn't able to get to it. thank you so much!!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

依然会出现任务执行过程中卡死的状态

2 participants