From 1e9d42c0538401a9e9c7540fa5329ca6811cdb55 Mon Sep 17 00:00:00 2001 From: Frank Song Date: Sat, 4 Jul 2026 22:30:40 +0800 Subject: [PATCH] Live Stream: harden compression recovery action Make focused compression-recovery continuation creation idempotent by marking child sessions and reusing the existing child on repeated start requests. Also surface malformed successful recovery responses in the UI instead of silently re-enabling the action. Co-Authored-By: Paperclip --- api/models.py | 88 +++++++++++++++++ api/routes.py | 110 +++++++++++++--------- docs/rfcs/session-sse-contract-v1.md | 30 +++--- static/ui.js | 19 ++-- tests/test_compression_recovery_action.py | 50 ++++++++++ 5 files changed, 225 insertions(+), 72 deletions(-) diff --git a/api/models.py b/api/models.py index 79f969c9..613cccee 100644 --- a/api/models.py +++ b/api/models.py @@ -1053,6 +1053,8 @@ class Session: last_prompt_tokens=None, compression_recovery=None, recommended_recovery_action=None, + compression_recovery_source_session_id=None, + compression_recovery_action=None, truncation_watermark=None, truncation_boundary=None, gateway_routing=None, gateway_routing_history=None, @@ -1106,6 +1108,16 @@ class Session: self.last_prompt_tokens = last_prompt_tokens self.compression_recovery = compression_recovery if isinstance(compression_recovery, dict) else {} self.recommended_recovery_action = recommended_recovery_action + self.compression_recovery_source_session_id = ( + str(compression_recovery_source_session_id).strip() + if compression_recovery_source_session_id + else None + ) + self.compression_recovery_action = ( + str(compression_recovery_action).strip() + if compression_recovery_action + else None + ) self.truncation_watermark = truncation_watermark self.truncation_boundary = truncation_boundary self.gateway_routing = gateway_routing if isinstance(gateway_routing, dict) else None @@ -1177,6 +1189,7 @@ class Session: 'compression_anchor_details', 'context_engine_state', 'context_length', 'threshold_tokens', 'last_prompt_tokens', 'compression_recovery', 'recommended_recovery_action', + 'compression_recovery_source_session_id', 'compression_recovery_action', 'truncation_watermark', 'truncation_boundary', 'gateway_routing', 'gateway_routing_history', 'llm_title_generated', 'manual_title', @@ -1426,6 +1439,10 @@ class Session: 'last_prompt_tokens': self.last_prompt_tokens, 'compression_recovery': self.compression_recovery, 'recommended_recovery_action': self.recommended_recovery_action, + **({ + 'compression_recovery_source_session_id': self.compression_recovery_source_session_id, + 'compression_recovery_action': self.compression_recovery_action, + } if (self.compression_recovery_source_session_id or self.compression_recovery_action) else {}), 'gateway_routing': self.gateway_routing, 'gateway_routing_history': self.gateway_routing_history, 'manual_title': self.manual_title, @@ -3273,6 +3290,77 @@ def get_session(sid, metadata_only=False): return s raise KeyError(sid) + +def _compression_recovery_child_matches(session, source_session_id: str, action: str) -> bool: + return ( + str(getattr(session, "compression_recovery_source_session_id", "") or "").strip() == source_session_id + and str(getattr(session, "compression_recovery_action", "") or "").strip() == action + ) + + +def find_compression_recovery_session(source_session_id: str, action: str): + """Return an existing focused recovery child for ``source_session_id``. + + The recovery-start endpoint is a retryable UI action. A persisted marker on + the child session makes double-clicks, repeated calls, and cache reloads + converge on the same continuation instead of creating duplicate siblings. + """ + + source_sid = str(source_session_id or "").strip() + recovery_action = str(action or "").strip() + if not source_sid or not recovery_action: + return None + + matches = [] + seen_ids: set[str] = set() + try: + with LOCK: + memory_sessions = list(SESSIONS.values()) + for session in memory_sessions: + sid = str(getattr(session, "session_id", "") or "").strip() + if sid: + seen_ids.add(sid) + if _compression_recovery_child_matches(session, source_sid, recovery_action): + matches.append(session) + except Exception: + logger.debug("Failed to scan cached compression recovery sessions", exc_info=True) + + try: + persisted_ids = _persisted_session_ids_snapshot() + except Exception: + persisted_ids = frozenset() + for sid in persisted_ids: + if sid in seen_ids: + continue + try: + meta = Session.load_metadata_only(sid) + except Exception: + logger.debug("Failed to inspect compression recovery session %s", sid, exc_info=True) + continue + if not meta or not _compression_recovery_child_matches(meta, source_sid, recovery_action): + continue + try: + matches.append(get_session(sid)) + except Exception: + matches.append(meta) + + if not matches: + return None + + def _sort_key(session): + try: + created_at = float(getattr(session, "created_at", 0) or 0) + except (TypeError, ValueError): + created_at = 0.0 + try: + updated_at = float(getattr(session, "updated_at", 0) or 0) + except (TypeError, ValueError): + updated_at = 0.0 + return (created_at, updated_at, str(getattr(session, "session_id", "") or "")) + + return sorted(matches, key=_sort_key)[0] + + def _profile_default_model_state(profile=None): """Return the default model/provider configured for *profile*.""" default_model = "" diff --git a/api/routes.py b/api/routes.py index 38d7fb72..46f61ac1 100644 --- a/api/routes.py +++ b/api/routes.py @@ -8581,6 +8581,7 @@ def _keep_latest_messaging_session_per_source( from api.models import ( Session, get_session, + find_compression_recovery_session, get_session_for_file_ops, new_session, all_sessions, @@ -8622,6 +8623,9 @@ from api.models import ( ) +_COMPRESSION_RECOVERY_START_LOCK = threading.Lock() + + def _pre_compression_continuation_session_id(session) -> str | None: """Return the newest visible descendant for a hidden compression snapshot. @@ -19756,53 +19760,61 @@ def _handle_session_compression_recovery_start(handler, body): if action != COMPRESSION_RECOVERY_ACTION_START_FOCUSED: return bad(handler, "Unsupported compression recovery action.", 409) - title = str(getattr(source, "title", None) or "Untitled").strip() or "Untitled" - if not title.endswith(" (focused continuation)"): - title = f"{title} (focused continuation)" - copied_session = Session( - session_id=uuid.uuid4().hex[:12], - title=title, - workspace=getattr(source, "workspace", get_last_workspace()), - model=getattr(source, "model", None), - model_provider=getattr(source, "model_provider", None), - messages=[], - tool_calls=[], - pinned=False, - archived=False, - project_id=getattr(source, "project_id", None), - profile=getattr(source, "profile", None), - personality=getattr(source, "personality", None), - enabled_toolsets=copy.deepcopy(getattr(source, "enabled_toolsets", None)), - context_length=getattr(source, "context_length", None), - threshold_tokens=getattr(source, "threshold_tokens", None), - gateway_routing=copy.deepcopy(getattr(source, "gateway_routing", None)), - gateway_routing_history=copy.deepcopy(getattr(source, "gateway_routing_history", None) or []), - parent_session_id=getattr(source, "session_id", sid), - worktree_path=getattr(source, "worktree_path", None), - worktree_branch=getattr(source, "worktree_branch", None), - worktree_repo_root=getattr(source, "worktree_repo_root", None), - worktree_created_at=getattr(source, "worktree_created_at", None), - ) - # Preserve the workspace/model/profile lane, but intentionally start with an - # empty model-facing transcript so a focused follow-up does not replay the - # exhausted state.db/context tail. - copied_session.context_messages = [] - copied_session.composer_draft = {"text": "", "files": []} - try: - copied_session.save() - except Exception as e: - logger.exception("failed to persist compression recovery session for %s", sid) - return bad(handler, f"Failed to start compression recovery: {_sanitize_error(e)}", 500) + created = False + with _COMPRESSION_RECOVERY_START_LOCK: + copied_session = find_compression_recovery_session(sid, action) + if copied_session is None: + title = str(getattr(source, "title", None) or "Untitled").strip() or "Untitled" + if not title.endswith(" (focused continuation)"): + title = f"{title} (focused continuation)" + copied_session = Session( + session_id=uuid.uuid4().hex[:12], + title=title, + workspace=getattr(source, "workspace", get_last_workspace()), + model=getattr(source, "model", None), + model_provider=getattr(source, "model_provider", None), + messages=[], + tool_calls=[], + pinned=False, + archived=False, + project_id=getattr(source, "project_id", None), + profile=getattr(source, "profile", None), + personality=getattr(source, "personality", None), + enabled_toolsets=copy.deepcopy(getattr(source, "enabled_toolsets", None)), + context_length=getattr(source, "context_length", None), + threshold_tokens=getattr(source, "threshold_tokens", None), + gateway_routing=copy.deepcopy(getattr(source, "gateway_routing", None)), + gateway_routing_history=copy.deepcopy(getattr(source, "gateway_routing_history", None) or []), + parent_session_id=getattr(source, "session_id", sid), + worktree_path=getattr(source, "worktree_path", None), + worktree_branch=getattr(source, "worktree_branch", None), + worktree_repo_root=getattr(source, "worktree_repo_root", None), + worktree_created_at=getattr(source, "worktree_created_at", None), + compression_recovery_source_session_id=sid, + compression_recovery_action=action, + ) + # Preserve the workspace/model/profile lane, but intentionally start with an + # empty model-facing transcript so a focused follow-up does not replay the + # exhausted state.db/context tail. + copied_session.context_messages = [] + copied_session.composer_draft = {"text": "", "files": []} + try: + copied_session.save() + except Exception as e: + logger.exception("failed to persist compression recovery session for %s", sid) + return bad(handler, f"Failed to start compression recovery: {_sanitize_error(e)}", 500) - with LOCK: - SESSIONS[copied_session.session_id] = copied_session - SESSIONS.move_to_end(copied_session.session_id) - _evict_sessions_over_cap() - publish_session_list_changed( - "session_compression_recovery", - profile=getattr(copied_session, "profile", None), - session_id=getattr(copied_session, "session_id", None), - ) + with LOCK: + SESSIONS[copied_session.session_id] = copied_session + SESSIONS.move_to_end(copied_session.session_id) + _evict_sessions_over_cap() + created = True + if created: + publish_session_list_changed( + "session_compression_recovery", + profile=getattr(copied_session, "profile", None), + session_id=getattr(copied_session, "session_id", None), + ) session_payload = redact_session_data(copied_session.compact() | {"messages": copied_session.messages}) return j( handler, @@ -19811,7 +19823,11 @@ def _handle_session_compression_recovery_start(handler, body): "session": session_payload, "source_session_id": sid, "recommended_recovery_action": action, - "message": "Started a focused continuation. Describe the next narrow task to continue.", + "message": ( + "Started a focused continuation. Describe the next narrow task to continue." + if created + else "Opened the existing focused continuation for this exhausted session." + ), }, ) diff --git a/docs/rfcs/session-sse-contract-v1.md b/docs/rfcs/session-sse-contract-v1.md index 7eefb9d2..1b4431fa 100644 --- a/docs/rfcs/session-sse-contract-v1.md +++ b/docs/rfcs/session-sse-contract-v1.md @@ -38,8 +38,8 @@ against current source before any route is added. - This RFC does **not** implement `GET /api/sessions/{session_id}/events`. No route, handler, or related code is added in this PR. - This RFC does **not** modify `GET /api/sessions/events` (the existing global - session-list invalidation stream routed at `api/routes.py:12351-12352` and - implemented by `_handle_session_events_stream()` at `api/routes.py:16186`). + session-list invalidation stream routed at `api/routes.py:12355-12356` and + implemented by `_handle_session_events_stream()` at `api/routes.py:16190`). - This RFC does **not** replace or modify existing streams: `/api/chat/stream`, `/api/approval/stream`, or `/api/clarify/stream`. - This RFC does **not** introduce Android, iOS, or PWA client code. @@ -53,15 +53,15 @@ against current source before any route is added. ### Existing global session-list stream `GET /api/sessions/events` is a **different endpoint** from the one this RFC -proposes. It is routed at `api/routes.py:12351-12352` and implemented by -`_handle_session_events_stream()` at `api/routes.py:16186`. It emits bare +proposes. It is routed at `api/routes.py:12355-12356` and implemented by +`_handle_session_events_stream()` at `api/routes.py:16190`. It emits bare `sessions_changed` events and keepalives for any change to the session list. It is a global invalidation signal, not a per-session lifecycle stream. The proposed `GET /api/sessions/{session_id}/events` is per-session and path-distinct. ### Heartbeat -`_SSE_HEARTBEAT_INTERVAL_SECONDS = 5` (`api/routes.py:1018-1030`) is the current +`_SSE_HEARTBEAT_INTERVAL_SECONDS = 5` (`api/routes.py:1028-1036`) is the current heartbeat interval for SSE streams. Phase 1 reuses this constant rather than adding a separate configurable knob. @@ -73,19 +73,19 @@ Line ranges in this inventory were verified against WebUI `master` when this RFC was written. Function, constant, and endpoint names are the stable anchors if source layout moves later. -- `_parse_run_journal_event_id()` (`api/routes.py:15682-15695`) and - `_parse_run_journal_after_seq()` (`api/routes.py:15697-15710`) parse the replay +- `_parse_run_journal_event_id()` (`api/routes.py:15686-15698`) and + `_parse_run_journal_after_seq()` (`api/routes.py:15701-15713`) parse the replay cursor from the `after_event_id` / `after_seq` **query params** (not the `Last-Event-ID` header — that header is the *proposed* new-endpoint contract below, §Reconnect). -- `_runner_event_id()` at `api/routes.py:15774-15781` constructs the event `id` +- `_runner_event_id()` at `api/routes.py:15778-15785` constructs the event `id` field as `stream_id:seq`. - SSE frames carry their `id:` via the `_sse_with_id()` helper, emitted on the - live `/api/chat/stream` path at `api/routes.py:15927`, on the runner-observe - path at `api/routes.py:15820`, and during journal replay at - `api/routes.py:15730` / `15743`. + live `/api/chat/stream` path at `api/routes.py:15931`, on the runner-observe + path at `api/routes.py:15824`, and during journal replay at + `api/routes.py:15734` / `15747`. - `_replay_run_journal()` reads events by `(session_id, stream_id)` at - `api/routes.py:15712-15744`. + `api/routes.py:15716-15748`. - `api/streaming.py:6265-6285` writes current live agent streams to `STREAMS[stream_id]`. - `api/streaming.py:6620-6634` appends SSE events to the run journal and carries @@ -165,7 +165,7 @@ position. **`event_id` is opaque to clients.** Its current source-compatible form is `stream_id:seq`, as constructed by `_runner_event_id()` at -`api/routes.py:15774-15781`. Clients must treat it as an opaque string and must +`api/routes.py:15778-15785`. Clients must treat it as an opaque string and must not parse or construct cursor values. **`seq` is monotonic within a stream/run.** It is not a session-global counter @@ -183,7 +183,7 @@ events. The live `STREAMS[stream_id]` queue (`api/streaming.py:6265-6285`) is not a reliable replay source because it holds only recent in-memory state. A future implementation must replay from the run journal via the existing -`_replay_run_journal()` path (`api/routes.py:15712-15744`) and fall back to the +`_replay_run_journal()` path (`api/routes.py:15716-15748`) and fall back to the snapshot mechanism when journal entries are unavailable for a given cursor. ## Snapshot fallback @@ -201,7 +201,7 @@ resync from the snapshot payload. ## Heartbeat -Phase 1 reuses `_SSE_HEARTBEAT_INTERVAL_SECONDS` (`api/routes.py:1018-1030`) for +Phase 1 reuses `_SSE_HEARTBEAT_INTERVAL_SECONDS` (`api/routes.py:1028-1036`) for heartbeat cadence. A new per-session configurable heartbeat knob is **not** added in Phase 1. The implementation PR must follow whatever value the constant holds at implementation time; it must not hard-code a separate interval. diff --git a/static/ui.js b/static/ui.js index 666fb557..69ee61e7 100644 --- a/static/ui.js +++ b/static/ui.js @@ -483,16 +483,15 @@ async function startCompressionRecovery(btn){ try{ const data=await api('/api/session/compression-recovery/start',{method:'POST',body:JSON.stringify({session_id:sourceSid})}); const sid=data&&data.session&&data.session.session_id; - if(sid){ - try{localStorage.setItem('hermes-webui-session',sid);}catch(_){} - if(typeof loadSession==='function') await loadSession(sid,{preserveActiveInput:false}); - else if(data.session){S.session=data.session;S.messages=data.session.messages||[];syncTopbar();renderMessages();} - if(typeof renderSessionList==='function') await renderSessionList(); - if(typeof _setActiveSessionUrl==='function') _setActiveSessionUrl(sid); - if(typeof showToast==='function') showToast((data&&data.message)||'Started focused continuation.',3000,'success'); - const composer=$('msg'); - if(composer&&typeof composer.focus==='function') composer.focus(); - } + if(!sid) throw new Error('Compression recovery did not return a session.'); + try{localStorage.setItem('hermes-webui-session',sid);}catch(_){} + if(typeof loadSession==='function') await loadSession(sid,{preserveActiveInput:false}); + else if(data.session){S.session=data.session;S.messages=data.session.messages||[];syncTopbar();renderMessages();} + if(typeof renderSessionList==='function') await renderSessionList(); + if(typeof _setActiveSessionUrl==='function') _setActiveSessionUrl(sid); + if(typeof showToast==='function') showToast((data&&data.message)||'Started focused continuation.',3000,'success'); + const composer=$('msg'); + if(composer&&typeof composer.focus==='function') composer.focus(); }catch(e){ if(typeof showToast==='function') showToast('Compression recovery failed: '+(e&&e.message||e),5000,'error'); }finally{ diff --git a/tests/test_compression_recovery_action.py b/tests/test_compression_recovery_action.py index 0e232f6f..0a3a7c5b 100644 --- a/tests/test_compression_recovery_action.py +++ b/tests/test_compression_recovery_action.py @@ -223,9 +223,58 @@ def test_recovery_start_creates_focused_linked_session(monkeypatch, tmp_path): saved = json.loads((session_dir / f"{new_session['session_id']}.json").read_text(encoding="utf-8")) assert saved["parent_session_id"] == sid assert saved["context_messages"] == [] + assert saved["compression_recovery_source_session_id"] == sid + assert saved["compression_recovery_action"] == "start_focused_continuation" assert compression_recovery_payload_for_session(session)["recommended_action"] == "start_focused_continuation" +def test_recovery_start_reuses_existing_focused_session(monkeypatch, tmp_path): + session_dir = _isolate_sessions(monkeypatch, tmp_path) + sid = "recoverysrc2" + session = Session( + session_id=sid, + title="Long task", + workspace=str(tmp_path), + model="gpt-4o", + messages=[{"role": "user", "content": "long task"}], + ) + stamp_compression_exhausted_recovery(session, message="Context length exceeded.") + session.save() + models.SESSIONS[sid] = session + routes.SESSIONS[sid] = session + + first_handler = _JSONHandler() + routes._handle_session_compression_recovery_start(first_handler, {"session_id": sid}) + first_payload = _payload(first_handler) + first_child_id = first_payload["session"]["session_id"] + + second_handler = _JSONHandler() + routes._handle_session_compression_recovery_start(second_handler, {"session_id": sid}) + second_payload = _payload(second_handler) + + assert second_handler.status == 200 + assert second_payload["session"]["session_id"] == first_child_id + assert second_payload["message"].startswith("Opened the existing") + + models.SESSIONS.clear() + routes.SESSIONS.clear() + third_handler = _JSONHandler() + routes._handle_session_compression_recovery_start(third_handler, {"session_id": sid}) + third_payload = _payload(third_handler) + + assert third_handler.status == 200 + assert third_payload["session"]["session_id"] == first_child_id + + recovery_children = [] + for path in session_dir.glob("*.json"): + if path.name.startswith("_"): + continue + data = json.loads(path.read_text(encoding="utf-8")) + if isinstance(data, dict) and data.get("compression_recovery_source_session_id") == sid: + recovery_children.append(data) + assert len(recovery_children) == 1 + + def test_recovery_metadata_is_persisted_and_exposed_in_compact_session(): session = Session(session_id="recovermeta1", title="x", messages=[]) recovery = stamp_compression_exhausted_recovery(session, message="Context length exceeded.") @@ -243,6 +292,7 @@ def test_compression_recovery_ui_wires_card_action_and_send_intercept(): assert "function _compressionRecoveryHtml" in ui assert "data-compression-recovery-card=\"1\"" in ui assert "api('/api/session/compression-recovery/start'" in ui + assert "Compression recovery did not return a session." in ui assert "function shouldInterceptCompressionRecoveryContinuation" in ui assert "shouldInterceptCompressionRecoveryContinuation(text,S.pendingFiles)" in messages assert "_compressionRecovery:recovery||undefined" in messages