turnstone/tests/test_storage_sqlite.py

"""Tests for the SQLite storage backend."""

from __future__ import annotations

from typing import Any

# -- Workstream registration ---------------------------------------------------


class TestRegisterWorkstream:
    def test_register_creates_workstream(self, backend):
        backend.register_workstream("s1", title="Test")
        name = backend.get_workstream_display_name("s1")
        assert name == "Test"

    def test_register_idempotent(self, backend):
        backend.register_workstream("s1", title="First")
        backend.register_workstream("s1", title="Second")
        name = backend.get_workstream_display_name("s1")
        assert name == "First"  # INSERT OR IGNORE preserves first


class TestSaveAndLoadMessages:
    def test_roundtrip(self, backend):
        backend.register_workstream("s1")
        backend.save_message("s1", "user", "hello")
        backend.save_message("s1", "assistant", "world")
        msgs = backend.load_messages("s1")
        assert len(msgs) == 2
        assert msgs[0]["role"] == "user"
        assert msgs[0]["content"] == "hello"
        assert msgs[1]["role"] == "assistant"
        assert msgs[1]["content"] == "world"

    def test_tool_call_grouping(self, backend):
        import json

        backend.register_workstream("s1")
        tc_json = json.dumps(
            [
                {
                    "id": "c1",
                    "type": "function",
                    "function": {"name": "bash", "arguments": '{"cmd":"ls"}'},
                }
            ]
        )
        backend.save_message("s1", "user", "do something")
        backend.save_message("s1", "assistant", None, tool_calls=tc_json)
        backend.save_message("s1", "tool", "file.txt", tool_call_id="c1")
        backend.save_message("s1", "assistant", "done")
        msgs = backend.load_messages("s1")
        assert len(msgs) == 4
        assert msgs[1]["role"] == "assistant"
        assert len(msgs[1]["tool_calls"]) == 1
        assert msgs[1]["tool_calls"][0]["id"] == "c1"
        assert msgs[2]["role"] == "tool"
        assert msgs[2]["content"] == "file.txt"

    def test_incomplete_turn_repair(self, backend):
        import json

        backend.register_workstream("s1")
        tc_json = json.dumps(
            [
                {
                    "id": "c1",
                    "type": "function",
                    "function": {"name": "bash", "arguments": '{"cmd":"ls"}'},
                },
                {
                    "id": "c2",
                    "type": "function",
                    "function": {"name": "read", "arguments": '{"path":"a"}'},
                },
            ]
        )
        backend.save_message("s1", "user", "do something")
        backend.save_message("s1", "assistant", None, tool_calls=tc_json)
        # Only 1 result for 2 calls — incomplete turn
        backend.save_message("s1", "tool", "ok", tool_call_id="c1")
        msgs = backend.load_messages("s1")
        # Incomplete turn should be stripped
        assert len(msgs) == 1  # only the user message remains

    def test_provider_data_preserved(self, backend):
        import json

        backend.register_workstream("s1")
        pd = json.dumps({"encrypted": True})
        backend.save_message("s1", "assistant", "hi", provider_data=pd)
        msgs = backend.load_messages("s1")
        assert msgs[0].get("_provider_content") == {"encrypted": True}

    def test_empty_workstream_returns_empty(self, backend):
        assert backend.load_messages("nonexistent") == []


class TestLoadMessagesLimit:
    """Phase 3 added ``limit=N`` so cluster-inspect can avoid reading
    thousands of rows to return a tail-20 preview.  The contract: fetch
    the last N conversation rows (DESC + LIMIT at the SQL layer), then
    reverse into chronological order for reconstruction.  Approximate
    tail-N — a tool-call group straddling the cut produces an
    incomplete turn that the existing repair step strips."""

    def test_limit_none_fetches_all(self, backend):
        backend.register_workstream("s1")
        for i in range(10):
            backend.save_message("s1", "user", f"msg-{i}")
        msgs = backend.load_messages("s1", limit=None)
        assert len(msgs) == 10

    def test_limit_fetches_tail_in_chronological_order(self, backend):
        backend.register_workstream("s1")
        for i in range(10):
            backend.save_message("s1", "user", f"msg-{i:02d}")
        msgs = backend.load_messages("s1", limit=3)
        assert len(msgs) == 3
        # Chronological order preserved even though SQL fetched DESC.
        assert msgs[0]["content"] == "msg-07"
        assert msgs[1]["content"] == "msg-08"
        assert msgs[2]["content"] == "msg-09"

    def test_limit_exceeds_total_returns_all(self, backend):
        backend.register_workstream("s1")
        for i in range(5):
            backend.save_message("s1", "user", f"msg-{i}")
        msgs = backend.load_messages("s1", limit=100)
        assert len(msgs) == 5

    def test_limit_zero_fetches_all(self, backend):
        """limit<=0 matches the ``None`` branch — the SQL LIMIT is
        skipped, full history returned.  Belt-and-suspenders against
        callers that pass the clamped ``max(0, limit)`` result."""
        backend.register_workstream("s1")
        for i in range(5):
            backend.save_message("s1", "user", f"msg-{i}")
        assert len(backend.load_messages("s1", limit=0)) == 5

    def test_limit_boundary_straddles_tool_call_group(self, backend):
        """Document the approximate-tail-N semantics the ``load_messages``
        docstring warns about: when the tail slice opens mid-tool-call-
        group, the orphaned ``role=tool`` row is returned verbatim
        (the incomplete-turn repair at ``_reconstruct_messages`` only
        strips incomplete *assistant-with-tool_calls* groups, not
        orphaned tool-response rows).

        Callers that need strict tail-N semantics (e.g. re-hydrating a
        session to resume generation) must either request more than
        they need and post-filter, or do a full load.  The cluster-
        inspect preview path tolerates orphan tool rows because the
        UI renders them as standalone tool-output blocks.

        Seed: [user, assistant w/ 1 tool_call, tool result, assistant].
        Fetch tail=2 → [tool result, assistant].  Orphan tool row
        survives; this is expected behavior, not a bug."""
        import json

        backend.register_workstream("s1")
        tc_json = json.dumps(
            [
                {
                    "id": "c1",
                    "type": "function",
                    "function": {"name": "bash", "arguments": '{"cmd":"ls"}'},
                }
            ]
        )
        backend.save_message("s1", "user", "do it")
        backend.save_message("s1", "assistant", None, tool_calls=tc_json)
        backend.save_message("s1", "tool", "output", tool_call_id="c1")
        backend.save_message("s1", "assistant", "done")

        # Full load: 4 messages (complete turn, assistant reply).
        assert len(backend.load_messages("s1")) == 4

        # Tail=2: orphan tool row + final assistant reply.
        tail = backend.load_messages("s1", limit=2)
        assert len(tail) == 2
        assert tail[0]["role"] == "tool"
        assert tail[0]["content"] == "output"
        assert tail[1]["role"] == "assistant"
        assert tail[1]["content"] == "done"

    def test_limit_keeps_complete_tool_call_group_when_fully_contained(self, backend):
        """Tool-call groups entirely inside the tail slice survive intact."""
        import json

        backend.register_workstream("s1")
        tc_json = json.dumps(
            [
                {
                    "id": "c1",
                    "type": "function",
                    "function": {"name": "read", "arguments": '{"p":"a"}'},
                }
            ]
        )
        backend.save_message("s1", "user", "older message")
        backend.save_message("s1", "assistant", None, tool_calls=tc_json)
        backend.save_message("s1", "tool", "contents", tool_call_id="c1")
        backend.save_message("s1", "assistant", "summarized")

        # Tail=3 captures the full group + assistant reply (drops
        # only the oldest user message).
        tail = backend.load_messages("s1", limit=3)
        assert len(tail) == 3
        assert tail[0]["role"] == "assistant"
        assert len(tail[0]["tool_calls"]) == 1
        assert tail[1]["role"] == "tool"
        assert tail[1]["content"] == "contents"
        assert tail[2]["content"] == "summarized"

    def test_limit_bounds_attachment_scan(self, backend):
        """When ``limit=N`` is set, ``load_attachments_for_messages``
        receives only the fetched message ids — the attachment query
        must not fall back to a full-workstream scan.  Otherwise the
        tail-N optimization on conversations is partly undone for
        workstreams with many attachments."""
        from unittest.mock import patch

        backend.register_workstream("s1")
        for i in range(20):
            backend.save_message("s1", "user", f"msg-{i:02d}")

        captured: dict[str, list[int] | None] = {}
        orig = backend.load_attachments_for_messages

        def _spy(ws_id, *, message_ids=None):
            captured["message_ids"] = list(message_ids) if message_ids is not None else None
            return orig(ws_id, message_ids=message_ids)

        with patch.object(backend, "load_attachments_for_messages", side_effect=_spy):
            backend.load_messages("s1", limit=5)
        # Tail-N request passed a bounded list of exactly 5 ids.
        assert captured["message_ids"] is not None
        assert len(captured["message_ids"]) == 5

        with patch.object(backend, "load_attachments_for_messages", side_effect=_spy):
            backend.load_messages("s1")
        # Full-load request passes None → backend scans all attachments.
        assert captured["message_ids"] is None


class TestSaveMessagesBulk:
    def test_bulk_roundtrip(self, backend):
        backend.register_workstream("s1")
        backend.save_messages_bulk(
            [
                {"ws_id": "s1", "role": "user", "content": "hello"},
                {"ws_id": "s1", "role": "assistant", "content": "hi there"},
                {"ws_id": "s1", "role": "user", "content": "bye"},
            ]
        )
        msgs = backend.load_messages("s1")
        assert len(msgs) == 3
        assert msgs[0]["content"] == "hello"
        assert msgs[2]["content"] == "bye"

    def test_bulk_preserves_tool_calls(self, backend):
        import json

        backend.register_workstream("s1")
        tc = json.dumps(
            [{"id": "c1", "type": "function", "function": {"name": "bash", "arguments": "{}"}}]
        )
        backend.save_messages_bulk(
            [
                {"ws_id": "s1", "role": "user", "content": "do it"},
                {"ws_id": "s1", "role": "assistant", "content": None, "tool_calls": tc},
                {"ws_id": "s1", "role": "tool", "content": "ok", "tool_call_id": "c1"},
            ]
        )
        msgs = backend.load_messages("s1")
        assert len(msgs) == 3
        assert msgs[1]["tool_calls"][0]["id"] == "c1"

    def test_bulk_empty_is_noop(self, backend):
        backend.save_messages_bulk([])

    def test_bulk_updates_workstream_timestamp(self, backend):
        backend.register_workstream("s1")
        # Save a message to establish an initial updated timestamp
        backend.save_message("s1", "user", "seed")
        rows_before = backend.list_workstreams_with_history()
        updated_before = rows_before[0][5]  # updated column

        backend.save_messages_bulk([{"ws_id": "s1", "role": "user", "content": "bulk"}])
        rows_after = backend.list_workstreams_with_history()
        updated_after = rows_after[0][5]
        assert updated_after >= updated_before


class TestListWorkstreamsWithHistory:
    def test_lists_workstreams_with_messages(self, backend):
        backend.register_workstream("s1")
        backend.save_message("s1", "user", "hi")
        backend.register_workstream("s2")  # no messages
        rows = backend.list_workstreams_with_history()
        assert len(rows) == 1
        assert rows[0][0] == "s1"

    def test_respects_limit(self, backend):
        for i in range(5):
            sid = f"s{i}"
            backend.register_workstream(sid)
            backend.save_message(sid, "user", f"msg {i}")
        rows = backend.list_workstreams_with_history(limit=3)
        assert len(rows) == 3

    def test_kind_filter_excludes_coordinators(self, backend):
        """The interactive 'saved workstreams' sidebar calls this with
        kind=INTERACTIVE so coordinator rows (which also persist
        conversation history) don't leak into the interactive UI."""
        from turnstone.core.workstream import WorkstreamKind

        backend.register_workstream("interactive-1", kind=WorkstreamKind.INTERACTIVE)
        backend.save_message("interactive-1", "user", "hi")
        backend.register_workstream("coord-1", kind=WorkstreamKind.COORDINATOR)
        backend.save_message("coord-1", "user", "plan something")

        # Default (no filter) returns both — preserves legacy behaviour.
        rows_all = backend.list_workstreams_with_history()
        assert {r[0] for r in rows_all} == {"interactive-1", "coord-1"}

        # kind=INTERACTIVE drops the coordinator row at the SQL layer.
        rows_i = backend.list_workstreams_with_history(kind=WorkstreamKind.INTERACTIVE)
        assert {r[0] for r in rows_i} == {"interactive-1"}

        # kind=COORDINATOR symmetric — for admin tooling that wants
        # the opposite view.
        rows_c = backend.list_workstreams_with_history(kind=WorkstreamKind.COORDINATOR)
        assert {r[0] for r in rows_c} == {"coord-1"}

    def test_kind_filter_accepts_string(self, backend):
        """String form (``"interactive"``) works too — matches how the
        memory.py helper forwards caller-supplied values."""
        from turnstone.core.workstream import WorkstreamKind

        backend.register_workstream("interactive-1", kind=WorkstreamKind.INTERACTIVE)
        backend.save_message("interactive-1", "user", "hi")
        backend.register_workstream("coord-1", kind=WorkstreamKind.COORDINATOR)
        backend.save_message("coord-1", "user", "plan")

        rows = backend.list_workstreams_with_history(kind="interactive")
        assert {r[0] for r in rows} == {"interactive-1"}


class TestDeleteWorkstream:
    def test_deletes_all_data(self, backend):
        backend.register_workstream("s1")
        backend.save_message("s1", "user", "hi")
        backend.save_workstream_config("s1", {"temp": "0.5"})
        assert backend.delete_workstream("s1")
        assert backend.load_messages("s1") == []
        assert backend.load_workstream_config("s1") == {}
        assert backend.get_workstream_display_name("s1") is None


class TestPruneWorkstreams:
    def test_orphan_removed(self, backend):
        backend.register_workstream("orphan")
        orphans, stale = backend.prune_workstreams()
        assert orphans == 1

    def test_stale_removed(self, backend):
        import sqlalchemy as sa

        backend.register_workstream("old")
        backend.save_message("old", "user", "hi")
        # Force old timestamp
        with backend._engine.connect() as conn:
            conn.execute(
                sa.text("UPDATE workstreams SET updated = '2020-01-01' WHERE ws_id = 'old'")
            )
            conn.commit()
        _, stale = backend.prune_workstreams(retention_days=30)
        assert stale == 1


class TestResolveWorkstream:
    def test_exact_alias(self, backend):
        backend.register_workstream("s1")
        backend.set_workstream_alias("s1", "myalias")
        assert backend.resolve_workstream("myalias") == "s1"

    def test_exact_id(self, backend):
        backend.register_workstream("abc-123-def")
        assert backend.resolve_workstream("abc-123-def") == "abc-123-def"

    def test_prefix_match(self, backend):
        backend.register_workstream("abc-123-def")
        assert backend.resolve_workstream("abc") == "abc-123-def"

    def test_not_found(self, backend):
        assert backend.resolve_workstream("nonexistent") is None


# -- Workstream config ---------------------------------------------------------


class TestWorkstreamConfig:
    def test_roundtrip(self, backend):
        backend.register_workstream("s1")
        backend.save_workstream_config("s1", {"temperature": "0.7", "effort": "high"})
        cfg = backend.load_workstream_config("s1")
        assert cfg == {"temperature": "0.7", "effort": "high"}

    def test_empty_config(self, backend):
        assert backend.load_workstream_config("nonexistent") == {}


# -- Workstream metadata ------------------------------------------------------


class TestWorkstreamMetadata:
    def test_alias(self, backend):
        backend.register_workstream("s1")
        assert backend.set_workstream_alias("s1", "my-session")
        assert backend.get_workstream_display_name("s1") == "my-session"

    def test_alias_conflict(self, backend):
        backend.register_workstream("s1")
        backend.register_workstream("s2")
        backend.set_workstream_alias("s1", "taken")
        assert not backend.set_workstream_alias("s2", "taken")

    def test_title(self, backend):
        backend.register_workstream("s1")
        backend.update_workstream_title("s1", "My Title")
        assert backend.get_workstream_display_name("s1") == "My Title"

    def test_alias_preferred_over_title(self, backend):
        backend.register_workstream("s1")
        backend.update_workstream_title("s1", "Title")
        backend.set_workstream_alias("s1", "Alias")
        assert backend.get_workstream_display_name("s1") == "Alias"


# -- Conversation search -------------------------------------------------------


class TestSearch:
    def test_search_history(self, backend):
        backend.register_workstream("s1")
        backend.save_message("s1", "user", "hello world")
        backend.save_message("s1", "user", "goodbye world")
        results = backend.search_history("hello")
        assert len(results) >= 1
        assert any("hello" in str(r[3]) for r in results)

    def test_search_recent(self, backend):
        backend.register_workstream("s1")
        backend.save_message("s1", "user", "msg1")
        backend.save_message("s1", "user", "msg2")
        results = backend.search_history_recent(limit=1)
        assert len(results) == 1


# -- Workstream operations -----------------------------------------------------


class TestWorkstreams:
    def test_register_and_list(self, backend):
        backend.register_workstream("ws1", node_id="node-a", name="first")
        backend.register_workstream("ws2", node_id="node-a", name="second")
        rows = backend.list_workstreams()
        assert len(rows) == 2
        ws_ids = {r[0] for r in rows}
        assert ws_ids == {"ws1", "ws2"}

    def test_register_idempotent(self, backend):
        backend.register_workstream("ws1", name="first")
        backend.register_workstream("ws1", name="overwrite")
        rows = backend.list_workstreams()
        assert len(rows) == 1
        assert rows[0][2] == "first"  # name preserved from first insert

    def test_update_state(self, backend):
        backend.register_workstream("ws1")
        backend.update_workstream_state("ws1", "running")
        rows = backend.list_workstreams()
        assert rows[0][3] == "running"

    def test_update_name(self, backend):
        backend.register_workstream("ws1", name="old")
        backend.update_workstream_name("ws1", "new")
        rows = backend.list_workstreams()
        assert rows[0][2] == "new"

    def test_delete(self, backend):
        backend.register_workstream("ws1")
        assert backend.delete_workstream("ws1") is True
        assert backend.list_workstreams() == []
        assert backend.delete_workstream("ws1") is False

    def test_list_by_node(self, backend):
        backend.register_workstream("ws1", node_id="node-a")
        backend.register_workstream("ws2", node_id="node-b")
        rows = backend.list_workstreams(node_id="node-a")
        assert len(rows) == 1
        assert rows[0][0] == "ws1"

    def test_workstream_with_messages_in_history(self, backend):
        backend.register_workstream("ws1", node_id="node-a")
        backend.save_message("ws1", "user", "hello")
        rows = backend.list_workstreams_with_history()
        assert len(rows) == 1
        # Columns: ws_id, alias, title, name, created, updated, count, node_id
        assert rows[0][0] == "ws1"
        assert rows[0][7] == "node-a"


# -- Structured memory touch ---------------------------------------------------


class TestTouchStructuredMemory:
    @staticmethod
    def _create_memory(
        backend: Any, name: str = "m1", scope: str = "global", scope_id: str = ""
    ) -> None:
        import uuid

        backend.create_structured_memory(
            memory_id=str(uuid.uuid4()),
            name=name,
            description="test desc",
            mem_type="project",
            scope=scope,
            scope_id=scope_id,
            content="test content",
        )

    def test_batch_touch_multiple(self, backend):
        self._create_memory(backend, name="a")
        self._create_memory(backend, name="b")
        self._create_memory(backend, name="c")

        count = backend.touch_structured_memories(
            [
                ("a", "global", ""),
                ("b", "global", ""),
                ("c", "global", ""),
            ]
        )
        assert count == 3

        for name in ("a", "b", "c"):
            mem = backend.get_structured_memory_by_name(name, "global", "")
            assert int(mem["access_count"]) == 1

    def test_batch_touch_empty_list(self, backend):
        assert backend.touch_structured_memories([]) == 0

    def test_batch_touch_partial_match(self, backend):
        self._create_memory(backend, name="exists")

        count = backend.touch_structured_memories(
            [
                ("exists", "global", ""),
                ("missing", "global", ""),
            ]
        )
        assert count == 1

        mem = backend.get_structured_memory_by_name("exists", "global", "")
        assert int(mem["access_count"]) == 1

    def test_batch_touch_with_duplicates(self, backend):
        """Duplicate keys in batch should each increment access_count once."""
        self._create_memory(backend, name="dup")

        # Two identical keys — storage gets called twice for the same row
        count = backend.touch_structured_memories([("dup", "global", ""), ("dup", "global", "")])
        assert count == 2

        mem = backend.get_structured_memory_by_name("dup", "global", "")
        assert int(mem["access_count"]) == 2


# -- Per-workstream usage aggregation -----------------------------------------


class TestSumWorkstreamTokens:
    """``sum_workstream_tokens`` powers the inspect-time token fallback for
    idle children — a regression here would surface as wrong tokens in
    the coordinator's inspect output rather than a focused test failure,
    so guard it directly."""

    def test_empty_ws_id_returns_zero(self, backend):
        assert backend.sum_workstream_tokens("") == 0

    def test_no_events_returns_zero(self, backend):
        assert backend.sum_workstream_tokens("never-seen") == 0

    def test_sums_prompt_and_completion_across_events(self, backend):
        backend.record_usage_event(
            event_id="e1", ws_id="ws-a", prompt_tokens=10, completion_tokens=5
        )
        backend.record_usage_event(
            event_id="e2", ws_id="ws-a", prompt_tokens=200, completion_tokens=80
        )
        assert backend.sum_workstream_tokens("ws-a") == 10 + 5 + 200 + 80

    def test_scoped_to_requested_ws_id(self, backend):
        """Other workstreams' usage events must not leak into the sum."""
        backend.record_usage_event(
            event_id="e1", ws_id="ws-a", prompt_tokens=100, completion_tokens=50
        )
        backend.record_usage_event(
            event_id="e2", ws_id="ws-b", prompt_tokens=999, completion_tokens=999
        )
        assert backend.sum_workstream_tokens("ws-a") == 150
        assert backend.sum_workstream_tokens("ws-b") == 1998


class TestBatchPrimitives:
    """``get_workstreams_batch`` and ``sum_workstream_tokens_batch`` power
    ``wait_for_workstream``'s per-tick polling.  Direct backend coverage
    here so a regression surfaces as a focused failure rather than as
    wrong tokens / spurious denied states in a coordinator session."""

    def test_get_workstreams_batch_empty_input(self, backend):
        assert backend.get_workstreams_batch([]) == {}

    def test_get_workstreams_batch_returns_row_per_id(self, backend):
        backend.register_workstream("a", title="A", kind="interactive")
        backend.register_workstream("b", title="B", kind="interactive", parent_ws_id="a")
        result = backend.get_workstreams_batch(["a", "b"])
        assert set(result.keys()) == {"a", "b"}
        assert result["a"]["ws_id"] == "a"
        assert result["b"]["parent_ws_id"] == "a"

    def test_get_workstreams_batch_missing_id_returns_none(self, backend):
        backend.register_workstream("a")
        result = backend.get_workstreams_batch(["a", "missing"])
        assert result["a"] is not None
        assert result["missing"] is None

    def test_get_workstreams_batch_drops_empty_strings(self, backend):
        """Empty / non-string ids must not pollute the IN clause."""
        backend.register_workstream("a")
        result = backend.get_workstreams_batch(["a", "", "  "])
        # Only the non-empty id is kept; whitespace-only strings are
        # passed through (the helper only strips truly-empty entries).
        assert "a" in result
        assert result["a"] is not None

    def test_sum_workstream_tokens_batch_empty_input(self, backend):
        assert backend.sum_workstream_tokens_batch([]) == {}

    def test_sum_workstream_tokens_batch_aggregates_per_id(self, backend):
        backend.record_usage_event(event_id="e1", ws_id="a", prompt_tokens=10, completion_tokens=5)
        backend.record_usage_event(event_id="e2", ws_id="a", prompt_tokens=20, completion_tokens=10)
        backend.record_usage_event(
            event_id="e3", ws_id="b", prompt_tokens=100, completion_tokens=50
        )
        result = backend.sum_workstream_tokens_batch(["a", "b", "c"])
        assert result == {"a": 45, "b": 150, "c": 0}

    def test_sum_workstream_tokens_batch_missing_id_defaults_zero(self, backend):
        result = backend.sum_workstream_tokens_batch(["never-seen"])
        assert result == {"never-seen": 0}


# -- Lifecycle -----------------------------------------------------------------


class TestLifecycle:
    def test_close(self, backend):
        backend.close()  # Should not raise

    def test_isinstance_check(self, backend):
        from turnstone.core.storage._protocol import StorageBackend

        assert isinstance(backend, StorageBackend)