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:
co-authored by
Frontend Developer 1
Backend Developer 1
Renn F
parent
1d1ec88aad
commit
547fe444f2
@@ -0,0 +1,129 @@
|
||||
"""
|
||||
Usage Event Publisher
|
||||
|
||||
Throttled helpers for publishing USAGE_UPDATE and USAGE_SNAPSHOT events
|
||||
to the StreamEventBus. Consumed by the orchestrator token sweep and
|
||||
forwarded to /ws/system WebSocket clients via the websocket_bridge.
|
||||
|
||||
Throttle window: one USAGE_UPDATE publish per agent per 5-second window.
|
||||
Subsequent calls within the window are silently dropped so a frequent
|
||||
sweep loop cannot flood the event bus or WebSocket clients.
|
||||
|
||||
USAGE_SNAPSHOT is always published (no per-agent throttle) — it is an
|
||||
aggregate and published at most once per sweep cycle.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from datetime import UTC, datetime
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from roboco.events.stream_bus import StreamEventBus
|
||||
|
||||
_THROTTLE_WINDOW_SECONDS: float = 5.0
|
||||
|
||||
|
||||
class _UsageThrottle:
|
||||
"""Per-agent last-publish timestamp tracker.
|
||||
|
||||
Uses ``time.monotonic()`` so clock adjustments (NTP, DST) don't cause
|
||||
spurious suppressions or double-fires.
|
||||
"""
|
||||
|
||||
def __init__(self, window: float = _THROTTLE_WINDOW_SECONDS) -> None:
|
||||
self._window = window
|
||||
self._last: dict[str, float] = {}
|
||||
|
||||
def should_publish(self, agent_id: str) -> bool:
|
||||
"""Return True if the agent is outside the throttle window.
|
||||
|
||||
Also records the current monotonic time as the new *last published*
|
||||
timestamp when it returns True, so the caller does not need to call a
|
||||
separate ``record()`` method.
|
||||
"""
|
||||
now = time.monotonic()
|
||||
if now - self._last.get(agent_id, 0.0) >= self._window:
|
||||
self._last[agent_id] = now
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
# Module-level singleton — shared across all callers in this process.
|
||||
_throttle = _UsageThrottle()
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class UsageUpdate:
|
||||
"""Per-agent cumulative token counts carried by a USAGE_UPDATE event.
|
||||
|
||||
Fields map directly onto the event payload the panel consumes (the
|
||||
backend emits these per active agent during the token sweep).
|
||||
"""
|
||||
|
||||
agent_id: str
|
||||
task_id: str | None
|
||||
input_tokens: int
|
||||
output_tokens: int
|
||||
model: str
|
||||
timestamp: datetime | None = None
|
||||
|
||||
def event_data(self) -> dict[str, Any]:
|
||||
"""Render the event payload, stamping ``timestamp`` if not supplied."""
|
||||
return {
|
||||
"agent_id": self.agent_id,
|
||||
"task_id": self.task_id,
|
||||
"input_tokens": self.input_tokens,
|
||||
"output_tokens": self.output_tokens,
|
||||
"model": self.model,
|
||||
"timestamp": (self.timestamp or datetime.now(UTC)).isoformat(),
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class UsageSnapshot:
|
||||
"""Aggregate token/cost totals carried by a USAGE_SNAPSHOT event.
|
||||
|
||||
``totals`` is ``{"input_tokens": int, "output_tokens": int}``; ``by_agent``
|
||||
is the per-agent breakdown (each item carries ``agent_id``, token counts,
|
||||
``model`` and ``cost_estimate``).
|
||||
"""
|
||||
|
||||
period: str
|
||||
totals: dict[str, int]
|
||||
cost_estimate: float
|
||||
by_agent: list[dict[str, Any]]
|
||||
timestamp: datetime | None = None
|
||||
|
||||
def event_data(self) -> dict[str, Any]:
|
||||
"""Render the event payload, stamping ``timestamp`` if not supplied."""
|
||||
return {
|
||||
"period": self.period,
|
||||
"totals": self.totals,
|
||||
"cost_estimate": self.cost_estimate,
|
||||
"by_agent": self.by_agent,
|
||||
"timestamp": (self.timestamp or datetime.now(UTC)).isoformat(),
|
||||
}
|
||||
|
||||
|
||||
async def publish_usage_update(bus: StreamEventBus, update: UsageUpdate) -> bool:
|
||||
"""Publish a USAGE_UPDATE event if the per-agent throttle window has elapsed.
|
||||
|
||||
Returns True if the event was published; False if suppressed by the throttle.
|
||||
"""
|
||||
if not _throttle.should_publish(update.agent_id):
|
||||
return False
|
||||
|
||||
from roboco.models.events import Event, EventType # lazy — avoids circular import
|
||||
|
||||
await bus.publish(Event(type=EventType.USAGE_UPDATE, data=update.event_data()))
|
||||
return True
|
||||
|
||||
|
||||
async def publish_usage_snapshot(bus: StreamEventBus, snapshot: UsageSnapshot) -> None:
|
||||
"""Publish a USAGE_SNAPSHOT aggregate event (no throttle)."""
|
||||
from roboco.models.events import Event, EventType # lazy — avoids circular import
|
||||
|
||||
await bus.publish(Event(type=EventType.USAGE_SNAPSHOT, data=snapshot.event_data()))
|
||||
Reference in New Issue
Block a user