From 7f20b1bc8417f7ef7401eb40e193ab5a9f0877eb Mon Sep 17 00:00:00 2001 From: Patrick Buckley Date: Thu, 2 Jul 2026 16:31:39 -0700 Subject: [PATCH] fix(session): durable shared-workstream state + fork sender persistence - _known_senders/_shared_workstream are now monotonic: union-only growth, latched shared flag, seeded once per workstream from a full-history distinct-sender read (new StorageBackend.list_message_senders) so compaction narrowing the resumable slice can no longer forget participants (duplicate join notes) or flip the banner back to single-user framing (prompt-prefix cache churn). - Recompute is memoized per turn (invalidated on stamped user-turn append); system-prompt composition no longer pays an O(n) trajectory scan on every recompose. - resume() resets the state: the monotonic guarantees are per workstream, not per session object. - resume(fork=True) bulk-persist now carries the user-turn sender stamp into the fork's meta column (was: _source_meta only, which dropped attribution for every forked user turn on reopen). --- tests/test_per_user_message_context.py | 137 ++++++++++++++++++++++++- tests/test_storage_sqlite.py | 29 ++++++ turnstone/core/session.py | 121 ++++++++++++++++++---- turnstone/core/storage/_postgresql.py | 17 +++ turnstone/core/storage/_protocol.py | 10 ++ turnstone/core/storage/_sqlite.py | 17 +++ turnstone/core/storage/_utils.py | 26 ++++- 7 files changed, 331 insertions(+), 26 deletions(-) diff --git a/tests/test_per_user_message_context.py b/tests/test_per_user_message_context.py index 43c35bd8..63be099c 100644 --- a/tests/test_per_user_message_context.py +++ b/tests/test_per_user_message_context.py @@ -219,15 +219,96 @@ def test_labels_render_resolved_usernames(): def test_recompute_shared_state_from_history(): s = make_session(user_id="owner") - s.messages.append(turn_from_dict({"role": "user", "content": "a", "_sender": "owner"})) - s._recompute_shared_state() - assert s._shared_workstream is False # owner alone is not shared - s.messages.append(turn_from_dict({"role": "user", "content": "b", "_sender": "alice"})) - s._recompute_shared_state() + with patch("turnstone.core.session.get_storage", return_value=None): + s.messages.append(turn_from_dict({"role": "user", "content": "a", "_sender": "owner"})) + s._invalidate_shared_state() # what _append_user_turn does for stamped turns + s._recompute_shared_state() + assert s._shared_workstream is False # owner alone is not shared + s.messages.append(turn_from_dict({"role": "user", "content": "b", "_sender": "alice"})) + s._invalidate_shared_state() + s._recompute_shared_state() assert s._shared_workstream is True assert s._known_senders == {"owner", "alice"} +def test_shared_state_latches_and_senders_never_shrink(): + # Compaction narrows self.messages to [summary]+[tail]; a participant whose + # turns were summarized away must stay known (no duplicate join note) and + # the workstream must stay shared (no banner flip, no prefix-cache churn). + s = make_session(user_id="owner") + with patch("turnstone.core.session.get_storage", return_value=None): + s.messages.append(turn_from_dict({"role": "user", "content": "a", "_sender": "alice"})) + s._invalidate_shared_state() + s._recompute_shared_state() + assert s._shared_workstream is True + # compaction-style narrowing: alice's turns vanish from the slice + s.messages = [turn_from_dict({"role": "user", "content": "s", "_sender": "owner"})] + s._invalidate_shared_state() + s._recompute_shared_state() + assert s._shared_workstream is True # latched + assert "alice" in s._known_senders # union, never overwrite + # ...so the returning participant does not re-fire the join note + n = len(s.messages) + s._maybe_note_new_participant("alice") + assert len(s.messages) == n + + +def test_recompute_unions_persisted_senders_once(): + # A rehydrating worker sees only the checkpointed slice; the one-time + # full-history read recovers participants summarized out of it. + s = make_session(user_id="owner") + s._reset_shared_state() # the state resume() leaves behind + fake = MagicMock() + fake.list_message_senders.return_value = ["alice"] + with patch("turnstone.core.session.get_storage", return_value=fake): + s._recompute_shared_state() + assert s._shared_workstream is True + assert "alice" in s._known_senders + s._invalidate_shared_state() + s._recompute_shared_state() # second turn: no second full-history read + fake.list_message_senders.assert_called_once() + + +def test_persisted_sender_read_retries_after_storage_error(): + # A transient storage error must not pin an incomplete participant set: + # the next recompute (next user turn) retries the full-history read. + s = make_session(user_id="owner") + s._reset_shared_state() + fake = MagicMock() + fake.list_message_senders.side_effect = [RuntimeError("storage down"), ["alice"]] + with patch("turnstone.core.session.get_storage", return_value=fake): + s._recompute_shared_state() # error -> degraded this turn, not cached + assert s._shared_workstream is False + s._invalidate_shared_state() # next user turn + s._recompute_shared_state() # retried, recovered + assert s._shared_workstream is True + assert fake.list_message_senders.call_count == 2 + + +def test_recompute_is_memoized_per_turn(): + # _init_system_messages fires many times within a turn; between user-turn + # appends the recompute is a no-op flag check, not an O(n) rescan. + s = make_session(user_id="owner") + with patch("turnstone.core.session.get_storage", return_value=None): + s._reset_shared_state() + s._recompute_shared_state() + s.messages.append(turn_from_dict({"role": "user", "content": "b", "_sender": "alice"})) + s._recompute_shared_state() # memoized: append not yet visible + assert s._shared_workstream is False + s._invalidate_shared_state() # what _append_user_turn does + s._recompute_shared_state() + assert s._shared_workstream is True + + +def test_append_user_turn_invalidates_shared_state(): + s = make_session(user_id="owner") + s._acting_user_id = "alice" + with patch("turnstone.core.session.save_message", return_value=1): + s._senders_dirty = False + s._append_user_turn("hello", ()) + assert s._senders_dirty is True + + def test_new_participant_flips_shared_and_emits_join_note_once(): s = make_session(user_id="owner") s._known_senders = {"owner"} @@ -255,6 +336,52 @@ def test_owner_only_never_shared(): recompose.assert_not_called() +# -- resume / fork carry attribution across the DB round-trip ----------------- + + +def test_resume_resets_shared_state(): + # resume() can point this session object at a different workstream's + # history; the monotonic shared-state guarantees are per workstream. + s = make_session(user_id="owner") + s._known_senders = {"alice"} + s._shared_workstream = True + turns = [turn_from_dict({"role": "user", "content": "x", "_sender": "owner"})] + with ( + patch("turnstone.core.session.load_message_turns", return_value=turns), + patch("turnstone.core.session.get_storage", return_value=None), + patch.object(s, "_reset_shared_state", wraps=s._reset_shared_state) as rst, + patch.object(s, "_save_config"), + patch.object(s, "_init_system_messages"), + ): + assert s.resume("ws-other") is True + rst.assert_called_once() + + +def test_fork_persists_sender_meta(): + # The fork bulk-persist must carry the user-turn sender stamp into the + # fork's rows (mirroring _append_user_turn), or the fork loses per-user + # attribution the first time it is reopened from the DB. + s = make_session(user_id="owner") + turns = [ + turn_from_dict({"role": "user", "content": "hi", "_sender": "alice"}), + turn_from_dict({"role": "user", "content": "wake", "_source": "wake"}), + turn_from_dict({"role": "assistant", "content": "yo"}), + ] + with ( + patch("turnstone.core.session.load_message_turns", return_value=turns), + patch("turnstone.core.session.save_messages_bulk") as bulk, + patch("turnstone.core.session.get_storage", return_value=None), + patch.object(s, "_save_config"), + patch.object(s, "_init_system_messages"), + ): + assert s.resume("src-ws", fork=True) is True + rows = bulk.call_args.args[0] + by_content = {r["content"]: r for r in rows} + assert json.loads(by_content["hi"]["meta"]) == {"sender": "alice"} + assert by_content["wake"]["meta"] is None # synthetic: no sender stamped + assert by_content["yo"]["meta"] is None # assistant rows carry no sender + + # -- Session Context banner (shared vs single-user) --------------------------- diff --git a/tests/test_storage_sqlite.py b/tests/test_storage_sqlite.py index 9494a9fc..cb80c616 100644 --- a/tests/test_storage_sqlite.py +++ b/tests/test_storage_sqlite.py @@ -132,6 +132,35 @@ class TestSaveAndLoadMessages: 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 diff --git a/turnstone/core/session.py b/turnstone/core/session.py index ce71a472..2c340574 100644 --- a/turnstone/core/session.py +++ b/turnstone/core/session.py @@ -1210,14 +1210,21 @@ class ChatSession: # the ``(user_id, callback)`` pair). self._acting_user_id: str = "" self._mcp_listener_user_id: str | None = user_id or None - # Shared-workstream context state (context-identity layer, atop the - # acting-user credential fix above): the model must be TOLD when more - # than one human is in the room. ``_shared_workstream`` flips True once a - # non-owner sender appears (live send OR rehydrated history) and drives - # the ``## Session Context`` banner; ``_known_senders`` tracks who has - # spoken so a first-time participant gets a one-time "has joined" note. + # Shared-workstream context state: the model must be TOLD when more + # than one human is in the room. ``_shared_workstream`` flips True once + # a non-owner sender appears (live send OR rehydrated history) and + # drives the ``## Session Context`` banner; it LATCHES — reverting + # would misframe a known-multi-user conversation and churn the + # provider-cached prompt prefix. ``_known_senders`` (everyone who has + # ever spoken here) gates the one-time "has joined" note and only + # grows: deriving it from the in-memory slice alone would forget + # participants once compaction narrows history. ``_senders_dirty`` + # memoizes recomputes per turn (the composer runs many times per + # turn); ``_db_senders_loaded`` marks the one-time full-history read. self._shared_workstream: bool = False self._known_senders: set[str] = set() + self._senders_dirty: bool = True + self._db_senders_loaded: bool = False # user_id -> display username cache for shared-workstream labels / join # notes, so senders read as usernames (like the owner banner) not raw # id hashes. Resolved lazily via storage; a handful of entries per ws. @@ -2924,6 +2931,9 @@ class ChatSession: if not fork: self._ws_id = ws_id self.messages = turns + # Shared-workstream state is per-workstream: this session object now + # points at (possibly different) history, so forget and re-derive. + self._reset_shared_state() self._read_files.clear() self._repeat_detector.clear() self._last_usage = None @@ -3033,11 +3043,21 @@ class ChatSession: except (TypeError, ValueError): pd_str = None src = msg.get("_source") - # Operator-context per-kind meta (``_source_meta`` dict) rides - # the fork too, so a forked watch-result keeps its structured - # card. Serialized to JSON for the ``conversations.meta`` column. + # The ``conversations.meta`` column rides the fork too, so a + # forked watch-result keeps its structured card and a forked + # user turn keeps its sender attribution. The two sources are + # role-exclusive (``_source_meta`` rides SYSTEM turns, the + # sender stamp rides USER turns — see ``reconstruct_turns``), + # mirroring the live save paths in ``_run_loop`` and + # ``_append_user_turn``. sm = msg.get("_source_meta") - meta_json = json.dumps(sm) if isinstance(sm, dict) and sm else None + sender = msg.get("_sender") + if isinstance(sm, dict) and sm: + meta_json = json.dumps(sm) + elif isinstance(sender, str) and sender: + meta_json = json.dumps({"sender": sender}) + else: + meta_json = None bulk_rows.append( { "ws_id": self._ws_id, @@ -3651,23 +3671,80 @@ class ChatSession: log.debug("display-name lookup failed for user=%s", user_id, exc_info=True) return name - def _recompute_shared_state(self) -> None: - """Recompute shared-workstream state from the current trajectory. + def _invalidate_shared_state(self) -> None: + """Mark shared-workstream state for recompute. - Scans user turns for recorded senders (the ``meta.extra["sender"]`` - stamped by :meth:`_append_user_turn`). The workstream is *shared* once a - sender other than the owner (``_mcp_user_id``) has spoken. Called at - system-prompt (re)composition so the ``## Session Context`` banner - reflects reality on a fresh compose AND on a worker rehydrating an - already-multi-user workstream from history.""" + The cheap flag half of the per-turn memo in + :meth:`_recompute_shared_state`; called when a sender-stamped user turn + is appended (the only live event that can change the participant set).""" + self._senders_dirty = True + + def _reset_shared_state(self) -> None: + """Forget shared-workstream state entirely. + + For :meth:`resume`, which points this session object at (possibly + different) history — the monotonic guarantees in + :meth:`_recompute_shared_state` hold per *workstream*, not per session + object, so carrying senders across a resume would leak one + workstream's participant set into another's framing.""" + self._shared_workstream = False + self._known_senders = set() + self._db_senders_loaded = False + self._senders_dirty = True + + def _load_persisted_senders(self) -> set[str]: + """One-time full-history sender read for :meth:`_recompute_shared_state`. + + Compaction narrows ``self.messages`` to a ``[summary] + [tail]`` slice, + so scanning it alone forgets participants whose turns were summarized + away. The persisted rows keep every sender ever stamped; read them once + per workstream. A storage error leaves ``_db_senders_loaded`` unset so + the next recompute (at most one per user turn, via the memo) retries + instead of pinning an incomplete participant set for the session's + lifetime.""" + try: + storage = get_storage() + if storage is not None: + senders = {s for s in storage.list_message_senders(self._ws_id) if s} + self._db_senders_loaded = True + return senders + # No storage configured (ephemeral session): nothing to read, ever. + self._db_senders_loaded = True + except Exception: + log.debug("persisted-sender load failed for ws=%s", self._ws_id, exc_info=True) + return set() + + def _recompute_shared_state(self) -> None: + """Refresh shared-workstream state from history — monotonically. + + ``_known_senders`` unions the current trajectory's recorded senders + (the ``meta.extra["sender"]`` stamped by :meth:`_append_user_turn`) + with a one-time read of the full persisted history; it never shrinks, + so a participant summarized out of the compacted slice stays known and + cannot re-trigger the one-time join note. ``_shared_workstream`` flips + True once any non-owner (``_mcp_user_id``) has spoken and then latches: + reverting would misattribute a known-multi-user conversation AND flip + the banner bytes, invalidating the provider prompt-prefix cache that + the hour-rounded timestamp above exists to protect. + + Memoized per turn via ``_senders_dirty``: system-prompt composition + runs many times within a turn (state transitions, MCP refresh, tool + results) but the sender set only changes on user-turn append and + history (re)load.""" + if not self._senders_dirty: + return owner = (self._mcp_user_id or "").strip() senders = { s for t in self.messages if t.role is Role.USER and (s := (t.meta.extra.get("sender") or "").strip()) } - self._known_senders = senders - self._shared_workstream = any(s != owner for s in senders) + if not self._db_senders_loaded: + senders |= self._load_persisted_senders() + self._known_senders |= senders + if not self._shared_workstream: + self._shared_workstream = any(s != owner for s in self._known_senders) + self._senders_dirty = False def _maybe_note_new_participant(self, sender_user_id: str | None) -> None: """Announce a first-time non-owner sender and flip the ws to shared. @@ -4498,6 +4575,10 @@ class ChatSession: for a in attachments ] self.messages.append(turn_from_dict(user_msg)) + if sender: + # A newly recorded sender can change shared-workstream state; let + # the next system-prompt compose re-derive it (memoized otherwise). + self._invalidate_shared_state() self._msg_tokens.append(max(1, int(self._msg_char_count(user_msg) / self._chars_per_token))) # DB row stores the raw text only; attachment bytes are written # content-addressed into workstream_attachments and the ordered id-list diff --git a/turnstone/core/storage/_postgresql.py b/turnstone/core/storage/_postgresql.py index 51582725..aade3eba 100644 --- a/turnstone/core/storage/_postgresql.py +++ b/turnstone/core/storage/_postgresql.py @@ -132,6 +132,7 @@ from turnstone.core.storage._utils import ( purge_orphan_conversations, release_attachment_refs, sanitize_text, + senders_from_user_meta, ) from turnstone.core.storage._utils import ( normalize_search_terms as _normalize_search_terms, @@ -359,6 +360,22 @@ class PostgreSQLBackend: conn.commit() return rowid + def list_message_senders(self, ws_id: str) -> list[str]: + # DISTINCT on the raw meta blob: a user row's meta carries only + # {"sender": ...}, so distinct blobs ≈ distinct senders and the JSON + # parse (shared, backend-neutral) runs on a handful of rows. + with self._conn() as conn: + rows = conn.execute( + sa.select(conversations.c.meta) + .distinct() + .where( + conversations.c.ws_id == ws_id, + conversations.c.role == "user", + conversations.c.meta.is_not(None), + ) + ).fetchall() + return senders_from_user_meta(meta for (meta,) in rows) + def save_messages_bulk(self, rows: list[dict[str, Any]]) -> None: if not rows: return diff --git a/turnstone/core/storage/_protocol.py b/turnstone/core/storage/_protocol.py index 7ae91b40..84d60846 100644 --- a/turnstone/core/storage/_protocol.py +++ b/turnstone/core/storage/_protocol.py @@ -241,6 +241,16 @@ class StorageBackend(Protocol): """ ... + def list_message_senders(self, ws_id: str) -> list[str]: + """Distinct sender user-ids recorded on a workstream's USER rows. + + Reads the full persisted history (``meta`` → ``{"sender": ...}``), not + a compaction-bounded view: the participant set drives shared-workstream + framing and the one-time join note, so it must survive compaction + narrowing the resumable ``[summary] + [tail]`` slice. + """ + ... + def get_max_event_id(self, ws_id: str) -> int | None: """Return the highest persisted ``event_id`` for ``ws_id``. diff --git a/turnstone/core/storage/_sqlite.py b/turnstone/core/storage/_sqlite.py index 5dd1f1d8..0a26e9bb 100644 --- a/turnstone/core/storage/_sqlite.py +++ b/turnstone/core/storage/_sqlite.py @@ -132,6 +132,7 @@ from turnstone.core.storage._utils import ( purge_orphan_conversations, release_attachment_refs, sanitize_text, + senders_from_user_meta, ) from turnstone.core.storage._utils import ( normalize_search_terms as _normalize_search_terms, @@ -407,6 +408,22 @@ class SQLiteBackend: conn.commit() return rowid + def list_message_senders(self, ws_id: str) -> list[str]: + # DISTINCT on the raw meta blob: a user row's meta carries only + # {"sender": ...}, so distinct blobs ≈ distinct senders and the JSON + # parse (shared, backend-neutral) runs on a handful of rows. + with self._conn() as conn: + rows = conn.execute( + sa.select(conversations.c.meta) + .distinct() + .where( + conversations.c.ws_id == ws_id, + conversations.c.role == "user", + conversations.c.meta.is_not(None), + ) + ).fetchall() + return senders_from_user_meta(meta for (meta,) in rows) + def save_messages_bulk(self, rows: list[dict[str, Any]]) -> None: if not rows: return diff --git a/turnstone/core/storage/_utils.py b/turnstone/core/storage/_utils.py index 58de9903..704ec4d5 100644 --- a/turnstone/core/storage/_utils.py +++ b/turnstone/core/storage/_utils.py @@ -6,10 +6,13 @@ import base64 import json import re from collections import Counter -from typing import Any +from typing import TYPE_CHECKING, Any import sqlalchemy as sa +if TYPE_CHECKING: + from collections.abc import Iterable + from turnstone.core.attachments import AUDIO_MIME_TO_FORMAT, unreadable_placeholder from turnstone.core.log import get_logger from turnstone.core.storage._schema import ( @@ -1141,3 +1144,24 @@ def reconstruct_turns_checkpointed( *marker_turns, *reconstruct_turns(tail, ws_id, attachments_by_msg), ] + + +def senders_from_user_meta(metas: Iterable[str | None]) -> list[str]: + """Distinct, stripped sender ids from USER-row ``meta`` JSON blobs. + + A user row's ``meta`` column carries only ``{"sender": ...}`` (the + role-exclusive routing in :func:`reconstruct_turns`); anything unparsable, + non-dict, or sender-less is skipped so one stray blob cannot poison the + participant set. Sorted for deterministic output across backends. + """ + senders: set[str] = set() + for raw in metas: + if not raw: + continue + try: + sender = json.loads(raw).get("sender") + except (ValueError, AttributeError): + continue + if isinstance(sender, str) and sender.strip(): + senders.add(sender.strip()) + return sorted(senders)