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()