mirror of
https://github.com/openswarm-ai/openswarm.git
synced 2026-08-23 21:12:22 +02:00
[aidan] fix/analytics-telemetry: emit final thinking snapshots once
This commit is contained in:
@@ -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
|
||||
|
||||
|
||||
@@ -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"
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user