mirror of
https://github.com/rennf93/roboco.git
synced 2026-08-03 07:23:24 +02:00
[4865ff8b] Add WebSocket support to the usage dashboard (#115)
* [e7349d84] feat(dashboard): WS usage store, hook extension, status badge, and smooth animations (#111) (#113) - Add src/store/usage-store.ts with typed UsageData interface, useUsageStore Zustand store, setUsageData, clearUsageData, and setWsState actions - Export useUsageStore and UsageData from store/index.ts - Extend use-rate-limit-websocket.ts: rename msg type to SystemWsMessage, add key_metrics field; add useEffect syncing wsState into useUsageStore; add USAGE_UPDATE/USAGE_SNAPSHOT handler dispatching to useUsageStore (RATE_LIMIT_HIT/LIFTED handling and onReconnect unchanged) - Update CommandCenter to read key_metrics from useUsageStore when wsState === 'connected' and usageData non-null; falls back to useCeoOverview() (refetchInterval: 60000) when WS disconnected - Update KeyMetricsPanel: add wsState prop, render connection status Badge matching AgentStreamViewer pattern (bg-green-500+Wifi / bg-yellow-500+ Loader2 spin / bg-gray-500+WifiOff); add transition-all duration-300 ease-in-out to metric value spans for smooth animated updates Co-authored-by: Frontend Developer 1 <fe-dev-1@agents.roboco.dev> * [c9745ee8] feat(events): add USAGE_UPDATE/SNAPSHOT event types, throttled publisher, /ws/system usage bridge (#112) (#114) - Add EventType.USAGE_UPDATE='usage.update' and EventType.USAGE_SNAPSHOT='usage.snapshot' to the EventType StrEnum in roboco/models/events.py - Create roboco/services/usage_events.py with _UsageThrottle class (5-second per-agent window using time.monotonic()) and publish_usage_update() / publish_usage_snapshot() helpers; lazy imports prevent circular dependency with roboco.events - Extend orchestrator._sweep_token_snapshots() to publish USAGE_UPDATE per active agent (throttled) and a USAGE_SNAPSHOT aggregate after each sweep cycle; wrapped in contextlib.suppress so event errors never abort DB snapshot operations - Add _handle_usage_event() to websocket_bridge.py following _handle_rate_limit_event pattern; register USAGE_UPDATE and USAGE_SNAPSHOT subscriptions in register_websocket_bridge_handlers() forwarding both to /ws/system via broadcast_system() - Add unit tests: test_usage_events.py (throttle suppression, publish helpers) and test_websocket_bridge.py extended with _handle_usage_event coverage and updated registration assertion to include USAGE_UPDATE/USAGE_SNAPSHOT Co-authored-by: Backend Developer 1 <be-dev-1@agents.roboco.dev> * fix(usage-ws): reconcile the realtime token/cost contract end-to-end The backend and frontend halves shipped mismatched contracts, so the usage dashboard never received live data: - The bridge forwarded the dotted event value ("usage.update") while the panel switched on "USAGE_UPDATE"; map both to the UPPER_SNAKE type string the same way the rate-limit handler does. - The backend emitted token/cost telemetry but the frontend read a key_metrics field and fed the org-metrics panel. Rewire the frontend to consume the USAGE_SNAPSHOT token/cost payload into the "Token Usage & Cost" panel — WS-first with polling fallback and a connection-status badge — and revert the unrelated KeyMetricsPanel / CommandCenter wiring. Backend cleanups in the same path: - Replace the multi-argument publish helpers with typed UsageUpdate / UsageSnapshot payloads, removing the too-many-arguments lint suppressions. - Extract _fetch_agent_tokens and _persist_token_snapshot from the token sweep, removing the too-many-statements suppression; label the live snapshot "live". Hardening uncovered while fixing the above: - _finalize_spawn_session pulled the full RAG stack into the session-finalization path through a transcript-parse import; move the pure parser into a dependency-light roboco.agent_sdk.transcript_usage module so finalization never imports the agent SDK server. - Reduce _finalize_spawn_session complexity by extracting _resolve_final_token_usage, and widen the transcript-fallback guard so a read error can never abort finalization. Also align KeyMetricsPanel with the metrics /dashboard/ceo actually returns: it read velocity_24h / avg_time_to_done / active_agents, none of which get_key_metrics() emits, so four of five rows rendered "—". Render velocity_weekly, completion_rate, documentation_coverage and active_blockers. * docs: note live usage push over /ws/system on the usage dashboard * fix(usage): finalize on self-exit and de-duplicate transcript token counts Two bugs left token capture broken even after the transcript-read fallback landed — surfaced by a live agent run: - Agents that self-exit (the normal i_am_idle -> container shutdown, exit 0) were never finalized. _finalize_spawn_session is only called from stop_agent(), but a graceful self-exit goes through _handle_stopped_container, which set the instance OFFLINE and returned without finalizing — leaving the spawn-session row open with zero tokens. Finalize there for both graceful (exit_reason="completed") and crash (exit_reason="crashed") exits. - sum_transcript_usage double-counted. Claude Code logs one assistant message as several JSONL lines (one per content block — thinking / text / tool_use), each repeating the same message.usage, so summing every line roughly doubled the totals. De-duplicate by message.id. Verified against a live agent transcript: the raw sum (12, 1068, 62502, 115828) vs the de-duped (6, 516, 62502, 63336), which matches the session's authoritative result.usage exactly. * feat(usage): fall back to the transcript in the live token sweep The 60s token sweep read only the agent SDK's /usage/status, which races container teardown and reports zero mid-run — so live usage (and the USAGE_SNAPSHOT pushed to /ws/system) stayed at zero for active agents. Extract _resolve_active_tokens: try the SDK, then fall back to the durable transcript (the same source finalize uses) so running agents report live. * feat(usage): add GET /usage/sessions for the dashboard's Recent Sessions The panel's Recent Sessions table was mock-only — the backend had no sessions endpoint, so production always showed 'No sessions recorded yet'. Add UsageService.get_recent_sessions + a /usage/sessions route returning the most recent spawn-session rows (token totals + cost), and point the panel client at it. --------- Co-authored-by: Frontend Developer 1 <fe-dev-1@agents.roboco.dev> Co-authored-by: Backend Developer 1 <be-dev-1@agents.roboco.dev> Co-authored-by: Renn F <rennf93@users.noreply.github.com>
This commit is contained in:
Renzo F
co-authored by
Frontend Developer 1
Backend Developer 1
Renn F
Frontend Developer 1
Backend Developer 1
Renn F
parent
1d1ec88aad
commit
547fe444f2
@@ -168,3 +168,51 @@ def test_parser_handles_entries_without_message(tmp_path: Path) -> None:
|
||||
exp["tokens_cache_read"],
|
||||
exp["tokens_cache_write"],
|
||||
)
|
||||
|
||||
|
||||
def _assistant_line_with_id(row: _UsageRow, message_id: str) -> str:
|
||||
"""An assistant line carrying a message id (for de-duplication tests)."""
|
||||
inp, out, cread, cwrite = row
|
||||
return json.dumps(
|
||||
{
|
||||
"type": "assistant",
|
||||
"message": {
|
||||
"id": message_id,
|
||||
"role": "assistant",
|
||||
"usage": {
|
||||
"input_tokens": inp,
|
||||
"output_tokens": out,
|
||||
"cache_read_input_tokens": cread,
|
||||
"cache_creation_input_tokens": cwrite,
|
||||
},
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def test_parser_dedupes_repeated_message_id(tmp_path: Path) -> None:
|
||||
"""One message logged across several lines (same id) is counted once.
|
||||
|
||||
Claude Code emits one transcript line per content block (thinking, text,
|
||||
tool_use), each repeating the message's ``usage``. Summing every line would
|
||||
roughly double the totals, so the parser must de-duplicate by message id.
|
||||
"""
|
||||
msg = (100, 20, 5, 3)
|
||||
other = (7, 2, 1, 0)
|
||||
transcript = tmp_path / "session.jsonl"
|
||||
_write(
|
||||
transcript,
|
||||
_assistant_line_with_id(msg, "msg_aaa"), # thinking block
|
||||
_assistant_line_with_id(msg, "msg_aaa"), # text block (same id)
|
||||
_assistant_line_with_id(msg, "msg_aaa"), # tool_use block (same id)
|
||||
_assistant_line_with_id(other, "msg_bbb"),
|
||||
)
|
||||
tin, tout, cread, cwrite = srv._sum_transcript_usage(transcript)
|
||||
# Counted once per id: msg + other, NOT msg * 3 + other.
|
||||
exp = _expected([msg, other])
|
||||
assert (tin, tout, cread, cwrite) == (
|
||||
exp["tokens_input"],
|
||||
exp["tokens_output"],
|
||||
exp["tokens_cache_read"],
|
||||
exp["tokens_cache_write"],
|
||||
)
|
||||
|
||||
@@ -17,6 +17,7 @@ from roboco.api.websocket_bridge import (
|
||||
_handle_notification_sent,
|
||||
_handle_rate_limit_event,
|
||||
_handle_session_event,
|
||||
_handle_usage_event,
|
||||
register_websocket_bridge_handlers,
|
||||
start_websocket_bridge,
|
||||
)
|
||||
@@ -291,13 +292,78 @@ async def test_handle_rate_limit_ignores_unrelated_event() -> None:
|
||||
mgr.broadcast_system.assert_not_called()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _handle_usage_event
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_usage_update_broadcasts_to_system() -> None:
|
||||
"""USAGE_UPDATE event → broadcast_system tagged USAGE_UPDATE + data fields."""
|
||||
expected_input = 100
|
||||
expected_output = 50
|
||||
event = _evt(
|
||||
EventType.USAGE_UPDATE,
|
||||
{
|
||||
"agent_id": "be-dev-1",
|
||||
"task_id": "task-abc",
|
||||
"input_tokens": expected_input,
|
||||
"output_tokens": expected_output,
|
||||
"model": "claude-sonnet-4-6",
|
||||
"timestamp": "2026-06-11T00:00:00+00:00",
|
||||
},
|
||||
)
|
||||
with patch("roboco.api.websocket_bridge.manager") as mgr:
|
||||
mgr.broadcast_system = AsyncMock()
|
||||
await _handle_usage_event(event)
|
||||
mgr.broadcast_system.assert_awaited_once()
|
||||
msg = mgr.broadcast_system.await_args.args[0]
|
||||
assert msg["type"] == "USAGE_UPDATE"
|
||||
assert msg["agent_id"] == "be-dev-1"
|
||||
assert msg["input_tokens"] == expected_input
|
||||
assert msg["output_tokens"] == expected_output
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_usage_snapshot_broadcasts_to_system() -> None:
|
||||
"""USAGE_SNAPSHOT event → broadcast_system tagged USAGE_SNAPSHOT + aggregate."""
|
||||
expected_input = 500
|
||||
event = _evt(
|
||||
EventType.USAGE_SNAPSHOT,
|
||||
{
|
||||
"period": "live",
|
||||
"totals": {"input_tokens": expected_input, "output_tokens": 200},
|
||||
"cost_estimate": 0.0025,
|
||||
"by_agent": [
|
||||
{
|
||||
"agent_id": "be-dev-1",
|
||||
"input_tokens": 500,
|
||||
"output_tokens": 200,
|
||||
"model": "sonnet",
|
||||
"cost_estimate": 0.0025,
|
||||
}
|
||||
],
|
||||
"timestamp": "2026-06-11T00:01:00+00:00",
|
||||
},
|
||||
)
|
||||
with patch("roboco.api.websocket_bridge.manager") as mgr:
|
||||
mgr.broadcast_system = AsyncMock()
|
||||
await _handle_usage_event(event)
|
||||
mgr.broadcast_system.assert_awaited_once()
|
||||
msg = mgr.broadcast_system.await_args.args[0]
|
||||
assert msg["type"] == "USAGE_SNAPSHOT"
|
||||
assert msg["period"] == "live"
|
||||
assert msg["totals"]["input_tokens"] == expected_input
|
||||
assert len(msg["by_agent"]) == 1
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Registration + start
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_register_websocket_bridge_handlers_subscribes_all_event_types() -> None:
|
||||
"""Registration wires up notification + session + agent event handlers."""
|
||||
"""Registration wires up all handler categories, including usage events."""
|
||||
|
||||
class _FakeBus:
|
||||
def __init__(self) -> None:
|
||||
@@ -310,7 +376,7 @@ def test_register_websocket_bridge_handlers_subscribes_all_event_types() -> None
|
||||
with patch("roboco.api.websocket_bridge.get_event_bus", return_value=fake):
|
||||
register_websocket_bridge_handlers()
|
||||
types = [t for t, _ in fake.subscribed]
|
||||
# All 10 expected event types appear at least once.
|
||||
# All 14 expected event types appear at least once.
|
||||
assert EventType.NOTIFICATION_SENT in types
|
||||
assert EventType.NOTIFICATION_ACKED in types
|
||||
assert EventType.SESSION_CREATED in types
|
||||
@@ -323,6 +389,9 @@ def test_register_websocket_bridge_handlers_subscribes_all_event_types() -> None
|
||||
assert EventType.AGENT_ERROR in types
|
||||
assert EventType.RATE_LIMIT_HIT in types
|
||||
assert EventType.RATE_LIMIT_LIFTED in types
|
||||
# Usage events forwarded to /ws/system.
|
||||
assert EventType.USAGE_UPDATE in types
|
||||
assert EventType.USAGE_SNAPSHOT in types
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
||||
@@ -265,6 +265,7 @@ async def test_finalize_spawn_session_http_error_uses_zero_tokens() -> None:
|
||||
|
||||
with (
|
||||
patch("roboco.runtime.orchestrator.httpx.AsyncClient", _client_cls),
|
||||
patch.object(orch, "_usage_from_transcript", return_value=(0, 0, 0, 0)),
|
||||
patch("roboco.db.base.get_session_factory", return_value=db_factory),
|
||||
patch("roboco.billing.pricing.calculate_cost", return_value=0.0) as mock_cost,
|
||||
):
|
||||
@@ -298,6 +299,7 @@ async def test_finalize_spawn_session_non_200_uses_zero_tokens() -> None:
|
||||
|
||||
with (
|
||||
patch("roboco.runtime.orchestrator.httpx.AsyncClient", _client_cls),
|
||||
patch.object(orch, "_usage_from_transcript", return_value=(0, 0, 0, 0)),
|
||||
patch("roboco.db.base.get_session_factory", return_value=db_factory),
|
||||
patch("roboco.billing.pricing.calculate_cost", return_value=0.0) as mock_cost,
|
||||
):
|
||||
@@ -383,6 +385,7 @@ async def test_sweep_token_snapshots_skips_zero_token_agents() -> None:
|
||||
|
||||
with (
|
||||
patch("roboco.runtime.orchestrator.httpx.AsyncClient", _client_cls),
|
||||
patch.object(orch, "_usage_from_transcript", return_value=(0, 0, 0, 0)),
|
||||
patch("roboco.db.base.get_session_factory", return_value=db_factory),
|
||||
):
|
||||
await orch._sweep_token_snapshots()
|
||||
@@ -525,7 +528,7 @@ async def test_stop_agent_finalizes_before_lock() -> None:
|
||||
|
||||
finalized: list[str] = []
|
||||
|
||||
async def _fake_finalize(agent_id: str, exit_reason: str = "stopped") -> None: # noqa: ARG001
|
||||
async def _fake_finalize(agent_id: str, **_kwargs: object) -> None:
|
||||
finalized.append(agent_id)
|
||||
|
||||
# Stub out the Docker subprocess so stop_agent doesn't actually run Docker
|
||||
@@ -541,3 +544,107 @@ async def test_stop_agent_finalizes_before_lock() -> None:
|
||||
|
||||
# _finalize_spawn_session must have been called exactly once with our agent id
|
||||
assert finalized == [_AGENT_ID]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _handle_stopped_container — self-exits finalize (stop_agent was not called)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def test_handle_stopped_container_graceful_finalizes() -> None:
|
||||
"""A graceful self-exit (exit 0) finalizes the spawn session.
|
||||
|
||||
The agent calls i_am_idle and its container exits 0 without stop_agent
|
||||
being invoked, so _handle_stopped_container must finalize to capture the
|
||||
token usage; otherwise the session row is left open with zero tokens.
|
||||
"""
|
||||
orch = _make_orchestrator()
|
||||
instance = _make_instance(_AGENT_ID)
|
||||
instance.container_id = "abc123def456"
|
||||
orch._instances[_AGENT_ID] = instance
|
||||
|
||||
calls: list[tuple[str, str]] = []
|
||||
|
||||
async def _fake_finalize(agent_id: str, exit_reason: str = "stopped") -> None:
|
||||
calls.append((agent_id, exit_reason))
|
||||
|
||||
with patch.object(orch, "_finalize_spawn_session", side_effect=_fake_finalize):
|
||||
await orch._handle_stopped_container(_AGENT_ID, instance, 0)
|
||||
|
||||
assert calls == [(_AGENT_ID, "completed")]
|
||||
assert instance.state is OrchestratorAgentState.OFFLINE
|
||||
|
||||
|
||||
async def test_handle_stopped_container_crash_finalizes_then_restarts() -> None:
|
||||
"""A non-zero exit finalizes (exit_reason='crashed') before auto-restart."""
|
||||
orch = _make_orchestrator()
|
||||
instance = _make_instance(_AGENT_ID)
|
||||
instance.container_id = "abc123def456"
|
||||
instance.error_count = 0
|
||||
orch._instances[_AGENT_ID] = instance
|
||||
|
||||
calls: list[tuple[str, str]] = []
|
||||
|
||||
async def _fake_finalize(agent_id: str, exit_reason: str = "stopped") -> None:
|
||||
calls.append((agent_id, exit_reason))
|
||||
|
||||
with (
|
||||
patch.object(orch, "_finalize_spawn_session", side_effect=_fake_finalize),
|
||||
patch.object(orch, "spawn_agent", AsyncMock()) as mock_spawn,
|
||||
):
|
||||
await orch._handle_stopped_container(_AGENT_ID, instance, 1)
|
||||
|
||||
assert calls == [(_AGENT_ID, "crashed")]
|
||||
mock_spawn.assert_awaited_once()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _resolve_active_tokens — SDK first, transcript fallback for live agents
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def test_resolve_active_tokens_falls_back_to_transcript() -> None:
|
||||
"""When the SDK reports all-zero, live resolution uses the transcript."""
|
||||
orch = _make_orchestrator()
|
||||
|
||||
def _handler(_url: str) -> Any:
|
||||
return _mock_response(
|
||||
200,
|
||||
{
|
||||
"tokens_input": 0,
|
||||
"tokens_output": 0,
|
||||
"tokens_cache_read": 0,
|
||||
"tokens_cache_write": 0,
|
||||
},
|
||||
)
|
||||
|
||||
client = _FakeHTTPClient(_handler)
|
||||
with patch.object(orch, "_usage_from_transcript", return_value=(6, 514, 100, 50)):
|
||||
tokens = await orch._resolve_active_tokens(client, _AGENT_ID)
|
||||
|
||||
assert tokens == (6, 514, 100, 50)
|
||||
|
||||
|
||||
async def test_resolve_active_tokens_prefers_sdk() -> None:
|
||||
"""A non-zero SDK response is used directly — no transcript fallback."""
|
||||
orch = _make_orchestrator()
|
||||
|
||||
def _handler(_url: str) -> Any:
|
||||
return _mock_response(
|
||||
200,
|
||||
{
|
||||
"tokens_input": 10,
|
||||
"tokens_output": 20,
|
||||
"tokens_cache_read": 0,
|
||||
"tokens_cache_write": 0,
|
||||
},
|
||||
)
|
||||
|
||||
client = _FakeHTTPClient(_handler)
|
||||
with patch.object(
|
||||
orch, "_usage_from_transcript", return_value=(999, 999, 999, 999)
|
||||
) as mock_tx:
|
||||
tokens = await orch._resolve_active_tokens(client, _AGENT_ID)
|
||||
|
||||
assert tokens == (10, 20, 0, 0)
|
||||
mock_tx.assert_not_called()
|
||||
|
||||
@@ -14,6 +14,7 @@ from __future__ import annotations
|
||||
|
||||
import datetime
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
from uuid import UUID
|
||||
|
||||
import pytest
|
||||
from roboco.services.usage import UsageService
|
||||
@@ -759,3 +760,77 @@ class TestGetCacheEfficiency:
|
||||
result = await svc.get_cache_efficiency("24h")
|
||||
for field in ("cache_hit_rate", "cost_saved_by_cache_usd"):
|
||||
assert field in result, f"Missing field: {field}"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# get_recent_sessions — maps spawn-session rows to the dashboard shape
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestGetRecentSessions:
|
||||
@pytest.mark.asyncio
|
||||
async def test_shapes_rows(self) -> None:
|
||||
"""Rows are mapped to id/agent/model/tokens/cache/total/cost fields."""
|
||||
exp_in, exp_out = 6, 514
|
||||
exp_cr, exp_cw = 111_032, 14_881
|
||||
exp_cost = 0.1614
|
||||
exp_count = 1
|
||||
sid = UUID("12345678-1234-5678-1234-567812345678")
|
||||
|
||||
row = MagicMock()
|
||||
row.id = sid
|
||||
row.agent_slug = "product-owner"
|
||||
row.model = "claude-opus-4-6"
|
||||
row.started_at = datetime.datetime(2026, 6, 11, 20, 41, tzinfo=datetime.UTC)
|
||||
row.ended_at = datetime.datetime(2026, 6, 11, 20, 42, tzinfo=datetime.UTC)
|
||||
row.tokens_input = exp_in
|
||||
row.tokens_output = exp_out
|
||||
row.tokens_cache_read = exp_cr
|
||||
row.tokens_cache_write = exp_cw
|
||||
row.estimated_cost_usd = exp_cost
|
||||
|
||||
scalars = MagicMock()
|
||||
scalars.all = MagicMock(return_value=[row])
|
||||
result = MagicMock()
|
||||
result.scalars = MagicMock(return_value=scalars)
|
||||
|
||||
svc = _service_with_execute(result)
|
||||
out = await svc.get_recent_sessions(limit=10)
|
||||
|
||||
assert len(out) == exp_count
|
||||
s = out[0]
|
||||
assert s["id"] == str(sid)
|
||||
assert s["agent_slug"] == "product-owner"
|
||||
assert s["model"] == "claude-opus-4-6"
|
||||
assert s["tokens_input"] == exp_in
|
||||
assert s["tokens_output"] == exp_out
|
||||
assert s["tokens_cache"] == exp_cr + exp_cw
|
||||
assert s["total_tokens"] == exp_in + exp_out + exp_cr + exp_cw
|
||||
assert s["cost"] == pytest.approx(exp_cost)
|
||||
assert s["ended_at"] is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_open_session_has_null_ended_at(self) -> None:
|
||||
"""A still-running session (ended_at None) serializes ended_at as None."""
|
||||
row = MagicMock()
|
||||
row.id = "00000000-0000-0000-0000-000000000001"
|
||||
row.agent_slug = "main-pm"
|
||||
row.model = "sonnet"
|
||||
row.started_at = datetime.datetime(2026, 6, 11, 20, 0, tzinfo=datetime.UTC)
|
||||
row.ended_at = None
|
||||
row.tokens_input = _ZERO
|
||||
row.tokens_output = _ZERO
|
||||
row.tokens_cache_read = _ZERO
|
||||
row.tokens_cache_write = _ZERO
|
||||
row.estimated_cost_usd = None
|
||||
|
||||
scalars = MagicMock()
|
||||
scalars.all = MagicMock(return_value=[row])
|
||||
result = MagicMock()
|
||||
result.scalars = MagicMock(return_value=scalars)
|
||||
|
||||
svc = _service_with_execute(result)
|
||||
out = await svc.get_recent_sessions()
|
||||
|
||||
assert out[0]["ended_at"] is None
|
||||
assert out[0]["cost"] == _ZERO
|
||||
|
||||
@@ -0,0 +1,262 @@
|
||||
"""Unit tests for roboco.services.usage_events.
|
||||
|
||||
Covers the _UsageThrottle class and the publish_usage_update /
|
||||
publish_usage_snapshot helpers. No real Redis or event bus is needed —
|
||||
we use AsyncMock to assert that bus.publish is called with the right
|
||||
payload and type.
|
||||
|
||||
The throttle suppression test is the acceptance-criterion gate:
|
||||
"Server-side throttle prevents more than 1 USAGE_UPDATE publish per
|
||||
agent per 5-second window."
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import UTC, datetime
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from roboco.services.usage_events import (
|
||||
UsageSnapshot,
|
||||
UsageUpdate,
|
||||
_UsageThrottle,
|
||||
publish_usage_snapshot,
|
||||
publish_usage_update,
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _UsageThrottle
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_throttle_allows_first_publish() -> None:
|
||||
"""A fresh agent has no prior timestamp — first publish is always allowed."""
|
||||
th = _UsageThrottle(window=5.0)
|
||||
assert th.should_publish("be-dev-1") is True
|
||||
|
||||
|
||||
def test_throttle_suppresses_second_publish_within_window() -> None:
|
||||
"""Second call within the 5-second window returns False (suppressed)."""
|
||||
th = _UsageThrottle(window=5.0)
|
||||
|
||||
with patch("roboco.services.usage_events.time") as mock_time:
|
||||
mock_time.monotonic.return_value = 100.0
|
||||
assert th.should_publish("be-dev-1") is True # first → allowed
|
||||
|
||||
mock_time.monotonic.return_value = 104.9 # 4.9 s later — still inside window
|
||||
assert th.should_publish("be-dev-1") is False # suppressed
|
||||
|
||||
|
||||
def test_throttle_allows_publish_after_window_expires() -> None:
|
||||
"""After the full window elapses, the next publish is allowed again."""
|
||||
th = _UsageThrottle(window=5.0)
|
||||
|
||||
with patch("roboco.services.usage_events.time") as mock_time:
|
||||
mock_time.monotonic.return_value = 100.0
|
||||
assert th.should_publish("be-dev-1") is True # first
|
||||
|
||||
mock_time.monotonic.return_value = 105.0 # exactly 5 s later
|
||||
assert th.should_publish("be-dev-1") is True # window elapsed → allowed
|
||||
|
||||
|
||||
def test_throttle_tracks_agents_independently() -> None:
|
||||
"""Different agents have independent throttle windows."""
|
||||
th = _UsageThrottle(window=5.0)
|
||||
|
||||
with patch("roboco.services.usage_events.time") as mock_time:
|
||||
mock_time.monotonic.return_value = 100.0
|
||||
|
||||
assert th.should_publish("be-dev-1") is True
|
||||
# be-dev-2 has never published, so it is always allowed.
|
||||
assert th.should_publish("be-dev-2") is True
|
||||
|
||||
mock_time.monotonic.return_value = 101.0
|
||||
# be-dev-1 is suppressed; be-dev-2 is also now suppressed.
|
||||
assert th.should_publish("be-dev-1") is False
|
||||
assert th.should_publish("be-dev-2") is False
|
||||
|
||||
|
||||
def test_throttle_records_timestamp_on_allow() -> None:
|
||||
"""should_publish records the current time when it returns True."""
|
||||
th = _UsageThrottle(window=5.0)
|
||||
recorded_at = 200.0
|
||||
|
||||
with patch("roboco.services.usage_events.time") as mock_time:
|
||||
mock_time.monotonic.return_value = recorded_at
|
||||
th.should_publish("be-dev-1")
|
||||
assert th._last["be-dev-1"] == recorded_at
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# publish_usage_update
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_publish_usage_update_calls_bus_publish() -> None:
|
||||
"""First call in a window publishes the event and returns True."""
|
||||
bus = MagicMock()
|
||||
bus.publish = AsyncMock()
|
||||
th = _UsageThrottle(window=5.0)
|
||||
expected_input = 100
|
||||
expected_output = 50
|
||||
|
||||
with patch("roboco.services.usage_events._throttle", th):
|
||||
result = await publish_usage_update(
|
||||
bus,
|
||||
UsageUpdate(
|
||||
agent_id="be-dev-1",
|
||||
task_id="task-abc",
|
||||
input_tokens=expected_input,
|
||||
output_tokens=expected_output,
|
||||
model="claude-sonnet-4-6",
|
||||
),
|
||||
)
|
||||
|
||||
assert result is True
|
||||
bus.publish.assert_awaited_once()
|
||||
event = bus.publish.await_args.args[0]
|
||||
assert event.type.value == "usage.update"
|
||||
assert event.data["agent_id"] == "be-dev-1"
|
||||
assert event.data["task_id"] == "task-abc"
|
||||
assert event.data["input_tokens"] == expected_input
|
||||
assert event.data["output_tokens"] == expected_output
|
||||
assert event.data["model"] == "claude-sonnet-4-6"
|
||||
assert "timestamp" in event.data
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_publish_usage_update_throttle_suppresses_second_call() -> None:
|
||||
"""Second publish within the throttle window is suppressed (returns False)."""
|
||||
bus = MagicMock()
|
||||
bus.publish = AsyncMock()
|
||||
th = _UsageThrottle(window=5.0)
|
||||
|
||||
with (
|
||||
patch("roboco.services.usage_events._throttle", th),
|
||||
patch("roboco.services.usage_events.time") as mock_time,
|
||||
):
|
||||
mock_time.monotonic.return_value = 100.0
|
||||
first = await publish_usage_update(
|
||||
bus,
|
||||
UsageUpdate(
|
||||
agent_id="be-dev-1",
|
||||
task_id=None,
|
||||
input_tokens=10,
|
||||
output_tokens=5,
|
||||
model="sonnet",
|
||||
),
|
||||
)
|
||||
|
||||
mock_time.monotonic.return_value = 102.0 # 2 s later — still suppressed
|
||||
second = await publish_usage_update(
|
||||
bus,
|
||||
UsageUpdate(
|
||||
agent_id="be-dev-1",
|
||||
task_id=None,
|
||||
input_tokens=20,
|
||||
output_tokens=10,
|
||||
model="sonnet",
|
||||
),
|
||||
)
|
||||
|
||||
assert first is True
|
||||
assert second is False
|
||||
# bus.publish should only have been called once.
|
||||
assert bus.publish.await_count == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_publish_usage_update_custom_timestamp() -> None:
|
||||
"""Custom timestamp is passed through to the event data."""
|
||||
bus = MagicMock()
|
||||
bus.publish = AsyncMock()
|
||||
ts = datetime(2026, 6, 11, 12, 0, 0, tzinfo=UTC)
|
||||
|
||||
# Use a fresh throttle so the first publish goes through.
|
||||
th = _UsageThrottle(window=5.0)
|
||||
with patch("roboco.services.usage_events._throttle", th):
|
||||
await publish_usage_update(
|
||||
bus,
|
||||
UsageUpdate(
|
||||
agent_id="be-dev-1",
|
||||
task_id=None,
|
||||
input_tokens=0,
|
||||
output_tokens=0,
|
||||
model="sonnet",
|
||||
timestamp=ts,
|
||||
),
|
||||
)
|
||||
|
||||
event = bus.publish.await_args.args[0]
|
||||
assert event.data["timestamp"] == ts.isoformat()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# publish_usage_snapshot
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_publish_usage_snapshot_always_publishes() -> None:
|
||||
"""publish_usage_snapshot has no throttle — always publishes."""
|
||||
bus = MagicMock()
|
||||
bus.publish = AsyncMock()
|
||||
expected_input = 500
|
||||
expected_cost = 0.0025
|
||||
expected_agents = 2
|
||||
|
||||
await publish_usage_snapshot(
|
||||
bus,
|
||||
UsageSnapshot(
|
||||
period="live",
|
||||
totals={"input_tokens": expected_input, "output_tokens": 200},
|
||||
cost_estimate=expected_cost,
|
||||
by_agent=[
|
||||
{
|
||||
"agent_id": "be-dev-1",
|
||||
"input_tokens": 300,
|
||||
"output_tokens": 100,
|
||||
"model": "sonnet",
|
||||
"cost_estimate": 0.0015,
|
||||
},
|
||||
{
|
||||
"agent_id": "be-dev-2",
|
||||
"input_tokens": 200,
|
||||
"output_tokens": 100,
|
||||
"model": "sonnet",
|
||||
"cost_estimate": 0.0010,
|
||||
},
|
||||
],
|
||||
),
|
||||
)
|
||||
|
||||
bus.publish.assert_awaited_once()
|
||||
event = bus.publish.await_args.args[0]
|
||||
assert event.type.value == "usage.snapshot"
|
||||
assert event.data["period"] == "live"
|
||||
assert event.data["totals"]["input_tokens"] == expected_input
|
||||
assert event.data["cost_estimate"] == expected_cost
|
||||
assert len(event.data["by_agent"]) == expected_agents
|
||||
assert "timestamp" in event.data
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_publish_usage_snapshot_twice_both_published() -> None:
|
||||
"""No throttle on snapshot: two rapid calls both publish."""
|
||||
bus = MagicMock()
|
||||
bus.publish = AsyncMock()
|
||||
expected_calls = 2
|
||||
|
||||
for _ in range(expected_calls):
|
||||
await publish_usage_snapshot(
|
||||
bus,
|
||||
UsageSnapshot(
|
||||
period="live",
|
||||
totals={"input_tokens": 0, "output_tokens": 0},
|
||||
cost_estimate=0.0,
|
||||
by_agent=[],
|
||||
),
|
||||
)
|
||||
|
||||
assert bus.publish.await_count == expected_calls
|
||||
Reference in New Issue
Block a user