From 3c493a9e37014d5283a7c2ddb233e4744456726a Mon Sep 17 00:00:00 2001 From: nesquena-hermes Date: Sat, 27 Jun 2026 16:43:37 +0000 Subject: [PATCH] fix(cancel): preserve partial turn payload via embedded cancel snapshot (#4647) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- CHANGELOG.md | 2 + api/streaming.py | 51 +++++++++++++++---- static/messages.js | 39 +++++++++----- tests/test_cancelled_turn_status.py | 29 +++++++++++ tests/test_issue1361_cancel_data_loss.py | 44 ++++++++++++++++ ...test_live_to_final_anchor_visible_order.py | 7 ++- .../test_session_lost_response_regression.py | 4 -- 7 files changed, 148 insertions(+), 28 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6f7c5536..95475fdc 100644 --- a/CHANGELOG.md +++ b/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) diff --git a/api/streaming.py b/api/streaming.py index 69a90410..ed523f96 100644 --- a/api/streaming.py +++ b/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") diff --git a/static/messages.js b/static/messages.js index 287efc91..0d4a41ec 100644 --- a/static/messages.js +++ b/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){ diff --git a/tests/test_cancelled_turn_status.py b/tests/test_cancelled_turn_status.py index ab847558..bb7c3118 100644 --- a/tests/test_cancelled_turn_status.py +++ b/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 diff --git a/tests/test_issue1361_cancel_data_loss.py b/tests/test_issue1361_cancel_data_loss.py index dbd59cea..c327403b 100644 --- a/tests/test_issue1361_cancel_data_loss.py +++ b/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" diff --git a/tests/test_live_to_final_anchor_visible_order.py b/tests/test_live_to_final_anchor_visible_order.py index b01d96e4..3fa78e3b 100644 --- a/tests/test_live_to_final_anchor_visible_order.py +++ b/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) diff --git a/tests/test_session_lost_response_regression.py b/tests/test_session_lost_response_regression.py index fa022bec..f37da008 100644 --- a/tests/test_session_lost_response_regression.py +++ b/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" - - - -