fix(cancel): preserve partial turn payload via embedded cancel snapshot (#4647)

Cancelling a running turn briefly collapsed the already-streamed reasoning,
tool cards, and text to a bare "Task cancelled" marker before a follow-up
GET /api/session repopulated them — a visible race between the terminal
cancel SSE event and the refetch (#4076, #1361).

The terminal cancel event now carries a canonical, redacted session snapshot
captured under the per-session agent lock right after the cancelled turn is
persisted (including the recovered partial-assistant message). The frontend
applies that embedded snapshot immediately via _applyCancelSessionPayload and
skips the GET round-trip. The snapshot is a fully detached deep copy (redaction
rebuilds every container) so it cannot alias or be staled by concurrent writes;
it is only applied when its session_id matches the active view; and a
missing/failed snapshot degrades cleanly to the existing GET fallback.

The seven other cancel call sites are standardized via _cancel_event_payload()
(message + type + status), backward-compatible (no session key when none).

Co-authored-by: franksong2702 <franksong2702@users.noreply.github.com>

CHANGELOG.md

@@ -5,6 +5,8 @@
 
 ### Fixed
 
+- **Cancelling a running turn now keeps the partial work it had already produced, with no flicker.** When you hit Stop mid-turn, the browser briefly collapsed the already-streamed reasoning, tool cards, and text down to a bare "Task cancelled" marker before a follow-up `GET /api/session` repopulated them — a visible race between the cancel event and the refetch. The terminal cancel event now carries a canonical, redacted session snapshot captured under the per-session agent lock right after the cancelled turn is persisted (including the recovered partial-assistant message), so the frontend renders the preserved partial reasoning/tool/text immediately and skips the GET round-trip entirely. The embedded snapshot is a fully detached deep copy (redaction rebuilds every container) so it can never alias or be staled by concurrent writes, it is only applied when its `session_id` matches the active view, and a missing/failed snapshot degrades cleanly to the existing GET fallback. Thanks @franksong2702. (#4647, fixes #4076, #1361)
+
 - **A WebUI turn whose stream lost its terminal event no longer loses already-produced text on refresh.** If a stream drops its `done`/`stream_end` while the agent keeps writing to `state.db`, the browser showed live output but a refresh reloaded the stale sidecar and the produced text appeared to vanish. WebUI now self-heals on read: when the durable `state.db` transcript is genuinely newer than the sidecar and the stream is truly gone (not in the live SSE/worker registries, and past a startup grace window), it reconciles the sidecar from `state.db` through the existing append-only merger (preserving compaction/truncation semantics) and clears the stale unfinished-run state. The reconcile + write run under the per-session lock against a freshly-reloaded session, so it can never race or clobber a concurrent worker writeback or a just-starting turn. Thanks @allenliang2022. (#5053)
 
 - **Concurrent stale session-visit model refreshes no longer each pay for a full live rebuild — and a slow/failed rebuild can't stall or hang the followers.** Building on #4798, overlapping stale session visits now coalesce behind a single in-flight live model-catalog rebuild instead of each launching their own. A forced-refresh follower waits within its configured live-rebuild budget (falling back to the stale/static catalog if the budget expires while a rebuild is still running) and, in the legacy unbounded mode, keeps coalescing behind the rebuild rather than duplicating it — and if that rebuild raises, followers are released to fall through to their own rebuild/fallback instead of blocking. Followers only adopt a freshly-published live rebuild (not a stale disk-cache publish) as the forced-refresh result. Thanks @rodboev. (#5051, fixes #5047)

api/streaming.py

@@ -93,6 +93,34 @@ def _strip_compact_echo_suffix(value: str, suffix: str, *, search_window: int =
     return raw, False
 
 
+def _redacted_session_payload_with_full_messages(session, *, tool_calls=None) -> dict | None:
+    """Best-effort terminal SSE session payload for already-persisted state."""
+    try:
+        return redact_session_data(
+            _session_payload_with_full_messages(session, tool_calls=tool_calls)
+        )
+    except Exception:
+        logger.debug("Failed to build redacted session payload", exc_info=True)
+        return None
+
+
+def _cancel_event_payload(
+    message: str = "Cancelled by user",
+    *,
+    session: dict | None = None,
+) -> dict:
+    """Return base cancel terminal event metadata."""
+    payload = {
+        'message': message,
+        'type': 'cancelled',
+        'status': 'cancelled',
+    }
+    if session:
+        payload['session'] = session
+        payload['session_id'] = session.get('session_id')
+    return payload
+
+
 # Global lock for os.environ writes. Per-session locks (_agent_lock) prevent
 # concurrent runs of the SAME session, but two DIFFERENT sessions can still
 # interleave their os.environ writes. This global lock serializes the env
@@ -6244,7 +6272,7 @@ def _run_agent_streaming(
         if cancel_event.is_set():
             with _agent_lock:
                 _finalize_cancelled_turn(s, ephemeral=ephemeral, message='Task cancelled before start.')
-            put('cancel', {'message': 'Cancelled before start'})
+            put('cancel', _cancel_event_payload('Cancelled before start'))
             return
 
         # Resolve profile home for this agent run — use the session's own profile
@@ -7426,7 +7454,7 @@ def _run_agent_streaming(
                         logger.debug("Failed to interrupt agent before start")
                     with _agent_lock:
                         _finalize_cancelled_turn(s, ephemeral=ephemeral, message='Task cancelled before start.')
-                    put('cancel', {'message': 'Cancelled by user'})
+                    put('cancel', _cancel_event_payload('Cancelled by user'))
                     return
 
             # Prepend workspace context so the agent always knows which directory
@@ -7584,7 +7612,7 @@ def _run_agent_streaming(
                             )
                         except Exception:
                             logger.debug("Failed to append cancelled turn journal event", exc_info=True)
-                put('cancel', {'message': 'Cancelled by user'})
+                put('cancel', _cancel_event_payload('Cancelled by user'))
                 return
             # ── Ephemeral mode (/btw): deliver answer, skip persistence, cleanup ──
             if ephemeral:
@@ -7626,7 +7654,7 @@ def _run_agent_streaming(
                         )
                     except Exception:
                         logger.debug("Failed to append cancelled turn journal event", exc_info=True)
-                put('cancel', {'message': 'Cancelled by user'})
+                put('cancel', _cancel_event_payload('Cancelled by user'))
                 return
             _writeback_timings = []
             _writeback_started = time.perf_counter()
@@ -7667,7 +7695,7 @@ def _run_agent_streaming(
                             )
                         except Exception:
                             logger.debug("Failed to append cancelled turn journal event", exc_info=True)
-                        put('cancel', {'message': 'Cancelled by user'})
+                        put('cancel', _cancel_event_payload('Cancelled by user'))
                         return
                     _next_context_messages = _restore_reasoning_metadata(
                         _previous_context_messages,
@@ -7874,7 +7902,7 @@ def _run_agent_streaming(
                                 )
                             except Exception:
                                 logger.debug("Failed to append cancelled turn journal event", exc_info=True)
-                        put('cancel', {'message': 'Cancelled by user'})
+                        put('cancel', _cancel_event_payload('Cancelled by user'))
                         return
                     _last_err = getattr(agent, '_last_error', None) or result.get('error') or ''
                     _err_str = str(_last_err) if _last_err else ''
@@ -8484,7 +8512,7 @@ def _run_agent_streaming(
                         )
                     except Exception:
                         logger.debug("Failed to append cancelled turn journal event", exc_info=True)
-                    put('cancel', {'message': 'Cancelled by user'})
+                    put('cancel', _cancel_event_payload('Cancelled by user'))
                     return
                 with _stream_writeback_stage(_writeback_timings, "session_save"):
                     s.save()
@@ -8502,7 +8530,7 @@ def _run_agent_streaming(
                         )
                     except Exception:
                         logger.debug("Failed to append cancelled turn journal event", exc_info=True)
-                    put('cancel', {'message': 'Cancelled by user'})
+                    put('cancel', _cancel_event_payload('Cancelled by user'))
                     return
                 if not ephemeral:
                     try:
@@ -8926,7 +8954,7 @@ def _run_agent_streaming(
                             )
                         except Exception:
                             logger.debug("Failed to append cancelled turn journal event", exc_info=True)
-            put('cancel', {'message': 'Cancelled by user'})
+            put('cancel', _cancel_event_payload('Cancelled by user'))
             return
         _exc_is_quota = _classification['type'] == 'quota_exhausted'
         # Exception quota text still includes: 'more credits' in _exc_lower, 'can only afford' in _exc_lower, 'fewer max_tokens' in _exc_lower.
@@ -9332,6 +9360,7 @@ def cancel_stream(stream_id: str) -> bool:
     _snap_tool_calls = None
     _snap_flag = None
     _snap_agent = None
+    _cancel_session_payload = None
 
     with streams_lock:
         stream_present = stream_id in streams
@@ -9632,6 +9661,7 @@ def cancel_stream(stream_id: str) -> bool:
                         'timestamp': int(time.time()),
                     })
                 _cs.save()
+                _cancel_session_payload = _redacted_session_payload_with_full_messages(_cs)
             except Exception:
                 logger.debug("Failed to clear session state on cancel for %s", _cancel_session_id)
 
@@ -9643,7 +9673,8 @@ def cancel_stream(stream_id: str) -> bool:
             except Exception:
                 logger.debug("Failed to note cancel event_id %s for stream %s", _cancel_event_id, stream_id, exc_info=True)
         try:
-            q.put_nowait(('cancel', {'message': 'Cancelled by user'}))
+            _payload = _cancel_event_payload('Cancelled by user', session=_cancel_session_payload)
+            q.put_nowait(('cancel', _payload))
         except Exception:
             logger.debug("Failed to put cancel event to queue")
 

static/messages.js

@@ -5194,22 +5194,35 @@ function attachLiveStream(activeSid, streamId, uploaded=[], options={}){
       if(S.session&&S.session.session_id===activeSid){
         S.activeStreamId=null;
       }
-      // Fetch latest session from server to get accurate message list (includes cancel status)
-      // This ensures messages stay in sync with server, fixing race condition where local
-      // "*Task cancelled.*" message gets lost when done event overwrites S.messages
+      const _applyCancelSessionPayload=(sessionPayload)=>{
+        if(!sessionPayload||typeof sessionPayload!=='object'||!S.session||S.session.session_id!==activeSid) return false;
+        // Belt-and-suspenders: the embedded cancel snapshot must be for THIS session.
+        // The GET path guarantees it via the URL; the embedded path via the stream→session
+        // binding — but reject a mismatched id so a stray payload can't overwrite the view.
+        if(sessionPayload.session_id&&sessionPayload.session_id!==activeSid) return false;
+        S.session=sessionPayload;
+        const _nextMsgs3018=(sessionPayload.messages||[]).filter(m=>m&&m.role);
+        _attachProjectedAnchorSceneToLastAssistant(_nextMsgs3018);
+        S.messages=_carryForwardEphemeralTurnFields(S.messages||[], _nextMsgs3018);
+        if(typeof _hydrateTodosFromSession==='function') _hydrateTodosFromSession(S.session);
+        clearLiveToolCards();if(!assistantText)removeThinking();
+        _markSessionViewed(activeSid, sessionPayload.message_count ?? S.messages.length);
+        renderMessages({preserveScroll:true});
+        return true;
+      };
+      // Prefer the canonical session snapshot embedded in the terminal cancel event.
+      // It includes _partial reasoning/tool rows captured by cancel_stream(), avoiding
+      // a second GET race where the visible cancelled work briefly collapses to only
+      // the fallback "Task cancelled" marker (#4076).
+      const _cancelSessionPayload=_cancelData&&typeof _cancelData.session==='object'?_cancelData.session:null;
       (async()=>{
         try{
+          if(_applyCancelSessionPayload(_cancelSessionPayload)) return;
+          // Fetch latest session from server to get accurate message list (includes cancel status)
+          // This ensures messages stay in sync with server, fixing race condition where local
+          // "*Task cancelled.*" message gets lost when done event overwrites S.messages
           const data=await api(`/api/session?session_id=${encodeURIComponent(activeSid)}`);
-          if(data&&data.session&&S.session&&S.session.session_id===activeSid){
-            S.session=data.session;
-            const _nextMsgs3018=(data.session.messages||[]).filter(m=>m&&m.role);
-            _attachProjectedAnchorSceneToLastAssistant(_nextMsgs3018);
-            S.messages=_carryForwardEphemeralTurnFields(S.messages||[], _nextMsgs3018);
-            if(typeof _hydrateTodosFromSession==='function') _hydrateTodosFromSession(S.session);
-            clearLiveToolCards();if(!assistantText)removeThinking();
-            _markSessionViewed(activeSid, data.session.message_count ?? S.messages.length);
-            renderMessages({preserveScroll:true});
-          }
+          if(data&&data.session) _applyCancelSessionPayload(data.session);
         }catch(_){
           // Fallback to local cancel message if API fails
           if(S.session&&S.session.session_id===activeSid){

tests/test_cancelled_turn_status.py

@@ -207,3 +207,32 @@ class TestCancelledTurnPersistenceGuards:
         assert "No response from provider" in block
         assert "Cancellation details" in block
         assert "Interruption details" in block
+
+    def test_frontend_cancel_prefers_embedded_session_payload(self):
+        src = _read("static/messages.js")
+        start = src.find("source.addEventListener('cancel'")
+        end = src.find("for(const _runJournalEventName", start)
+        assert start != -1 and end != -1, "cancel handler not found"
+        block = src[start:end]
+
+        assert "const _applyCancelSessionPayload=(sessionPayload)=>" in block
+        assert "const _cancelSessionPayload=_cancelData&&typeof _cancelData.session==='object'?_cancelData.session:null;" in block
+        assert "if(_applyCancelSessionPayload(_cancelSessionPayload)) return;" in block
+        assert "const data=await api(`/api/session?session_id=${encodeURIComponent(activeSid)}`);" in block
+        assert block.index("if(_applyCancelSessionPayload(_cancelSessionPayload)) return;") < block.index("const data=await api("), (
+            "Cancel handler must apply the terminal SSE session payload before falling back "
+            "to /api/session so captured _partial reasoning/tool rows are visible immediately."
+        )
+
+    def test_worker_cancel_events_do_not_embed_session_payload(self):
+        src = _read("api/streaming.py")
+        worker_start = src.find("def _run_agent_streaming(")
+        cancel_stream_start = src.find("def cancel_stream(", worker_start)
+        assert worker_start != -1 and cancel_stream_start != -1, "streaming worker/cancel_stream not found"
+        worker_block = src[worker_start:cancel_stream_start]
+        cancel_stream_block = src[cancel_stream_start:]
+
+        assert "_cancel_event_payload('Cancelled by user', s)" not in worker_block
+        assert "_cancel_event_payload('Cancelled by user', session=" not in worker_block
+        assert "None if ephemeral else s" not in worker_block
+        assert "_cancel_event_payload('Cancelled by user', session=_cancel_session_payload)" in cancel_stream_block

tests/test_issue1361_cancel_data_loss.py

@@ -297,6 +297,50 @@ class TestCancelWithReasoningOnlyNoText:
         assert len(partial_msgs) > 0, \
             f"Expected at least one partial assistant msg for tools-only cancel. Got: {assistant_msgs}"
 
+    def test_cancel_event_payload_includes_partial_session_snapshot(self):
+        """#4076: the terminal cancel SSE must carry the persisted partial turn.
+
+        The browser should not need a second /api/session request before it can
+        render reasoning/tool-only partial work captured by cancel_stream().
+        """
+        sid = "test_4076_cancel_payload"
+        stream_id = "stream_4076_payload"
+        _make_session(session_id=sid)
+        _setup_cancel_state(sid, stream_id)
+        q = config.STREAMS[stream_id]
+
+        config.STREAM_PARTIAL_TEXT[stream_id] = ""
+        assert hasattr(config, "STREAM_REASONING_TEXT"), (
+            "cancel payload test requires reasoning capture"
+        )
+        assert hasattr(config, "STREAM_LIVE_TOOL_CALLS"), (
+            "cancel payload test requires live tool capture"
+        )
+        config.STREAM_REASONING_TEXT[stream_id] = "Important cancelled reasoning"
+        config.STREAM_LIVE_TOOL_CALLS[stream_id] = [
+            {"name": "terminal", "args": {"command": "pytest -q"}, "done": False},
+        ]
+
+        cancel_stream(stream_id)
+
+        event_type, payload = q.get_nowait()
+        assert event_type == "cancel"
+        assert payload["type"] == "cancelled"
+        assert payload["status"] == "cancelled"
+        assert payload["session_id"] == sid
+        session_payload = payload["session"]
+        partials = [
+            m for m in session_payload["messages"]
+            if isinstance(m, dict) and m.get("_partial")
+        ]
+        assert partials, f"Expected persisted partial in cancel SSE payload: {session_payload}"
+        assert partials[-1]["reasoning"] == "Important cancelled reasoning"
+        assert partials[-1]["_partial_tool_calls"][0]["name"] == "terminal"
+        assert any(
+            isinstance(m, dict) and m.get("_error") and "Task cancelled" in str(m.get("content") or "")
+            for m in session_payload["messages"]
+        )
+
     def test_no_reasoning_no_tools_no_partial(self):
         """Cancel with no reasoning and no tools and no text = only cancel marker (no change)."""
         sid = "test_1361_c3"

tests/test_live_to_final_anchor_visible_order.py

@@ -620,12 +620,17 @@ def test_stream_end_restore_attaches_projected_anchor_scene_before_render():
 def test_cancel_settlement_attaches_projected_anchor_scene_before_render():
     cancel = _event_listener_body(MESSAGES_JS, "cancel")
 
-    fetch_idx = cancel.index("const _nextMsgs3018=(data.session.messages||[]).filter(m=>m&&m.role);")
+    fetch_idx = cancel.index("const _nextMsgs3018=(sessionPayload.messages||[]).filter(m=>m&&m.role);")
     attach_idx = cancel.index("_attachProjectedAnchorSceneToLastAssistant(_nextMsgs3018);")
     carry_idx = cancel.index("S.messages=_carryForwardEphemeralTurnFields(S.messages||[], _nextMsgs3018);")
     render_idx = cancel.index("renderMessages({preserveScroll:true});")
     assert fetch_idx < attach_idx < carry_idx < render_idx
 
+    embedded_idx = cancel.index("if(_applyCancelSessionPayload(_cancelSessionPayload)) return;")
+    fallback_get_idx = cancel.index("const data=await api(`/api/session?session_id=${encodeURIComponent(activeSid)}`);")
+    fallback_apply_idx = cancel.index("if(data&&data.session) _applyCancelSessionPayload(data.session);")
+    assert embedded_idx < fallback_get_idx < fallback_apply_idx
+
     fallback_push_idx = cancel.index("S.messages.push({role:'assistant',content:`**Task cancelled:**")
     fallback_attach_idx = cancel.index("_attachProjectedAnchorSceneToLastAssistant(S.messages);", fallback_push_idx)
     fallback_render_idx = cancel.index("renderMessages({preserveScroll:true});", fallback_attach_idx)

tests/test_session_lost_response_regression.py

@@ -1157,7 +1157,3 @@ def test_sync_revalidates_against_concurrent_disk_write_under_lock(monkeypatch):
     # On-disk record keeps the concurrent writer's value — not overwritten.
     reloaded = Session.load(sid)
     assert reloaded.active_stream_id == "rotated_stream_after_compression"
-
-
-
-