From 828a58c4b46b8a67e88f0c033f9b3fcd80eb7416 Mon Sep 17 00:00:00 2001 From: abccodes Date: Thu, 16 Jul 2026 17:14:12 -0700 Subject: [PATCH] [aidan] fix/analytics-telemetry: emit final thinking snapshots once --- backend/apps/agents/core/ws_manager.py | 6 ++- .../streaming/handle_result_message.py | 1 + .../apps/agents/manager/streaming/thinking.py | 12 ++++- .../apps/service/analytics/agent_bridge.py | 12 ++++- backend/tests/test_analytics_agent_bridge.py | 47 +++++++++++++++++++ backend/tests/test_result_message.py | 25 ++++++++++ 6 files changed, 98 insertions(+), 5 deletions(-) create mode 100644 backend/tests/test_analytics_agent_bridge.py diff --git a/backend/apps/agents/core/ws_manager.py b/backend/apps/agents/core/ws_manager.py index 5b8022ff..443e339a 100644 --- a/backend/apps/agents/core/ws_manager.py +++ b/backend/apps/agents/core/ws_manager.py @@ -113,7 +113,11 @@ class ConnectionManager: if event == "agent:message": try: from backend.apps.service.analytics.agent_bridge import bridge_agent_message, BroadcastMessage - bridge_agent_message(session_id, BroadcastMessage.model_validate(data.get("message") or {})) + bridge_agent_message( + session_id, + BroadcastMessage.model_validate(data.get("message") or {}), + final=bool(data.get("analytics_final")), + ) except Exception: logger.debug("agent:message analytics bridge failed", exc_info=True) diff --git a/backend/apps/agents/manager/streaming/handle_result_message.py b/backend/apps/agents/manager/streaming/handle_result_message.py index b9f72259..4587fac6 100644 --- a/backend/apps/agents/manager/streaming/handle_result_message.py +++ b/backend/apps/agents/manager/streaming/handle_result_message.py @@ -74,6 +74,7 @@ async def handle_result_message( await thinking_mod.emit_consolidated_thinking( thinking, turn, session, session_id, sessions, force_provider_unavailable=route_strips_reasoning, + analytics_final=True, ) except Exception: pass diff --git a/backend/apps/agents/manager/streaming/thinking.py b/backend/apps/agents/manager/streaming/thinking.py index aa788e34..4de55ce9 100644 --- a/backend/apps/agents/manager/streaming/thinking.py +++ b/backend/apps/agents/manager/streaming/thinking.py @@ -18,7 +18,15 @@ logger = logging.getLogger(__name__) @typechecked -async def emit_consolidated_thinking(thinking: ThinkingState, turn: TurnState, session: AgentSession, session_id: str, sessions: Dict[str, AgentSession], force_provider_unavailable: bool = False) -> None: +async def emit_consolidated_thinking( + thinking: ThinkingState, + turn: TurnState, + session: AgentSession, + session_id: str, + sessions: Dict[str, AgentSession], + force_provider_unavailable: bool = False, + analytics_final: bool = False, +) -> None: """Build the running aggregate Message and broadcast it. Safe to call multiple times, uses a stable per-turn id so the frontend dedupes by id and updates the bubble in @@ -156,6 +164,7 @@ async def emit_consolidated_thinking(thinking: ThinkingState, turn: TurnState, s await ws_manager.send_to_session(session_id, "agent:message", { "session_id": session_id, "message": consolidated.model_dump(mode="json"), + "analytics_final": analytics_final, }) except Exception: logger.exception("Failed to emit consolidated thinking message") @@ -174,4 +183,3 @@ async def ticker_loop(thinking: ThinkingState, turn: TurnState, session: AgentSe await emit_consolidated_thinking(thinking, turn, session, session_id, sessions) except asyncio.CancelledError: pass - diff --git a/backend/apps/service/analytics/agent_bridge.py b/backend/apps/service/analytics/agent_bridge.py index 60a31d20..671f8fc9 100644 --- a/backend/apps/service/analytics/agent_bridge.py +++ b/backend/apps/service/analytics/agent_bridge.py @@ -1,7 +1,8 @@ """Bridge a broadcast `agent:message` into the typed `events.agent.message`. Called from ws_manager.send_to_session, the single chokepoint every agent message -flows through. Best-effort: never raises into the broadcast path. +flows through. Live thinking snapshots are skipped until the finalized message. +Best-effort: never raises into the broadcast path. """ from __future__ import annotations @@ -55,10 +56,17 @@ def p_branch_version(session: AgentSession, message: BroadcastMessage) -> int: @typechecked -def bridge_agent_message(session_id: str, message: BroadcastMessage) -> None: +def bridge_agent_message( + session_id: str, + message: BroadcastMessage, + *, + final: bool = False, +) -> None: # seq is the message's stable index in the persisted history (survives close -> reopen -> restart); transient messages with no anchor are skipped. if not message.id or not message.role: return + if message.role == "thinking" and not final: + return try: from backend.apps.agents.agent_manager import agent_manager sess = agent_manager.sessions.get(session_id) diff --git a/backend/tests/test_analytics_agent_bridge.py b/backend/tests/test_analytics_agent_bridge.py new file mode 100644 index 00000000..ad910fce --- /dev/null +++ b/backend/tests/test_analytics_agent_bridge.py @@ -0,0 +1,47 @@ +import sys +from types import ModuleType, SimpleNamespace +from unittest.mock import Mock + +from backend.apps.agents.core.models import AgentSession, Message +from backend.apps.service.analytics import agent_bridge + + +def p_session_module(session_id: str, session: AgentSession) -> ModuleType: + module = ModuleType("backend.apps.agents.agent_manager") + module.agent_manager = SimpleNamespace(sessions={session_id: session}) + return module + + +def test_thinking_telemetry_only_tracks_final_snapshot(monkeypatch): + session_id = "session-1" + thought = Message(id="thought-1", role="thinking", content="final") + session = AgentSession(name="test", messages=[thought]) + track = Mock() + monkeypatch.setitem(sys.modules, "backend.apps.agents.agent_manager", p_session_module(session_id, session)) + monkeypatch.setattr(agent_bridge, "track_agent_message", track) + broadcast = agent_bridge.BroadcastMessage.model_validate(thought.model_dump(mode="json")) + + agent_bridge.bridge_agent_message(session_id, broadcast) + track.assert_not_called() + + agent_bridge.bridge_agent_message(session_id, broadcast, final=True) + track.assert_called_once() + assert track.call_args.kwargs["id"] == "thought-1" + assert track.call_args.kwargs["content"] == "final" + + +def test_non_thinking_message_tracks_without_final_marker(monkeypatch): + session_id = "session-1" + message = Message(id="user-1", role="user", content="hello") + session = AgentSession(name="test", messages=[message]) + track = Mock() + monkeypatch.setitem(sys.modules, "backend.apps.agents.agent_manager", p_session_module(session_id, session)) + monkeypatch.setattr(agent_bridge, "track_agent_message", track) + + agent_bridge.bridge_agent_message( + session_id, + agent_bridge.BroadcastMessage.model_validate(message.model_dump(mode="json")), + ) + + track.assert_called_once() + assert track.call_args.kwargs["role"] == "user" diff --git a/backend/tests/test_result_message.py b/backend/tests/test_result_message.py index 916ab2a4..04cd73e0 100644 --- a/backend/tests/test_result_message.py +++ b/backend/tests/test_result_message.py @@ -70,3 +70,28 @@ async def test_resets_per_turn_state_at_completion(): assert turn.tool_count == 0 assert thinking.total_ms == 0 assert thinking.block_starts == {} + + +@pytest.mark.asyncio +async def test_final_thinking_emit_is_marked_for_analytics(): + session, turn, thinking = p_fixt() + thinking.text_parts.append("final thought") + emit = AsyncMock() + with ( + patch.object(result_message.thinking_mod, "emit_consolidated_thinking", new=emit), + patch.object(result_message.ws_manager, "send_to_session", new=AsyncMock()), + ): + await result_message.handle_result_message( + p_result(), + session, + session.id, + turn, + thinking, + {}, + "sonnet", + "anthropic", + load_settings(), + ) + + emit.assert_awaited_once() + assert emit.call_args.kwargs["analytics_final"] is True