[aidan] fix/analytics-telemetry: emit final thinking snapshots once

This commit is contained in:
abccodes
2026-07-16 17:22:03 -07:00
parent 24d197eefd
commit 828a58c4b4
6 changed files with 98 additions and 5 deletions
+5 -1
View File
@@ -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)
@@ -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
@@ -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
+10 -2
View File
@@ -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)
@@ -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"
+25
View File
@@ -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