From a00a75b547cc48fd0478c5e759d1257548d7c76b Mon Sep 17 00:00:00 2001 From: Phil Merrell Date: Sun, 6 Sep 2026 15:48:31 -0600 Subject: [PATCH] fix(sessions): store displayText when the user's turn lands, not at turn end MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `displayText` is the clean copy of what the user typed, and what the UI renders whenever the prompt sent to the model was augmented — RAG context, attachment guidance, an MCP App context block, or the `` addressed to the model after an interrupted turn. It was written by a single call site in the stream coordinator's success path, as the last thing a turn does. Nothing on the Stop, disconnect, or error paths wrote it, so any turn that didn't reach that line left the raw augmented prompt as the only thing the UI could render. Two effects, both reproduced on dev: * transiently, for every augmented turn — reload while it is still streaming and the note is in the user's own chat bubble; * permanently, for any turn that never completes. Stop it and the `D#` record is never written, so the note stays in the bubble on every load. The permanent case lands where it hurts most: a turn only carries an interruption note because the PREVIOUS turn was interrupted, so the note is disproportionately likely to be on a turn that gets interrupted too. Not model-specific — reproduced on GPT-5.6 Luna and Claude Haiku 4.5. `DisplayTextHook` writes the record on `MessageAddedEvent` instead — the moment the user's message enters history, before the model call — so every later exit path already has it. Writing at request start instead would risk a record keyed to an index the failed turn never filled, which a later message would inherit; anchoring to the append keeps index and record together. Two subtleties the hook has to respect: * Not every role-`user` message is the user. Tool results carry that role, and Strands prepends a SYNTHETIC tool-result message ahead of the prompt when history ends on a dangling `toolUse` — exactly what an interrupted tool turn leaves behind. Content carrying `toolResult`/`toolUse` is skipped so the arm is spent on the user's own message. * The agent instance is cached across turns (#741/#751), so the hook is armed unconditionally at the head of every turn, including to None. The coordinator's write stays as a backstop for wrappers that carry no hook (voice, tests) and for a failed hook write, gated on `wrote_this_turn` so the normal path still makes exactly one put. Co-Authored-By: Claude Opus 5 --- backend/src/agents/main_agent/base_agent.py | 13 + .../main_agent/session/hooks/__init__.py | 2 + .../main_agent/session/hooks/display_text.py | 147 ++++++++++++ .../streaming/stream_coordinator.py | 67 +++++- .../session/test_display_text_hook.py | 226 ++++++++++++++++++ .../streaming/test_display_text_write.py | 165 +++++++++++++ 6 files changed, 618 insertions(+), 2 deletions(-) create mode 100644 backend/src/agents/main_agent/session/hooks/display_text.py create mode 100644 backend/tests/agents/main_agent/session/test_display_text_hook.py create mode 100644 backend/tests/agents/main_agent/streaming/test_display_text_write.py diff --git a/backend/src/agents/main_agent/base_agent.py b/backend/src/agents/main_agent/base_agent.py index e87d4ba1..15a29fe0 100644 --- a/backend/src/agents/main_agent/base_agent.py +++ b/backend/src/agents/main_agent/base_agent.py @@ -13,6 +13,7 @@ from agents.main_agent.core import ModelConfig, SystemPromptBuilder, AgentFactory from agents.main_agent.session import SessionFactory from agents.main_agent.session.hooks import ( + DisplayTextHook, SteeringHook, StopHook, OAuthConsentHook, @@ -276,6 +277,9 @@ def _create_hooks(self) -> List: - StopHook: Always enabled, cancels tool execution on user stop - SteeringHook: Injects a follow-up queued mid-turn at the next tool boundary + - DisplayTextHook: Stores the user's original message for UI display + as soon as their turn is appended, so an augmented prompt is never + what the UI renders for an interrupted turn - OAuthConsentHook: Pauses the agent (Strands interrupt) when an OAuth-gated MCP tool is about to run without a cached token - Approval hooks: Gate dangerous operations for user confirmation @@ -300,6 +304,15 @@ def _create_hooks(self) -> List: self.steering_hook = SteeringHook(self.session_manager) hooks.append(self.steering_hook) + # Persist the user's own words (`displayText`) the moment their turn + # enters history, so an augmented prompt — RAG context, attachment + # guidance, an `` — never becomes what the UI shows + # for a turn that doesn't finish. Held on the wrapper so the stream + # coordinator can arm it per turn and skip its own end-of-turn write + # once this has done it. + self.display_text_hook = DisplayTextHook() + hooks.append(self.display_text_hook) + # OAuth consent gate for external MCP tools. Registered unconditionally; # the hook is a no-op for tools that don't have a registered provider. hooks.append(self._build_oauth_consent_hook()) diff --git a/backend/src/agents/main_agent/session/hooks/__init__.py b/backend/src/agents/main_agent/session/hooks/__init__.py index 840944ff..0251a49a 100644 --- a/backend/src/agents/main_agent/session/hooks/__init__.py +++ b/backend/src/agents/main_agent/session/hooks/__init__.py @@ -1,6 +1,7 @@ """Hooks for Main Agent""" from agents.main_agent.session.hooks.context_attribution import ContextAttributionHook +from agents.main_agent.session.hooks.display_text import DisplayTextHook from agents.main_agent.session.hooks.oauth_consent import OAuthConsentHook from agents.main_agent.session.hooks.prefix_fingerprint import PrefixFingerprintHook from agents.main_agent.session.hooks.steering import SteeringHook @@ -9,6 +10,7 @@ __all__ = [ "ContextAttributionHook", + "DisplayTextHook", "OAuthConsentHook", "PrefixFingerprintHook", "SteeringHook", diff --git a/backend/src/agents/main_agent/session/hooks/display_text.py b/backend/src/agents/main_agent/session/hooks/display_text.py new file mode 100644 index 00000000..268aeb90 --- /dev/null +++ b/backend/src/agents/main_agent/session/hooks/display_text.py @@ -0,0 +1,147 @@ +"""Persist the user's own words as soon as their turn enters history. + +The prompt that reaches the model is often not the prompt the user typed. RAG +prepends retrieved context, attachments add guidance, an embedded MCP App +pushes a context block, and an interrupted previous turn prepends an +```` addressed to the model. All of that is deliberately +kept in persisted history — it is an honest record of what the model actually +read — and the UI is supposed to show the clean original instead, via the +``displayText`` (``D#``) record this hook writes. + +**Why a hook, and why this event.** That write used to live at the very end of +``stream_coordinator.stream_response``, in the success path. Nothing on the +Stop, disconnect, or error paths wrote it, so any turn that did not reach that +final line left the raw augmented prompt as the only thing the UI could +render — and a turn is at its most likely to be interrupted precisely when it +is carrying an interruption note, because the note only exists because the +*previous* turn was interrupted. The visible result was the model-directed +note sitting in the user's own chat bubble, permanently. It also showed +transiently on any reload mid-turn, for every augmentation. + +``MessageAddedEvent`` fires from ``Agent._append_messages``, the moment the +user's turn is really in the conversation — before the model call, so every +later exit path (completion, Stop, cancellation, error, container death) +already has the record written. That is the whole point: the write no longer +depends on how the turn ends. + +**Why not simply write at request start.** If the turn died before the user +message was appended, a record keyed to that index would be inherited by +whatever message later takes the index — showing one turn's clean text on a +different turn's bubble. Anchoring to the actual append makes the index and +the record land together. + +**One-shot per turn, and why role alone is not enough.** Tool-result messages +are also role ``user`` under Bedrock Converse, and mid-turn steering appends +into them. The hook is armed once per turn and disarms on the first user-role +message it writes, which is the user's prompt — tool results only exist after +the first model call. + +Armed unconditionally at the head of every turn, *including to ``None``*, for +the same reason ``turn_lease`` is stamped unconditionally: the agent instance +is cached and outlives the turn (#741/#751), so an arm left behind by a +previous turn would fire against the wrong one. +""" + +from __future__ import annotations + +import logging +from typing import Any, Optional + +from strands.hooks import HookProvider, HookRegistry, MessageAddedEvent + +logger = logging.getLogger(__name__) + + +class DisplayTextHook(HookProvider): + """Write the turn's ``displayText`` when the user message is appended. + + Best-effort in every direction: ``displayText`` is a UI nicety, and a + failure here must never break a turn. When it does fail, the stream + coordinator's end-of-turn write is still there as a backstop for turns + that complete. + """ + + def __init__(self) -> None: + self._armed: Optional[dict] = None + self._written = False + + def register_hooks(self, registry: HookRegistry, **kwargs: Any) -> None: + registry.add_callback(MessageAddedEvent, self.write_display_text) + + @property + def wrote_this_turn(self) -> bool: + """Whether this turn's record is already stored. + + Read by the stream coordinator so its end-of-turn backstop doesn't + repeat a write this hook already made. + """ + return self._written + + def arm( + self, + *, + session_id: str, + user_id: str, + message_index: int, + display_text: Optional[str], + ) -> None: + """Prime the hook for one turn, or clear it when there's nothing to write. + + ``display_text`` is the user's original message, passed only when the + prompt was modified before reaching the model. A turn that sends the + user's text verbatim (and a resume / continuation, which sends no new + user turn at all) passes ``None`` and disarms. + """ + self._written = False + if not display_text: + self._armed = None + return + self._armed = { + "session_id": session_id, + "user_id": user_id, + "message_id": message_index, + "display_text": display_text, + } + + async def write_display_text(self, event: MessageAddedEvent) -> None: + """Store the clean text once the user's message is in history.""" + armed = self._armed + if armed is None: + return + + message = getattr(event, "message", None) or {} + if message.get("role") != "user": + return + # Not every role-`user` message is the user speaking. Tool results + # carry that role under Bedrock Converse, and Strands prepends a + # SYNTHETIC tool-result message ahead of the prompt when history ends + # on a dangling `toolUse` (`Agent._run_loop`, "appending a toolResult + # message to have valid conversation") — which is precisely the shape + # an interrupted tool turn leaves behind, i.e. the case this hook + # exists for. Consuming the arm there would stamp the clean text onto + # the repair message instead of the user's own. + if any( + isinstance(block, dict) and ("toolResult" in block or "toolUse" in block) + for block in message.get("content") or [] + ): + return + + # One-shot: consume before the await so a tool-result message later in + # the same turn can never re-enter this. + self._armed = None + + try: + from apis.shared.sessions.metadata import store_user_display_text + + await store_user_display_text(**armed) + self._written = True + logger.info( + "💾 Stored displayText for user message %s at append time", + armed["message_id"], + ) + except Exception: # noqa: BLE001 - a UI nicety must never break a turn + logger.error( + "Failed to store displayText for user message %s", + armed["message_id"], + exc_info=True, + ) diff --git a/backend/src/agents/main_agent/streaming/stream_coordinator.py b/backend/src/agents/main_agent/streaming/stream_coordinator.py index ac644d2e..6b954c04 100644 --- a/backend/src/agents/main_agent/streaming/stream_coordinator.py +++ b/backend/src/agents/main_agent/streaming/stream_coordinator.py @@ -344,6 +344,22 @@ async def stream_response( initial_message_count = self._get_initial_message_count(session_manager) logger.info(f"📊 Initial message count before streaming: {initial_message_count}") + # Arm the displayText write for this turn. The hook stores the user's + # original message on `MessageAddedEvent` — i.e. before the model + # call — so a turn that is stopped, dropped, or errors still has the + # clean text to render instead of the augmented prompt the model was + # sent. Armed UNCONDITIONALLY, including to None: the agent instance + # is cached across turns (#741/#751), so an arm left by a previous + # turn would otherwise fire against this one. See the end-of-turn + # backstop below for wrappers that carry no hook. + self._arm_display_text( + main_agent_wrapper, + session_id=session_id, + user_id=user_id, + message_index=initial_message_count, + display_text=original_message, + ) + # MCP Apps PR #5: subscribe this conversation stream to the # app-initiated tool-event broker so a `tools/call` proxied from an # embedded MCP App surfaces as a tool_use/tool_result card in the @@ -1186,8 +1202,13 @@ async def stream_response( logger.info(f"✅ Message metadata stored for {len(message_ids_to_store)} assistant messages (sequential)") - # Store displayText for user message if original_message differs from augmented - if original_message: + # displayText backstop. `DisplayTextHook` normally wrote this at + # append time, which is the write that matters — it is the only + # one an interrupted turn ever reaches. This runs only when that + # didn't happen: a wrapper with no hook (voice, tests), or a + # failed write. Skipped otherwise, so the normal path still makes + # exactly one put. + if original_message and not self._display_text_written(main_agent_wrapper): user_message_index = initial_message_count # User message is first in this turn try: from apis.shared.sessions.metadata import store_user_display_text @@ -2136,6 +2157,48 @@ def _emit_tool_input_partial( logger.warning("Failed to emit ui_tool_input_partial event: %s", e) return [] + def _arm_display_text( + self, + main_agent_wrapper: Any, + *, + session_id: str, + user_id: str, + message_index: int, + display_text: Optional[str], + ) -> None: + """Prime this turn's ``displayText`` write on the agent's hook. + + No-op for a wrapper that carries no hook (voice, tests) — those fall + through to the coordinator's end-of-turn backstop, which is exactly + the behaviour they had before the hook existed. + """ + hook = getattr(main_agent_wrapper, "display_text_hook", None) + if hook is None: + return + try: + hook.arm( + session_id=session_id, + user_id=user_id, + message_index=message_index, + display_text=display_text, + ) + except Exception: # noqa: BLE001 - never break a turn on a UI nicety + logger.warning("Could not arm displayText hook", exc_info=True) + + def _display_text_written(self, main_agent_wrapper: Any) -> bool: + """Whether the hook already stored this turn's ``displayText``. + + False whenever we can't tell, so the backstop runs — a duplicate put + of an identical record is harmless, a missing one is the bug. + """ + hook = getattr(main_agent_wrapper, "display_text_hook", None) + if hook is None: + return False + try: + return bool(hook.wrote_this_turn) + except Exception: # noqa: BLE001 + return False + def _drain_steering_events( self, main_agent_wrapper: Any, session_id: str ) -> List[str]: diff --git a/backend/tests/agents/main_agent/session/test_display_text_hook.py b/backend/tests/agents/main_agent/session/test_display_text_hook.py new file mode 100644 index 00000000..8a5d6dc3 --- /dev/null +++ b/backend/tests/agents/main_agent/session/test_display_text_hook.py @@ -0,0 +1,226 @@ +"""Tests for DisplayTextHook — persist the user's own words at append time. + +The bug this exists to close: `displayText` used to be written only by the +stream coordinator's success path, so any turn that was stopped, dropped, or +errored left the *augmented* prompt as the only thing the UI could render. +That put a model-directed `` in the user's own chat bubble, +permanently — and it landed most often on exactly the turns carrying such a +note, since the note only exists because the previous turn was interrupted. + +So the properties under test, in order of how expensive they are to get wrong: + +1. **The write happens on append, before the model call.** That is what makes + it independent of how the turn ends. +2. **One-shot per turn.** Tool-result messages are role `user` too; a second + write would relabel the wrong message index. +3. **No stale arm.** The agent instance is cached across turns (#741/#751), so + an un-armed turn must never inherit the previous turn's text. +4. Fail-soft: a storage failure never propagates into the turn. +""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from strands.hooks import MessageAddedEvent + +from agents.main_agent.session.hooks.display_text import DisplayTextHook + + +def _user_message(text: str = "hi"): + return {"role": "user", "content": [{"text": text}]} + + +def _assistant_message(): + return {"role": "assistant", "content": [{"text": "answer"}]} + + +def _tool_result_message(): + return { + "role": "user", + "content": [{"toolResult": {"toolUseId": "t1", "content": [{"text": "ok"}]}}], + } + + +def _event(message): + return MessageAddedEvent(agent=MagicMock(), message=message) + + +@pytest.fixture +def armed_hook(): + hook = DisplayTextHook() + hook.arm( + session_id="s1", + user_id="u1", + message_index=4, + display_text="what the user actually typed", + ) + return hook + + +@pytest.fixture +def store(): + with patch( + "apis.shared.sessions.metadata.store_user_display_text", new_callable=AsyncMock + ) as mock: + yield mock + + +class TestWritesOnAppend: + @pytest.mark.asyncio + async def test_stores_the_original_text_when_the_user_message_lands( + self, armed_hook, store + ): + await armed_hook.write_display_text(_event(_user_message())) + + store.assert_awaited_once_with( + session_id="s1", + user_id="u1", + message_id=4, + display_text="what the user actually typed", + ) + assert armed_hook.wrote_this_turn is True + + @pytest.mark.asyncio + async def test_ignores_assistant_messages(self, armed_hook, store): + await armed_hook.write_display_text(_event(_assistant_message())) + + store.assert_not_awaited() + # Still armed — the user turn hasn't landed yet. + assert armed_hook.wrote_this_turn is False + + @pytest.mark.asyncio + async def test_writes_once_even_though_tool_results_are_role_user( + self, armed_hook, store + ): + """Under Bedrock Converse a tool-result message is role `user` too. + + A second write would stamp this turn's clean text onto a message index + that isn't the user's prompt. + """ + await armed_hook.write_display_text(_event(_user_message())) + await armed_hook.write_display_text(_event(_tool_result_message())) + await armed_hook.write_display_text(_event(_tool_result_message())) + + assert store.await_count == 1 + + @pytest.mark.asyncio + async def test_a_synthetic_tool_result_repair_does_not_consume_the_arm( + self, armed_hook, store + ): + """Strands prepends a role-`user` tool-result message ahead of the + prompt when history ends on a dangling `toolUse` (agent.py, "appending + a toolResult message to have valid conversation"). + + That is exactly the shape an interrupted tool turn leaves behind — the + case this hook exists for — so consuming the arm there would stamp the + clean text onto the repair message and leave the user's own prompt + showing the augmented text. + """ + await armed_hook.write_display_text(_event(_tool_result_message())) + store.assert_not_awaited() + + await armed_hook.write_display_text(_event(_user_message())) + + store.assert_awaited_once_with( + session_id="s1", + user_id="u1", + message_id=4, + display_text="what the user actually typed", + ) + + @pytest.mark.asyncio + async def test_a_tool_use_message_does_not_consume_the_arm(self, armed_hook, store): + await armed_hook.write_display_text( + _event( + { + "role": "user", + "content": [{"toolUse": {"toolUseId": "t1", "name": "x", "input": {}}}], + } + ) + ) + + store.assert_not_awaited() + + +class TestArming: + @pytest.mark.asyncio + async def test_unarmed_hook_writes_nothing(self, store): + hook = DisplayTextHook() + + await hook.write_display_text(_event(_user_message())) + + store.assert_not_awaited() + assert hook.wrote_this_turn is False + + @pytest.mark.asyncio + async def test_arming_with_no_text_disarms(self, store): + """A turn that sends the user's text verbatim — and every resume / + continuation, which sends no new user turn at all — passes None.""" + hook = DisplayTextHook() + hook.arm(session_id="s1", user_id="u1", message_index=4, display_text="orig") + hook.arm(session_id="s1", user_id="u1", message_index=6, display_text=None) + + await hook.write_display_text(_event(_user_message())) + + store.assert_not_awaited() + + @pytest.mark.asyncio + async def test_re_arming_replaces_the_previous_turns_state(self, store): + """The agent instance is cached across turns (#741/#751). + + A second turn must write ITS text at ITS index, never the first's. + """ + hook = DisplayTextHook() + hook.arm(session_id="s1", user_id="u1", message_index=4, display_text="first") + hook.arm(session_id="s1", user_id="u1", message_index=6, display_text="second") + + await hook.write_display_text(_event(_user_message())) + + store.assert_awaited_once_with( + session_id="s1", user_id="u1", message_id=6, display_text="second" + ) + + @pytest.mark.asyncio + async def test_re_arming_clears_the_written_flag(self, armed_hook, store): + """`wrote_this_turn` gates the coordinator's backstop, so a stale True + from last turn would suppress a write this turn genuinely needs.""" + await armed_hook.write_display_text(_event(_user_message())) + assert armed_hook.wrote_this_turn is True + + armed_hook.arm( + session_id="s1", user_id="u1", message_index=6, display_text="next turn" + ) + + assert armed_hook.wrote_this_turn is False + + +class TestFailSoft: + @pytest.mark.asyncio + async def test_a_storage_failure_never_reaches_the_turn(self, armed_hook): + with patch( + "apis.shared.sessions.metadata.store_user_display_text", + new_callable=AsyncMock, + side_effect=RuntimeError("dynamo down"), + ): + await armed_hook.write_display_text(_event(_user_message())) + + # Not marked written, so the coordinator's end-of-turn backstop still + # runs for a turn that completes. + assert armed_hook.wrote_this_turn is False + + @pytest.mark.asyncio + async def test_a_message_without_a_role_is_ignored(self, armed_hook, store): + await armed_hook.write_display_text(_event({"content": []})) + + store.assert_not_awaited() + + +class TestRegistration: + def test_registers_for_message_added(self): + hook = DisplayTextHook() + registry = MagicMock() + + hook.register_hooks(registry) + + registered = {call.args[0] for call in registry.add_callback.call_args_list} + assert MessageAddedEvent in registered diff --git a/backend/tests/agents/main_agent/streaming/test_display_text_write.py b/backend/tests/agents/main_agent/streaming/test_display_text_write.py new file mode 100644 index 00000000..0177d547 --- /dev/null +++ b/backend/tests/agents/main_agent/streaming/test_display_text_write.py @@ -0,0 +1,165 @@ +"""The coordinator's half of the displayText fix: arm early, back stop late. + +`displayText` is what the UI renders in place of a prompt the model saw but +the user never typed — RAG context, attachment guidance, an +``. It used to be written only here, at the end of a +successful turn, so a stopped or dropped turn left the augmented prompt as the +only renderable text. `DisplayTextHook` now writes it at append time instead. + +What stays the coordinator's job, and is tested here: + +1. **Arm the hook every turn, unconditionally — including to None.** The agent + instance is cached across turns (#741/#751); an arm left by a previous turn + would stamp its text onto this turn's message index. Same discipline as the + `turn_lease` stamp next to it. +2. **Back stop only what the hook didn't do.** A wrapper with no hook (voice, + tests) must keep the old end-of-turn write, and a hook whose write failed + must not silence it — but the normal path must not put twice. + +Driven through the real `stream_response`, like the steering-events suite. +""" + +from typing import Any, AsyncIterator, Dict, List, Optional +from unittest.mock import AsyncMock, patch + +import pytest + +from agents.main_agent.streaming.stream_coordinator import StreamCoordinator + + +class _FakeAgent: + def __init__(self) -> None: + self.messages = [{"role": "user", "content": [{"text": "hi"}]}] + + def stream_async(self, prompt: Any) -> AsyncIterator[Dict[str, Any]]: + async def _gen() -> AsyncIterator[Dict[str, Any]]: + return + yield # pragma: no cover - empty stream + + return _gen() + + +class _SessionManager: + def __init__(self) -> None: + self.cancelled = False + self.turn_lease = None + + async def update_after_turn(self, input_tokens, current_messages=None): + return None + + +class _RecordingHook: + """Stands in for DisplayTextHook — records arming, reports its result.""" + + def __init__(self, wrote: bool = False) -> None: + self.arms: List[dict] = [] + self._wrote = wrote + + def arm(self, **kwargs) -> None: + self.arms.append(kwargs) + + @property + def wrote_this_turn(self) -> bool: + return self._wrote + + +class _Wrapper: + def __init__(self, hook=None) -> None: + if hook is not None: + self.display_text_hook = hook + + +async def _run(wrapper=None, original_message: Optional[str] = None) -> None: + coordinator = StreamCoordinator() + async for _ in coordinator.stream_response( + agent=_FakeAgent(), + prompt="augmented prompt the model saw", + session_manager=_SessionManager(), + session_id="sess-1", + user_id="user-1", + main_agent_wrapper=wrapper, + original_message=original_message, + ): + pass + + +@pytest.fixture +def store(): + with patch( + "apis.shared.sessions.metadata.store_user_display_text", new_callable=AsyncMock + ) as mock: + yield mock + + +class TestArming: + @pytest.mark.asyncio + async def test_arms_the_hook_with_this_turns_text_and_index(self, store): + hook = _RecordingHook() + + await _run(_Wrapper(hook), original_message="what the user typed") + + assert hook.arms == [ + { + "session_id": "sess-1", + "user_id": "user-1", + "message_index": 0, + "display_text": "what the user typed", + } + ] + + @pytest.mark.asyncio + async def test_arms_to_none_when_the_prompt_was_not_modified(self, store): + """Unconditional arming is the point: a cached agent whose previous + turn was augmented must not write that turn's text against this one.""" + hook = _RecordingHook() + + await _run(_Wrapper(hook), original_message=None) + + assert hook.arms == [ + { + "session_id": "sess-1", + "user_id": "user-1", + "message_index": 0, + "display_text": None, + } + ] + + @pytest.mark.asyncio + async def test_a_wrapper_without_the_hook_is_not_an_error(self, store): + await _run(_Wrapper(), original_message="what the user typed") + await _run(None, original_message="what the user typed") + + +class TestBackstop: + @pytest.mark.asyncio + async def test_skipped_once_the_hook_has_written(self, store): + """The hook's write is the one that matters; repeating it at turn end + would put the same record twice on every augmented turn.""" + await _run(_Wrapper(_RecordingHook(wrote=True)), original_message="typed") + + store.assert_not_awaited() + + @pytest.mark.asyncio + async def test_runs_when_the_hook_write_failed(self, store): + """`wrote_this_turn` stays False on a storage failure, so a turn that + completes still gets its record.""" + await _run(_Wrapper(_RecordingHook(wrote=False)), original_message="typed") + + store.assert_awaited_once_with( + session_id="sess-1", user_id="user-1", message_id=0, display_text="typed" + ) + + @pytest.mark.asyncio + async def test_runs_for_a_wrapper_that_carries_no_hook(self, store): + """Voice and tests keep exactly the behaviour they had before.""" + await _run(_Wrapper(), original_message="typed") + + store.assert_awaited_once_with( + session_id="sess-1", user_id="user-1", message_id=0, display_text="typed" + ) + + @pytest.mark.asyncio + async def test_nothing_written_when_the_prompt_was_not_modified(self, store): + await _run(_Wrapper(), original_message=None) + + store.assert_not_awaited()