mirror of
https://github.com/turnstonelabs/turnstone.git
synced 2026-08-23 12:24:46 -06:00
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:
@@ -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
|
||||
|
||||
@@ -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",
|
||||
|
||||
Reference in New Issue
Block a user