fix(session): preserve _source/_reminders on fork + cap persisted reminder text

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 91e7f2daca)
This commit is contained in:
Patrick Buckley
2026-05-06 22:14:02 -07:00
parent 2e393d76b4
commit 0c58910c4b
2 changed files with 117 additions and 8 deletions
+56
View File
@@ -3362,6 +3362,62 @@ class TestReminderSidechannelIsolation:
assert "HISTORICAL_REMINDER_BODY" not in rendered
assert "<system-reminder>" 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
+61 -8
View File
@@ -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",