From f879fd6bc39dc08f4857f168cfe527ccb8e5c1b4 Mon Sep 17 00:00:00 2001
From: ai-ag2026 <261867348+ai-ag2026@users.noreply.github.com>
Date: Thu, 28 May 2026 08:19:49 +0200
Subject: [PATCH] fix: add dry-run discoverability safe repair

---
 CHANGELOG.md                                 |   3 +
 api/session_discoverability.py               | 196 ++++++++++++++++++-
 tests/test_session_discoverability_repair.py | 175 +++++++++++++++++
 3 files changed, 372 insertions(+), 2 deletions(-)
 create mode 100644 tests/test_session_discoverability_repair.py

diff --git a/CHANGELOG.md b/CHANGELOG.md
index a990d204abc..3ac7e23e1d5 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -3,6 +3,9 @@
 
 ## [Unreleased]
 
+### Added
+- Session discoverability audit now has a default-dry-run `--repair-safe` routine for deterministic cleanup: stale persisted WebUI-as-CLI flags can be cleared from sidecars/index entries, and messageful WebUI rows present only in `state.db` can be materialized into sidecars/index entries when `--apply --backup-dir <dir>` is explicitly provided.
+
 ## [v0.51.152] — 2026-05-28 — Release DX (stage-batch34 — single-PR optional gateway-backed browser chat)
 
 ### Added
diff --git a/api/session_discoverability.py b/api/session_discoverability.py
index 0f218561bc5..a04ca753517 100644
--- a/api/session_discoverability.py
+++ b/api/session_discoverability.py
@@ -12,6 +12,8 @@
 
 import argparse
 import json
+import os
+import shutil
 import sqlite3
 from collections import Counter
 from pathlib import Path
@@ -373,6 +375,184 @@ def audit_session_discoverability(
     }
 
 
+def _atomic_write_json(path: Path, payload) -> None:
+    tmp = path.with_suffix(path.suffix + f".tmp.{os.getpid()}")
+    tmp.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
+    os.replace(tmp, path)
+
+
+def _backup_file(path: Path, backup_dir: Path, backed_up: dict[Path, str]) -> str | None:
+    if not path.exists():
+        return None
+    resolved = path.resolve()
+    if resolved in backed_up:
+        return backed_up[resolved]
+    backup_dir.mkdir(parents=True, exist_ok=True)
+    target = backup_dir / path.name
+    if target.exists():
+        stem = target.name
+        i = 1
+        while (backup_dir / f"{stem}.{i}").exists():
+            i += 1
+        target = backup_dir / f"{stem}.{i}"
+    shutil.copy2(path, target)
+    backed_up[resolved] = str(target)
+    return str(target)
+
+
+def _plan_discoverability_repairs(report: dict) -> list[dict]:
+    actions: list[dict] = []
+    for item in report.get("items") or []:
+        sid = str(item.get("session_id") or "")
+        if not sid:
+            continue
+        if item.get("kind") == "persisted_source_flag_stale":
+            if item.get("sidecar_is_cli_session") is True:
+                actions.append({"session_id": sid, "action": "clear_sidecar_cli_flag"})
+            if item.get("index_is_cli_session") is True:
+                actions.append({"session_id": sid, "action": "clear_index_cli_flag"})
+        elif item.get("kind") == "state_db_messageful_missing_sidecar":
+            actions.append({"session_id": sid, "action": "materialize_sidecar_from_state_db"})
+    return actions
+
+
+def _clear_sidecar_cli_flag(session_dir: Path, sid: str, backup_dir: Path, backed_up: dict[Path, str]) -> dict:
+    path = session_dir / f"{sid}.json"
+    payload = _read_json(path)
+    if not isinstance(payload, dict):
+        return {"session_id": sid, "action": "clear_sidecar_cli_flag", "applied": False, "error": "sidecar_unreadable"}
+    if not _webui_origin(payload):
+        return {"session_id": sid, "action": "clear_sidecar_cli_flag", "applied": False, "skipped": "not_webui_origin"}
+    if payload.get("is_cli_session") is not True:
+        return {"session_id": sid, "action": "clear_sidecar_cli_flag", "applied": False, "skipped": "already_clear"}
+    backup = _backup_file(path, backup_dir, backed_up)
+    payload["is_cli_session"] = False
+    _atomic_write_json(path, payload)
+    return {"session_id": sid, "action": "clear_sidecar_cli_flag", "applied": True, "backup": backup}
+
+
+def _clear_index_cli_flag(session_dir: Path, sid: str, backup_dir: Path, backed_up: dict[Path, str]) -> dict:
+    path = session_dir / "_index.json"
+    payload = _read_json(path)
+    if not isinstance(payload, list):
+        return {"session_id": sid, "action": "clear_index_cli_flag", "applied": False, "error": "index_unreadable"}
+    changed = False
+    for entry in payload:
+        if not isinstance(entry, dict):
+            continue
+        if str(entry.get("session_id") or "") != sid:
+            continue
+        if not _webui_origin(entry):
+            continue
+        if entry.get("is_cli_session") is True:
+            entry["is_cli_session"] = False
+            changed = True
+    if not changed:
+        return {"session_id": sid, "action": "clear_index_cli_flag", "applied": False, "skipped": "already_clear_or_missing"}
+    backup = _backup_file(path, backup_dir, backed_up)
+    _atomic_write_json(path, payload)
+    return {"session_id": sid, "action": "clear_index_cli_flag", "applied": True, "backup": backup}
+
+
+def _materialize_sidecar_from_state_db(session_dir: Path, state_db_path: Path | None, sid: str, backup_dir: Path, backed_up: dict[Path, str]) -> dict:
+    if state_db_path is None:
+        return {"session_id": sid, "action": "materialize_sidecar_from_state_db", "applied": False, "error": "state_db_required"}
+    target = session_dir / f"{sid}.json"
+    if target.exists():
+        return {"session_id": sid, "action": "materialize_sidecar_from_state_db", "applied": False, "skipped": "sidecar_exists"}
+    try:
+        from api.session_recovery import _read_state_db_missing_sidecar_rows, _state_db_row_to_sidecar
+    except Exception as exc:
+        return {"session_id": sid, "action": "materialize_sidecar_from_state_db", "applied": False, "error": f"recovery_import_failed:{exc}"}
+    rows = {str(row.get("id") or ""): row for row in _read_state_db_missing_sidecar_rows(session_dir, state_db_path)}
+    row = rows.get(sid)
+    if not row:
+        return {"session_id": sid, "action": "materialize_sidecar_from_state_db", "applied": False, "skipped": "state_row_not_repairable"}
+    payload = _state_db_row_to_sidecar(row)
+    _backup_file(state_db_path, backup_dir, backed_up)
+    session_dir.mkdir(parents=True, exist_ok=True)
+    tmp = target.with_suffix(target.suffix + f".tmp.{os.getpid()}")
+    tmp.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
+    try:
+        os.link(str(tmp), str(target))
+    except FileExistsError:
+        return {"session_id": sid, "action": "materialize_sidecar_from_state_db", "applied": False, "skipped": "sidecar_appeared_during_repair"}
+    finally:
+        try:
+            tmp.unlink(missing_ok=True)
+        except OSError:
+            pass
+    index_updated = False
+    index_path = session_dir / "_index.json"
+    index_payload = _read_json(index_path)
+    if not isinstance(index_payload, list):
+        index_payload = []
+    if not any(isinstance(entry, dict) and str(entry.get("session_id") or "") == sid for entry in index_payload):
+        _backup_file(index_path, backup_dir, backed_up)
+        index_entry = {key: value for key, value in payload.items() if key not in {"messages", "tool_calls"}}
+        index_payload.append(index_entry)
+        _atomic_write_json(index_path, index_payload)
+        index_updated = True
+    return {
+        "session_id": sid,
+        "action": "materialize_sidecar_from_state_db",
+        "applied": True,
+        "messages": len(payload.get("messages") or []),
+        "index_updated": index_updated,
+        "backup": str((backup_dir / state_db_path.name)) if (backup_dir / state_db_path.name).exists() else None,
+    }
+
+
+def repair_session_discoverability(
+    session_dir: Path,
+    state_db_path: Path | None = None,
+    *,
+    api_sessions: Iterable[dict] | None = None,
+    dry_run: bool = True,
+    backup_dir: Path | None = None,
+) -> dict:
+    """Plan or apply deterministic discoverability repairs.
+
+    Default mode is read-only. Applying mutations requires ``backup_dir`` and is
+    limited to stale persisted WebUI-as-CLI flags plus materializing WebUI
+    messageful sidecars from canonical state.db rows.
+    """
+    before = audit_session_discoverability(session_dir, state_db_path=state_db_path, api_sessions=api_sessions)
+    planned = _plan_discoverability_repairs(before)
+    if dry_run:
+        return {"ok": True, "dry_run": True, "planned": planned, "applied": [], "before": before, "after": before}
+    if backup_dir is None:
+        return {"ok": False, "dry_run": False, "error": "backup_dir_required_for_apply", "planned": planned, "applied": [], "before": before}
+
+    session_dir = Path(session_dir)
+    backup_dir = Path(backup_dir)
+    backed_up: dict[Path, str] = {}
+    applied: list[dict] = []
+    for action in planned:
+        sid = str(action.get("session_id") or "")
+        name = action.get("action")
+        try:
+            if name == "clear_sidecar_cli_flag":
+                applied.append(_clear_sidecar_cli_flag(session_dir, sid, backup_dir, backed_up))
+            elif name == "clear_index_cli_flag":
+                applied.append(_clear_index_cli_flag(session_dir, sid, backup_dir, backed_up))
+            elif name == "materialize_sidecar_from_state_db":
+                applied.append(_materialize_sidecar_from_state_db(session_dir, state_db_path, sid, backup_dir, backed_up))
+        except Exception as exc:
+            applied.append({"session_id": sid, "action": name, "applied": False, "error": str(exc)})
+    after = audit_session_discoverability(session_dir, state_db_path=state_db_path, api_sessions=api_sessions)
+    errors = [item for item in applied if item.get("error")]
+    return {
+        "ok": not errors,
+        "dry_run": False,
+        "planned": planned,
+        "applied": applied,
+        "backups": sorted(set(backed_up.values())),
+        "before": before,
+        "after": after,
+    }
+
+
 def render_discoverability_markdown(report: dict) -> str:
     lines = [
         "# WebUI Session Discoverability Audit",
@@ -416,11 +596,23 @@ def _main() -> int:
     parser.add_argument("--session-dir", type=Path, required=True)
     parser.add_argument("--state-db", type=Path, default=None)
     parser.add_argument("--format", choices=("json", "markdown"), default="json")
+    parser.add_argument("--repair-safe", action="store_true", help="Plan/apply deterministic discoverability repairs")
+    parser.add_argument("--apply", action="store_true", help="Apply --repair-safe changes; default is dry-run")
+    parser.add_argument("--backup-dir", type=Path, default=None, help="Required with --repair-safe --apply")
     parser.add_argument("--out", type=Path, default=None)
     args = parser.parse_args()
 
-    report = audit_session_discoverability(args.session_dir, state_db_path=args.state_db)
-    text = render_discoverability_markdown(report) if args.format == "markdown" else json.dumps(report, sort_keys=True)
+    if args.repair_safe:
+        report = repair_session_discoverability(
+            args.session_dir,
+            state_db_path=args.state_db,
+            dry_run=not args.apply,
+            backup_dir=args.backup_dir,
+        )
+        text = json.dumps(report, sort_keys=True)
+    else:
+        report = audit_session_discoverability(args.session_dir, state_db_path=args.state_db)
+        text = render_discoverability_markdown(report) if args.format == "markdown" else json.dumps(report, sort_keys=True)
     if args.out:
         args.out.parent.mkdir(parents=True, exist_ok=True)
         args.out.write_text(text, encoding="utf-8")
diff --git a/tests/test_session_discoverability_repair.py b/tests/test_session_discoverability_repair.py
new file mode 100644
index 00000000000..c43c9fc11b3
--- /dev/null
+++ b/tests/test_session_discoverability_repair.py
@@ -0,0 +1,175 @@
+import json
+import sqlite3
+import subprocess
+import sys
+from pathlib import Path
+
+from api.session_discoverability import repair_session_discoverability
+
+
+def _write_sidecar(session_dir: Path, sid: str, *, messages=1, **metadata):
+    payload = {
+        "session_id": sid,
+        "id": sid,
+        "title": metadata.pop("title", sid),
+        "messages": [{"role": "user", "content": f"message {i}"} for i in range(messages)],
+        **metadata,
+    }
+    path = session_dir / f"{sid}.json"
+    path.write_text(json.dumps(payload), encoding="utf-8")
+    return path
+
+
+def _write_index(session_dir: Path, *entries):
+    (session_dir / "_index.json").write_text(json.dumps(list(entries)), encoding="utf-8")
+
+
+def _state_db(session_dir: Path, rows, message_counts=None):
+    db = session_dir / "state.db"
+    message_counts = message_counts or {}
+    with sqlite3.connect(db) as conn:
+        conn.execute(
+            """
+            create table sessions (
+                id text primary key,
+                source text,
+                title text,
+                parent_session_id text,
+                message_count integer,
+                started_at real,
+                model text,
+                workspace text
+            )
+            """
+        )
+        conn.execute("create table messages (id integer primary key, session_id text, role text, content text, timestamp real)")
+        for row in rows:
+            conn.execute(
+                """
+                insert into sessions (id, source, title, parent_session_id, message_count, started_at, model, workspace)
+                values (?, ?, ?, ?, ?, ?, ?, ?)
+                """,
+                (
+                    row["id"],
+                    row.get("source"),
+                    row.get("title") or row["id"],
+                    row.get("parent_session_id"),
+                    row.get("message_count", message_counts.get(row["id"], 0)),
+                    row.get("started_at", 10.0),
+                    row.get("model", "gpt-test"),
+                    row.get("workspace", "/tmp/workspace"),
+                ),
+            )
+            for i in range(message_counts.get(row["id"], 0)):
+                conn.execute(
+                    "insert into messages (session_id, role, content, timestamp) values (?, 'user', ?, ?)",
+                    (row["id"], f"message {i}", 10.0 + i),
+                )
+    return db
+
+
+def test_repair_discoverability_dry_run_plans_without_mutating_files(tmp_path):
+    stale = "webui-stale-cli-flag"
+    missing = "state-only-messageful"
+    _write_sidecar(tmp_path, stale, messages=3, source_tag="webui", session_source="webui", is_cli_session=True)
+    _write_index(tmp_path, {"session_id": stale, "message_count": 3, "source_tag": "webui", "session_source": "webui", "is_cli_session": True})
+    db = _state_db(
+        tmp_path,
+        [
+            {"id": stale, "source": "webui", "message_count": 3},
+            {"id": missing, "source": "webui", "message_count": 2},
+        ],
+        {stale: 3, missing: 2},
+    )
+
+    result = repair_session_discoverability(tmp_path, state_db_path=db, dry_run=True, backup_dir=tmp_path / "backup")
+
+    assert result["dry_run"] is True
+    assert result["applied"] == []
+    assert {action["action"] for action in result["planned"]} == {
+        "clear_sidecar_cli_flag",
+        "clear_index_cli_flag",
+        "materialize_sidecar_from_state_db",
+    }
+    assert json.loads((tmp_path / f"{stale}.json").read_text())["is_cli_session"] is True
+    assert json.loads((tmp_path / "_index.json").read_text())[0]["is_cli_session"] is True
+    assert not (tmp_path / f"{missing}.json").exists()
+    assert not (tmp_path / "backup").exists()
+
+
+def test_repair_discoverability_apply_requires_backup_dir(tmp_path):
+    sid = "webui-stale-cli-flag"
+    _write_sidecar(tmp_path, sid, messages=1, source_tag="webui", session_source="webui", is_cli_session=True)
+    _write_index(tmp_path, {"session_id": sid, "message_count": 1, "source_tag": "webui", "session_source": "webui", "is_cli_session": True})
+    db = _state_db(tmp_path, [{"id": sid, "source": "webui", "message_count": 1}], {sid: 1})
+
+    result = repair_session_discoverability(tmp_path, state_db_path=db, dry_run=False)
+
+    assert result["ok"] is False
+    assert result["error"] == "backup_dir_required_for_apply"
+    assert json.loads((tmp_path / f"{sid}.json").read_text())["is_cli_session"] is True
+
+
+def test_repair_discoverability_apply_backs_up_and_repairs_safe_findings(tmp_path):
+    stale = "webui-stale-cli-flag"
+    missing = "state-only-messageful"
+    _write_sidecar(tmp_path, stale, messages=3, source_tag="webui", session_source="webui", is_cli_session=True)
+    _write_index(tmp_path, {"session_id": stale, "message_count": 3, "source_tag": "webui", "session_source": "webui", "is_cli_session": True})
+    db = _state_db(
+        tmp_path,
+        [
+            {"id": stale, "source": "webui", "message_count": 3},
+            {"id": missing, "source": "webui", "message_count": 2, "title": "Recovered From State"},
+        ],
+        {stale: 3, missing: 2},
+    )
+
+    result = repair_session_discoverability(tmp_path, state_db_path=db, dry_run=False, backup_dir=tmp_path / "backup")
+
+    assert result["ok"] is True
+    assert result["dry_run"] is False
+    assert {action["action"] for action in result["applied"]} == {
+        "clear_sidecar_cli_flag",
+        "clear_index_cli_flag",
+        "materialize_sidecar_from_state_db",
+    }
+    assert json.loads((tmp_path / f"{stale}.json").read_text())["is_cli_session"] is False
+    assert json.loads((tmp_path / "_index.json").read_text())[0]["is_cli_session"] is False
+    index_rows = json.loads((tmp_path / "_index.json").read_text())
+    assert {row["session_id"] for row in index_rows} == {stale, missing}
+    recovered = json.loads((tmp_path / f"{missing}.json").read_text())
+    assert recovered["title"] == "Recovered From State"
+    assert recovered["message_count"] == 2
+    assert len(recovered["messages"]) == 2
+    backed_up = {p.name for p in (tmp_path / "backup").iterdir()}
+    assert f"{stale}.json" in backed_up
+    assert "_index.json" in backed_up
+    assert "state.db" in backed_up
+
+
+def test_repair_discoverability_cli_defaults_to_dry_run(tmp_path):
+    sid = "webui-stale-cli-flag"
+    _write_sidecar(tmp_path, sid, messages=1, source_tag="webui", session_source="webui", is_cli_session=True)
+    _write_index(tmp_path, {"session_id": sid, "message_count": 1, "source_tag": "webui", "session_source": "webui", "is_cli_session": True})
+    db = _state_db(tmp_path, [{"id": sid, "source": "webui", "message_count": 1}], {sid: 1})
+
+    completed = subprocess.run(
+        [
+            sys.executable,
+            "-m",
+            "api.session_discoverability",
+            "--repair-safe",
+            "--session-dir",
+            str(tmp_path),
+            "--state-db",
+            str(db),
+        ],
+        check=True,
+        text=True,
+        capture_output=True,
+    )
+
+    result = json.loads(completed.stdout)
+    assert result["dry_run"] is True
+    assert [action["action"] for action in result["planned"]] == ["clear_sidecar_cli_flag", "clear_index_cli_flag"]
+    assert json.loads((tmp_path / f"{sid}.json").read_text())["is_cli_session"] is True