ewilliams28/turnstone
1
0
Fork 0
mirror of https://github.com/turnstonelabs/turnstone.git synced 2026-08-12 23:12:23 -06:00
Code Issues Packages Projects Releases Wiki Activity
Files
stable/1.7
turnstone/tests/test_storage_sqlite.py
T
Patrick Buckley aeab2535b1 feat(preview): rich preview pane + open_preview tool 
Tool results only ever rendered as plain text in the transcript. This
adds the model-driven rich-preview lane every comparable surface has,
in turnstone's developer-tool idiom: a preview pane that opens BESIDE
the conversation, keyboard-operable, sandboxed, never replacing the
transcript that spawned it.

Backend
- New built-in open_preview(target, kind?, title?): resolves an http(s)
  URL, a file path, or attachment:<id> to bytes; classifies into
  web/pdf/image/table/text/markdown (magic bytes > MIME hint >
  extension > UTF-8 fallback, legacy-charset pages transcoded); caps
  size per kind; persists content-addressed with kind="preview" —
  refcounted and GC'd with the workstream, skipped by trajectory
  reconstruction so preview bytes can never materialize onto the wire.
  URL targets gate like web_fetch (network egress); paths/attachments
  run unprompted like read_file.
- New core.web.fetch_with_ssrf_guard: manual redirect walk that
  SSRF-screens every hop BEFORE requesting it (follow_redirects=True
  checked nothing between hops); adopted by both open_preview and
  web_fetch. URL userinfo is stripped before the descriptor or the
  stored bytes see it; <base href> is injected doctype-safely so
  relative assets resolve without quirks mode.
- The preview descriptor rides the tool turn's meta side channel with
  ONE shape on every boundary: the live tool_result SSE event, the
  conversations.meta column, and the /history projection. Cancelled
  batches commit an already-announced preview (blob + meta) instead of
  stranding the open pane on a permanent 404.
- New GET {ws}/attachments/{id}/preview (read scope, same ownership
  gate as /content) serves the STORED type with per-MIME hardening:
  bare CSP sandbox for text/html (renderable, scriptless, opaque
  origin), no CSP for application/pdf (Chromium's viewer refuses
  sandboxed contexts), full default-src 'none' otherwise; filenames
  fold to latin-1-safe ASCII. The console /node proxy now forwards
  CSP/nosniff/disposition/cache-control instead of dropping them.
- History loads exclude preview blobs from the bulk content fetch at
  the query (they were read and discarded on every load).

Frontend
- New "preview" pane type registered in the shared shell (server +
  console): openPaneBeside placement, per-kind renderers — fully
  sandboxed iframe for pages, browser PDF viewer, sortable tables
  (CSV/TSV/JSON, ragged-file safe, 5k-row cap), rendered markdown,
  text — plus back/forward history with arrow keys, reload persistence
  via pane meta, and backoff auto-retry (0.9s..7.2s) bridging the gap
  between the live descriptor and the batch fold that commits its blob.
- Tool results carrying a descriptor render a credential-redacted
  preview chip (the reopen + replay affordance); live results auto-open
  the pane only while the originating pane holds focus.

Docs: docs/tools.md + prompts/tools.md. Tests: policy unit tests, tool
prepare/exec (mocked fetch), serving route + proxy header pass-through,
storage exclusion on both backends, cancel-path commit, JS static
guards; a headless-Chrome harness drives the real module graph (32 DOM
assertions).

(cherry picked from commit e010124008)
2026-07-08 17:30:55 -07:00

1169 lines
49 KiB
Python
Raw Permalink Blame History

"""Tests for the SQLite storage backend."""
from __future__ import annotations
from typing import Any
import pytest
import sqlalchemy as sa
from turnstone.core.storage._schema import workstreams
# -- 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_system_turn_meta_roundtrip(self, backend):
# The structured per-kind operator meta saved on a ``system`` row
# round-trips back as ``_source_meta`` through the real insert + SELECT
# (the ``conversations.meta`` column, migration 060).
import json
backend.register_workstream("s1")
backend.save_message("s1", "user", "go")
backend.save_message(
"s1",
"system",
"ci failed",
source="watch_triggered",
meta=json.dumps({"watch_name": "ci", "command": "make test", "poll_count": 3}),
)
msgs = backend.load_messages("s1")
assert msgs[1] == {
"role": "system",
"content": "ci failed",
"_source": "watch_triggered",
"_source_meta": {"watch_name": "ci", "command": "make test", "poll_count": 3},
}
def test_ordinary_row_has_null_meta(self, backend):
# A non-operator row carries no meta — no ``_source_meta`` on reload.
backend.register_workstream("s1")
backend.save_message("s1", "user", "hello")
msgs = backend.load_messages("s1")
assert "_source_meta" not in msgs[0]
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")
# The native lane is a block list (or the {producer, blocks} envelope);
# it round-trips verbatim through _provider_content.
blocks = [{"type": "thinking", "thinking": "secret", "signature": "s"}]
backend.save_message("s1", "assistant", "hi", provider_data=json.dumps(blocks))
msgs = backend.load_messages("s1")
assert msgs[0].get("_provider_content") == blocks
def test_empty_workstream_returns_empty(self, backend):
assert backend.load_messages("nonexistent") == []
class TestListMessageSenders:
def test_distinct_senders_from_user_rows_only(self, backend):
import json
backend.register_workstream("s1")
backend.save_message("s1", "user", "a", meta=json.dumps({"sender": "alice"}))
backend.save_message("s1", "user", "b", meta=json.dumps({"sender": "bob"}))
backend.save_message("s1", "user", "c", meta=json.dumps({"sender": "alice"}))
backend.save_message("s1", "user", "plain") # unstamped: meta is NULL
# A system row's meta rides the source_meta channel; even a stray
# "sender" key there must never count as a participant.
backend.save_message(
"s1", "system", "note", source="watch_triggered", meta=json.dumps({"sender": "evil"})
)
backend.register_workstream("s2")
backend.save_message("s2", "user", "x", meta=json.dumps({"sender": "carol"}))
assert backend.list_message_senders("s1") == ["alice", "bob"]
assert backend.list_message_senders("s2") == ["carol"] # ws-scoped
assert backend.list_message_senders("nope") == []
def test_garbage_meta_is_skipped(self, backend):
backend.register_workstream("s1")
backend.save_message("s1", "user", "a", meta="not json{")
backend.save_message("s1", "user", "b", meta='"just a string"')
backend.save_message("s1", "user", "c", meta='{"sender": " "}')
backend.save_message("s1", "user", "d", meta='{"sender": 7}')
assert backend.list_message_senders("s1") == []
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):
"""The content-addressed attachment resolution only fetches blobs the
*fetched* rows reference — so a tail-N load that doesn't include the
attachment-bearing row issues no blob fetch at all, and a full load
fetches exactly the referenced ids. This keeps the tail-N conversations
LIMIT from being undone by a full-workstream attachment scan."""
from unittest.mock import patch
backend.register_workstream("s1")
# Oldest row carries the attachment; then 20 plain rows after it.
aid = "a" * 64
att_msg_id = backend.save_message("s1", "user", "see attachment")
backend.save_attachment(aid, "x.txt", "text/plain", 1, "text", b"x")
backend.set_message_attachments("s1", att_msg_id, [aid])
for i in range(20):
backend.save_message("s1", "user", f"msg-{i:02d}")
captured: list[list[str]] = []
orig = backend.get_attachments
def _spy(ids, exclude_kinds=()):
captured.append(sorted(ids))
return orig(ids, exclude_kinds=exclude_kinds)
# Tail-N=5 fetches only the 5 newest rows (all plain) — the
# attachment row is excluded, so NO blob fetch is issued.
with patch.object(backend, "get_attachments", side_effect=_spy):
backend.load_messages("s1", limit=5)
assert captured == []
# Full load resolves exactly the one referenced id (not a full scan).
captured.clear()
with patch.object(backend, "get_attachments", side_effect=_spy):
backend.load_messages("s1")
assert captured == [[aid]]
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"}
def test_enriched_columns(self, backend):
"""The saved-list query carries the enrichment trailing columns:
node_id, state, model_alias + launch_skill (workstream_config),
child_count (parent_ws_id), context_tokens (latest usage_events row)
and context_window (model_definitions join)."""
from turnstone.core.workstream import WorkstreamKind
backend.register_workstream(
"parent", node_id="n1", state="error", kind=WorkstreamKind.COORDINATOR
)
backend.save_message("parent", "user", "orchestrate")
backend.save_workstream_config("parent", {"model_alias": "m1", "skill": "news"})
# child via parent_ws_id → child_count = 1 (no message needed; the
# child_count subquery doesn't gate on EXISTS conversation)
backend.register_workstream("child", kind=WorkstreamKind.INTERACTIVE, parent_ws_id="parent")
backend.record_usage_event("e-parent", ws_id="parent", prompt_tokens=250)
# a usage event on a different ws must not bleed into parent's tokens
backend.record_usage_event("e-other", ws_id="child", prompt_tokens=999)
backend.create_model_definition("d1", alias="m1", model="m1-model", context_window=1000)
rows = backend.list_workstreams_with_history()
row = next(r for r in rows if r[0] == "parent")
# (ws_id, alias, title, name, created, updated, message_count,
# node_id, state, kind, model_alias, launch_skill, child_count,
# context_tokens, context_window)
assert row[7] == "n1" # node_id
assert row[8] == "error" # state
assert row[10] == "m1" # model_alias
assert row[11] == "news" # launch_skill
assert row[12] == 1 # child_count
assert row[13] == 250 # context_tokens — parent's event, not child's 999
assert row[14] == 1000 # context_window from model_definitions
def test_enriched_columns_null_when_absent(self, backend):
"""A bare workstream (no config / usage / model_def / children) leaves
the enrichment columns NULL / zero — the LEFT JOIN + subquery misses
degrade gracefully (the handler coerces these to defaults)."""
backend.register_workstream("bare")
backend.save_message("bare", "user", "hi")
rows = backend.list_workstreams_with_history()
row = next(r for r in rows if r[0] == "bare")
assert row[10] is None # model_alias — no config row
assert row[11] is None # launch_skill
assert row[12] == 0 # child_count
assert row[13] is None # context_tokens — no usage events
assert row[14] is None # context_window — no model def
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
def test_search_history_survives_oversized_row(self, backend):
# A multi-MB row of mostly-unique words: on PostgreSQL its full
# tsvector exceeds the 1MB hard limit, which used to abort every
# search_history scan ("string is too long for tsvector") — one
# giant tool dump silently killed history recall entirely.
backend.register_workstream("s1")
giant = "gargantuan beacon " + " ".join(f"w{i}" for i in range(300_000))
assert len(giant) > 2_000_000
backend.save_message("s1", "tool", giant)
backend.save_message("s1", "user", "hello world")
results = backend.search_history("hello")
assert any("hello" in str(r[3]) for r in results)
# The oversized row itself stays findable by its head.
results = backend.search_history("gargantuan beacon")
assert any("gargantuan" in str(r[3]) for r in results)
def test_search_history_fts_error_falls_back_to_ilike(self, request, backend, monkeypatch):
# PostgreSQL only: a failed FTS statement aborts the connection's
# autobegun transaction, and the ILIKE fallback runs on that same
# connection — without a rollback first it dies with
# InFailedSqlTransaction instead of returning results.
if request.config.getoption("--storage-backend") != "postgresql":
pytest.skip("exercises PostgreSQL aborted-transaction fallback")
backend.register_workstream("s1")
backend.save_message("s1", "user", "hello fallback world")
real_execute = sa.engine.Connection.execute
def failing_fts_execute(self, statement, *args, **kwargs):
if "to_tsvector" in str(statement):
# A genuine server-side error, so the transaction is aborted
# exactly as when to_tsvector rejects a row.
return real_execute(self, sa.text("SELECT 1/0"))
return real_execute(self, statement, *args, **kwargs)
monkeypatch.setattr(sa.engine.Connection, "execute", failing_fts_execute)
results = backend.search_history("fallback")
assert any("fallback" in str(r[3]) for r in results)
# -- 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="general",
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}
# -- bulk_close_stale_orphans --------------------------------------------------
def _force_updated(backend: Any, ws_id: str, updated: str) -> None:
"""Stamp a workstream row's ``updated`` column directly.
The public surface only sets ``updated`` to ``now``, which makes it
impossible to fabricate a stale row through register/update calls.
Reaches into ``backend._engine`` — same access pattern conftest uses
for cross-backend cleanup.
"""
with backend._engine.connect() as conn:
conn.execute(
sa.update(workstreams).where(workstreams.c.ws_id == ws_id).values(updated=updated)
)
conn.commit()
class TestBulkCloseStaleOrphans:
def test_closes_stale_non_terminal_rows_of_kind(self, backend):
backend.register_workstream("stale-idle", kind="interactive")
backend.register_workstream("stale-thinking", kind="interactive")
backend.update_workstream_state("stale-thinking", "thinking")
backend.register_workstream("fresh-idle", kind="interactive")
_force_updated(backend, "stale-idle", "2020-01-01T00:00:00")
_force_updated(backend, "stale-thinking", "2020-01-01T00:00:00")
# fresh-idle stays at registration time (effectively now)
closed = backend.bulk_close_stale_orphans(
"interactive", cutoff="2024-01-01T00:00:00", exclude_ws_ids=[]
)
assert set(closed) == {"stale-idle", "stale-thinking"}
rows = backend.get_workstreams_batch(["stale-idle", "stale-thinking", "fresh-idle"])
assert rows["stale-idle"]["state"] == "closed"
assert rows["stale-thinking"]["state"] == "closed"
assert rows["fresh-idle"]["state"] == "idle"
def test_skips_already_closed(self, backend):
backend.register_workstream("already-closed", kind="interactive")
backend.update_workstream_state("already-closed", "closed")
_force_updated(backend, "already-closed", "2020-01-01T00:00:00")
closed = backend.bulk_close_stale_orphans(
"interactive", cutoff="2024-01-01T00:00:00", exclude_ws_ids=[]
)
assert closed == []
def test_filters_by_kind(self, backend):
backend.register_workstream("interactive-stale", kind="interactive")
backend.register_workstream("coord-stale", kind="coordinator")
_force_updated(backend, "interactive-stale", "2020-01-01T00:00:00")
_force_updated(backend, "coord-stale", "2020-01-01T00:00:00")
closed = backend.bulk_close_stale_orphans(
"interactive", cutoff="2024-01-01T00:00:00", exclude_ws_ids=[]
)
assert closed == ["interactive-stale"]
rows = backend.get_workstreams_batch(["interactive-stale", "coord-stale"])
assert rows["interactive-stale"]["state"] == "closed"
assert rows["coord-stale"]["state"] == "idle"
def test_excludes_loaded_ws_ids(self, backend):
backend.register_workstream("ws-keep", kind="interactive")
backend.register_workstream("ws-close", kind="interactive")
_force_updated(backend, "ws-keep", "2020-01-01T00:00:00")
_force_updated(backend, "ws-close", "2020-01-01T00:00:00")
closed = backend.bulk_close_stale_orphans(
"interactive", cutoff="2024-01-01T00:00:00", exclude_ws_ids=["ws-keep"]
)
assert closed == ["ws-close"]
rows = backend.get_workstreams_batch(["ws-keep", "ws-close"])
assert rows["ws-keep"]["state"] == "idle"
assert rows["ws-close"]["state"] == "closed"
def test_empty_exclude_list_does_not_break_sql(self, backend):
backend.register_workstream("orphan", kind="interactive")
_force_updated(backend, "orphan", "2020-01-01T00:00:00")
closed = backend.bulk_close_stale_orphans(
"interactive", cutoff="2024-01-01T00:00:00", exclude_ws_ids=[]
)
assert closed == ["orphan"]
def test_no_orphans_returns_empty(self, backend):
backend.register_workstream("fresh", kind="interactive")
closed = backend.bulk_close_stale_orphans(
"interactive", cutoff="2024-01-01T00:00:00", exclude_ws_ids=[]
)
assert closed == []
def test_closes_all_non_terminal_states(self, backend):
for ws_id, state in [
("o-idle", "idle"),
("o-thinking", "thinking"),
("o-attention", "attention"),
("o-running", "running"),
]:
backend.register_workstream(ws_id, kind="interactive")
if state != "idle":
backend.update_workstream_state(ws_id, state)
_force_updated(backend, ws_id, "2020-01-01T00:00:00")
closed = backend.bulk_close_stale_orphans(
"interactive", cutoff="2024-01-01T00:00:00", exclude_ws_ids=[]
)
assert set(closed) == {"o-idle", "o-thinking", "o-attention", "o-running"}
def test_bumps_updated_on_close(self, backend):
stale_updated = "2020-01-01T00:00:00"
backend.register_workstream("orphan", kind="interactive")
_force_updated(backend, "orphan", stale_updated)
backend.bulk_close_stale_orphans(
"interactive", cutoff="2024-01-01T00:00:00", exclude_ws_ids=[]
)
# ``updated`` must change away from the forced stale value. Asserting
# inequality from the seed (rather than ``> "2024-01-01..."``) keeps
# the test independent of wall-clock date.
with backend._engine.connect() as conn:
row = conn.execute(
sa.select(workstreams.c.updated).where(workstreams.c.ws_id == "orphan")
).one()
assert row[0] != stale_updated
def test_protects_rows_owned_by_live_services(self, backend):
"""Liveness scoping (post-#384 rendezvous-routing world): rows
whose ``node_id`` matches a heartbeating service must NOT be
reaped, because that owner may legitimately have them loaded on
another worker. Rows whose ``node_id`` matches a dead service
ARE eligible — that's how dead-pod orphans get reclaimed in
containerized deployments with dynamic hostnames."""
backend.register_workstream("dead-node", node_id="dead-pod-x4k2", kind="interactive")
backend.register_workstream("alive-node", node_id="alive-pod-y9p3", kind="interactive")
_force_updated(backend, "dead-node", "2020-01-01T00:00:00")
_force_updated(backend, "alive-node", "2020-01-01T00:00:00")
closed = backend.bulk_close_stale_orphans(
"interactive",
cutoff="2024-01-01T00:00:00",
exclude_ws_ids=[],
live_node_ids=["alive-pod-y9p3"],
)
assert closed == ["dead-node"]
rows = backend.get_workstreams_batch(["dead-node", "alive-node"])
assert rows["dead-node"]["state"] == "closed"
assert rows["alive-node"]["state"] == "idle"
def test_null_node_id_always_eligible(self, backend):
"""A row with NULL ``node_id`` has no owner identity — age alone
gates the reap. Belt-and-suspenders against ``NULL NOT IN (...)``
evaluating to NULL (not TRUE) and silently protecting orphans
forever."""
backend.register_workstream("no-owner", node_id=None, kind="interactive")
_force_updated(backend, "no-owner", "2020-01-01T00:00:00")
closed = backend.bulk_close_stale_orphans(
"interactive",
cutoff="2024-01-01T00:00:00",
exclude_ws_ids=[],
live_node_ids=["some-other-node"],
)
assert closed == ["no-owner"]
def test_live_node_ids_none_skips_filter(self, backend):
"""``live_node_ids=None`` is the single-process / operator-backfill
mode — all rows of *kind* are eligible regardless of node_id."""
backend.register_workstream("node-a", node_id="node-a", kind="interactive")
backend.register_workstream("node-b", node_id="node-b", kind="interactive")
_force_updated(backend, "node-a", "2020-01-01T00:00:00")
_force_updated(backend, "node-b", "2020-01-01T00:00:00")
closed = backend.bulk_close_stale_orphans(
"interactive", cutoff="2024-01-01T00:00:00", exclude_ws_ids=[]
)
assert set(closed) == {"node-a", "node-b"}
def test_empty_live_node_ids_treats_all_as_dead(self, backend):
"""Empty list ``live_node_ids=[]`` means "no nodes alive" — every
row's owner is unprotected. Useful for operator scripts that
want to reap regardless of liveness."""
backend.register_workstream("any", node_id="node-a", kind="interactive")
_force_updated(backend, "any", "2020-01-01T00:00:00")
closed = backend.bulk_close_stale_orphans(
"interactive",
cutoff="2024-01-01T00:00:00",
exclude_ws_ids=[],
live_node_ids=[],
)
assert closed == ["any"]
def test_combines_live_node_ids_and_exclude_ws_ids(self, backend):
"""Both filters stack as AND clauses on the UPDATE. Covers the
full 2x2 matrix to catch a future edit that replaces an AND with
an OR or drops one of the filters: only the (orphan + dead-node)
cell should be reaped."""
# All four registered with the same stale ``updated``.
for ws_id, node in [
("loaded-alive", "alive-node"),
("loaded-dead", "dead-node"),
("orphan-alive", "alive-node"),
("orphan-dead", "dead-node"),
]:
backend.register_workstream(ws_id, node_id=node, kind="interactive")
_force_updated(backend, ws_id, "2020-01-01T00:00:00")
closed = backend.bulk_close_stale_orphans(
"interactive",
cutoff="2024-01-01T00:00:00",
exclude_ws_ids=["loaded-alive", "loaded-dead"],
live_node_ids=["alive-node"],
)
# Only orphan-dead is unprotected by both filters.
assert closed == ["orphan-dead"]
rows = backend.get_workstreams_batch(
["loaded-alive", "loaded-dead", "orphan-alive", "orphan-dead"]
)
assert rows["loaded-alive"]["state"] == "idle"
assert rows["loaded-dead"]["state"] == "idle"
assert rows["orphan-alive"]["state"] == "idle"
assert rows["orphan-dead"]["state"] == "closed"
# -- touch_workstream ----------------------------------------------------------
class TestTouchWorkstream:
def test_bumps_updated_only(self, backend):
"""Used by ``open()`` on rehydrate to defend against the orphan
reaper clobbering a freshly-loaded row. Must not change ``state``
(the open() path explicitly avoids state writes to dodge a race
with concurrent close())."""
stale_updated = "2020-01-01T00:00:00"
backend.register_workstream("ws-touch", kind="interactive")
backend.update_workstream_state("ws-touch", "closed") # simulate prior close
_force_updated(backend, "ws-touch", stale_updated)
backend.touch_workstream("ws-touch")
with backend._engine.connect() as conn:
row = conn.execute(
sa.select(workstreams.c.state, workstreams.c.updated).where(
workstreams.c.ws_id == "ws-touch"
)
).one()
assert row[0] == "closed", "state must not be modified by touch"
# Compare against the forced stale value rather than a fixed calendar
# date so the test is independent of wall-clock time.
assert row[1] != stale_updated, "updated must be bumped"
def test_unknown_id_is_noop(self, backend):
"""Touch on a missing id must not raise — open()'s exception
handler is best-effort."""
backend.touch_workstream("nonexistent") # must not raise
# -- MCP OAuth columns ---------------------------------------------------------
class TestMcpServerOauthColumns:
def test_mcp_servers_oauth_columns_round_trip(self, backend: Any) -> None:
"""An oauth_user row round-trips through create -> get with all
seven OAuth text columns intact."""
sid = "oauth-srv-1"
backend.create_mcp_server(
server_id=sid,
name="oauth-srv",
transport="streamable-http",
url="https://mcp.example.com/sse",
auth_type="oauth_user",
oauth_client_id="cli_abc123",
oauth_scopes="openid profile",
oauth_audience="https://mcp.example.com",
oauth_registration_mode="preregistered",
oauth_authorization_server_url="https://auth.example.com",
oauth_as_issuer_cached="https://auth.example.com",
)
s = backend.get_mcp_server(sid)
assert s is not None
assert s["auth_type"] == "oauth_user"
assert s["oauth_client_id"] == "cli_abc123"
assert s["oauth_scopes"] == "openid profile"
assert s["oauth_audience"] == "https://mcp.example.com"
assert s["oauth_registration_mode"] == "preregistered"
assert s["oauth_authorization_server_url"] == "https://auth.example.com"
assert s["oauth_as_issuer_cached"] == "https://auth.example.com"
# Phase 2 leaves the ciphertext slot NULL even when other oauth
# fields are populated; Phase 3 wires the encryption write path.
assert s["oauth_client_secret_ct"] is None
def test_update_auth_type_static_to_oauth(self, backend: Any) -> None:
sid = "oauth-srv-2"
backend.create_mcp_server(
server_id=sid,
name="static-then-oauth",
transport="streamable-http",
url="https://mcp.example.com/sse",
)
assert backend.get_mcp_server(sid)["auth_type"] == "static"
ok = backend.update_mcp_server(
sid,
auth_type="oauth_user",
oauth_client_id="cli_after",
oauth_audience="https://mcp.example.com",
)
assert ok is True
s = backend.get_mcp_server(sid)
assert s is not None
assert s["auth_type"] == "oauth_user"
assert s["oauth_client_id"] == "cli_after"
assert s["oauth_audience"] == "https://mcp.example.com"
# -- 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)
Reference in New Issue View Git Blame Copy Permalink