From dc35cbc7bf162f29c2c4ffe22267c4bb3ed9e4cd Mon Sep 17 00:00:00 2001 From: Patrick Buckley Date: Thu, 7 May 2026 16:27:28 -0700 Subject: [PATCH] fix(replay): seam 1 splice + storage symmetry for queued user messages MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reverses the seam-2-only design from the prior commits on this branch. Queued user messages arriving DURING a tool batch (Seam 1) splice into the last tool result's envelope as ``UserInterjection`` advisories via ``wrap_tool_result``. Messages arriving BETWEEN turns (Seam 2) drain as a single trailing user row via ``_flush_queued_messages`` with ``user_feedback`` (operator text alongside an approval, e.g. "y, use full path") folded in as a prefix. Cancel/exception drains (Seam 3) keep the existing ``_flush_queued_messages()`` call unchanged. Why all three seams: * Strict-template providers (Mistral, Llama via vLLM with stock chat templates) reject role-alternation violations. A literal ``user`` row mid-tool-batch breaks ``assistant(tool_calls) → tool → ... → assistant``; back-to-back ``user → user`` rows on the wire also fail. * The seam-2-only design produced back-to-back ``user`` whenever ``user_feedback`` and queued items both fired — bug-1 from the round-1 review. Folding ``user_feedback`` as a prefix to the queue-drain collapses the two into one row. * During-batch arrivals couldn't ride seam 2 — the splice was the only way to deliver same-turn without violating role alternation. Storage symmetry: Tool DB rows now store the wrapped ``output`` (envelope + advisories) unconditionally — ``self.messages[i]['content']`` and ``conversations.content`` match exactly. List-typed output (image / structured MCP results) uses ``wrap_tool_result(raw_joined_text, advisories)`` at save time so the persisted string is anchored on ``\n`` for the replay parser. ``TOOL_RESULT_STORAGE_CAP`` is removed entirely; tools are responsible for bounding their own output, storage faithfully represents in-memory. Removing the cap also simplifies the parser — no truncated-envelope edge case. Replay extraction: ``decorate_history_messages`` (REST ``/history``) and ``_build_history`` (SSE replay, resume, rewind, retry, post-load, rename re-replay) both call the public ``extract_advisories_from_tool_envelope`` helper to pull the envelope back into structured ``advisories`` for JS replay. Both string content and list-typed content (image+queued-message combo) covered. JS renders extracted advisories as normal user bubbles after the tool block via the shared ``replayAdvisoriesAfterTool`` helper in ``shared_static/utils.js``. Wrapper-tag escape and provider splice: ``escape_wrapper_tags`` now encodes pre-existing ``&`` first using an ``&`` sentinel so tool output containing literal entity strings (documentation viewers, code analyzers, web scrapers returning entity- encoded markup) round-trips correctly. Both encode and decode helpers short-circuit on absence of ``<`` / ``&``. ``_apply_reminders_for_provider`` detects already-wrapped content (string body and list text-part) by ``startswith("\n")`` and skips re-escape so existing envelopes survive intact when a tool message also carries ``_reminders`` (the queued-message + tool-error co-occurrence case is now common). ``decorate_history_messages`` runs in ``asyncio.to_thread`` to keep MB-scale string work off the event loop. Other cleanup: * ``_collect_advisories`` delegates the queue drain to a named helper ``_drain_queued_messages_to_advisories`` so the swap-and-clear pattern lives next to ``_flush_queued_messages``'s identical pattern and the side-effect is documented at the call site. * Preamble strings + body marker for ``UserInterjection`` round-trip detection moved to module-level constants in ``tool_advisory.py``; imported by ``history_decoration.py`` so a producer-side rephrase can't silently desync the parser. * ``_send_with_mocks`` ctxmgr extracted in ``test_session.py`` — the six new send-driven tests share an 8-deep ``patch.object`` block. * ``replayAdvisoriesAfterTool`` shared helper in ``shared_static/utils.js``; ``app.js`` and ``coordinator.js`` both invoke it. * Dead truncation-pill CSS removed (``.tool-output-truncated`` and ``.coord-tool-truncated``); the JS that added these elements went away with ``TOOL_RESULT_STORAGE_CAP``. * Tautological tests (``TestBuildHistoryAdvisoryPropagation``) replaced with production-realistic round-trip tests built from ``wrap_tool_result(...)`` envelopes — REST and SSE-replay surfaces pinned to the same wire shape; full DB round-trip pinned end-to-end. Negative-tested: * Reverting the prefix-merge in ``_flush_queued_messages`` produces back-to-back ``user`` rows, breaking ``test_user_feedback_and_queued_coexistence_single_row_with_prefix``. * Reverting the ``extract_advisories_from_tool_envelope`` call in ``_build_history``'s tool branch leaves the envelope verbatim in wire content, breaking the round-trip tests. * Reverting the wrapper-detection in ``_apply_reminders_for_provider`` entity-encodes the existing envelope's literal tags, breaking both the string-content and list-content envelope-preservation tests. * Reverting the ``wrap_tool_result(raw_text, advisories)`` projection at the DB save site produces a string starting with the original raw text, breaking ``test_tool_db_row_round_trips_list_output_with_advisories``. Tests: 5918 passed, 3 deselected. Lint + format + mypy clean on touched files. (cherry picked from commit eca4bb79e4a5fc381f85218f6b8e4bfd4b982f44) --- tests/test_app_js.py | 60 ++ tests/test_coordinator_page.py | 40 ++ tests/test_history_decoration.py | 207 +++++- tests/test_session.py | 654 ++++++++++++++++-- tests/test_tool_advisory.py | 58 ++ tests/test_workstream_endpoints.py | 182 +++++ .../static/coordinator/coordinator.css | 19 - .../console/static/coordinator/coordinator.js | 40 +- turnstone/core/history_decoration.py | 142 +++- turnstone/core/session.py | 234 +++++-- turnstone/core/session_routes.py | 9 +- turnstone/core/tool_advisory.py | 63 +- turnstone/server.py | 94 ++- turnstone/shared_static/utils.js | 19 + turnstone/ui/static/app.js | 29 +- turnstone/ui/static/style.css | 31 +- 16 files changed, 1593 insertions(+), 288 deletions(-) diff --git a/tests/test_app_js.py b/tests/test_app_js.py index 2416f949..4ed2fe99 100644 --- a/tests/test_app_js.py +++ b/tests/test_app_js.py @@ -169,6 +169,66 @@ def test_replay_history_renders_persisted_verdict_badge() -> None: ) +def test_shared_utils_defines_replay_advisories_after_tool() -> None: + """The shared ``replayAdvisoriesAfterTool`` helper in + ``shared_static/utils.js`` is the single source of advisory-walk + + type-filter logic for both ``app.js`` (interactive) and + ``coordinator.js`` (coord). A refactor that drops the helper + breaks both surfaces, so guard its definition + filter shape here. + """ + utils_js = Path(__file__).resolve().parent.parent / "turnstone/shared_static/utils.js" + body = utils_js.read_text(encoding="utf-8") + assert "function replayAdvisoriesAfterTool" in body, ( + "shared/utils.js must define replayAdvisoriesAfterTool — " + "interactive and coord both invoke it." + ) + # The type filter — ``adv.type !== 'user_interjection'`` — must + # remain in the helper so a future advisory shape (output_guard, + # metacognitive nudge, etc.) doesn't silently render as a user + # bubble. + assert 'adv.type !== "user_interjection"' in body, ( + "replayAdvisoriesAfterTool must filter by advisory type so a " + "future non-user_interjection advisory shape doesn't silently " + "render as a user bubble." + ) + + +def test_replay_renders_user_interjection_advisory_after_tool_block() -> None: + """Queued user messages spliced into the last tool-result envelope + of a batch (Seam 1) persist on the tool DB row as a wrapped + ```` envelope. ``decorate_history_messages`` extracts + the advisory back out and the wire layer projects it onto + ``msg.advisories``; ``replayHistory`` must invoke the shared + ``replayAdvisoriesAfterTool`` helper (defined in + ``shared/utils.js``) so each ``user_interjection`` renders through + ``addUserMessage`` and the bubble looks identical to a Seam 2/3 + user row. + + This test pins the call site so a refactor that drops the helper + invocation regresses the queued-during-batch replay shape + silently.""" + body = _APP_JS.read_text(encoding="utf-8") + start = body.index("Pane.prototype.replayHistory = function") + end = body.index("Pane.prototype._attachRetryToLastAssistant", start) + fn = body[start:end] + # The replay loop must invoke the shared helper, passing + # ``msg.advisories`` and a renderer that routes through + # ``addUserMessage``. The helper itself filters on + # ``adv.type !== "user_interjection"``; that branch lives in + # ``shared/utils.js`` (test_shared_utils_js or runtime smoke covers + # the helper's body). + assert "replayAdvisoriesAfterTool(msg.advisories" in fn, ( + "replayHistory must invoke replayAdvisoriesAfterTool with " + "msg.advisories so queued messages spliced into the tool " + "envelope render as user bubbles after the tool block." + ) + assert "addUserMessage(text" in fn, ( + "replayHistory's renderer callback must route the extracted " + "advisory text through addUserMessage so the rendered bubble " + "matches a normal user-row replay." + ) + + # --------------------------------------------------------------------------- # Phase 8 — Chunk D: MCP error embed + settings panel UX # --------------------------------------------------------------------------- diff --git a/tests/test_coordinator_page.py b/tests/test_coordinator_page.py index c9e23060..8c06166a 100644 --- a/tests/test_coordinator_page.py +++ b/tests/test_coordinator_page.py @@ -302,3 +302,43 @@ def test_coordinator_js_handle_child_state_no_longer_reads_sse_pending_approval_ "cycles) — without this, the second bulk-poll after an SSE " "transition silently clobbers." ) + + +def test_coord_history_renders_user_interjection_advisory_after_tool_block(): + """Queued user messages spliced into the last tool-result envelope + of a batch (Seam 1) persist on the tool DB row as a wrapped + ```` envelope. ``decorate_history_messages`` extracts + the advisory back out and the wire layer projects it onto + ``m.advisories``; the coord history loop must invoke the shared + ``replayAdvisoriesAfterTool`` helper (defined in + ``shared/utils.js``) so each ``user_interjection`` renders through + ``appendUserMessageWithAttachments`` and the bubble looks identical + to a Seam 2/3 user row. + + This test pins the call site so a refactor that drops the helper + invocation regresses the queued-during-batch replay shape silently. + Mirrors ``test_app_js.py``'s same-shape pin on interactive's + ``replayHistory``.""" + import re + from pathlib import Path + + coord_js = Path(__file__).resolve().parent.parent / ( + "turnstone/console/static/coordinator/coordinator.js" + ) + body = coord_js.read_text(encoding="utf-8") + + assert "replayAdvisoriesAfterTool(m.advisories" in body, ( + "Coord history loop must invoke replayAdvisoriesAfterTool with " + "m.advisories so queued messages spliced into the tool envelope " + "render as user bubbles after the tool block." + ) + # The renderer callback routes through appendUserMessageWithAttachments + # so the bubble matches a normal user-row replay. + assert re.search( + r"appendUserMessageWithAttachments\(\s*text", + body, + ), ( + "Coord history loop's renderer callback must route the extracted " + "advisory text through appendUserMessageWithAttachments so the " + "rendered bubble matches a normal user-row replay." + ) diff --git a/tests/test_history_decoration.py b/tests/test_history_decoration.py index c843e8f3..d0f44612 100644 --- a/tests/test_history_decoration.py +++ b/tests/test_history_decoration.py @@ -167,7 +167,7 @@ class TestDecorateHistoryMessages: """End-to-end mutation of a /history-shaped message list — covers the full transform applied by ``make_history_handler``.""" - def test_decorates_tool_calls_and_marks_truncated(self) -> None: + def test_decorates_tool_calls_with_verdict_and_assessment(self) -> None: verdicts = { "call_a": { "risk_level": "high", @@ -181,13 +181,6 @@ class TestDecorateHistoryMessages: assessments = { "call_a": {"risk_level": "high", "flags": '["secret"]', "redacted": 1}, } - # Tool result content of exactly TOOL_RESULT_STORAGE_CAP chars - # hits the storage cap (longer is impossible — storage clamps - # at the cap). Reference the constant rather than a literal so - # this test stays correct if the cap moves again. - from turnstone.core.history_decoration import TOOL_RESULT_STORAGE_CAP - - truncated_content = "x" * TOOL_RESULT_STORAGE_CAP messages: list[dict[str, object]] = [ {"role": "user", "content": "hi"}, { @@ -200,7 +193,7 @@ class TestDecorateHistoryMessages: } ], }, - {"role": "tool", "tool_call_id": "call_a", "content": truncated_content}, + {"role": "tool", "tool_call_id": "call_a", "content": "long output"}, {"role": "tool", "tool_call_id": "call_b", "content": "short"}, ] decorate_history_messages(messages, verdicts, assessments) @@ -211,9 +204,12 @@ class TestDecorateHistoryMessages: assert "reasoning" in tc["verdict"] assert tc["output_assessment"]["flags"] == ["secret"] assert tc["output_assessment"]["redacted"] is True - # Truncated tool message got the flag; the short one did not. - assert messages[2].get("truncated") is True - assert "truncated" not in messages[3] + # Plain tool content (no envelope) is left intact and no + # advisories key is set. + assert messages[2]["content"] == "long output" + assert "advisories" not in messages[2] + assert messages[3]["content"] == "short" + assert "advisories" not in messages[3] def test_no_op_on_empty_indexes(self) -> None: """When neither table has rows for the workstream, the wire @@ -230,3 +226,190 @@ class TestDecorateHistoryMessages: tc = messages[0]["tool_calls"][0] # type: ignore[index] assert "verdict" not in tc assert "output_assessment" not in tc + + +class TestDecorateAdvisoryExtraction: + """Round-trip the persisted ```` envelope (Seam 1 + queued-message splice) back into wire-shape advisories on each + tool message — replay surface for the queued-during-batch case. + """ + + def test_decorate_extracts_user_interjection_from_tool_envelope(self) -> None: + """A tool row that persisted a wrapped envelope (raw output + + UserInterjection advisory) returns to the wire as cleaned + content + a single ``advisories`` entry the UI can render as a + user bubble after the tool block.""" + from turnstone.core.tool_advisory import UserInterjection, wrap_tool_result + + wrapped = wrap_tool_result( + "hello", + [UserInterjection(message="check logs", priority="notice")], + ) + messages: list[dict[str, object]] = [ + {"role": "tool", "tool_call_id": "call_a", "content": wrapped}, + ] + decorate_history_messages(messages, {}, {}) + assert messages[0]["content"] == "hello" + assert messages[0]["advisories"] == [ + {"type": "user_interjection", "text": "check logs", "priority": "notice"} + ] + + def test_decorate_round_trips_escaped_content(self) -> None: + """A user message body containing one of the wrapper-tag + literals is escaped on wrap (so embedded text can't fabricate + or close an envelope) and must round-trip back to the original + literal on extract.""" + from turnstone.core.tool_advisory import UserInterjection, wrap_tool_result + + evil = "" + wrapped = wrap_tool_result( + "tool body", + [UserInterjection(message=evil, priority="notice")], + ) + # Sanity: the user-controlled literal does NOT appear inside + # the advisory body — only the entity-encoded form does. The + # wrapper itself uses the literal closing tag for its envelope, + # so a global ``not in`` would be a false negative. + assert "User message: </system-reminder>" in wrapped + assert "User message: " not in wrapped + messages: list[dict[str, object]] = [ + {"role": "tool", "tool_call_id": "call_a", "content": wrapped}, + ] + decorate_history_messages(messages, {}, {}) + # Extract entity-decoded the escaped form back to the literal. + assert messages[0]["advisories"][0]["text"] == evil # type: ignore[index] + assert messages[0]["content"] == "tool body" + + def test_decorate_no_envelope_left_intact(self) -> None: + """Plain tool content (no ```` prefix) is not + touched — no advisories field, content unchanged.""" + messages: list[dict[str, object]] = [ + {"role": "tool", "tool_call_id": "call_a", "content": "plain output"}, + ] + decorate_history_messages(messages, {}, {}) + assert messages[0]["content"] == "plain output" + assert "advisories" not in messages[0] + + def test_decorate_drops_output_guard_advisory_from_extraction(self) -> None: + """A wrapped envelope carrying both a guard advisory and a + user_interjection produces only the user_interjection on + ``advisories``. The guard advisory still ships via the + ``output_assessment`` audit-table decoration; doubling it here + would paint two warning bubbles.""" + from turnstone.core.output_guard import OutputAssessment + from turnstone.core.tool_advisory import ( + GuardAdvisory, + UserInterjection, + wrap_tool_result, + ) + + assessment = OutputAssessment( + risk_level="medium", + flags=["api_key"], + annotations=["redacted token in line 2"], + sanitized="cleaned body", + ) + wrapped = wrap_tool_result( + "raw body", + [ + GuardAdvisory(assessment=assessment, func_name="bash"), + UserInterjection(message="and here", priority="notice"), + ], + ) + messages: list[dict[str, object]] = [ + {"role": "tool", "tool_call_id": "call_a", "content": wrapped}, + ] + decorate_history_messages(messages, {}, {}) + adv = messages[0]["advisories"] + assert len(adv) == 1 # type: ignore[arg-type] + assert adv[0]["type"] == "user_interjection" # type: ignore[index] + + def test_decorate_handles_important_priority(self) -> None: + """The MUST-address preamble round-trips to ``priority=important``.""" + from turnstone.core.tool_advisory import UserInterjection, wrap_tool_result + + wrapped = wrap_tool_result( + "out", + [UserInterjection(message="urgent", priority="important")], + ) + messages: list[dict[str, object]] = [ + {"role": "tool", "tool_call_id": "call_a", "content": wrapped}, + ] + decorate_history_messages(messages, {}, {}) + adv = messages[0]["advisories"][0] # type: ignore[index] + assert adv["priority"] == "important" + assert adv["text"] == "urgent" + + def test_wrap_extract_round_trips_preexisting_entities(self) -> None: + """A user message body containing literal HTML-entity references + matching the wrapper-escape forms must round-trip identically + through ``wrap_tool_result + extract_advisories_from_tool_envelope``. + Without escaping ``&`` first in the encode step, encode→decode + would produce the bare wrapper tag, fabricating an envelope the + wrapper layer never produced. + """ + from turnstone.core.history_decoration import ( + extract_advisories_from_tool_envelope, + ) + from turnstone.core.tool_advisory import UserInterjection, wrap_tool_result + + tricky = "I describe XML tags like <tool_output> in my docs." + wrapped = wrap_tool_result( + "tool body", + [UserInterjection(message=tricky, priority="notice")], + ) + result = extract_advisories_from_tool_envelope(wrapped) + assert result is not None + cleaned, advisories = result + assert cleaned == "tool body" + assert len(advisories) == 1 + # The original literal entity-reference text round-trips + # identically — the parser does not silently turn it into a + # bare wrapper tag. + assert advisories[0]["text"] == tricky + + def test_save_load_decorate_round_trips_envelope(self, backend) -> None: + """End-to-end round-trip pinning the persisted-envelope + contract. Persists a wrapped tool-output envelope via + ``save_message``, loads via ``load_messages``, runs + ``decorate_history_messages``, asserts the wire shape carries + the extracted advisory + cleaned content. Pins the contract + every component in the chain participates in (persistence + layer ↔ in-memory replay ↔ wire projection) so a schema drift, + an envelope-format change, or a parser regression surfaces + here rather than only in production. + """ + from turnstone.core.tool_advisory import UserInterjection, wrap_tool_result + + wrapped = wrap_tool_result( + "command output", + [UserInterjection(message="check the logs", priority="notice")], + ) + backend.register_workstream("ws_rt_1") + backend.save_message("ws_rt_1", "user", "go") + backend.save_message( + "ws_rt_1", + "assistant", + None, + tool_calls='[{"id":"call_a","type":"function","function":{"name":"bash","arguments":"{}"}}]', + ) + backend.save_message( + "ws_rt_1", + "tool", + wrapped, + tool_call_id="call_a", + ) + msgs = backend.load_messages("ws_rt_1") + # Persisted shape — content survives the storage layer + # untouched. Symmetry with in-memory ``self.messages[i]['content']`` + # is what makes envelope extraction lossless on replay. + tool_msg = next(m for m in msgs if m["role"] == "tool") + assert tool_msg["content"] == wrapped + # Decorate (the /history shared transform) — extracts the + # advisory and strips the envelope. + decorate_history_messages(msgs, {}, {}) + tool_msg = next(m for m in msgs if m["role"] == "tool") + assert tool_msg["content"] == "command output" + assert tool_msg["advisories"] == [ + {"type": "user_interjection", "text": "check the logs", "priority": "notice"} + ] diff --git a/tests/test_session.py b/tests/test_session.py index 94317ab2..55c308ca 100644 --- a/tests/test_session.py +++ b/tests/test_session.py @@ -86,6 +86,48 @@ def _make_session( return ChatSession(**defaults) +@contextlib.contextmanager +def _send_with_mocks(session, responses, mock_execute, **extra_patches): + """Stand up the mock context that the queued-message ``send()`` tests share. + + Six tests in ``TestMetacognitiveBuffers`` previously inlined the + same nine ``patch.object`` / ``patch`` declarations. Extracting + the ctxmgr keeps each test focused on its scenario (responses + + execute behaviour + assertions) rather than re-asserting the + common mock surface. + + Yields the ``save_message`` MagicMock so callers that need to + assert on persistence can ``... as save_msg`` over the helper. + Extra per-test patches (e.g. wrapping ``_collect_advisories``) ride + via ``**extra_patches`` — keyword name maps to attribute on the + session, value is the ``side_effect`` to inject. + """ + from unittest.mock import patch as _patch + + def mock_stream(_msgs): + return iter([]) + + def mock_response(_stream, _gen): + return responses.pop(0) + + with contextlib.ExitStack() as stack: + stack.enter_context( + _patch.object(session, "_create_stream_with_retry", side_effect=mock_stream) + ) + stack.enter_context(_patch.object(session, "_stream_response", side_effect=mock_response)) + stack.enter_context(_patch.object(session, "_execute_tools", side_effect=mock_execute)) + for attr, side_effect in extra_patches.items(): + stack.enter_context(_patch.object(session, attr, side_effect=side_effect)) + stack.enter_context(_patch.object(session, "_full_messages", return_value=[])) + stack.enter_context(_patch.object(session, "_update_token_table")) + stack.enter_context(_patch.object(session, "_print_status_line")) + stack.enter_context(_patch.object(session, "_emit_state")) + stack.enter_context(_patch.object(session, "_visible_memory_count", return_value=0)) + stack.enter_context(_patch.object(session, "_apply_post_execute_advisories")) + save_msg = stack.enter_context(_patch("turnstone.core.session.save_message")) + yield save_msg + + def _user_pending(session) -> list[tuple[str, str]]: """Return user-channel queued nudges as ``(type, text)`` tuples. @@ -2127,50 +2169,73 @@ class TestMetacognitiveBuffers: assert metacog == [] assert len(_tool_pending(session)) == 1 - def test_collect_advisories_does_not_drain_queued_messages(self, tmp_db): - """Queued user messages are NOT drained inside the tool result - envelope — they're appended as a separate user turn after the - full batch completes (see the ``_flush_queued_messages`` call - in ``send`` after the tool-result loop). The pre-fix shape - wrapped them as ``UserInterjection`` advisories that rode - inside ``wrap_tool_result``; the model saw them inline but no - user row landed in storage, so the optimistic queued bubble - vanished on page reload (and on cross-tab replay). + def test_collect_advisories_drains_queued_messages_on_last_result(self, tmp_db): + """Queued user messages drain into a ``UserInterjection`` + persistent advisory on the LAST result of a batch (Seam 1 in + the queued-message architecture). The wrapped tool-result + envelope persists on the tool row; replay extracts the + advisory back via ``decorate_history_messages``. - Persistent + metacog drains still happen here for the output - guard finding and tool-channel nudges; only the queued-message - drain moved out. + Splice-on-last-result mirrors the metacognitive tool-channel + drain — single envelope per batch, no duplication across N + tool results. The queue is cleared so the next-turn flow + doesn't double-deliver. """ + from turnstone.core.tool_advisory import UserInterjection + session = _make_session() pre_count = len(session.messages) session.queue_message("hows it going?", queue_msg_id="q1") persistent, metacog = session._collect_advisories( assessment=None, func_name="bash", is_last_in_batch=True ) - # Neither list carries the queued message anymore. assert metacog == [] - assert persistent == [] - # Queue still holds the message — it'll drain via - # ``_flush_queued_messages`` after the tool batch finishes. - assert "q1" in session._queued_messages - # No separate user turn appended yet (that happens post-batch). + assert len(persistent) == 1 + adv = persistent[0] + assert isinstance(adv, UserInterjection) + assert adv.message == "hows it going?" + assert adv.priority == "notice" + # Queue cleared by the drain. + assert session._queued_messages == {} + # No separate user turn appended (the splice rides inside the + # tool envelope and persists via the wrapped DB row). assert len(session.messages) == pre_count - def test_queued_message_persists_as_user_row_after_tool_batch(self, tmp_db): - """A queued message arriving during a tool batch lands in - ``self.messages`` as a separate user turn after the batch - completes — and gets persisted to storage so a reconnect / - page-reload sees the bubble survive. + def test_collect_advisories_does_not_drain_queued_when_not_last(self, tmp_db): + """Mid-batch results must NOT splice the queued message — the + drain is bound to the last result so a parallel fan-out doesn't + paint the same advisory N times. Queue stays intact until the + last result fires (or until cancel/exception/no-tool-call paths + flush it as Seams 2/3).""" + session = _make_session() + session.queue_message("hows it going?", queue_msg_id="q1") + persistent, metacog = session._collect_advisories( + assessment=None, func_name="bash", is_last_in_batch=False + ) + assert persistent == [] + assert metacog == [] + # Queue intact — the next call (with is_last_in_batch=True) + # will drain it. + assert "q1" in session._queued_messages - Pre-fix shape wrapped the queue into ``wrap_tool_result``'s - ```` envelope; the model saw the text but no - user row was written, so reconnecting tabs lost the bubble. + def test_queued_message_persists_as_user_row_after_tool_batch(self, tmp_db): + """A queued message arriving during a tool batch rides as a + ``UserInterjection`` advisory spliced into the LAST tool result + envelope (Seam 1). The wrapped ```` envelope + persists on the tool DB row; replay's + ``decorate_history_messages`` extracts the advisory back as a + wire-shape user bubble that survives reconnect / page reload. + + Asserts: + - Last tool message in ``session.messages`` carries the + ```` envelope containing the queued text. + - No extra user row landed between the tool batch and the next + assistant (role sequence stays ``assistant -> tool``). + - ``save_message`` was called for the tool row with the + WRAPPED envelope as its content (DB↔memory symmetry). + - Queue cleared post-batch. """ session = _make_session() - # Two-iteration stream: first call returns tool_calls (drives - # the batch path); second call has a queued message arriving - # in between (the post-batch drain seam). Third call returns - # no tool_calls, ending the loop. responses = [ { "role": "assistant", @@ -2185,59 +2250,400 @@ class TestMetacognitiveBuffers: }, {"role": "assistant", "content": "ack"}, ] - stream_idx = 0 - - def mock_stream(_msgs): - nonlocal stream_idx - stream_idx += 1 - return iter([]) - - def mock_response(_stream, _gen): - return responses.pop(0) def mock_execute(_tool_calls): + # Queue arrives DURING the tool batch — Seam 1 fires on + # the last result of the batch. session.queue_message("typed during tool", queue_msg_id="q1") return [("call_x", "ok")], None - with ( - patch.object(session, "_create_stream_with_retry", side_effect=mock_stream), - patch.object(session, "_stream_response", side_effect=mock_response), - patch.object(session, "_execute_tools", side_effect=mock_execute), - patch.object(session, "_full_messages", return_value=[]), - patch.object(session, "_update_token_table"), - patch.object(session, "_print_status_line"), - patch.object(session, "_emit_state"), - patch.object(session, "_visible_memory_count", return_value=0), - patch.object(session, "_apply_post_execute_advisories"), - patch("turnstone.core.session.save_message") as save_msg, + with _send_with_mocks(session, responses, mock_execute) as save_msg: + session._title_generated = True + session.send("first") + + # Role sequence: user(first) -> assistant(tool_calls) -> + # tool(call_x with envelope) -> assistant(ack). No extra + # trailing user row from a post-batch flush — Seam 1 splice + # is the only delivery path for the queued text. + roles = [m.get("role") for m in session.messages] + assert roles == ["user", "assistant", "tool", "assistant"], ( + f"expected user->assistant->tool->assistant, got {roles!r}" + ) + # Last tool message carries the spliced advisory. + tool_msg = session.messages[2] + tool_content = tool_msg["content"] + assert isinstance(tool_content, str) + assert "" in tool_content + assert "typed during tool" in tool_content + assert tool_content.startswith("\n") + # The DB save for the tool row carried the WRAPPED envelope — + # storage matches in-memory shape, so replay extraction is + # complete and lossless. + tool_saves = [ + c for c in save_msg.call_args_list if len(c.args) >= 3 and c.args[1] == "tool" + ] + assert len(tool_saves) == 1 + saved_text = tool_saves[0].args[2] + assert "" in saved_text + assert "typed during tool" in saved_text + # Queue empty after drain. + assert session._queued_messages == {} + + def test_user_feedback_only_creates_single_user_row_via_flush(self, tmp_db): + """When ``_execute_tools`` returns a non-empty ``user_feedback`` + and the queue is empty, the post-batch flush still produces + exactly one trailing user row (the feedback alone). Seam 2 in + the queued-message architecture: flush absorbs the feedback as + a prefix-only call.""" + session = _make_session() + responses = [ + { + "role": "assistant", + "content": "calling", + "tool_calls": [ + { + "id": "call_x", + "type": "function", + "function": {"name": "echo", "arguments": "{}"}, + } + ], + }, + {"role": "assistant", "content": "ack"}, + ] + + def mock_execute(_tool_calls): + return [("call_x", "ok")], "y, use full path" + + with _send_with_mocks(session, responses, mock_execute) as save_msg: + session._title_generated = True + session.send("first") + + roles = [m.get("role") for m in session.messages] + # Single trailing user row carrying the feedback before the + # final assistant ack. + assert roles == ["user", "assistant", "tool", "user", "assistant"], ( + f"expected feedback-as-user-row sequence, got {roles!r}" + ) + assert session.messages[3]["content"] == "y, use full path" + # Persisted via _append_user_turn -> save_message("user", ...). + user_saves = [ + c for c in save_msg.call_args_list if len(c.args) >= 3 and c.args[1] == "user" + ] + assert any(c.args[2] == "y, use full path" for c in user_saves), ( + f"feedback must persist as a user row; saw: {user_saves!r}" + ) + + def test_user_feedback_and_queued_coexistence_single_row_with_prefix(self, tmp_db): + """Seam 1 + Seam 2 coexistence — a queued message that lands + AFTER ``_collect_advisories`` already drained the queue for + the last tool result (Seam 1 closed) but BEFORE + ``_flush_queued_messages`` ran (Seam 2). In production this + race is operator-typing during the approval prompt narrowly + crossing the boundary; here we simulate it by wrapping + ``_collect_advisories`` with a pass-through that queues a + new message AFTER the original returned. + + The queued text rides Seam 2's flush as the suffix of a + single trailing user row, with ``user_feedback`` as the + prefix. Crucially: NO back-to-back user rows (the strict- + template hazard the prefix-merge logic was added to fix). + Reverting the prefix-merge in ``_flush_queued_messages`` + breaks this test.""" + session = _make_session() + responses = [ + { + "role": "assistant", + "content": "calling", + "tool_calls": [ + { + "id": "call_x", + "type": "function", + "function": {"name": "echo", "arguments": "{}"}, + } + ], + }, + {"role": "assistant", "content": "ack"}, + ] + + def mock_execute(_tool_calls): + return [("call_x", "ok")], "y, use full path" + + # Wrap _collect_advisories so we can queue a message AFTER + # the original ran (Seam 1 already closed for this batch). + # The next stop in the call chain is _flush_queued_messages + # — which is Seam 2 and must fold the late arrival in + # alongside user_feedback. + original_collect = session._collect_advisories + + def collect_then_queue_late(*args, **kwargs): + persistent, metacog = original_collect(*args, **kwargs) + # Only queue once, AFTER the last-in-batch drain ran so + # the queue is genuinely empty when we fill it. + if kwargs.get("is_last_in_batch") or (len(args) >= 3 and args[2]): + session.queue_message("late arrival", queue_msg_id="q-late") + return persistent, metacog + + with _send_with_mocks( + session, + responses, + mock_execute, + _collect_advisories=collect_then_queue_late, ): session._title_generated = True session.send("first") - # The queued user message landed in self.messages as a user - # turn after the tool result. roles = [m.get("role") for m in session.messages] - # Expect: user(first) → assistant(tool_calls) → tool(call_x) - # → user(typed during tool) → assistant(ack) - assert roles[-2] == "user", f"expected user before final assistant, got {roles!r}" - assert "typed during tool" in str(session.messages[-2].get("content", "")) - # And persisted: at least one save_message call carried "user" - # role with the queued text. - user_saves = [ - c for c in save_msg.call_args_list if len(c.args) >= 3 and c.args[1] == "user" - ] - assert any("typed during tool" in str(c.args[2]) for c in user_saves), ( - f"queued message must be persisted as a user row; saw: {user_saves!r}" + # NO back-to-back user rows. Pre-fix the user_feedback + # appended a separate user row and the queue drained another: + # roles == [..., "user", "user", ...] which broke strict + # vLLM-template providers (Mistral / Llama). + for i in range(1, len(roles)): + assert not (roles[i] == "user" and roles[i - 1] == "user"), ( + f"back-to-back user rows at idx {i - 1}/{i} in {roles!r}" + ) + # Single trailing user row containing prefix + queued. + assert roles == ["user", "assistant", "tool", "user", "assistant"], ( + f"expected single-trailing-user shape, got {roles!r}" ) - # Queue is empty after drain. + flushed_content = session.messages[3]["content"] + # The two pieces are joined by the canonical separator. + assert flushed_content == "y, use full path\n\nlate arrival" + # Queue cleared. assert session._queued_messages == {} - # Two model turns: one for the tool-call iteration, one for the - # follow-up after the post-batch flush appended the queued user - # row. Without the second stream the model would never see / - # respond to the queued message — pinning this guards against a - # future regression where the post-batch flush runs but the - # send-loop short-circuits before the next iteration. - assert stream_idx == 2, f"expected 2 stream calls (tool-iter + follow-up), got {stream_idx}" + + def test_flush_queued_messages_with_prefix_only(self, tmp_db): + """Empty queue + non-empty prefix produces one user row + carrying the prefix verbatim. Returns True so the caller + knows a turn was appended.""" + session = _make_session() + pre_count = len(session.messages) + appended = session._flush_queued_messages(prefix="hello") + assert appended is True + assert len(session.messages) == pre_count + 1 + last = session.messages[-1] + assert last["role"] == "user" + assert last["content"] == "hello" + + def test_flush_queued_messages_with_prefix_and_items(self, tmp_db): + """Both prefix and queued items produce ONE user row joining + prefix + items with the canonical ``\\n\\n`` separator. Queue + is cleared on drain so a re-entry doesn't double-deliver.""" + session = _make_session() + session.queue_message("a", queue_msg_id="q-a") + session.queue_message("b", queue_msg_id="q-b") + appended = session._flush_queued_messages(prefix="approve") + assert appended is True + last = session.messages[-1] + assert last["role"] == "user" + assert last["content"] == "approve\n\na\n\nb" + assert session._queued_messages == {} + + def test_tool_db_row_stores_wrapped_output_when_advisories_present(self, tmp_db): + """A tool row whose batch produced a ``UserInterjection`` + advisory persists the WRAPPED envelope to storage, not the + bare raw output. Symmetry between + ``self.messages[i]['content']`` and the DB row is what makes + replay's envelope-extraction lossless.""" + session = _make_session() + responses = [ + { + "role": "assistant", + "content": "calling", + "tool_calls": [ + { + "id": "call_x", + "type": "function", + "function": {"name": "echo", "arguments": "{}"}, + } + ], + }, + {"role": "assistant", "content": "ack"}, + ] + + def mock_execute(_tool_calls): + session.queue_message("during", queue_msg_id="q-d") + return [("call_x", "raw output")], None + + with _send_with_mocks(session, responses, mock_execute) as save_msg: + session._title_generated = True + session.send("first") + + tool_saves = [ + c for c in save_msg.call_args_list if len(c.args) >= 3 and c.args[1] == "tool" + ] + assert len(tool_saves) == 1 + saved_text = tool_saves[0].args[2] + assert saved_text.startswith("\n") + assert "" in saved_text + assert "during" in saved_text + # And the raw text is still inside the envelope. + assert "raw output" in saved_text + + def test_tool_db_row_stores_raw_output_when_no_advisories(self, tmp_db): + """When ``_collect_advisories`` returns empty (no guard + finding, no queued message), ``wrap_tool_result`` no-ops and + the DB row gets the bare raw output — symmetry with + ``self.messages[i]['content']`` is preserved without forcing + every tool through the envelope path.""" + session = _make_session() + responses = [ + { + "role": "assistant", + "content": "calling", + "tool_calls": [ + { + "id": "call_x", + "type": "function", + "function": {"name": "echo", "arguments": "{}"}, + } + ], + }, + {"role": "assistant", "content": "ack"}, + ] + + def mock_execute(_tool_calls): + return [("call_x", "raw output")], None + + with _send_with_mocks(session, responses, mock_execute) as save_msg: + session._title_generated = True + session.send("first") + + tool_saves = [ + c for c in save_msg.call_args_list if len(c.args) >= 3 and c.args[1] == "tool" + ] + assert len(tool_saves) == 1 + saved_text = tool_saves[0].args[2] + # Bare raw output — no envelope at all. + assert saved_text == "raw output" + assert "" not in saved_text + assert "" not in saved_text + + def test_tool_db_row_stores_wrapped_output_for_list_content(self, tmp_db): + """Image / structured tool output (list-typed) projects to a + proper envelope-anchored-at-start string in the TEXT column — + ``wrap_tool_result(joined_raw_text, advisories)`` rebuilds the + envelope from the PRE-wrap text parts so the saved row starts + with ``\\n``. The previous join-of-post-wrap-parts + produced a `` ...`` shape that + ``extract_advisories_from_tool_envelope``'s ``startswith`` + prefix check rejected on replay. Replacing ``raw_output`` with + ``output`` (the post-wrap list) in the persistence projection + breaks this test.""" + session = _make_session() + responses = [ + { + "role": "assistant", + "content": "calling", + "tool_calls": [ + { + "id": "call_x", + "type": "function", + "function": {"name": "view_image", "arguments": "{}"}, + } + ], + }, + {"role": "assistant", "content": "ack"}, + ] + + def mock_execute(_tool_calls): + session.queue_message("about that image", queue_msg_id="q-i") + return [ + ( + "call_x", + [ + {"type": "text", "text": "raw text part"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,xxx"}}, + ], + ) + ], None + + with _send_with_mocks(session, responses, mock_execute) as save_msg: + session._title_generated = True + session.send("first") + + tool_saves = [ + c for c in save_msg.call_args_list if len(c.args) >= 3 and c.args[1] == "tool" + ] + assert len(tool_saves) == 1 + saved_text = tool_saves[0].args[2] + # Envelope anchored at the start so replay's ``startswith`` + # prefix check matches and extracts the advisory cleanly. + assert saved_text.startswith("\n") + assert "raw text part" in saved_text + assert "" in saved_text + assert "about that image" in saved_text + + def test_tool_db_row_round_trips_list_output_with_advisories(self, tmp_db): + """Round-trip the persistence projection for list-typed output: + the saved string must extract back to the original raw text + plus the user_interjection advisory — same contract the string + case satisfies via the symmetry between + ``self.messages[i]['content']`` and the DB row. + + Replacing the persistence projection's + ``wrap_tool_result(raw_text, persistent_advisories)`` with the + previous join-of-post-wrap-parts shape produces a string that + does NOT start with ``\\n``; + ``extract_advisories_from_tool_envelope`` returns ``None`` and + the raw envelope XML leaks into the rendered tool bubble on + replay. This test fails in that broken-shape branch.""" + from turnstone.core.history_decoration import ( + extract_advisories_from_tool_envelope, + ) + + session = _make_session() + responses = [ + { + "role": "assistant", + "content": "calling", + "tool_calls": [ + { + "id": "call_x", + "type": "function", + "function": {"name": "view_image", "arguments": "{}"}, + } + ], + }, + {"role": "assistant", "content": "ack"}, + ] + + def mock_execute(_tool_calls): + session.queue_message("inspect the histogram", queue_msg_id="q-i") + return [ + ( + "call_x", + [ + {"type": "text", "text": "the chart shows X"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,xxx"}}, + ], + ) + ], None + + with _send_with_mocks(session, responses, mock_execute) as save_msg: + session._title_generated = True + session.send("first") + + tool_saves = [ + c for c in save_msg.call_args_list if len(c.args) >= 3 and c.args[1] == "tool" + ] + assert len(tool_saves) == 1 + saved_text = tool_saves[0].args[2] + # Saved row is anchored on the envelope prefix replay's + # extractor checks for. + assert saved_text.startswith("\n") + # And the extractor recovers the cleaned text + the queued- + # message advisory it carried. + result = extract_advisories_from_tool_envelope(saved_text) + assert result is not None + cleaned, advisories = result + assert cleaned == "the chart shows X" + assert advisories == [ + { + "type": "user_interjection", + "text": "inspect the histogram", + "priority": "notice", + } + ] def test_start_nudge_fires_through_send(self, tmp_db): """Pin the +1 count-shift invariant — `start` must still fire on the @@ -2822,6 +3228,112 @@ class TestApplyRemindersForProvider: session._apply_reminders_for_provider([msg]) assert msg["_reminders"] == [{"type": "correction", "text": "WATCH"}] + def test_apply_reminders_preserves_existing_wrapper_envelope(self, tmp_db): + """A tool message whose ``content`` already carries a + ```` envelope (queued-message ``UserInterjection`` + spliced via ``wrap_tool_result`` on Seam 1, or a pre-existing + output_guard advisory) must not be entity-encoded by the + provider splice — the wrapper has already escaped the raw body + and re-escaping would turn the envelope's literal + ```` and ```` tags into + ``<tool_output>...`` so the model sees mangled text + instead of a parseable envelope. + + Removing the ``content.startswith("\\n")`` + wrapper-detection branch in ``_apply_reminders_for_provider`` + (so the splice always re-escapes) breaks this test.""" + from turnstone.core.tool_advisory import UserInterjection, wrap_tool_result + + session = _make_session() + wrapped = wrap_tool_result( + "raw body", + [UserInterjection(message="check logs", priority="notice")], + ) + msg = { + "role": "tool", + "tool_call_id": "call_a", + "content": wrapped, + "_reminders": [{"type": "tool_error", "text": "retry suggestion"}], + } + out = session._apply_reminders_for_provider([msg]) + wire = out[0]["content"] + assert isinstance(wire, str) + # Envelope still anchored at the start — replay's prefix check + # would still match this projection. + assert wire.startswith("\n") + # The original envelope's literal tags survive (NOT entity- + # encoded). ```` should appear at least + # twice: once for the queued-message advisory inside the + # original envelope, once for the freshly appended tool_error + # block. + assert wire.count("") >= 2 + assert wire.count("") >= 2 + # The wrapper's literal tags must NOT have been entity-encoded + # by escape_wrapper_tags. + assert "<tool_output>" not in wire + assert "<system-reminder>" not in wire + # Both the original advisory body and the appended reminder + # text are present. + assert "check logs" in wire + assert "retry suggestion" in wire + + def test_apply_reminders_preserves_wrapper_envelope_in_list_text_part(self, tmp_db): + """Parallel of the string-content case for list-typed tool + output (image / structured MCP results). The Seam 1 splice + for list content appends the wrap envelope as a separate + text part (see ``session.py`` tool-result loop where + ``wrap_tool_result("", advisories)`` lands as + ``{"type": "text", "text": }``). When the same + message also carries ``_reminders``, the per-text-part + ``escape_wrapper_tags`` loop in ``_apply_reminders_for_provider`` + must skip parts whose text already starts with + ``\\n`` — re-escaping would mangle the + envelope's literal tags so the model sees opaque entity- + encoded text instead of a parseable envelope. + + Removing the ``text.startswith("\\n")`` skip in + the list-branch escape loop breaks this test.""" + from turnstone.core.tool_advisory import UserInterjection, wrap_tool_result + + session = _make_session() + wrap_text = wrap_tool_result( + "", + [UserInterjection(message="check logs", priority="notice")], + ) + msg = { + "role": "tool", + "tool_call_id": "call_a", + "content": [ + {"type": "text", "text": "the chart shows X"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,xxx"}}, + {"type": "text", "text": wrap_text}, + ], + "_reminders": [{"type": "tool_error", "text": "retry suggestion"}], + } + out = session._apply_reminders_for_provider([msg]) + wire_parts = out[0]["content"] + assert isinstance(wire_parts, list) + # The wrap text-part must survive intact — its literal + # ```` and ```` tags are + # not entity-encoded. + wrap_part_text = next( + p["text"] + for p in wire_parts + if isinstance(p, dict) + and p.get("type") == "text" + and p.get("text", "").startswith("\n") + ) + assert "<tool_output>" not in wrap_part_text + assert "<system-reminder>" not in wrap_part_text + assert "check logs" in wrap_part_text + # The non-envelope text part WAS escape-walked (it contains + # no wrapper-tag literals so the escape is a no-op, but the + # appended reminder block lands here per the existing + # contract). + last_text = wire_parts[-1] + assert isinstance(last_text, dict) and last_text.get("type") == "text" + assert "retry suggestion" in last_text.get("text", "") + class TestMarkRemindersDelivered: """``_mark_reminders_delivered`` flips the wire-suppression flag on diff --git a/tests/test_tool_advisory.py b/tests/test_tool_advisory.py index dc525bc4..764a7cb9 100644 --- a/tests/test_tool_advisory.py +++ b/tests/test_tool_advisory.py @@ -7,6 +7,7 @@ from turnstone.core.tool_advisory import ( GuardAdvisory, MetacognitiveAdvisory, UserInterjection, + escape_wrapper_tags, parse_priority, render_system_reminder, wrap_tool_result, @@ -224,6 +225,63 @@ class TestMetacognitiveAdvisory: assert "don't repeat tool calls" in result +class TestEscapeWrapperTags: + """``escape_wrapper_tags`` must round-trip through + ``_entity_decode_wrapper_tags`` for any input — not just text that + happens to contain only wrapper tags. + """ + + def test_short_circuit_passes_through_plain_text(self) -> None: + """No ``<`` and no ``&`` — ``escape_wrapper_tags`` must avoid + the four ``replace`` chains. Common case for most tool outputs; + the short-circuit keeps wrap_tool_result's overhead near zero.""" + text = "plain text without any markup" + assert escape_wrapper_tags(text) == text + + def test_escape_wrapper_tags_round_trips_preexisting_entities(self) -> None: + """Asymmetry guard — a tool output that happens to contain the + literal string ``<tool_output>`` (e.g. documentation + describing the wrapper format) must round-trip identically. + Without escaping ``&`` first, encode→decode would produce the + bare ```` tag, fabricating an envelope the wrapper + layer never produced.""" + from turnstone.core.history_decoration import _entity_decode_wrapper_tags + + text = "I describe XML tags like <tool_output> in my docs." + encoded = escape_wrapper_tags(text) + # Sanity: the original literal got escaped to a sentinel form + # that can't collide with our wrapper-tag escapes. + assert "&lt;tool_output&gt;" in encoded + assert "<tool_output>" not in encoded + # Round-trip back to the literal source. + assert _entity_decode_wrapper_tags(encoded) == text + + def test_escape_wrapper_tags_round_trips_real_wrapper_tag(self) -> None: + """A literal ```` in source text round-trips back + correctly — encoding produces ``<tool_output>`` (no + ``&`` prefix because there was no pre-existing entity), and + decoding restores the literal.""" + from turnstone.core.history_decoration import _entity_decode_wrapper_tags + + text = "Here is a literal tag in my doc." + encoded = escape_wrapper_tags(text) + assert "" not in encoded + assert "<tool_output>" in encoded + assert _entity_decode_wrapper_tags(encoded) == text + + def test_escape_wrapper_tags_round_trips_mixed_content(self) -> None: + """Mixed: literal wrapper tags AND pre-existing entity + references — both round-trip.""" + from turnstone.core.history_decoration import _entity_decode_wrapper_tags + + text = ( + "Mixed: literal next to escaped <system-reminder> " + "and a stray & on its own." + ) + encoded = escape_wrapper_tags(text) + assert _entity_decode_wrapper_tags(encoded) == text + + class TestRenderSystemReminder: """render_system_reminder builds a standalone envelope.""" diff --git a/tests/test_workstream_endpoints.py b/tests/test_workstream_endpoints.py index 3804f53f..7727fbdf 100644 --- a/tests/test_workstream_endpoints.py +++ b/tests/test_workstream_endpoints.py @@ -1031,6 +1031,188 @@ class TestBuildHistoryReminderPropagation: assert history[0]["content"] == content +class TestBuildHistoryAdvisoryRoundTrip: + """``_build_history`` must round-trip the persisted + ```` envelope (Seam 1 queued-message splice) to + cleaned content + a wire-shape ``advisories`` array. + + Production realism note: ``session.messages`` never carries an + ``advisories`` key — only ``decorate_history_messages`` mutates + dicts to add it for the REST ``/history`` path, and the SSE replay + surface bypasses that decoration entirely. The earlier + ``TestBuildHistoryAdvisoryPropagation`` class pre-populated + ``advisories`` directly on the session messages, which tested a + passthrough that doesn't exist in production — the SSE replay code + path silently dropped queued messages despite the green tests. + These round-trip tests exercise the production shape (wrapped + envelope on the tool row's ``content``) so a regression in the + inline ``extract_advisories_from_tool_envelope`` call inside + ``_build_history`` surfaces here. + """ + + def _session_with_messages(self, messages: list[dict]) -> MagicMock: + session = MagicMock() + session.messages = messages + return session + + def test_build_history_round_trips_envelope_to_advisories(self): + """The production-realistic shape: a tool row whose ``content`` + is the wrapped ```` envelope (no ``advisories`` + key set — that's the bug-1 footprint). ``_build_history`` + must extract the advisory back out and ship it on the wire as + cleaned content + ``advisories``. + + Reverting the inline ``extract_advisories_from_tool_envelope`` + call in ``server._build_history``'s tool-message branch breaks + this test. + """ + from turnstone.core.tool_advisory import UserInterjection, wrap_tool_result + from turnstone.server import _build_history + + wrapped = wrap_tool_result( + "tool body", + [UserInterjection(message="check logs", priority="notice")], + ) + session = self._session_with_messages( + [ + { + "role": "tool", + "tool_call_id": "call_a", + "content": wrapped, + } + ] + ) + history = _build_history(session) + # Cleaned content rides on the wire — envelope stripped. + assert history[0]["content"] == "tool body" + # Advisory survives as a wire-shape entry the JS can render + # as a user bubble after the tool block. + assert history[0]["advisories"] == [ + {"type": "user_interjection", "text": "check logs", "priority": "notice"} + ] + + def test_build_history_round_trips_important_priority(self): + """The ``important`` priority preamble round-trips — pin both + the priority detection in the parser and the projection through + to the wire shape.""" + from turnstone.core.tool_advisory import UserInterjection, wrap_tool_result + from turnstone.server import _build_history + + wrapped = wrap_tool_result( + "out", + [UserInterjection(message="urgent", priority="important")], + ) + session = self._session_with_messages( + [{"role": "tool", "tool_call_id": "call_a", "content": wrapped}] + ) + history = _build_history(session) + assert history[0]["content"] == "out" + assert history[0]["advisories"] == [ + {"type": "user_interjection", "text": "urgent", "priority": "important"} + ] + + def test_build_history_no_envelope_passes_through_unchanged(self): + """Plain tool content (no ```` prefix) — no + advisories field, content unchanged.""" + from turnstone.server import _build_history + + session = self._session_with_messages( + [{"role": "tool", "tool_call_id": "call_a", "content": "plain output"}] + ) + history = _build_history(session) + assert history[0]["content"] == "plain output" + assert "advisories" not in history[0] + + def test_build_history_round_trip_through_full_decoration_chain(self): + """End-to-end pin: persist a wrapped envelope into ``messages``, + run the full decoration chain (``decorate_history_messages`` + followed by ``_build_history``), assert the wire shape carries + the advisory. This pins the contract every component in the + chain participates in — REST ``/history`` callers go through + ``decorate_history_messages``, and SSE replay goes through + ``_build_history`` — both must produce the same wire shape. + """ + from turnstone.core.history_decoration import decorate_history_messages + from turnstone.core.tool_advisory import UserInterjection, wrap_tool_result + from turnstone.server import _build_history + + wrapped = wrap_tool_result( + "raw", + [UserInterjection(message="hi", priority="notice")], + ) + # Decorate first — REST /history shape. + rest_messages: list[dict] = [{"role": "tool", "tool_call_id": "call_a", "content": wrapped}] + decorate_history_messages(rest_messages, {}, {}) + # And separately drive _build_history with a fresh undecorated + # message — SSE replay shape. + session = self._session_with_messages( + [{"role": "tool", "tool_call_id": "call_a", "content": wrapped}] + ) + sse_history = _build_history(session) + # Both surfaces produce the same advisory + cleaned content. + assert rest_messages[0]["content"] == "raw" + assert rest_messages[0]["advisories"] == [ + {"type": "user_interjection", "text": "hi", "priority": "notice"} + ] + assert sse_history[0]["content"] == "raw" + assert sse_history[0]["advisories"] == [ + {"type": "user_interjection", "text": "hi", "priority": "notice"} + ] + + def test_build_history_extracts_advisories_from_list_content_text_part(self): + """List-typed tool output (image / structured MCP results) + with a Seam 1 splice carries the wrap envelope as a separate + text part (``session.py``'s tool-result loop appends + ``{"type": "text", "text": wrap_tool_result("", advisories)}`` + when ``output`` is a list). ``_build_history`` must walk the + list parts, extract advisories from any wrap-envelope text + part, and DROP that text part from the projected list — the + cleaned inner content is empty by construction, and leaving + the part would cause the JS replay to render the literal + envelope text as a chunk inside the tool block AND fail to + render the queued message as a user bubble. + + Removing the list-content branch in ``_build_history``'s tool- + message advisory extraction breaks this test. + """ + from turnstone.core.tool_advisory import UserInterjection, wrap_tool_result + from turnstone.server import _build_history + + wrap_text = wrap_tool_result( + "", + [UserInterjection(message="inspect histogram", priority="notice")], + ) + list_content = [ + {"type": "text", "text": "the chart shows X"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,xxx"}}, + {"type": "text", "text": wrap_text}, + ] + session = self._session_with_messages( + [{"role": "tool", "tool_call_id": "call_a", "content": list_content}] + ) + history = _build_history(session) + # Wire-shape content keeps the original text + image parts but + # has the wrap text-part dropped. + wire_content = history[0]["content"] + assert isinstance(wire_content, list) + assert len(wire_content) == 2 + assert wire_content[0] == {"type": "text", "text": "the chart shows X"} + assert wire_content[1] == { + "type": "image_url", + "image_url": {"url": "data:image/png;base64,xxx"}, + } + # Advisory rides on the wire so JS replay renders the user + # bubble after the tool block — same contract as the string- + # content path. + assert history[0]["advisories"] == [ + { + "type": "user_interjection", + "text": "inspect histogram", + "priority": "notice", + } + ] + + class TestDetailInteractive: """Interactive parity for the lifted ``GET /v1/api/workstreams/{ws_id}``. diff --git a/turnstone/console/static/coordinator/coordinator.css b/turnstone/console/static/coordinator/coordinator.css index 2025152a..e93ac1b6 100644 --- a/turnstone/console/static/coordinator/coordinator.css +++ b/turnstone/console/static/coordinator/coordinator.css @@ -439,25 +439,6 @@ font-size: 10px; } -/* Storage-truncation indicator — same convention as the interactive - UI's `.tool-output-truncated` pill (transparent bg, dim border, - small font) so the operator reads the affordance the same way on - both surfaces. Sibling node next to .coord-tool-row-result rather - than text-in-content so a future "best-effort JSON repair" pass - on the result body doesn't have to strip a marker string. */ -.coord-tool-truncated { - display: inline-block; - margin-top: 4px; - margin-left: 6px; - padding: 1px 6px; - font-size: 10px; - font-family: var(--font-mono); - color: var(--ink-3); - background: transparent; - border: 1px solid var(--ink-3); - border-radius: 3px; -} - /* memory/recall calls are background metadata — the audit trail is useful but they crowd the tree on workstreams with heavy memory usage. Dim the row by default; full opacity on hover so they diff --git a/turnstone/console/static/coordinator/coordinator.js b/turnstone/console/static/coordinator/coordinator.js index 95b1142f..ff2ab267 100644 --- a/turnstone/console/static/coordinator/coordinator.js +++ b/turnstone/console/static/coordinator/coordinator.js @@ -949,10 +949,6 @@ if (!row) return; const existing = row.querySelector(".coord-tool-row-result"); if (existing) existing.remove(); - // Re-fires (cancel + rerun, error + retry) clear any prior - // truncation pill so it doesn't stack on the new result. - const existingTrunc = row.querySelector(".coord-tool-truncated"); - if (existingTrunc) existingTrunc.remove(); if (isError) { row.classList.add("error"); // Lift the row's error onto the enclosing batch so the left @@ -1008,19 +1004,6 @@ body.textContent = pretty; block.appendChild(body); row.appendChild(block); - // Storage-truncation indicator — sibling pill (not text inside - // the result body) so renderers / parsers / copy-as-text paths - // see the unmodified output. Same convention as interactive's - // .tool-output-truncated; styled by .coord-tool-truncated in - // coordinator.css. - if (opts && opts.truncated) { - const pill = document.createElement("span"); - pill.className = "coord-tool-truncated"; - pill.textContent = "… truncated in storage"; - pill.title = - "Full tool output was sent to the model live; only the first 10000 characters are persisted to the conversation row."; - row.appendChild(pill); - } } function _makeActionButton(label, role, kbdHint, ariaLabel) { @@ -4088,20 +4071,27 @@ const toolName = (callId && toolNameByCallId.get(callId)) || m.tool_name || "tool"; const isError = callOutcomes.get(callId) === "error"; - // Storage truncation surfaces as a sibling pill next to - // the result (see _appendResultToRow's opts.truncated - // branch) rather than as text inside the result body — a - // future "best-effort JSON repair" pass would otherwise - // need to strip a marker string before parsing. - appendToolResult(toolName, callId, content || "", isError, { - truncated: !!m.truncated, - }); + appendToolResult(toolName, callId, content || "", isError); // Tool-channel metacog reminders ride the same _reminders // side-channel as the user channel; surface as a themed // bubble below the .coord-tool-batch construct. if (Array.isArray(m.reminders) && m.reminders.length) { appendToolReminderLive(m.reminders, callId); } + // Queued user messages spliced into the last tool-result + // envelope of a batch (Seam 1) replay as proper user bubbles + // after the tool block. ``decorate_history_messages`` + // extracts the user_interjection advisory from the persisted + // envelope and the wire layer projects it onto + // ``m.advisories``; rendering through + // ``appendUserMessageWithAttachments`` matches the live shape + // a Seam 2/3 message would produce. The walk/filter is + // shared via ``replayAdvisoriesAfterTool`` in + // ``shared/utils.js`` so coord and interactive can never drift + // on advisory-shape filtering. + replayAdvisoriesAfterTool(m.advisories, function (text) { + appendUserMessageWithAttachments(text, [], { label: "user" }); + }); } else if (role === "assistant") { // Render content BEFORE the tool batch so DOM order matches // chronological order (the model emits text first, then diff --git a/turnstone/core/history_decoration.py b/turnstone/core/history_decoration.py index c8fccccf..79a09df8 100644 --- a/turnstone/core/history_decoration.py +++ b/turnstone/core/history_decoration.py @@ -22,24 +22,14 @@ import json from typing import Any from turnstone.core.log import get_logger +from turnstone.core.tool_advisory import ( + _USER_INTERJECTION_BODY_MARKER, + _USER_INTERJECTION_IMPORTANT_PREAMBLE, + _USER_INTERJECTION_NOTICE_PREAMBLE, +) log = get_logger(__name__) -# Tool results are clamped at this length per row at storage time -# (see ``session.py``'s ``store_text = raw_output[:TOOL_RESULT_STORAGE_CAP]``). -# Keeping the constant here lets the truncation flag detection in -# ``decorate_history_messages`` stay in sync without a magic number -# duplicated across server.py / session.py. -# -# Raised from 2000 → 10000 because a 2000-char clip routinely cut -# the body of a single grep / file read mid-line, leaving the -# historical record useless for retrospective debugging. FTS5 -# index + row size grow proportionally; the per-tool upper bound is -# still bounded upstream by ``_truncate_output``'s context-budget -# clamp (so a single huge result can't blow past the live context -# window). -TOOL_RESULT_STORAGE_CAP = 10000 - def load_verdict_indexes( ws_id: str, @@ -174,6 +164,103 @@ def decorate_tool_call( tc["output_assessment"] = assessment +def _entity_decode_wrapper_tags(text: str) -> str: + """Reverse :func:`tool_advisory.escape_wrapper_tags` on extraction. + + The wrap layer escapes the four wrapper-tag forms to HTML entities + so embedded user / advisory text cannot fabricate or close an + envelope. When the replay decorator pulls advisories back out of + the persisted envelope, the inner text needs to be returned to its + literal form for UI rendering. + + Decodes ``&`` last so a tool output that contains the literal + string ``<tool_output>`` round-trips identically to its + source: encode produces ``&lt;tool_output&gt;`` (no + collision with wrapper-tag escapes), decode walks the wrapper + escapes first, then strips the ``&`` sentinel back to ``&``. + The short-circuit on ``"&" not in text`` covers the common case + where no escaped entities are present. + """ + if "&" not in text: + return text + return ( + text.replace("</tool_output>", "") + .replace("<tool_output>", "") + .replace("<system-reminder>", "") + .replace("</system-reminder>", "") + .replace("&", "&") + ) + + +def _classify_advisory(render_text: str) -> dict[str, str] | None: + """Map a ```` body back to a wire-shape advisory. + + Returns a dict with ``type`` / ``text`` / optional ``priority`` for + advisory shapes the UI knows how to render, or ``None`` to suppress + the advisory entirely (output-guard findings already render via the + ``output_assessment`` audit-table decoration; doubling them would + paint two warning bubbles). Unknown advisory shapes fall through + to ``None`` rather than rendering an opaque envelope blob. + """ + if render_text.startswith("Output guard:"): + return None + if _USER_INTERJECTION_BODY_MARKER in render_text: + # UserInterjection is the only producer that uses this marker. + # The preamble disambiguates priority: "important" gets the + # MUST-address framing, "notice" gets the incorporate-if-relevant + # framing. The body sits after the marker. Preamble + marker + # constants are imported from ``tool_advisory`` so the parser + # and producer can never drift on wording. + if render_text.startswith(_USER_INTERJECTION_IMPORTANT_PREAMBLE): + priority = "important" + elif render_text.startswith(_USER_INTERJECTION_NOTICE_PREAMBLE): + priority = "notice" + else: + # Marker present but preamble drifted — still render as a + # notice rather than dropping the user's text. + priority = "notice" + body = render_text.split(_USER_INTERJECTION_BODY_MARKER, 1)[1] + return {"type": "user_interjection", "text": body, "priority": priority} + return None + + +def extract_advisories_from_tool_envelope( + content: str, +) -> tuple[str, list[dict[str, str]]] | None: + """Strip a ```` envelope and return ``(clean, advisories)``. + + Returns ``None`` when *content* doesn't look like a wrapped tool + result — caller should leave the message unchanged. When the + envelope parses but no advisories survive classification (e.g. + only an output_guard advisory rode along), returns the cleaned + output with an empty advisories list — the caller still needs to + strip the envelope from the rendered content. + """ + if not content.startswith("\n"): + return None + close = content.find("\n") + if close == -1: + return None + inner = content[len("\n") : close] + rest = content[close + len("\n") :] + advisories: list[dict[str, str]] = [] + cursor = 0 + while True: + open_idx = rest.find("\n", cursor) + if open_idx == -1: + break + close_idx = rest.find("\n", open_idx) + if close_idx == -1: + break + body = rest[open_idx + len("\n") : close_idx] + decoded = _entity_decode_wrapper_tags(body) + classified = _classify_advisory(decoded) + if classified is not None: + advisories.append(classified) + cursor = close_idx + len("\n") + return _entity_decode_wrapper_tags(inner), advisories + + def decorate_history_messages( messages: list[dict[str, Any]], verdicts_by_call_id: dict[str, dict[str, Any]], @@ -184,8 +271,12 @@ def decorate_history_messages( Used by the ``/history`` REST endpoint after ``load_messages`` returns. For each assistant message with ``tool_calls``, runs :func:`decorate_tool_call` on every entry. For each tool message - whose content hits the storage cap, sets ``truncated: True`` so - the client can render the "… truncated in storage" pill. + whose ``content`` carries a ```` envelope (queued + user message spliced via :class:`UserInterjection` during a tool + batch), strips the envelope, restores literal wrapper tags inside + the body, and surfaces the extracted advisories on + ``msg["advisories"]`` so the wire layer can replay them as user + bubbles after the tool result. Pure transform — no I/O. Async callers should pre-load the indexes via :func:`load_verdict_indexes` (in ``to_thread``) and @@ -201,5 +292,18 @@ def decorate_history_messages( decorate_tool_call(tc, verdicts_by_call_id, assessments_by_call_id) elif role == "tool": content = msg.get("content") - if isinstance(content, str) and len(content) >= TOOL_RESULT_STORAGE_CAP: - msg["truncated"] = True + if not isinstance(content, str): + continue + try: + extracted = extract_advisories_from_tool_envelope(content) + except Exception: + # Defensive — on any unexpected parse failure leave the + # message untouched rather than crashing the replay. + log.debug("advisory extraction failed; leaving content intact", exc_info=True) + continue + if extracted is None: + continue + cleaned, advisories = extracted + msg["content"] = cleaned + if advisories: + msg["advisories"] = advisories diff --git a/turnstone/core/session.py b/turnstone/core/session.py index 99638c08..3a6c493c 100644 --- a/turnstone/core/session.py +++ b/turnstone/core/session.py @@ -45,7 +45,6 @@ from turnstone.core.attachments import ( ) from turnstone.core.config import get_tavily_key from turnstone.core.edit import find_occurrences, pick_nearest -from turnstone.core.history_decoration import TOOL_RESULT_STORAGE_CAP from turnstone.core.log import get_logger from turnstone.core.memory import ( count_structured_memories, @@ -555,10 +554,9 @@ _WATCH_QUEUE_SOFT_CAP = 50 # producers (a watch streaming unbounded shell output, a corruption-class # steering payload). The in-memory side-channel keeps the full body — # only the persisted JSON is clamped — so the live splice and UI render -# see the same shape. Mirrors ``TOOL_RESULT_STORAGE_CAP`` on tool -# result rows. Cap is in UTF-8 bytes so non-ASCII payloads (CJK, emoji) -# can't blow the byte ceiling 4x by living entirely under the codepoint -# count. +# see the same shape. Cap is in UTF-8 bytes so non-ASCII payloads +# (CJK, emoji) can't blow the byte ceiling 4x by living entirely under +# the codepoint count. REMINDER_TEXT_STORAGE_CAP = 8192 @@ -2255,7 +2253,21 @@ class ChatSession: copy = dict(msg) content = copy.get("content") if isinstance(content, str): - copy["content"] = escape_wrapper_tags(content) + block + # Tool messages whose content already carries a + # ```` envelope (queued-message advisories + # via ``wrap_tool_result``, output_guard findings) must + # not be re-escaped — the wrapper has already escaped + # the raw body, and re-escaping would entity-encode the + # envelope's literal ```` and + # ```` tags so the model sees + # ``<tool_output>...`` instead of a parseable + # envelope. Append the metacog block as a trailing + # ```` (same shape ``wrap_tool_result`` + # produces when stacking multiple advisories). + if content.startswith("\n"): + copy["content"] = content + block + else: + copy["content"] = escape_wrapper_tags(content) + block elif isinstance(content, list): # Shallow-copy the parts list and any text parts we'll # mutate so the original list/dicts in self.messages stay @@ -2268,7 +2280,21 @@ class ChatSession: p for p in new_parts if isinstance(p, dict) and p.get("type") == "text" ] for part in text_parts: - part["text"] = escape_wrapper_tags(part.get("text", "")) + text = part.get("text", "") + # Skip escape on text parts that already carry a + # ```` envelope — those are produced by + # ``wrap_tool_result`` (Seam 1 splice for list-typed + # tool output appends the envelope as its own text + # part). The wrapper has already escaped the inner + # body; re-escaping here would entity-encode the + # envelope's literal ```` and + # ```` tags so the model sees + # ``<tool_output>...`` instead of a parseable + # envelope. Mirrors the string-branch detection + # above. + if text.startswith("\n"): + continue + part["text"] = escape_wrapper_tags(text) if text_parts: text_parts[-1]["text"] = text_parts[-1]["text"] + block else: @@ -3067,23 +3093,20 @@ class ChatSession: budget = self._remaining_token_budget() output = self._truncate_output(output, remaining_budget_tokens=budget) - # Capture raw output for DB storage before advisory wrapping - raw_output = output - - # Advisory injection: persistent advisories (output - # guard findings) wrap into the tool-result envelope - # and stay in self.messages. Metacognitive - # tool-channel reminders (tool_error / repeat) ride - # a side-channel — never inside content — so the - # model sees the splice only at the wire boundary - # via _apply_reminders_for_provider, while UI/replay + # Advisory injection: persistent advisories — output + # guard findings AND queued user messages (Seam 1) + # — wrap into the tool-result envelope and stay in + # self.messages. Metacognitive tool-channel + # reminders (tool_error / repeat) ride a side- + # channel — never inside content — so the model + # sees the splice only at the wire boundary via + # _apply_reminders_for_provider, while UI/replay # surfaces them as a themed bubble below the tool - # result. Queued user messages drain via - # ``_flush_queued_messages`` AFTER this loop, so - # they no longer ride the persistent-advisory path. + # result. persistent_advisories, metacog_reminders = self._collect_advisories( assessment, _tc_names.get(tc_id, ""), _ri == _last_idx ) + raw_output = output if isinstance(output, str): output = wrap_tool_result(output, persistent_advisories) elif isinstance(output, list) and persistent_advisories: @@ -3129,20 +3152,44 @@ class ChatSession: tok_est = max(1, int(len(output) / self._chars_per_token)) self._msg_tokens.append(tok_est) - # Log tool result. Use raw_output (pre-advisory-wrap) - # so the DB stores clean tool output without ephemeral - # advisory XML. memory/recall persist alongside every - # other tool: replays show the full audit trail, and - # output already passes through _truncate_output above - # so size is bounded by the same budget every other - # tool uses. + # Log tool result. Store the WRAPPED output so the + # persisted row matches ``self.messages[i]['content']`` + # exactly — when ``persistent_advisories`` is empty + # ``wrap_tool_result`` no-ops and the row carries + # bare text, when non-empty the row carries the + # ```` envelope so replay's + # ``decorate_history_messages`` can extract the + # advisory back out. Size is already bounded by + # ``_truncate_output`` above (per-turn context + # budget); no second cap needed. + # + # List-typed output (image / structured MCP results) + # has no native string projection, so we rebuild the + # envelope from the joined PRE-wrap text parts plus + # the same advisories. Joining the POST-wrap parts + # would put the original raw text first and the + # already-wrapped advisory text second, producing a + # `` ...`` shape that + # ``extract_advisories_from_tool_envelope``'s + # ``startswith("\n")`` prefix check + # rejects on replay — the raw envelope XML would + # leak into the rendered tool bubble. Round-trip + # works because ``wrap_tool_result(text, [])`` + # returns ``text`` unchanged when there are no + # advisories. _tname = _tc_names.get(tc_id, "") if isinstance(raw_output, list): - store_text = " ".join( - p.get("text", "") for p in raw_output if p.get("type") == "text" - )[:TOOL_RESULT_STORAGE_CAP] + raw_text = " ".join( + p.get("text", "") + for p in raw_output + if isinstance(p, dict) and p.get("type") == "text" + ) + store_text: str = wrap_tool_result(raw_text, persistent_advisories) else: - store_text = raw_output[:TOOL_RESULT_STORAGE_CAP] + # ``raw_output`` was a string, so the post-wrap + # ``output`` is also a string (envelope or raw). + assert isinstance(output, str) + store_text = output # Mirror the user-channel persistence: tool-channel # reminders (``tool_error`` / ``repeat``) ride the same # ``_reminders`` JSON column so a tab reconnecting via @@ -3157,23 +3204,17 @@ class ChatSession: tool_call_id=tc_id, reminders=tool_reminders_json, ) - # Inject user feedback from approval prompt (e.g. "y, use full path") - if user_feedback: - self.messages.append({"role": "user", "content": user_feedback}) - self._msg_tokens.append(max(1, int(len(user_feedback) / self._chars_per_token))) - - # Drain queued user messages as a separate user turn - # AFTER the assistant→tool batch is complete. The - # role sequence becomes assistant(tool_calls) → tool … - # tool → user, which is valid for strict providers - # (Mistral, Anthropic) and gives the user a real DB - # row that survives reconnect. Pre-fix the queued - # messages rode inside the tool envelope as - # ``UserInterjection`` advisories — visible to the - # model on the same turn, but with no persisted user - # row, so the optimistic queued bubble vanished on - # page reload (and on cross-tab replay). - self._flush_queued_messages() + # Fold ``user_feedback`` (text typed alongside an approval, + # e.g. "y, use full path") and any queued messages that + # raced past Seam 1's drain into a single trailing user + # row. Seam 2 in the queued-message architecture: the + # safety net for items that landed in ``_queued_messages`` + # *after* ``_collect_advisories`` ran for the last result + # but *before* ``_execute_tools`` returned. Common case + # (no feedback, queue empty) no-ops; coexistence case + # produces one user turn with feedback as the prefix + # joined to queued items by ``\n\n``. + self._flush_queued_messages(prefix=user_feedback or "") # Mid-turn compaction: prevent context overflow during long # tool chains. Uses local estimates since _last_usage reflects @@ -4375,7 +4416,7 @@ class ChatSession: popped = self._queued_messages.pop(msg_id, None) return popped is not None - def _flush_queued_messages(self) -> bool: + def _flush_queued_messages(self, prefix: str = "") -> bool: """Drain queued messages into a single combined user turn. Queued items are always text-only (attachments are rejected at @@ -4383,20 +4424,55 @@ class ChatSession: so a single combined turn avoids back-to-back user messages that some models handle poorly. - Returns ``True`` when any items drained, ``False`` otherwise. + ``prefix`` (when non-empty) is the ``user_feedback`` string from + the post-tool-batch path: text the user typed alongside an + approval prompt (e.g. "y, use full path"). Folding it into + the same flush keeps the role sequence to exactly one trailing + user row whether feedback alone, queued items alone, or both + are present. When both are present the rendered content is + ``prefix + "\\n\\n" + "\\n\\n".join(parts)``. + + Returns ``True`` when any user row was appended (prefix or + items), ``False`` when both were empty. """ from turnstone.core.tool_advisory import PRIORITY_IMPORTANT with self._queued_lock: items = list(self._queued_messages.values()) self._queued_messages.clear() - if not items: + if not items and not prefix: return False parts = [f"[IMPORTANT] {msg}" if pri == PRIORITY_IMPORTANT else msg for msg, pri in items] - self._append_user_turn("\n\n".join(parts), ()) + if prefix and parts: + content = prefix + "\n\n" + "\n\n".join(parts) + elif prefix: + content = prefix + else: + content = "\n\n".join(parts) + self._append_user_turn(content, ()) return True + def _drain_queued_messages_to_advisories(self) -> list[ToolAdvisory]: + """Drain ``_queued_messages`` into ``UserInterjection`` advisories. + + Mirrors the swap-and-clear pattern in + :meth:`_flush_queued_messages` but produces wire-bound + ``UserInterjection`` instances rather than appending a user + turn — the splice path used by Seam 1 (queued message rides + inside the last tool-result envelope of a batch). Held under + ``_queued_lock`` to keep the read+clear atomic against a + concurrent ``queue_message`` from the SSE worker. + """ + from turnstone.core.tool_advisory import UserInterjection + + with self._queued_lock: + queued_items = list(self._queued_messages.values()) + self._queued_messages.clear() + return [ + UserInterjection(message=text, priority=priority) for text, priority in queued_items + ] + def _collect_advisories( self, assessment: OutputAssessment | None, @@ -4407,9 +4483,10 @@ class ChatSession: Returns ``(persistent, metacog_reminders)``: - - ``persistent`` — guard findings that ride inside the - tool-result envelope via ``wrap_tool_result``. These are - conversation history and must persist in ``self.messages``. + - ``persistent`` — guard findings and queued user messages + (Seam 1) that ride inside the tool-result envelope via + ``wrap_tool_result``. These are conversation history and + must persist in ``self.messages``. - ``metacog_reminders`` — list of ``{"type", "text", ...optional}`` dicts for ``tool_error`` / ``repeat`` nudges that the caller attaches to the tool message dict's ``_reminders`` side-channel. @@ -4420,19 +4497,21 @@ class ChatSession: surfaced separately on the UI as a themed bubble below the tool result. - Queued user messages are NOT drained here — they're appended - as a separate user turn after the batch completes (see - ``_flush_queued_messages`` invocations in ``send``). The prior - UserInterjection splice into ``wrap_tool_result`` left no DB - row for the queued message, so it vanished on reconnect (the - tool result content carried the text inside a transient - ```` envelope). Appending after the batch - keeps the assistant → tool sequence intact for strict providers - (Mistral) while persisting a real user row for replay parity. + Queued user messages drain into a :class:`UserInterjection` + advisory on the LAST result of a batch — Seam 1 in the + queued-message architecture. The wrapped envelope persists + on the tool row; replay extracts it back out via + :func:`history_decoration.decorate_history_messages` and the + wire layer ships it as ``advisories`` so the UI renders a + proper user bubble after the tool block. Cancel / exception / + no-tool-call paths drain the queue as a real user row instead + (Seams 2 and 3), since there is no envelope to splice into. Both lists are empty when no advisories apply (common case). - Guard advisories attach per-result; metacognitive nudges drain - on the last result in the batch only. + Guard advisories attach per-result; queued messages and + metacognitive nudges drain on the last result in the batch + only — single splice envelope per batch, no duplication + across N tool results. """ from turnstone.core.tool_advisory import GuardAdvisory @@ -4443,13 +4522,24 @@ class ChatSession: if assessment is not None: persistent.append(GuardAdvisory(assessment=assessment, func_name=func_name)) - # Metacognitive tool-channel drain — fires once per batch on the - # last result. Queued by _queue_tool_advisory from the - # tool_error / repeat detection paths just before this loop. - # Lands on the tool message dict's ``_reminders`` side-channel - # (caller's responsibility) so it stays out of persisted content - # and rides the wire only via the transient-copy splice. + # Last-result-in-batch drain seams: queued user messages (Seam 1) + # and tool-channel metacog nudges (tool_error / repeat). Both + # fire once per batch so a parallel fan-out doesn't paint the + # same advisory N times. Queue-drain is delegated to a named + # helper so the swap-and-clear pattern lives next to + # ``_flush_queued_messages``'s identical pattern, and the + # side-effect is documented at the call site rather than buried + # mid-function. if is_last_in_batch: + persistent.extend(self._drain_queued_messages_to_advisories()) + + # Metacognitive tool-channel drain. Queued by + # ``_queue_tool_advisory`` from the tool_error / repeat + # detection paths just before this loop. Lands on the + # tool message dict's ``_reminders`` side-channel + # (caller's responsibility) so it stays out of persisted + # content and rides the wire only via the transient-copy + # splice. drained = self._nudge_queue.drain(TOOL_DRAIN) for nt, text, meta in drained: entry: dict[str, Any] = {"type": nt, "text": text} diff --git a/turnstone/core/session_routes.py b/turnstone/core/session_routes.py index f6379c37..811a9965 100644 --- a/turnstone/core/session_routes.py +++ b/turnstone/core/session_routes.py @@ -2287,7 +2287,14 @@ def make_history_handler(cfg: SessionEndpointConfig) -> Handler: ) indexes = await asyncio.to_thread(load_verdict_indexes, ws_id) - decorate_history_messages(messages, indexes[0], indexes[1]) + # Pure transform but iterates every message and every + # tool_call dict — for a long workstream the pass takes + # tens of milliseconds and would otherwise block the + # event loop's hot path on the request handler. + # ``decorate_history_messages`` is thread-safe (no + # shared mutable state beyond the per-call message + # list) so the off-loop hop is free. + await asyncio.to_thread(decorate_history_messages, messages, indexes[0], indexes[1]) except Exception: # Operationally interesting: a persistent decoration # failure (missing migration, driver mismatch, schema diff --git a/turnstone/core/tool_advisory.py b/turnstone/core/tool_advisory.py index c905a06b..ec55fd1d 100644 --- a/turnstone/core/tool_advisory.py +++ b/turnstone/core/tool_advisory.py @@ -23,6 +23,22 @@ PRIORITY_IMPORTANT: Final = "important" PRIORITY_NOTICE: Final = "notice" +# UserInterjection preamble strings — shared between the producer +# (``UserInterjection.render``) and the replay parser +# (``history_decoration._classify_advisory``). Keeping them in one +# place ensures the parser and producer can never drift on the exact +# wording the parser uses to disambiguate priority. Marker is the +# fixed substring that separates the preamble from the user's body. +_USER_INTERJECTION_NOTICE_PREAMBLE: Final = ( + "The user sent additional context while you were working. " + "Incorporate if relevant, otherwise continue." +) +_USER_INTERJECTION_IMPORTANT_PREAMBLE: Final = ( + "The user sent a message while you were working. You MUST address this before continuing." +) +_USER_INTERJECTION_BODY_MARKER: Final = "\n\nUser message: " + + # -- Protocol ----------------------------------------------------------------- @@ -68,14 +84,16 @@ class GuardAdvisory: class UserInterjection: """Advisory for a message the user sent while the model was executing. - Currently no production producer — pre-PR-#487 the queued-message - drain in ``_collect_advisories`` instantiated this and rode it - through ``wrap_tool_result``. PR #487 dropped the splice and now - drains queued messages as a real user turn after the tool batch. - Class + tests retained pending the back-to-back-user role-ordering - decision (see review finding bug-1) — one viable fix path resumes - the splice for the ``user_feedback`` + queue-drain coexistence - case, in which case this advisory shape stays load-bearing. + Produced by ``ChatSession._collect_advisories`` on the LAST tool + result of a batch when ``_queued_messages`` is non-empty (Seam 1 in + the queued-message architecture). Splicing inside the tool-result + envelope keeps the role sequence ``assistant -> tool`` intact for + strict-template providers and delivers the user's text on the + same turn as the tool batch. The persisted DB row is the wrapped + envelope — replay extracts the advisory back out via + :func:`history_decoration.decorate_history_messages` and renders + it as a proper user bubble in the wire/UI layer, so the queued + message survives reconnect / page reload. """ message: str @@ -87,16 +105,10 @@ class UserInterjection: def render(self) -> str: if self.priority == PRIORITY_IMPORTANT: - preamble = ( - "The user sent a message while you were working. " - "You MUST address this before continuing." - ) + preamble = _USER_INTERJECTION_IMPORTANT_PREAMBLE else: - preamble = ( - "The user sent additional context while you were working. " - "Incorporate if relevant, otherwise continue." - ) - return f"{preamble}\n\nUser message: {self.message}" + preamble = _USER_INTERJECTION_NOTICE_PREAMBLE + return f"{preamble}{_USER_INTERJECTION_BODY_MARKER}{self.message}" @dataclass(frozen=True) @@ -132,9 +144,24 @@ def escape_wrapper_tags(text: str) -> str: cannot fabricate or close one of the wrapper blocks. Use this on any untrusted content that is glued next to a wrapper tag — tool output, user message bodies, and (defense-in-depth) advisory render output. + + Round-trip symmetry with :func:`history_decoration._entity_decode_wrapper_tags` + requires escaping ``&`` first. Otherwise a tool output that contains + the literal string ``<tool_output>`` (e.g. documentation + describing the wrapper format) would round-trip to the bare + ```` tag, fabricating an envelope the wrapper layer + never produced. The short-circuit on ``"<" not in text and "&" not + in text`` covers the common case where neither wrapper tag nor any + pre-existing entity is present — most tool outputs. """ + if "<" not in text and "&" not in text: + return text + # Encode ``&`` first so a pre-existing literal like ``<tool_output>`` + # in the source text becomes ``&lt;tool_output&gt;`` and + # cannot collide with our wrapper-tag escape strings. return ( - text.replace("", "</tool_output>") + text.replace("&", "&") + .replace("", "</tool_output>") .replace("", "<tool_output>") .replace("", "<system-reminder>") .replace("", "</system-reminder>") diff --git a/turnstone/server.py b/turnstone/server.py index ba8f1279..0e103d3b 100644 --- a/turnstone/server.py +++ b/turnstone/server.py @@ -53,10 +53,10 @@ from turnstone.core.auth import ( jwt_version_slot, ) from turnstone.core.history_decoration import ( - TOOL_RESULT_STORAGE_CAP, + decorate_tool_call as _decorate_tool_call, ) from turnstone.core.history_decoration import ( - decorate_tool_call as _decorate_tool_call, + extract_advisories_from_tool_envelope, ) from turnstone.core.history_decoration import ( load_verdict_indexes as _load_verdict_indexes, @@ -587,15 +587,60 @@ def _build_history( result_call_id = msg.get("tool_call_id") if result_call_id: entry["tool_call_id"] = str(result_call_id) - # Tool results are clamped to TOOL_RESULT_STORAGE_CAP - # chars per row at storage time (session.py). Surface - # that on replay so the user knows the visible output is - # a clipped view of what the live session saw, rather - # than the full result. Reference the shared constant - # rather than a literal so the UI pill logic can't - # silently desync if the cap ever changes. - if isinstance(content, str) and len(content) >= TOOL_RESULT_STORAGE_CAP: - entry["truncated"] = True + # Extract advisories from a wrapped ```` envelope + # (Seam 1 queued-message splice). ``session.messages`` never + # carries an ``advisories`` key on its own — only + # ``decorate_history_messages`` mutates dicts to add it for + # the REST ``/history`` path, and the SSE replay surface + # bypasses that decoration entirely. Calling the same + # idempotent extraction helper here pins both surfaces to + # the same wire shape: cleaned content + extracted advisories + # ride as a user bubble after the tool block. No-ops cleanly + # for plain (unwrapped) content. + extracted_advisories: list[dict[str, str]] = [] + if isinstance(content, str): + try: + extracted = extract_advisories_from_tool_envelope(content) + except Exception: + extracted = None + if extracted is not None: + cleaned, extracted_advisories = extracted + content = cleaned + entry["content"] = cleaned + elif isinstance(content, list): + # List-typed tool output (image / structured MCP + # results) carries any Seam 1 splice as an appended + # text part produced by ``wrap_tool_result("", ...)`` + # — the inner cleaned content is empty by construction, + # so the part exists only to carry advisories. Walk + # the parts, extract advisories from any wrap + # envelope, and drop those parts from the projected + # list. Without this, the JS replay would render the + # raw envelope as a text bubble inside the tool block + # AND fail to render the queued message as a user + # bubble (no ``advisories`` array). + new_parts: list[Any] = [] + changed = False + for part in content: + text = ( + part.get("text") + if isinstance(part, dict) and part.get("type") == "text" + else None + ) + if isinstance(text, str) and text.startswith("\n"): + try: + extracted = extract_advisories_from_tool_envelope(text) + except Exception: + extracted = None + if extracted is not None: + _cleaned_text, advisories_from_part = extracted + extracted_advisories.extend(advisories_from_part) + changed = True + continue + new_parts.append(part) + if changed: + content = new_parts + entry["content"] = new_parts if isinstance(content, str): if content.startswith("Denied by user") or content.startswith("Blocked"): entry["denied"] = True @@ -612,6 +657,33 @@ def _build_history( or content.startswith("MCP prompt error") ): entry["is_error"] = True + # Surface advisories on the wire so the JS replay can render + # them as user bubbles after the tool block. Project on a + # known set of keys — narrows the blast radius if a future + # producer stuffs sensitive fields into the dict. Mirrors + # the ``reminders`` filter above for the same reason. The + # extracted-from-envelope list is the production-realistic + # source: ``decorate_history_messages`` populates the + # ``advisories`` key only on the REST ``/history`` path, but + # ``session.messages`` (the SSE-replay source) never has it. + if extracted_advisories: + clean_advisories: list[dict[str, Any]] = [] + for a in extracted_advisories: + if not isinstance(a, dict): + continue + atype = str(a.get("type") or "") + atext = str(a.get("text") or "") + if not atype or not atext: + continue + clean_advisories.append( + { + "type": atype, + "text": atext, + "priority": str(a.get("priority") or "notice"), + } + ) + if clean_advisories: + entry["advisories"] = clean_advisories history.append(entry) # Propagate denial from tool results to their parent assistant entry. diff --git a/turnstone/shared_static/utils.js b/turnstone/shared_static/utils.js index a98acdb6..d86c2d95 100644 --- a/turnstone/shared_static/utils.js +++ b/turnstone/shared_static/utils.js @@ -70,3 +70,22 @@ function cssEscape(s) { } return str.replace(/["\\]/g, "\\$&"); } + +// Replay the ``advisories`` array attached to a tool history message. +// Queued user messages spliced into the last tool-result envelope of +// a batch (Seam 1) persist on the tool DB row as a wrapped +// ```` envelope; the wire layer projects them onto +// ``msg.advisories``. Both interactive's ``replayHistory`` and the +// coord history loop walk the array, filter on ``user_interjection``, +// and route ``adv.text`` through their own renderer (interactive uses +// ``addUserMessage``; coord uses ``appendUserMessageWithAttachments``). +// This shared helper centralises the walk + filter so a future +// advisory shape change lands once. +function replayAdvisoriesAfterTool(advisories, renderUserText) { + if (!Array.isArray(advisories) || !advisories.length) return; + for (var ai = 0; ai < advisories.length; ai++) { + var adv = advisories[ai]; + if (!adv || adv.type !== "user_interjection") continue; + renderUserText(adv.text || ""); + } +} diff --git a/turnstone/ui/static/app.js b/turnstone/ui/static/app.js index 7998eeaa..2ec7f966 100644 --- a/turnstone/ui/static/app.js +++ b/turnstone/ui/static/app.js @@ -1292,8 +1292,8 @@ Pane.prototype.replayHistory = function (messages) { // bug where calling resultTarget.after(node) twice put the // second node BETWEEN resultTarget and the first (the second // .after call was always relative to the same anchor). - // Resulting order with all three present: - // [tool div][output][truncation pill][output-warning] + // Resulting order with all present: + // [tool div][output][output-warning] var insertCursor = resultTarget; var insertChained = function (node) { if (insertCursor) { @@ -1316,17 +1316,6 @@ Pane.prototype.replayHistory = function (messages) { } insertChained(out); } - // Truncation pill — server marks this when the stored row - // hit the 2000-char cap. Live tool_result events carry full - // output so they don't need the indicator. - if (msg.truncated) { - var pill = document.createElement("span"); - pill.className = "tool-output-truncated"; - pill.textContent = "… truncated in storage"; - pill.title = - "The full tool output was sent to the model live; only the first 10000 characters are persisted to the conversation row."; - insertChained(pill); - } } if (isToolError && !lastToolBlock.classList.contains("denied")) { lastToolBlock.classList.add("error"); @@ -1353,6 +1342,20 @@ Pane.prototype.replayHistory = function (messages) { if (Array.isArray(msg.reminders) && msg.reminders.length) { this.addToolReminder(msg.reminders, ""); } + // Queued user messages spliced into the last tool-result envelope + // (Seam 1) replay as proper user bubbles after the tool block. + // ``decorate_history_messages`` extracts the user_interjection + // advisory from the persisted envelope and the wire layer projects + // it onto ``msg.advisories``; rendering through ``addUserMessage`` + // matches the live shape a Seam 2/3 message would produce. The + // walk/filter is shared via ``replayAdvisoriesAfterTool`` in + // ``shared/utils.js`` so coord and interactive can never drift on + // advisory-shape filtering. + var self = this; + replayAdvisoriesAfterTool(msg.advisories, function (text) { + self.addUserMessage(text, null); + lastToolBlock = null; + }); } } // Flush any output_assessments left in the map — these correspond diff --git a/turnstone/ui/static/style.css b/turnstone/ui/static/style.css index a7bac418..23a5f14c 100644 --- a/turnstone/ui/static/style.css +++ b/turnstone/ui/static/style.css @@ -1655,9 +1655,7 @@ body { .ts-approval-tool[data-func-name="memory"] ~ .media-embed, .ts-approval-tool[data-func-name="recall"] ~ .media-embed, .ts-approval-tool[data-func-name="memory"] ~ .output-warning, -.ts-approval-tool[data-func-name="recall"] ~ .output-warning, -.ts-approval-tool[data-func-name="memory"] ~ .tool-output-truncated, -.ts-approval-tool[data-func-name="recall"] ~ .tool-output-truncated { +.ts-approval-tool[data-func-name="recall"] ~ .output-warning { opacity: 0.55; transition: opacity 120ms ease-out; } @@ -1665,8 +1663,8 @@ body { subtree. Without :focus-within on the siblings, a keyboard user tabbing into a link or collapsible toggle inside .tool-output sees the content remain dimmed — a11y regression. Cover the - warning + truncation pills too so they fully reveal alongside - the result they decorate. */ + warning pill too so it fully reveals alongside the result it + decorates. */ .ts-approval-tool[data-func-name="memory"] ~ .tool-output:hover, .ts-approval-tool[data-func-name="recall"] ~ .tool-output:hover, .ts-approval-tool[data-func-name="memory"] ~ .tool-output:focus-within, @@ -1678,30 +1676,9 @@ body { .ts-approval-tool[data-func-name="memory"] ~ .output-warning:hover, .ts-approval-tool[data-func-name="recall"] ~ .output-warning:hover, .ts-approval-tool[data-func-name="memory"] ~ .output-warning:focus-within, -.ts-approval-tool[data-func-name="recall"] ~ .output-warning:focus-within, -.ts-approval-tool[data-func-name="memory"] ~ .tool-output-truncated:hover, -.ts-approval-tool[data-func-name="recall"] ~ .tool-output-truncated:hover, -.ts-approval-tool[data-func-name="memory"] ~ .tool-output-truncated:focus-within, -.ts-approval-tool[data-func-name="recall"] ~ .tool-output-truncated:focus-within { +.ts-approval-tool[data-func-name="recall"] ~ .output-warning:focus-within { opacity: 1; } -/* Truncation indicator — the persisted tool result is clamped at - 2000 chars per row in storage; surface that on replay so users - know they're seeing a clipped view rather than the full output - the live session saw. Aligns with .output-warning's left gutter - (margin-left: 16px) and uses transparent background + dim border - so it reads as quiet metadata rather than a foreign element. */ -.tool-output-truncated { - display: inline-block; - margin-top: 4px; - margin-left: 16px; - padding: 1px 6px; - font-size: 10px; - color: var(--fg-dim); - background: transparent; - border-radius: 3px; - border: 1px solid var(--fg-dim); -} /* .ts-approval (chat.css) stacks its children with flex gap, so a border-top on the body would float above a strip of container background instead of sitting flush against the previous tool row.