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
main
turnstone/tests/test_rewind_retry.py
T
Patrick Buckley 480a1426b3 Fail-closed history-commit handoff (#1005) 
* fix(session): fail-closed history-commit handoff (#981)

The deleted-workstream discovery is now a terminal, ws_id-keyed latch:
keyed conversation commits refuse admission once the durable parent is
gone (convergence finalizers and force-abandon are exempt), history
handoff refuses to mint a proof token so /history fails closed with a
503 instead of silently wiping the pane, and the SSE stream carries a
workstream_gone resync reason. Discarded commits leave a forensic log
of commit keys and roles, never content.

Conversation rows gain a commit_key (migration 071): keyed saves are
idempotent under retry, validated against the full commit identity, and
refused when they would cross a workstream deletion. The prune orphan
category now requires a NULL alias plus a two-hour updated grace, with
cutoffs computed at discovery time and carried into both dialects'
rechecks.

The mid-turn interjection queue is owner-partitioned with no per-site
mode flags: pops take the acting principal's and unowned rows, other
participants' rows are structurally retained, and enforcement lives at
queue admission plus the shared before_spawn gates. The retraction
ledger is bounded by open pop windows: pops open a window atomically
with the queue delete, restores close their ids atomically with the
ledger consume, every other exit closes through one helper, and misses
for unheld ids record nothing. The workstream-gone latch refuses
unattended wakes at all three gates (watcher spawn, claim, delivery
pre-pop), and the retry dispatcher regained its pre-envelope
cancel/error convergence net.

Persistence-state reporting derives through the session bound to each
UI instead of a registry lookup by id that failed open to healthy
during tombstone retention. The dashboard roster no longer re-inserts
ghost entries from trailing activity events, the history tool-outcome
scan tolerates interleaved non-turn rows, and the shared
handoff-deadline handle owns its own retirement.

Single-sourced across call sites: keyed-commit row values, attachment
save wrappers, tail-truncation and conflict-resolution bodies for both
storage dialects; worker-slot lifecycle field sets; the direct-commit
admission frame; queued-row layout accessors; the string-aware comment
stripper shared by every JS harness suite.

Refs #981 #964

* fix(session): sweep handoff fixes to their sibling surfaces

The interactive replay loop treated a system row as a tool-batch
boundary, so every tool result after an interleaved row vanished from
that pane while the coordinator rendered the same history correctly.
Only a conversational turn ends the batch window now, matching the
shared outcome index.

Accepted user turns clear the composer's attachment chips on the same
viewer policy that settles optimistic bubbles rather than on having
matched a local bubble, so a workstream created with an upload no
longer keeps a chip for an attachment the create dispatch already
consumed. The coordinator's raced-Stop arm emits the stream-end hook it
inherits alongside the idle state, leaving no unfinalized bubble or
unflushed tool output. Ending a session surfaces a failure toast when
the request never lands or answers with a non-JSON body.

The per-second persistence reconcile now probes each session without
blocking: a workstream whose generation and handoff locks are held is
skipped until the next pass instead of contending the locks every
commit needs. The one-shot repair that gates workstream creation at
capacity keeps a definite probe — it has no next pass, and the sessions
likeliest to be contended are the ones whose unresolved journals
emptied its candidate list.

Single-sourced: the attachment lane builds its conversation row through
the shared commit-identity builder; the ordinary worker exit releases
its slot through the lifecycle owner; both operator surfaces snapshot
their counters through one non-consuming helper; the replay preamble
loses its per-kind wrappers and its config hook; the browser harness
suites share one brace walker; and each in-flight history attempt is
one record carrying both its abort controller and its deadline.

Refs #981 #964
2026-08-11 04:18:36 -07:00

514 lines
17 KiB
Python
Raw Permalink Blame History

"""Tests for conversation rewind and retry functionality."""
from __future__ import annotations
from unittest.mock import MagicMock
from turnstone.core.memory import register_workstream
from turnstone.core.session import ChatSession
from turnstone.core.trajectory import turn_to_dict, turns_from_dicts
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
class NullUI:
"""UI adapter that discards all output."""
def on_turn_start(self):
pass
def on_turn_committed(self):
pass
def on_thinking_start(self):
pass
def on_thinking_stop(self):
pass
def on_reasoning_token(self, text):
pass
def on_content_token(self, text):
pass
def on_stream_end(self):
pass
def approve_tools(self, items):
return True, None
def on_tool_result(self, call_id, name, output, **kwargs):
pass
def on_tool_output_chunk(self, call_id, chunk):
pass
def on_status(self, usage, context_window, effort):
pass
def on_info(self, message):
pass
def on_error(self, message):
pass
def on_state_change(self, state):
pass
def on_rename(self, name):
pass
def on_output_warning(self, call_id, assessment):
pass
def record_output_assessment(
self,
call_id,
assessment,
*,
tier="heuristic",
reasoning="",
judge_model="",
latency_ms=0,
confidence=0.0,
):
pass
def _make_session(tmp_db, ws_id: str | None = None) -> ChatSession:
session = ChatSession(
client=MagicMock(),
model="test-model",
ui=NullUI(),
instructions="",
temperature=0.5,
max_tokens=4096,
tool_timeout=30,
ws_id=ws_id,
)
# Production creates the durable parent before exposing a live session.
# Strict tail truncation shares that parent boundary with keyed commits.
register_workstream(session.ws_id, user_id="test-user")
return session
def _populate_simple(session: ChatSession) -> None:
"""Populate with 2 simple turns (no tool calls)."""
session.messages = turns_from_dicts(
[
{"role": "user", "content": "Hello"},
{"role": "assistant", "content": "Hi there!"},
{"role": "user", "content": "How are you?"},
{"role": "assistant", "content": "I'm fine."},
]
)
session._msg_tokens = [10, 20, 10, 20]
def _populate_with_tools(session: ChatSession) -> None:
"""Populate with 2 turns, first has tool calls."""
session.messages = turns_from_dicts(
[
{"role": "user", "content": "Write a test"},
{
"role": "assistant",
"content": "",
"tool_calls": [
{"id": "tc1", "function": {"name": "bash", "arguments": '{"cmd":"echo hi"}'}}
],
},
{"role": "tool", "tool_call_id": "tc1", "content": "hi"},
{"role": "assistant", "content": "Done."},
{"role": "user", "content": "Fix the import"},
{"role": "assistant", "content": "Fixed."},
]
)
session._msg_tokens = [10, 20, 10, 20, 10, 20]
# ---------------------------------------------------------------------------
# _find_turn_boundaries
# ---------------------------------------------------------------------------
class TestFindTurnBoundaries:
def test_empty_messages(self, tmp_db):
session = _make_session(tmp_db)
assert session._find_turn_boundaries() == []
def test_single_turn(self, tmp_db):
session = _make_session(tmp_db)
session.messages = turns_from_dicts(
[
{"role": "user", "content": "Hello"},
{"role": "assistant", "content": "Hi!"},
]
)
assert session._find_turn_boundaries() == [0]
def test_multi_turn(self, tmp_db):
session = _make_session(tmp_db)
_populate_simple(session)
assert session._find_turn_boundaries() == [0, 2]
def test_with_tool_calls(self, tmp_db):
session = _make_session(tmp_db)
_populate_with_tools(session)
assert session._find_turn_boundaries() == [0, 4]
# ---------------------------------------------------------------------------
# rewind
# ---------------------------------------------------------------------------
class TestRewind:
def test_rewind_zero(self, tmp_db):
session = _make_session(tmp_db)
_populate_simple(session)
assert session.rewind(0) == 0
assert len(session.messages) == 4
def test_rewind_one_turn(self, tmp_db):
session = _make_session(tmp_db)
_populate_simple(session)
removed = session.rewind(1)
assert removed == 2 # user + assistant
assert len(session.messages) == 2
assert turn_to_dict(session.messages[0])["content"] == "Hello"
assert turn_to_dict(session.messages[1])["content"] == "Hi there!"
assert len(session._msg_tokens) == 2
def test_rewind_all_turns(self, tmp_db):
session = _make_session(tmp_db)
_populate_simple(session)
removed = session.rewind(2)
assert removed == 4
assert len(session.messages) == 0
assert len(session._msg_tokens) == 0
def test_rewind_clamped(self, tmp_db):
"""Rewinding more turns than exist should clamp to available."""
session = _make_session(tmp_db)
_populate_simple(session)
removed = session.rewind(999)
assert removed == 4
assert len(session.messages) == 0
def test_rewind_empty(self, tmp_db):
session = _make_session(tmp_db)
assert session.rewind(1) == 0
def test_rewind_with_tools(self, tmp_db):
"""Rewinding 1 turn on a multi-sub-turn conversation."""
session = _make_session(tmp_db)
_populate_with_tools(session)
removed = session.rewind(1)
assert removed == 2 # user "Fix the import" + assistant "Fixed."
assert len(session.messages) == 4
assert turn_to_dict(session.messages[-1])["content"] == "Done."
def test_rewind_tokens_sync(self, tmp_db):
"""_msg_tokens stays in sync with messages."""
session = _make_session(tmp_db)
_populate_simple(session)
session.rewind(1)
assert len(session._msg_tokens) == len(session.messages)
# ---------------------------------------------------------------------------
# retry
# ---------------------------------------------------------------------------
class TestRetry:
def test_retry_returns_user_message(self, tmp_db):
session = _make_session(tmp_db)
_populate_simple(session)
msg = session.retry()
assert msg == "How are you?"
# Only Turn 1 remains, without the second user message
assert len(session.messages) == 2
assert turn_to_dict(session.messages[-1])["content"] == "Hi there!"
def test_retry_empty(self, tmp_db):
session = _make_session(tmp_db)
assert session.retry() is None
def test_retry_with_tools(self, tmp_db):
session = _make_session(tmp_db)
_populate_with_tools(session)
msg = session.retry()
assert msg == "Fix the import"
# Only Turn 1 remains (user + assistant w/tools + tool result + assistant)
assert len(session.messages) == 4
def test_retry_sets_pending(self, tmp_db):
"""handle_command for /retry should set _pending_retry."""
session = _make_session(tmp_db)
_populate_simple(session)
session.handle_command("/retry")
assert session._pending_retry == "How are you?"
def test_retry_tokens_sync(self, tmp_db):
session = _make_session(tmp_db)
_populate_simple(session)
session.retry()
assert len(session._msg_tokens) == len(session.messages)
def test_retry_multipart_content_returns_none(self, tmp_db):
"""retry() should refuse multipart (vision/image) messages."""
session = _make_session(tmp_db)
session.messages = turns_from_dicts(
[
{
"role": "user",
"content": [
{"type": "text", "text": "describe this"},
# Canonical non-text content is a by-reference placeholder;
# the turn stays multipart, so retry refuses it.
{"type": "image", "attachment_id": "sha256:abc"},
],
},
{"role": "assistant", "content": "It's an image."},
]
)
session._msg_tokens = [10, 20]
assert session.retry() is None
# Messages should be unchanged
assert len(session.messages) == 2
def test_retry_none_content_returns_none(self, tmp_db):
"""retry() should handle content=None gracefully."""
session = _make_session(tmp_db)
session.messages = turns_from_dicts(
[
{"role": "user", "content": None},
{"role": "assistant", "content": "Ok."},
]
)
session._msg_tokens = [10, 20]
assert session.retry() is None
# ---------------------------------------------------------------------------
# handle_command integration
# ---------------------------------------------------------------------------
class TestHandleCommand:
def test_rewind_command(self, tmp_db):
session = _make_session(tmp_db)
_populate_simple(session)
session.handle_command("/rewind 1")
assert len(session.messages) == 2
def test_rewind_no_arg(self, tmp_db):
session = _make_session(tmp_db)
ui = session.ui
ui.on_info = MagicMock()
session.handle_command("/rewind")
ui.on_info.assert_called_once()
assert "Usage" in ui.on_info.call_args[0][0]
def test_rewind_invalid_arg(self, tmp_db):
session = _make_session(tmp_db)
ui = session.ui
ui.on_info = MagicMock()
session.handle_command("/rewind abc")
ui.on_info.assert_called_once()
assert "integer" in ui.on_info.call_args[0][0]
def test_retry_nothing_to_retry(self, tmp_db):
session = _make_session(tmp_db)
ui = session.ui
ui.on_info = MagicMock()
session.handle_command("/retry")
ui.on_info.assert_called_once()
assert "Nothing" in ui.on_info.call_args[0][0]
def test_rewind_refusal_reports_instead_of_killing_the_repl(self, tmp_db):
"""Round-4 review pin: rewind()'s raising contract (GenerationCancelled
on a refused admission, e.g. the workstream-gone latch) converts to an
on_error message in the CLI arm — the REPL dispatches commands
uncaught, so a raise here previously killed the whole process."""
session = _make_session(tmp_db)
_populate_simple(session)
ui = session.ui
ui.on_error = MagicMock()
session._workstream_gone_ws = session._ws_id
session.handle_command("/rewind 1")
ui.on_error.assert_called_once()
assert "refused" in ui.on_error.call_args[0][0].lower()
def test_retry_storage_failure_reports_instead_of_killing_the_repl(self, tmp_db):
"""Round-4 review pin: a storage error escaping retry() (the durable
truncation batch is raising now) becomes an on_error message, never an
uncaught REPL death."""
from unittest.mock import patch
session = _make_session(tmp_db)
_populate_simple(session)
ui = session.ui
ui.on_error = MagicMock()
with patch.object(session, "retry", side_effect=RuntimeError("database is locked")):
session.handle_command("/retry")
ui.on_error.assert_called_once()
assert "Retry failed" in ui.on_error.call_args[0][0]
# ---------------------------------------------------------------------------
# Storage integration — delete_messages_after
# ---------------------------------------------------------------------------
class TestDeleteMessagesAfter:
def test_delete_truncates_db(self, tmp_db):
from turnstone.core.memory import (
delete_messages_after,
load_messages,
register_workstream,
save_message,
)
ws_id = "test-ws-delete"
register_workstream(ws_id)
save_message(ws_id, "user", "Hello")
save_message(ws_id, "assistant", "Hi!")
save_message(ws_id, "user", "Bye")
save_message(ws_id, "assistant", "Goodbye!")
deleted = delete_messages_after(ws_id, 2)
assert deleted == 2
msgs = load_messages(ws_id)
assert len(msgs) == 2
assert msgs[0]["content"] == "Hello"
assert msgs[1]["content"] == "Hi!"
def test_delete_nothing(self, tmp_db):
from turnstone.core.memory import (
delete_messages_after,
register_workstream,
save_message,
)
ws_id = "test-ws-noop"
register_workstream(ws_id)
save_message(ws_id, "user", "Hello")
deleted = delete_messages_after(ws_id, 10)
assert deleted == 0
def test_delete_all(self, tmp_db):
from turnstone.core.memory import (
delete_messages_after,
load_messages,
register_workstream,
save_message,
)
ws_id = "test-ws-all"
register_workstream(ws_id)
save_message(ws_id, "user", "Hello")
save_message(ws_id, "assistant", "Hi!")
deleted = delete_messages_after(ws_id, 0)
assert deleted == 2
assert load_messages(ws_id) == []
# ---------------------------------------------------------------------------
# End-to-end: rewind + DB sync
# ---------------------------------------------------------------------------
class TestRewindDBSync:
def test_rewind_persists_to_db(self, tmp_db):
from turnstone.core.memory import load_messages, register_workstream, save_message
session = _make_session(tmp_db)
ws_id = session.ws_id
register_workstream(ws_id)
# Persist messages to DB and set in-memory state
save_message(ws_id, "user", "Hello")
save_message(ws_id, "assistant", "Hi!")
save_message(ws_id, "user", "Bye")
save_message(ws_id, "assistant", "Goodbye!")
session.messages = turns_from_dicts(
[
{"role": "user", "content": "Hello"},
{"role": "assistant", "content": "Hi!"},
{"role": "user", "content": "Bye"},
{"role": "assistant", "content": "Goodbye!"},
]
)
session._msg_tokens = [5, 5, 5, 5]
session.rewind(1)
# Verify DB matches in-memory state
db_msgs = load_messages(ws_id)
assert len(db_msgs) == 2
assert db_msgs[0]["content"] == "Hello"
assert db_msgs[1]["content"] == "Hi!"
def test_truncation_bumps_history_generation(tmp_db) -> None:
"""#894: the /history single-flight keys on _history_generation, so a
truncation that DELETED storage rows must bump it — and a truncation
that could not delete (storage count unavailable/empty — the
error-path early returns) must NOT, because flights rebuild from
storage and an unbumped generation on error is the safe direction (a
stale-keyed flight reading post-delete rows is spuriously fresh; a
new-keyed flight reading pre-delete rows would be wrongly joinable).
The endpoint-level flight test mocks the counter, so this is the
PRODUCER pin: delete the bump in _persist_truncation and the first
arm fails while everything else stays green."""
from turnstone.core.storage import get_storage
storage = get_storage()
storage.register_workstream("ws-gen-pin", kind="interactive", user_id="test-user")
session = _make_session(tmp_db, ws_id="ws-gen-pin")
_populate_simple(session)
for role, content in (
("user", "Hello"),
("assistant", "Hi there!"),
("user", "How are you?"),
("assistant", "I'm fine."),
):
storage.save_message("ws-gen-pin", role, content)
g0 = session._history_generation
assert session.rewind(1) > 0
assert session._history_generation == g0 + 1, (
"a storage-deleting rewind must bump the history generation"
)
# retry: re-persist the tail the rewind removed, then drop the last
# assistant turn.
storage.save_message("ws-gen-pin", "user", "How are you?")
storage.save_message("ws-gen-pin", "assistant", "I'm fine.")
_populate_simple(session)
g1 = session._history_generation
assert session.retry() is not None
assert session._history_generation == g1 + 1, (
"a storage-deleting retry must bump the history generation"
)
# A registered workstream with no durable rows is still a successful
# atomic cut of the live trajectory, so its handoff revision advances.
bare = _make_session(tmp_db)
_populate_simple(bare)
g2 = bare._history_generation
assert bare.rewind(1) > 0
assert bare._history_generation == g2 + 1, (
"a successful live truncation must invalidate history even when the "
"durable tail is already empty"
)
Reference in New Issue View Git Blame Copy Permalink