From 81d5e897e95460bdc39e577ccb6f4b7f07f27832 Mon Sep 17 00:00:00 2001 From: Frank Song Date: Sun, 5 Jul 2026 07:46:43 +0800 Subject: [PATCH] fix: isolate compression recovery child sessions Co-Authored-By: Paperclip --- api/models.py | 28 ++++++-- api/routes.py | 9 ++- static/ui.js | 2 +- tests/test_compression_recovery_action.py | 79 +++++++++++++++++++++++ 4 files changed, 112 insertions(+), 6 deletions(-) diff --git a/api/models.py b/api/models.py index ad2b8bba..ad8308c4 100644 --- a/api/models.py +++ b/api/models.py @@ -3291,14 +3291,34 @@ def get_session(sid, metadata_only=False): raise KeyError(sid) -def _compression_recovery_child_matches(session, source_session_id: str, action: str) -> bool: +_COMPRESSION_RECOVERY_PROFILE_UNSET = object() + + +def _compression_recovery_child_matches( + session, + source_session_id: str, + action: str, + source_profile=_COMPRESSION_RECOVERY_PROFILE_UNSET, +) -> bool: + if source_profile is not _COMPRESSION_RECOVERY_PROFILE_UNSET: + try: + from api.profiles import _profiles_match + if not _profiles_match(getattr(session, "profile", None), source_profile): + return False + except Exception: + logger.debug("Failed to profile-check compression recovery session", exc_info=True) + return False 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): +def find_compression_recovery_session( + source_session_id: str, + action: str, + source_profile=_COMPRESSION_RECOVERY_PROFILE_UNSET, +): """Return an existing focused recovery child for ``source_session_id``. The recovery-start endpoint is a retryable UI action. A persisted marker on @@ -3320,7 +3340,7 @@ def find_compression_recovery_session(source_session_id: str, action: str): sid = str(getattr(session, "session_id", "") or "").strip() if sid: seen_ids.add(sid) - if _compression_recovery_child_matches(session, source_sid, recovery_action): + if _compression_recovery_child_matches(session, source_sid, recovery_action, source_profile): matches.append(session) except Exception: logger.debug("Failed to scan cached compression recovery sessions", exc_info=True) @@ -3337,7 +3357,7 @@ def find_compression_recovery_session(source_session_id: str, action: str): 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): + if not meta or not _compression_recovery_child_matches(meta, source_sid, recovery_action, source_profile): continue try: matches.append(get_session(sid)) diff --git a/api/routes.py b/api/routes.py index 97bd67e4..55250cf5 100644 --- a/api/routes.py +++ b/api/routes.py @@ -8123,6 +8123,11 @@ def _merged_webui_lineage_messages_for_display(session, messages=None) -> list: parent_id = str(getattr(session, "parent_session_id", "") or "").strip() if not parent_id: return primary_messages + if ( + str(getattr(session, "compression_recovery_source_session_id", "") or "").strip() + and str(getattr(session, "compression_recovery_action", "") or "").strip() + ): + return primary_messages source = str(getattr(session, "session_source", "") or "").strip().lower() relationship = str(getattr(session, "relationship_type", "") or "").strip().lower() if source == "fork" or relationship == "child_session": @@ -19851,7 +19856,8 @@ def _handle_session_compression_recovery_start(handler, body): created = False with _COMPRESSION_RECOVERY_START_LOCK: - copied_session = find_compression_recovery_session(sid, action) + source_profile = getattr(source, "profile", None) + copied_session = find_compression_recovery_session(sid, action, source_profile=source_profile) if copied_session is None: title = str(getattr(source, "title", None) or "Untitled").strip() or "Untitled" if not title.endswith(" (focused continuation)"): @@ -19868,6 +19874,7 @@ def _handle_session_compression_recovery_start(handler, body): archived=False, project_id=getattr(source, "project_id", None), profile=getattr(source, "profile", None), + session_source="fork", personality=getattr(source, "personality", None), enabled_toolsets=copy.deepcopy(getattr(source, "enabled_toolsets", None)), context_length=getattr(source, "context_length", None), diff --git a/static/ui.js b/static/ui.js index 535fbe5c..83097544 100644 --- a/static/ui.js +++ b/static/ui.js @@ -422,7 +422,7 @@ function _compressionRecoveryHtml(recovery, sessionId){ if(String(recovery.terminal_state||'')!=='compression_exhausted') return ''; const action=String(recovery.recommended_action||''); if(action!=='start_focused_continuation') return ''; - const sid=String(sessionId||recovery.source_session_id||''); + const sid=String(recovery.source_session_id||sessionId||''); const title=String(recovery.title||'Context compression exhausted'); const summary=String(recovery.summary||'Start a focused continuation, then describe the next narrow task.'); const actionLabel=String(recovery.action_label||'Start focused continuation'); diff --git a/tests/test_compression_recovery_action.py b/tests/test_compression_recovery_action.py index 9c536d93..502bc7d0 100644 --- a/tests/test_compression_recovery_action.py +++ b/tests/test_compression_recovery_action.py @@ -221,15 +221,45 @@ def test_recovery_start_creates_focused_linked_session(monkeypatch, tmp_path): assert new_session["model"] == "gpt-4o" assert new_session["model_provider"] == "openai" assert new_session["messages"] == [] + assert new_session["session_source"] == "fork" saved = json.loads((session_dir / f"{new_session['session_id']}.json").read_text(encoding="utf-8")) assert saved["parent_session_id"] == sid + assert saved["session_source"] == "fork" 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_child_does_not_merge_parent_transcript(monkeypatch, tmp_path): + _isolate_sessions(monkeypatch, tmp_path) + sid = "recoverysrcisolate" + session = Session( + session_id=sid, + title="Long task", + workspace=str(tmp_path), + model="gpt-4o", + messages=[ + {"role": "user", "content": "long task"}, + {"role": "assistant", "content": "compression exhausted"}, + ], + ) + stamp_compression_exhausted_recovery(session, message="Context length exceeded.") + session.save() + models.SESSIONS[sid] = session + routes.SESSIONS[sid] = session + + handler = _JSONHandler() + routes._handle_session_compression_recovery_start(handler, {"session_id": sid}) + payload = _payload(handler) + child_id = payload["session"]["session_id"] + child = models.SESSIONS[child_id] + + assert child.messages == [] + assert routes._merged_webui_lineage_messages_for_display(child) == [] + + def test_recovery_start_reuses_existing_focused_session(monkeypatch, tmp_path): session_dir = _isolate_sessions(monkeypatch, tmp_path) sid = "recoverysrc2" @@ -277,6 +307,54 @@ def test_recovery_start_reuses_existing_focused_session(monkeypatch, tmp_path): assert len(recovery_children) == 1 +def test_recovery_start_ignores_existing_child_from_other_profile(monkeypatch, tmp_path): + session_dir = _isolate_sessions(monkeypatch, tmp_path) + sid = "recoverysrcprofile" + source = Session( + session_id=sid, + title="Long task", + workspace=str(tmp_path), + model="gpt-4o", + profile="default", + messages=[{"role": "user", "content": "long task"}], + ) + stamp_compression_exhausted_recovery(source, message="Context length exceeded.") + source.save() + foreign_child = Session( + session_id="foreignchild1", + title="Foreign focused continuation", + workspace=str(tmp_path), + model="gpt-4o", + profile="other-profile", + messages=[], + parent_session_id=sid, + compression_recovery_source_session_id=sid, + compression_recovery_action="start_focused_continuation", + ) + foreign_child.save() + models.SESSIONS.clear() + routes.SESSIONS.clear() + models.SESSIONS[sid] = source + routes.SESSIONS[sid] = source + + handler = _JSONHandler() + routes._handle_session_compression_recovery_start(handler, {"session_id": sid}) + payload = _payload(handler) + + assert handler.status == 200 + assert payload["session"]["session_id"] != "foreignchild1" + assert payload["session"]["profile"] == "default" + + 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 {child["profile"] for child in recovery_children} == {"default", "other-profile"} + + 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.") @@ -320,6 +398,7 @@ def test_compression_recovery_ui_wires_card_action_and_send_intercept(): 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 "const sid=String(recovery.source_session_id||sessionId||'')" in ui assert "function shouldInterceptCompressionRecoveryContinuation" in ui assert "shouldInterceptCompressionRecoveryContinuation(text,S.pendingFiles)" in messages assert "_compressionRecovery:recovery||undefined" in messages