diff --git a/tests/test_orphan_cleanup.py b/tests/test_orphan_cleanup.py new file mode 100644 index 00000000..c0a454ae --- /dev/null +++ b/tests/test_orphan_cleanup.py @@ -0,0 +1,207 @@ +"""Orphan-conversation scan + purge (the ``turnstone-admin orphan-conversations`` verb). + +Orphans are conversation rows whose ``workstreams`` row is gone — written by +historical unregistered paths or by the delete-during-inflight race (a late +tool-result save re-creating rows after ``delete_workstream``). +""" + +from __future__ import annotations + +import argparse +import hashlib + +import pytest + +from turnstone.admin import _cmd_orphan_conversations + + +def _orphan(backend, ws_id: str, n: int = 2) -> None: + """Persist *n* conversation rows for *ws_id* WITHOUT registering it.""" + for i in range(n): + backend.save_message(ws_id, "user" if i % 2 == 0 else "assistant", f"m{i}") + + +def _blob(backend, payload: bytes, origin: str = "upload") -> str: + """Save a content-addressed attachment; each save bumps the refcount.""" + aid = hashlib.sha256(payload).hexdigest() + backend.save_attachment(aid, "f.txt", "text/plain", len(payload), "text", payload, origin) + return aid + + +class TestOrphanScan: + def test_clean_db_has_no_orphans(self, backend): + backend.register_workstream("live1") + backend.save_message("live1", "user", "hello") + assert backend.list_orphan_conversations() == [] + + def test_orphan_reported_with_stats(self, backend): + _orphan(backend, "ghost1", n=3) + scan = backend.list_orphan_conversations() + assert len(scan) == 1 + entry = scan[0] + assert entry["ws_id"] == "ghost1" + assert entry["rows"] == 3 + assert entry["first"] <= entry["last"] + assert entry["attachment_refs"] == 0 + + def test_scan_counts_attachment_refs(self, backend): + _orphan(backend, "ghost2", n=1) + msg_id = backend.save_message("ghost2", "user", "with attachment") + aid = _blob(backend, b"orphan-bytes") + backend.set_message_attachments("ghost2", msg_id, [aid]) + scan = backend.list_orphan_conversations() + assert scan[0]["attachment_refs"] == 1 + + def test_scan_is_oldest_first(self, backend): + _orphan(backend, "newer") + _orphan(backend, "older") + # Timestamps are insertion-ordered ISO text; rewrite to force ordering. + import sqlalchemy as sa + + from turnstone.core.storage._schema import conversations + + with backend._conn() as conn: + conn.execute( + sa.update(conversations) + .where(conversations.c.ws_id == "older") + .values(timestamp="2020-01-01T00:00:00") + ) + conn.commit() + scan = backend.list_orphan_conversations() + assert [o["ws_id"] for o in scan] == ["older", "newer"] + + +class TestOrphanPurge: + def test_purge_deletes_only_orphans(self, backend): + backend.register_workstream("live1") + backend.save_message("live1", "user", "keep me") + _orphan(backend, "ghost1", n=4) + result = backend.delete_orphan_conversations(["ghost1"]) + assert result == {"workstreams": 1, "rows": 4, "released_refs": 0, "skipped": 0} + assert backend.list_orphan_conversations() == [] + assert len(backend.load_messages("live1")) == 1 + + def test_purge_skips_reregistered_ws(self, backend): + """A ws_id that gained a workstreams row between scan and purge survives.""" + _orphan(backend, "ghost1", n=2) + scan = [o["ws_id"] for o in backend.list_orphan_conversations()] + backend.register_workstream("ghost1") + result = backend.delete_orphan_conversations(scan) + assert result["skipped"] == 1 + assert result["workstreams"] == 0 + assert result["rows"] == 0 + assert len(backend.load_messages("ghost1")) == 2 + + def test_purge_releases_refcounts_and_prunes_at_zero(self, backend): + _orphan(backend, "ghost1", n=1) + msg_id = backend.save_message("ghost1", "user", "img") + aid = _blob(backend, b"only-orphan-referenced") + backend.set_message_attachments("ghost1", msg_id, [aid]) + result = backend.delete_orphan_conversations(["ghost1"]) + assert result["released_refs"] == 1 + assert backend.get_attachment(aid) is None + + def test_purge_keeps_blob_shared_with_live_ws(self, backend): + payload = b"shared-bytes" + backend.register_workstream("live1") + live_msg = backend.save_message("live1", "user", "live ref") + aid_live = _blob(backend, payload) + backend.set_message_attachments("live1", live_msg, [aid_live]) + + _orphan(backend, "ghost1", n=1) + ghost_msg = backend.save_message("ghost1", "user", "ghost ref") + aid_ghost = _blob(backend, payload) # same content hash; refcount -> 2 + backend.set_message_attachments("ghost1", ghost_msg, [aid_ghost]) + assert aid_live == aid_ghost + + result = backend.delete_orphan_conversations(["ghost1"]) + assert result["released_refs"] == 1 + row = backend.get_attachment(aid_live) + assert row is not None + assert row["refcount"] == 1 + + def test_purge_sweeps_config_rows(self, backend): + _orphan(backend, "ghost1", n=1) + backend.save_workstream_config("ghost1", {"model": "x"}) + backend.delete_orphan_conversations(["ghost1"]) + import sqlalchemy as sa + + from turnstone.core.storage._schema import workstream_config + + with backend._conn() as conn: + left = conn.execute( + sa.select(sa.func.count()).where(workstream_config.c.ws_id == "ghost1") + ).scalar() + assert left == 0 + + def test_purge_empty_list_is_noop(self, backend): + assert backend.delete_orphan_conversations([]) == { + "workstreams": 0, + "rows": 0, + "released_refs": 0, + "skipped": 0, + } + + +class TestAdminVerb: + """The CLI handler over a real (ephemeral) backend.""" + + def _args(self, **kw) -> argparse.Namespace: + return argparse.Namespace(delete=False, yes=False, **kw) + + def test_scan_reports_and_does_not_delete(self, backend, monkeypatch, capsys): + _orphan(backend, "ghost1", n=2) + monkeypatch.setattr("turnstone.admin._get_storage", lambda args: backend) + _cmd_orphan_conversations(self._args()) + out = capsys.readouterr().out + assert "ghost1" in out + assert "--delete" in out + assert len(backend.load_messages("ghost1")) == 2 + + def test_delete_yes_purges(self, backend, monkeypatch, capsys): + _orphan(backend, "ghost1", n=2) + monkeypatch.setattr("turnstone.admin._get_storage", lambda args: backend) + ns = self._args() + ns.delete = True + ns.yes = True + _cmd_orphan_conversations(ns) + out = capsys.readouterr().out + assert "Purged 2" in out + assert backend.list_orphan_conversations() == [] + + def test_delete_confirmation_abort(self, backend, monkeypatch, capsys): + _orphan(backend, "ghost1", n=1) + monkeypatch.setattr("turnstone.admin._get_storage", lambda args: backend) + monkeypatch.setattr("builtins.input", lambda prompt: "n") + ns = self._args() + ns.delete = True + with pytest.raises(SystemExit): + _cmd_orphan_conversations(ns) + assert len(backend.load_messages("ghost1")) == 1 + + def test_delete_summary_reports_partial_skip(self, backend, monkeypatch, capsys): + """Mixed batch: summary shows ACTUAL purge counts plus the skipped clause.""" + _orphan(backend, "ghost1", n=2) + _orphan(backend, "ghost2", n=3) + real_list = backend.list_orphan_conversations + + def list_then_register(): + scan = real_list() + backend.register_workstream("ghost2") # wins the scan-to-purge race + return scan + + monkeypatch.setattr(backend, "list_orphan_conversations", list_then_register) + monkeypatch.setattr("turnstone.admin._get_storage", lambda args: backend) + ns = self._args() + ns.delete = True + ns.yes = True + _cmd_orphan_conversations(ns) + out = capsys.readouterr().out + assert "Purged 2 row(s) across 1 workstream(s)" in out + assert "skipped 1 re-registered" in out + assert len(backend.load_messages("ghost2")) == 3 + + def test_clean_db_message(self, backend, monkeypatch, capsys): + monkeypatch.setattr("turnstone.admin._get_storage", lambda args: backend) + _cmd_orphan_conversations(self._args()) + assert "No orphan conversation rows." in capsys.readouterr().out diff --git a/turnstone/admin.py b/turnstone/admin.py index a4c5ceba..4249aefb 100644 --- a/turnstone/admin.py +++ b/turnstone/admin.py @@ -439,6 +439,44 @@ def _discover_console_url() -> str: # --------------------------------------------------------------------------- +def _cmd_orphan_conversations(args: argparse.Namespace) -> None: + """Scan for (and with --delete purge) conversation rows whose workstream is gone.""" + storage = _get_storage(args) + orphans = storage.list_orphan_conversations() + if not orphans: + print("No orphan conversation rows.") + return + width = max(len(o["ws_id"]) for o in orphans) + print(f"{'ws_id':<{width}} {'rows':>5} {'refs':>4} first last") + for o in orphans: + print( + f"{o['ws_id']:<{width}} {o['rows']:>5} {o['attachment_refs']:>4} " + f"{(o['first'] or '')[:10]} {(o['last'] or '')[:10]}" + ) + total_rows = sum(o["rows"] for o in orphans) + total_refs = sum(o["attachment_refs"] for o in orphans) + print( + f"\n{len(orphans)} orphan workstream(s), {total_rows} conversation row(s), " + f"{total_refs} attachment ref(s)." + ) + if not args.delete: + print("Re-run with --delete to purge them.") + return + if not args.yes: + reply = input(f"Delete {total_rows} row(s) across {len(orphans)} workstream(s)? [y/N] ") + if reply.strip().lower() not in ("y", "yes"): + print("Aborted.", file=sys.stderr) + sys.exit(1) + result = storage.delete_orphan_conversations([o["ws_id"] for o in orphans]) + summary = ( + f"Purged {result['rows']} row(s) across {result['workstreams']} workstream(s); " + f"released {result['released_refs']} attachment ref(s)" + ) + if result["skipped"]: + summary += f"; skipped {result['skipped']} re-registered workstream(s)" + print(summary + ".") + + def _cmd_rerank_calibrate(args: argparse.Namespace) -> None: from turnstone.core.config import get_rerank_instruction from turnstone.core.config_store import ConfigStore @@ -619,6 +657,15 @@ def main() -> None: help="Write the calibration onto the model's capabilities", ) + p_orph = sub.add_parser( + "orphan-conversations", + help="Scan (and with --delete purge) conversation rows whose workstream row is gone", + ) + p_orph.add_argument( + "--delete", action="store_true", help="Purge the orphans after the scan report" + ) + p_orph.add_argument("--yes", action="store_true", help="Skip the interactive confirmation") + args = parser.parse_args() if not args.command: parser.print_help() @@ -638,6 +685,7 @@ def main() -> None: "set-node-metadata": _cmd_set_node_metadata, "delete-node-metadata": _cmd_delete_node_metadata, "export": _cmd_export, + "orphan-conversations": _cmd_orphan_conversations, "rerank-calibrate": _cmd_rerank_calibrate, } dispatch[args.command](args) diff --git a/turnstone/core/storage/_postgresql.py b/turnstone/core/storage/_postgresql.py index 48e5b6b5..5e1ddc28 100644 --- a/turnstone/core/storage/_postgresql.py +++ b/turnstone/core/storage/_postgresql.py @@ -114,17 +114,19 @@ from turnstone.core.storage._utils import ( from turnstone.core.storage._utils import ( escape_like as _escape_like, ) +from turnstone.core.storage._utils import ( + find_orphan_conversations, + prepare_provider_data_for_save, + purge_orphan_conversations, + release_attachment_refs, + sanitize_text, +) from turnstone.core.storage._utils import ( normalize_search_terms as _normalize_search_terms, ) from turnstone.core.storage._utils import ( parse_attachment_refs as _parse_attachment_refs, ) -from turnstone.core.storage._utils import ( - prepare_provider_data_for_save, - release_attachment_refs, - sanitize_text, -) from turnstone.core.storage._utils import ( reconstruct_messages as _reconstruct_messages, ) @@ -909,6 +911,16 @@ class PostgreSQLBackend: conn.commit() return result.rowcount > 0 + def list_orphan_conversations(self) -> list[dict[str, Any]]: + with self._conn() as conn: + return find_orphan_conversations(conn) + + def delete_orphan_conversations(self, ws_ids: list[str]) -> dict[str, int]: + with self._conn() as conn: + result = purge_orphan_conversations(conn, ws_ids) + conn.commit() + return result + # -- Workstream attachments (content-addressed, refcounted) ---------------- def save_attachment( diff --git a/turnstone/core/storage/_protocol.py b/turnstone/core/storage/_protocol.py index 7b054dd8..a3381d87 100644 --- a/turnstone/core/storage/_protocol.py +++ b/turnstone/core/storage/_protocol.py @@ -633,6 +633,27 @@ class StorageBackend(Protocol): """Delete a workstream and all its conversations + config.""" ... + def list_orphan_conversations(self) -> list[dict[str, Any]]: + """Conversation ws_ids with no ``workstreams`` row. + + One dict per orphan workstream — keys ``ws_id``, ``rows``, ``first``, + ``last`` (ISO text timestamps), ``attachment_refs`` — ordered + oldest-first. Read-only; feeds the ``turnstone-admin + orphan-conversations`` maintenance verb. + """ + ... + + def delete_orphan_conversations(self, ws_ids: list[str]) -> dict[str, int]: + """Purge conversation rows for the *ws_ids* that are STILL orphaned. + + Re-verifies against ``workstreams`` in-transaction (a re-registered + ws_id is skipped, never deleted), releases the deleted rows' + attachment refcounts, and sweeps matching ``workstream_config`` / + ``workstream_overrides`` rows. Returns counts keyed ``workstreams``, + ``rows``, ``released_refs``, ``skipped``. + """ + ... + def list_workstreams( self, node_id: str | None = None, diff --git a/turnstone/core/storage/_sqlite.py b/turnstone/core/storage/_sqlite.py index a913d5e5..a60540df 100644 --- a/turnstone/core/storage/_sqlite.py +++ b/turnstone/core/storage/_sqlite.py @@ -114,17 +114,19 @@ from turnstone.core.storage._utils import ( from turnstone.core.storage._utils import ( escape_like as _escape_like, ) +from turnstone.core.storage._utils import ( + find_orphan_conversations, + prepare_provider_data_for_save, + purge_orphan_conversations, + release_attachment_refs, + sanitize_text, +) from turnstone.core.storage._utils import ( normalize_search_terms as _normalize_search_terms, ) from turnstone.core.storage._utils import ( parse_attachment_refs as _parse_attachment_refs, ) -from turnstone.core.storage._utils import ( - prepare_provider_data_for_save, - release_attachment_refs, - sanitize_text, -) from turnstone.core.storage._utils import ( reconstruct_messages as _reconstruct_messages, ) @@ -1055,6 +1057,16 @@ class SQLiteBackend: conn.commit() return result.rowcount > 0 + def list_orphan_conversations(self) -> list[dict[str, Any]]: + with self._conn() as conn: + return find_orphan_conversations(conn) + + def delete_orphan_conversations(self, ws_ids: list[str]) -> dict[str, int]: + with self._conn() as conn: + result = purge_orphan_conversations(conn, ws_ids) + conn.commit() + return result + # -- Workstream attachments (content-addressed, refcounted) ---------------- def save_attachment( diff --git a/turnstone/core/storage/_utils.py b/turnstone/core/storage/_utils.py index 3b258455..744f8cc9 100644 --- a/turnstone/core/storage/_utils.py +++ b/turnstone/core/storage/_utils.py @@ -12,7 +12,13 @@ import sqlalchemy as sa from turnstone.core.attachments import unreadable_placeholder from turnstone.core.log import get_logger -from turnstone.core.storage._schema import workstream_attachments +from turnstone.core.storage._schema import ( + conversations, + workstream_attachments, + workstream_config, + workstream_overrides, + workstreams, +) from turnstone.core.trajectory import ( AttachmentRef, ContentBlock, @@ -173,6 +179,100 @@ def release_attachment_refs(conn: Any, attachment_ids: list[str]) -> None: ) +def find_orphan_conversations(conn: Any) -> list[dict[str, Any]]: + """Conversation ws_ids that have no ``workstreams`` row, with row stats. + + Orphans come from writers that persisted without a registered workstream: + historically the pre-unification CLI/server paths, and the + delete-during-inflight race (a late tool-result save re-creating rows + after ``delete_workstream``). Read-only; ordered oldest-first. Each + entry carries the attachment-ref count so a purge's refcount release is + visible before it happens. + """ + rows = conn.execute( + sa.select( + conversations.c.ws_id, + sa.func.count().label("row_count"), + sa.func.min(conversations.c.timestamp).label("first"), + sa.func.max(conversations.c.timestamp).label("last"), + ) + .select_from( + conversations.outerjoin(workstreams, conversations.c.ws_id == workstreams.c.ws_id) + ) + .where(workstreams.c.ws_id.is_(None)) + .group_by(conversations.c.ws_id) + .order_by(sa.func.min(conversations.c.timestamp)) + ).fetchall() + orphans: list[dict[str, Any]] = [] + for ws_id, row_count, first, last in rows: + referenced = conn.execute( + sa.select(conversations.c.attachments).where( + sa.and_( + conversations.c.ws_id == ws_id, + conversations.c.attachments.is_not(None), + ) + ) + ).fetchall() + ref_count = sum(len(parse_attachment_refs(refs)) for (refs,) in referenced) + orphans.append( + { + "ws_id": ws_id, + "rows": int(row_count), + "first": first, + "last": last, + "attachment_refs": ref_count, + } + ) + return orphans + + +def purge_orphan_conversations(conn: Any, ws_ids: list[str]) -> dict[str, int]: + """Delete conversation rows for the *ws_ids* that are STILL orphans. + + Orphan-ness is re-verified here, inside the caller's transaction — a + ws_id that gained a ``workstreams`` row between scan and purge is counted + in ``skipped`` and left untouched, so a stale scan can never delete a + live workstream's history. Mirrors ``delete_workstream``'s cascade for + rows with no owning workstream: release the deleted rows' attachment + refcounts, then sweep the matching ``workstream_config`` / + ``workstream_overrides`` rows. Caller owns commit. + """ + if not ws_ids: + return {"workstreams": 0, "rows": 0, "released_refs": 0, "skipped": 0} + registered = { + row[0] + for row in conn.execute( + sa.select(workstreams.c.ws_id).where(workstreams.c.ws_id.in_(ws_ids)) + ).fetchall() + } + targets = [w for w in ws_ids if w not in registered] + if not targets: + return {"workstreams": 0, "rows": 0, "released_refs": 0, "skipped": len(registered)} + referenced = conn.execute( + sa.select(conversations.c.attachments).where( + sa.and_( + conversations.c.ws_id.in_(targets), + conversations.c.attachments.is_not(None), + ) + ) + ).fetchall() + ref_ids: list[str] = [] + for (refs,) in referenced: + ref_ids.extend(parse_attachment_refs(refs)) + release_attachment_refs(conn, ref_ids) + deleted = conn.execute( + sa.delete(conversations).where(conversations.c.ws_id.in_(targets)) + ).rowcount + conn.execute(sa.delete(workstream_config).where(workstream_config.c.ws_id.in_(targets))) + conn.execute(sa.delete(workstream_overrides).where(workstream_overrides.c.ws_id.in_(targets))) + return { + "workstreams": len(targets), + "rows": int(deleted or 0), + "released_refs": len(ref_ids), + "skipped": len(registered), + } + + def attachment_to_content_part(att: dict[str, Any]) -> dict[str, Any] | None: """Convert a stored attachment row into an OpenAI-style content part.