From 0c58910c4bedf36a613700555ff6c19ec3d1ec84 Mon Sep 17 00:00:00 2001 From: Patrick Buckley Date: Wed, 6 May 2026 22:14:02 -0700 Subject: [PATCH] fix(session): preserve _source/_reminders on fork + cap persisted reminder text MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes round-1 review findings bug-2 (major), perf-1 (minor), perf-6 (nit). * **bug-2:** ``ChatSession.resume(..., fork=True)``'s bulk-row builder silently dropped the ``_source`` and ``_reminders`` side-channel data the source workstream had persisted via ``_append_user_turn``. Both backends' ``save_messages_bulk`` already accept these keys (the columns exist post-migration 050) — the bulk builder just didn't supply them. The fork's resumed transcript would then look like the assistant turn answered out of nowhere: every wake marker and every reminder bubble that survived to disk on the source got dropped on the fork. New regression test ``test_fork_preserves_source_and_reminders`` pins the contract. * **perf-6:** Extracts ``_encode_reminders(reminders) -> str | None`` near ``_apply_reminders_for_provider`` so the user-turn save path, the tool-turn save path, and the new fork bulk builder share one encoder. Eliminates the drift risk between three near-identical ``json.dumps(..., separators=(",", ":")) if X else None`` patterns. * **perf-1:** The new helper clamps each entry's ``text`` field at ``REMINDER_TEXT_STORAGE_CAP = 8192`` characters before encoding so a single rogue producer (a watch streaming unbounded shell output, a corruption-class steering payload) can't blow the conversations row width or the FTS5 index. The in-memory side-channel keeps the full body — only the persisted JSON is clamped. Mirrors ``TOOL_RESULT_STORAGE_CAP`` on tool result rows. 5734 non-live tests pass; ruff + mypy clean. (cherry picked from commit 91e7f2dacaeff329f78ec1b08dcad02b2a701a35) --- tests/test_session.py | 56 +++++++++++++++++++++++++++++++ turnstone/core/session.py | 69 ++++++++++++++++++++++++++++++++++----- 2 files changed, 117 insertions(+), 8 deletions(-) diff --git a/tests/test_session.py b/tests/test_session.py index 2e15b4f6..34c1cc9f 100644 --- a/tests/test_session.py +++ b/tests/test_session.py @@ -3362,6 +3362,62 @@ class TestReminderSidechannelIsolation: assert "HISTORICAL_REMINDER_BODY" not in rendered assert "" not in rendered + def test_fork_preserves_source_and_reminders(self, tmp_db): + """``ChatSession.resume(..., fork=True)`` bulk-inserts the source + workstream's messages into the fork's own ws_id. The bulk-row + builder must carry the persisted side-channels (``_source`` / + ``_reminders``) — both backends' ``save_messages_bulk`` accept + them post-migration 050. Dropping them was the original bug: + the fork's resumed transcript lost every wake marker and every + reminder bubble that survived to disk on the source, so the + resumed transcript looked like the assistant turn answered out + of nowhere. + """ + from turnstone.core.memory import register_workstream, save_message + + # Stage a source workstream with both a wake row (``_source = + # system_nudge``) and a reminders-bearing row. + register_workstream("fork_source") + save_message("fork_source", "user", "real turn") + save_message("fork_source", "user", "", source="system_nudge") + save_message( + "fork_source", + "user", + "advised turn", + reminders=json.dumps( + [{"type": "denial", "text": "FORKED_REMINDER_BODY"}], + separators=(",", ":"), + ), + ) + save_message("fork_source", "assistant", "ok") + + # Fork into a fresh session — keeps its own ws_id, copies the + # messages. Then resume the fork (no fork=True) into a second + # fresh session and assert the side-channels round-tripped. + forking_session = _make_session() + fork_ws_id = forking_session._ws_id + assert forking_session.resume("fork_source", fork=True) is True + + resumed_fork = _make_session() + assert resumed_fork.resume(fork_ws_id) is True + + # Wake row's ``_source`` survived. + wake_msgs = [ + m + for m in resumed_fork.messages + if m.get("role") == "user" and m.get("_source") == "system_nudge" + ] + assert len(wake_msgs) == 1 + assert wake_msgs[0].get("content") == "" + + # Reminders survived. + advised_msg = next( + m + for m in resumed_fork.messages + if m.get("role") == "user" and m.get("content") == "advised turn" + ) + assert advised_msg.get("_reminders") == [{"type": "denial", "text": "FORKED_REMINDER_BODY"}] + class TestSessionUIBaseUserReminderHook: """``on_user_reminder`` enqueues a ``user_reminder`` SSE event with diff --git a/turnstone/core/session.py b/turnstone/core/session.py index f540a244..f8496611 100644 --- a/turnstone/core/session.py +++ b/turnstone/core/session.py @@ -547,6 +547,18 @@ _TEMPLATE_VAR_RE = re.compile(r"\{\{(\w+)\}\}") # not its earliest. _WATCH_QUEUE_SOFT_CAP = 50 +# Per-reminder ``text`` field clamp at storage time. The conversations +# row's ``_reminders`` JSON column persists every metacog nudge a tab +# rendered live so reconnecting tabs see the same bubbles, but a single +# pathological producer — a watch streaming unbounded shell output, a +# corrupted steering payload — should not blow the row width or the +# FTS5 index. 8 KiB matches ``docs/design/watch-card-ux-briefing.md``'s +# spec for the reminder body. Mirrors ``TOOL_RESULT_STORAGE_CAP`` on +# tool result rows; the in-memory side-channel keeps the full body so +# the live splice and UI render see the same shape, only the persisted +# JSON is clamped. +REMINDER_TEXT_STORAGE_CAP = 8192 + def _without_tool(tools: list[dict[str, Any]], name: str) -> list[dict[str, Any]]: """Return *tools* with the named tool removed.""" @@ -1839,6 +1851,15 @@ class ChatSession: pd_str = json.dumps(pd) if pd and not isinstance(pd, str) else pd except (TypeError, ValueError): pd_str = None + # Carry the persisted side-channels (``_source`` / + # ``_reminders``) onto the fork's rows. Both backends' + # ``save_messages_bulk`` accept them post-migration 050; + # without them, a fork dropped every wake marker and + # every reminder bubble that survived to disk on the + # source workstream — the resumed fork's transcript + # would then look like the assistant turn answered out + # of nowhere. + src = msg.get("_source") bulk_rows.append( { "ws_id": self._ws_id, @@ -1848,6 +1869,8 @@ class ChatSession: "tool_call_id": msg.get("tool_call_id"), "tool_calls": tc_json, "provider_data": pd_str, + "source": src if isinstance(src, str) and src else None, + "reminders": self._encode_reminders(msg.get("_reminders")), } ) save_messages_bulk(bulk_rows) @@ -2101,6 +2124,42 @@ class ChatSession: """System messages + conversation history.""" return self.system_messages + self.messages + @staticmethod + def _encode_reminders( + reminders: list[dict[str, Any]] | None, + ) -> str | None: + """JSON-encode a ``_reminders`` list for the storage column. + + Returns ``None`` when *reminders* is empty / falsy so callers can + feed the result straight into ``save_message(..., reminders=...)`` + without a separate empty check. + + **Per-reminder text cap.** Each entry's ``text`` field is + clamped to :data:`REMINDER_TEXT_STORAGE_CAP` characters before + encoding so a single rogue producer (a watch streaming an + unbounded shell command, a corruption-class steering payload) + can't blow the conversations row width or the FTS5 index. The + in-memory dict on the message side-channel keeps the full body + — only the persisted JSON is clamped. Mirrors + ``TOOL_RESULT_STORAGE_CAP``'s row-level clamp on tool results. + """ + if not reminders: + return None + capped: list[dict[str, Any]] = [] + for r in reminders: + if not isinstance(r, dict): + continue + text = r.get("text") + if isinstance(text, str) and len(text) > REMINDER_TEXT_STORAGE_CAP: + clamped = dict(r) + clamped["text"] = text[:REMINDER_TEXT_STORAGE_CAP] + capped.append(clamped) + else: + capped.append(r) + if not capped: + return None + return json.dumps(capped, separators=(",", ":")) + def _apply_reminders_for_provider( self, messages: list[dict[str, Any]], @@ -2740,9 +2799,7 @@ class ChatSession: # nowhere. source = user_msg.get("_source") reminders_payload = user_msg.get("_reminders") - reminders_json = ( - json.dumps(reminders_payload, separators=(",", ":")) if reminders_payload else None - ) + reminders_json = self._encode_reminders(reminders_payload) message_id = save_message( self._ws_id, "user", @@ -3081,11 +3138,7 @@ class ChatSession: # ``_reminders`` JSON column so a tab reconnecting via # /history sees the below-the-tool bubble the # originating tab rendered live. - tool_reminders_json = ( - json.dumps(metacog_reminders, separators=(",", ":")) - if metacog_reminders - else None - ) + tool_reminders_json = self._encode_reminders(metacog_reminders) save_message( self._ws_id, "tool",