Live Stream: add compression recovery action

tests/test_compression_recovery_action.py

@@ -0,0 +1,248 @@
+import io
+import json
+from pathlib import Path
+
+from api import models, routes
+from api.compression_recovery import (
+    compression_recovery_payload_for_session,
+    is_generic_continuation_intent,
+    stamp_compression_exhausted_recovery,
+)
+from api.models import Session
+
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+class _JSONHandler:
+    headers = {}
+
+    def __init__(self):
+        self.status = None
+        self.wfile = io.BytesIO()
+        self.headers_sent = {}
+
+    def send_response(self, status):
+        self.status = status
+
+    def send_header(self, key, value):
+        self.headers_sent[key] = value
+
+    def end_headers(self):
+        pass
+
+
+def _payload(handler):
+    raw = handler.wfile.getvalue().decode("utf-8")
+    return json.loads(raw) if raw else {}
+
+
+def _isolate_sessions(monkeypatch, tmp_path):
+    session_dir = tmp_path / "sessions"
+    session_dir.mkdir()
+    monkeypatch.setattr(models, "SESSION_DIR", session_dir)
+    monkeypatch.setattr(models, "SESSION_INDEX_FILE", session_dir / "_index.json")
+    models.SESSIONS.clear()
+    routes.SESSIONS.clear()
+    return session_dir
+
+
+def test_generic_continuation_intent_is_scoped_to_empty_continue_words():
+    assert is_generic_continuation_intent("continue")
+    assert is_generic_continuation_intent("继续吧。")
+    assert is_generic_continuation_intent("go on")
+    assert not is_generic_continuation_intent("continue by summarizing the workspace changes")
+    assert not is_generic_continuation_intent("继续修复 4685 的恢复卡")
+
+
+def test_chat_start_blocks_generic_continue_after_compression_exhausted(monkeypatch, tmp_path):
+    _isolate_sessions(monkeypatch, tmp_path)
+    sid = "recoverychat1"
+    session = Session(
+        session_id=sid,
+        title="Recovery",
+        workspace=str(tmp_path),
+        model="gpt-4o",
+        messages=[{"role": "user", "content": "long task"}],
+    )
+    stamp_compression_exhausted_recovery(session, message="Context length exceeded.")
+    models.SESSIONS[sid] = session
+    routes.SESSIONS[sid] = session
+
+    handler = _JSONHandler()
+    routes._handle_chat_start(handler, {"session_id": sid, "message": "继续"})
+    payload = _payload(handler)
+
+    assert handler.status == 409
+    assert payload["type"] == "compression_recovery_required"
+    assert payload["recommended_recovery_action"] == "start_focused_continuation"
+    assert payload["compression_recovery"]["terminal_state"] == "compression_exhausted"
+
+
+def test_chat_start_keeps_recovery_when_substantive_prompt_fails_validation(monkeypatch, tmp_path):
+    session_dir = _isolate_sessions(monkeypatch, tmp_path)
+    sid = "recoverychat2"
+    session = Session(
+        session_id=sid,
+        title="Recovery",
+        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
+    monkeypatch.setattr(
+        routes,
+        "_resolve_chat_workspace_with_recovery",
+        lambda *_args, **_kwargs: (_ for _ in ()).throw(ValueError("bad workspace")),
+    )
+
+    handler = _JSONHandler()
+    routes._handle_chat_start(handler, {"session_id": sid, "message": "continue by checking the repo"})
+    saved = json.loads((session_dir / f"{sid}.json").read_text(encoding="utf-8"))
+
+    assert handler.status == 400
+    assert saved["recommended_recovery_action"] == "start_focused_continuation"
+    assert saved["compression_recovery"]["terminal_state"] == "compression_exhausted"
+
+
+def test_chat_start_clears_recovery_when_substantive_prompt_starts(monkeypatch, tmp_path):
+    session_dir = _isolate_sessions(monkeypatch, tmp_path)
+    sid = "recoverychat3"
+    session = Session(
+        session_id=sid,
+        title="Recovery",
+        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
+    monkeypatch.setattr(routes, "_resolve_chat_workspace_with_recovery", lambda *_args, **_kwargs: str(tmp_path))
+    monkeypatch.setattr(routes, "_read_profile_model_config", lambda *_args, **_kwargs: (None, None, {}))
+    monkeypatch.setattr(
+        routes,
+        "_resolve_compatible_session_model_state",
+        lambda requested_model, requested_provider, **_kwargs: (requested_model or "gpt-4o", requested_provider or "openai", False),
+    )
+    monkeypatch.setattr(routes, "webui_gateway_chat_enabled", lambda _config: False)
+
+    def _fake_start_run(run_session, **_kwargs):
+        assert compression_recovery_payload_for_session(run_session) is None
+        run_session.save()
+        return {"session_id": sid, "stream_id": "stream1", "_status": 200}
+
+    monkeypatch.setattr(routes, "_start_run", _fake_start_run)
+
+    handler = _JSONHandler()
+    routes._handle_chat_start(handler, {"session_id": sid, "message": "continue by checking the repo"})
+    payload = _payload(handler)
+    saved = json.loads((session_dir / f"{sid}.json").read_text(encoding="utf-8"))
+
+    assert handler.status == 200
+    assert payload["stream_id"] == "stream1"
+    assert saved["compression_recovery"] == {}
+    assert saved["recommended_recovery_action"] is None
+
+
+def test_chat_start_restores_recovery_when_substantive_prompt_start_is_rejected(monkeypatch, tmp_path):
+    session_dir = _isolate_sessions(monkeypatch, tmp_path)
+    sid = "recoverychat4"
+    session = Session(
+        session_id=sid,
+        title="Recovery",
+        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
+    monkeypatch.setattr(routes, "_resolve_chat_workspace_with_recovery", lambda *_args, **_kwargs: str(tmp_path))
+    monkeypatch.setattr(routes, "_read_profile_model_config", lambda *_args, **_kwargs: (None, None, {}))
+    monkeypatch.setattr(
+        routes,
+        "_resolve_compatible_session_model_state",
+        lambda requested_model, requested_provider, **_kwargs: (requested_model or "gpt-4o", requested_provider or "openai", False),
+    )
+    monkeypatch.setattr(routes, "webui_gateway_chat_enabled", lambda _config: False)
+
+    def _fake_start_run(run_session, **_kwargs):
+        assert compression_recovery_payload_for_session(run_session) is None
+        run_session.save()
+        return {"error": "session already has an active stream", "_status": 409}
+
+    monkeypatch.setattr(routes, "_start_run", _fake_start_run)
+
+    handler = _JSONHandler()
+    routes._handle_chat_start(handler, {"session_id": sid, "message": "continue by checking the repo"})
+    saved = json.loads((session_dir / f"{sid}.json").read_text(encoding="utf-8"))
+
+    assert handler.status == 409
+    assert saved["recommended_recovery_action"] == "start_focused_continuation"
+    assert saved["compression_recovery"]["terminal_state"] == "compression_exhausted"
+
+
+def test_recovery_start_creates_focused_linked_session(monkeypatch, tmp_path):
+    session_dir = _isolate_sessions(monkeypatch, tmp_path)
+    sid = "recoverysrc1"
+    session = Session(
+        session_id=sid,
+        title="Long task",
+        workspace=str(tmp_path),
+        model="gpt-4o",
+        model_provider="openai",
+        profile="default",
+        project_id="proj_1",
+        messages=[{"role": "user", "content": "long task"}],
+        context_messages=[{"role": "user", "content": "large context"}],
+    )
+    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)
+
+    assert handler.status == 200
+    new_session = payload["session"]
+    assert new_session["session_id"] != sid
+    assert new_session["parent_session_id"] == sid
+    assert new_session["workspace"] == str(tmp_path)
+    assert new_session["model"] == "gpt-4o"
+    assert new_session["model_provider"] == "openai"
+    assert new_session["messages"] == []
+
+    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 compression_recovery_payload_for_session(session)["recommended_action"] == "start_focused_continuation"
+
+
+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.")
+    compact = session.compact()
+
+    assert recovery["terminal_state"] == "compression_exhausted"
+    assert compact["recommended_recovery_action"] == "start_focused_continuation"
+    assert compact["compression_recovery"]["recommended_action"] == "start_focused_continuation"
+
+
+def test_compression_recovery_ui_wires_card_action_and_send_intercept():
+    ui = (ROOT / "static/ui.js").read_text(encoding="utf-8")
+    messages = (ROOT / "static/messages.js").read_text(encoding="utf-8")
+
+    assert "function _compressionRecoveryHtml" in ui
+    assert "data-compression-recovery-card=\"1\"" in ui
+    assert "api('/api/session/compression-recovery/start'" in ui
+    assert "function shouldInterceptCompressionRecoveryContinuation" in ui
+    assert "shouldInterceptCompressionRecoveryContinuation(text,S.pendingFiles)" in messages
+    assert "_compressionRecovery:recovery||undefined" in messages