[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>
2026-08-03 07:23:24 +02:00 · 2026-06-11 23:19:50 +02:00
parent 1d1ec88aad
commit 547fe444f2
22 changed files with 1300 additions and 187 deletions
@@ -2,7 +2,7 @@

 import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
 import { Skeleton } from "@/components/ui/skeleton";
-import { TrendingUp, Clock, CheckCircle, Users, BarChart3 } from "lucide-react";
+import { TrendingUp, CheckCircle, BarChart3, AlertTriangle } from "lucide-react";

 interface KeyMetricsProps {
  metrics: Record<string, unknown> | undefined;
@@ -16,17 +16,14 @@ interface MetricItem {
  format?: (value: number) => string;
 }

+// Keys must match DashboardService.get_key_metrics() — the shape /dashboard/ceo
+// returns. (velocity_weekly + completion_rate + documentation_coverage are the
+// 7-day rollups; active_blockers is the live blocked-task count.)
 const METRIC_CONFIG: MetricItem[] = [
  {
-    key: "velocity_24h",
-    label: "Velocity (24h)",
-    icon: <TrendingUp className="h-4 w-4" />,
-    format: (v) => `${v} tasks`,
-  },
-  {
-    key: "velocity_7d",
+    key: "velocity_weekly",
    label: "Velocity (7d)",
-    icon: <BarChart3 className="h-4 w-4" />,
+    icon: <TrendingUp className="h-4 w-4" />,
    format: (v) => `${v} tasks`,
  },
  {
@@ -36,15 +33,15 @@ const METRIC_CONFIG: MetricItem[] = [
    format: (v) => `${Math.round(v * 100)}%`,
  },
  {
-    key: "avg_time_to_done",
-    label: "Avg. Time to Done",
-    icon: <Clock className="h-4 w-4" />,
-    format: (v) => `${(typeof v === "number" ? v : 0).toFixed(1)}h`,
+    key: "documentation_coverage",
+    label: "Documentation Coverage",
+    icon: <BarChart3 className="h-4 w-4" />,
+    format: (v) => `${Math.round(v * 100)}%`,
  },
  {
-    key: "active_agents",
-    label: "Active Agents",
-    icon: <Users className="h-4 w-4" />,
+    key: "active_blockers",
+    label: "Active Blockers",
+    icon: <AlertTriangle className="h-4 w-4" />,
    format: (v) => `${v}`,
  },
 ];
@@ -1,9 +1,21 @@
 "use client";

 import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
+import { Badge } from "@/components/ui/badge";
 import { Skeleton } from "@/components/ui/skeleton";
 import { useUsageSummary } from "@/hooks/use-usage";
-import { Coins, TrendingUp, TrendingDown, Zap, Activity } from "lucide-react";
+import { useUsageStore } from "@/store/usage-store";
+import type { ConnectionState } from "@/lib/websocket/connection";
+import {
+  Coins,
+  TrendingUp,
+  TrendingDown,
+  Zap,
+  Activity,
+  Wifi,
+  WifiOff,
+  Loader2,
+} from "lucide-react";

 function fmt(n: number, decimals = 0): string {
  if (n >= 1_000_000) return (n / 1_000_000).toFixed(1) + "M";
@@ -30,28 +42,85 @@ function MetricRow({ icon, label, value, sub }: MetricRowProps) {
        {label}
      </div>
      <div className="flex items-center gap-1">
-        <span className="font-semibold text-sm">{value}</span>
+        <span className="font-semibold text-sm transition-all duration-300 ease-in-out">
+          {value}
+        </span>
        {sub}
      </div>
    </div>
  );
 }

+/**
+ * Badge props mirroring the AgentStreamViewer connection-status pattern:
+ * green/Live when the /ws/system stream is connected, yellow while
+ * (re)connecting, gray/Polling when it is down and the panel uses HTTP polling.
+ */
+function getConnectionBadge(wsState: ConnectionState) {
+  switch (wsState) {
+    case "connected":
+      return {
+        className: "bg-green-500 text-white",
+        icon: <Wifi className="h-3 w-3 mr-1" />,
+        label: "Live",
+      };
+    case "connecting":
+    case "reconnecting":
+      return {
+        className: "bg-yellow-500 text-white",
+        icon: <Loader2 className="h-3 w-3 mr-1 animate-spin" />,
+        label: wsState === "reconnecting" ? "Reconnecting..." : "Connecting...",
+      };
+    case "disconnected":
+    default:
+      return {
+        className: "bg-gray-500 text-white",
+        icon: <WifiOff className="h-3 w-3 mr-1" />,
+        label: "Polling",
+      };
+  }
+}
+
 export function UsageOverviewPanel() {
+  // The polling summary always runs in the background: it is the fallback when
+  // the WS is down, and it supplies the trend (which the live snapshot, being a
+  // point-in-time aggregate, does not compute).
  const { data: summary, isLoading } = useUsageSummary("24h");

-  const trendUp = (summary?.trend_pct ?? 0) >= 0;
+  // Live token/cost pushed over /ws/system (USAGE_SNAPSHOT). Prefer it whenever
+  // the stream is connected; `live` is non-null only then, so it narrows safely.
+  const { wsState, usageData } = useUsageStore();
+  const live = wsState === "connected" ? usageData : null;
+
+  const tokensInput = live ? live.tokens_input : summary?.tokens_input;
+  const tokensOutput = live ? live.tokens_output : summary?.tokens_output;
+  const totalCost = live ? live.total_cost_usd : summary?.total_cost_usd;
+  const periodLabel = live ? live.period : summary?.period;
+
+  // Trend always comes from the polling summary (needs a prior-period baseline).
+  const trendPct = summary?.trend_pct ?? 0;
+  const trendUp = trendPct >= 0;
+
+  const badge = getConnectionBadge(wsState);
+  // Only show skeletons on first load with no data from either source.
+  const showSkeleton = isLoading && !live;

  return (
    <Card>
      <CardHeader className="pb-3">
-        <CardTitle className="text-lg flex items-center gap-2">
-          <Coins className="h-5 w-5" />
-          Token Usage &amp; Cost
-        </CardTitle>
+        <div className="flex items-center justify-between">
+          <CardTitle className="text-lg flex items-center gap-2">
+            <Coins className="h-5 w-5" />
+            Token Usage &amp; Cost
+          </CardTitle>
+          <Badge className={badge.className}>
+            {badge.icon}
+            {badge.label}
+          </Badge>
+        </div>
      </CardHeader>
      <CardContent>
-        {isLoading ? (
+        {showSkeleton ? (
          <div className="space-y-3">
            {Array.from({ length: 5 }).map((_, i) => (
              <Skeleton key={i} className="h-6" />
@@ -62,17 +131,17 @@ export function UsageOverviewPanel() {
            <MetricRow
              icon={<Zap className="h-4 w-4" />}
              label="Tokens (input)"
-              value={summary ? fmt(summary.tokens_input) : "—"}
+              value={tokensInput != null ? fmt(tokensInput) : "—"}
            />
            <MetricRow
              icon={<Zap className="h-4 w-4 text-muted-foreground" />}
              label="Tokens (output)"
-              value={summary ? fmt(summary.tokens_output) : "—"}
+              value={tokensOutput != null ? fmt(tokensOutput) : "—"}
            />
            <MetricRow
              icon={<Coins className="h-4 w-4" />}
              label="Total cost"
-              value={summary ? fmtCost(summary.total_cost_usd) : "—"}
+              value={totalCost != null ? fmtCost(totalCost) : "—"}
            />
            <MetricRow
              icon={
@@ -83,7 +152,7 @@ export function UsageOverviewPanel() {
                )
              }
              label="Trend vs prior period"
-              value={summary ? (trendUp ? "+" : "") + summary.trend_pct.toFixed(1) + "%" : "—"}
+              value={summary ? (trendUp ? "+" : "") + trendPct.toFixed(1) + "%" : "—"}
              sub={
                summary ? (
                  <span className={"text-xs " + (trendUp ? "text-red-500" : "text-green-500")}>
@@ -95,7 +164,7 @@ export function UsageOverviewPanel() {
            <MetricRow
              icon={<Activity className="h-4 w-4 text-blue-500" />}
              label="Period"
-              value={summary?.period ?? "—"}
+              value={periodLabel ?? "—"}
            />
          </div>
        )}
@@ -3,13 +3,24 @@
 import { useEffect, useRef } from "react";
 import { useWebSocket } from "./use-websocket";
 import { useRateLimitStore } from "@/store/rate-limit-store";
+import { useUsageStore } from "@/store/usage-store";
 import type { RateLimitHitEvent, RateLimitLiftedEvent } from "@/types/rate-limits";

-interface RateLimitWsMessage {
+/**
+ * Unified shape for all messages arriving on the /ws/system endpoint.
+ * Fields are optional because different message types use different subsets.
+ */
+interface SystemWsMessage {
  type: string;
+  // Rate-limit fields (RATE_LIMIT_HIT / RATE_LIMIT_LIFTED)
  provider?: string;
  affectedAgents?: string[];
  retryAfterSeconds?: number;
+  // Usage fields (USAGE_SNAPSHOT — aggregate token/cost across active agents)
+  totals?: { input_tokens?: number; output_tokens?: number };
+  cost_estimate?: number;
+  period?: string;
+  // Shared
  timestamp?: string;
 }

@@ -19,9 +30,18 @@ interface UseRateLimitWebSocketOptions {
 }

 /**
- * Subscribes to RATE_LIMIT_HIT and RATE_LIMIT_LIFTED WebSocket events and
- * dispatches them to the useRateLimitStore. Accepts an optional onReconnect
- * callback that fires when the connection recovers from a reconnecting state.
+ * Subscribes to the /ws/system WebSocket (single shared instance mounted in
+ * RateLimitBanner). Handles:
+ *
+ *  - RATE_LIMIT_HIT / RATE_LIMIT_LIFTED → dispatched to useRateLimitStore
+ *  - USAGE_SNAPSHOT                     → dispatched to useUsageStore
+ *
+ * Also syncs the live WebSocket connection state into useUsageStore so that
+ * other components (e.g. UsageOverviewPanel) can read it without creating a
+ * second /ws/system connection.
+ *
+ * Accepts an optional onReconnect callback that fires when the connection
+ * recovers from a reconnecting state.
 */
 export function useRateLimitWebSocket(options: UseRateLimitWebSocketOptions = {}) {
  const { onReconnect } = options;
@@ -30,12 +50,18 @@ export function useRateLimitWebSocket(options: UseRateLimitWebSocketOptions = {}
  // getWebSocketUrl() already supplies the "/ws" base, so the endpoint is just
  // the path (matching the agents/channels/notifications hooks). Passing
  // "/ws/system" here produced the doubled "/ws/ws/system" URL.
-  const { state, lastMessage } = useWebSocket<RateLimitWsMessage>(
+  const { state, lastMessage } = useWebSocket<SystemWsMessage>(
    "/system",
    undefined,
    true
  );

+  // Sync WS connection state into useUsageStore for cross-component visibility.
+  // This is the ONLY place wsState is written; no second useWebSocket call is needed.
+  useEffect(() => {
+    useUsageStore.getState().setWsState(state);
+  }, [state]);
+
  // Fire onReconnect when state transitions from reconnecting → connected
  useEffect(() => {
    if (prevStateRef.current === "reconnecting" && state === "connected") {
@@ -49,6 +75,7 @@ export function useRateLimitWebSocket(options: UseRateLimitWebSocketOptions = {}
    if (!lastMessage) return;

    const { hitRateLimit, liftRateLimit } = useRateLimitStore.getState();
+    const { setUsageData } = useUsageStore.getState();

    if (lastMessage.type === "RATE_LIMIT_HIT") {
      const event: RateLimitHitEvent = {
@@ -66,6 +93,14 @@ export function useRateLimitWebSocket(options: UseRateLimitWebSocketOptions = {}
        timestamp: lastMessage.timestamp ?? new Date().toISOString(),
      };
      liftRateLimit(event);
+    } else if (lastMessage.type === "USAGE_SNAPSHOT") {
+      setUsageData({
+        tokens_input: lastMessage.totals?.input_tokens ?? 0,
+        tokens_output: lastMessage.totals?.output_tokens ?? 0,
+        total_cost_usd: lastMessage.cost_estimate ?? 0,
+        period: lastMessage.period ?? "live",
+        timestamp: lastMessage.timestamp,
+      });
    }
  }, [lastMessage]);

@@ -240,15 +240,14 @@ export const usageApi = {
  },

  /**
-   * Recent inference sessions — mock-mode only.
-   *
-   * The backend has no /usage/sessions endpoint.  In production this
-   * returns an empty array so SessionsTable shows a graceful "no data"
-   * state instead of throwing a 404.
+   * Recent spawn sessions — the raw per-session rows behind the aggregate
+   * panels, served by GET /usage/sessions.
   */
-  // eslint-disable-next-line @typescript-eslint/no-unused-vars
-  getUsageSessions: async (_limit: number = 100): Promise<UsageSession[]> => {
+  getUsageSessions: async (limit: number = 100): Promise<UsageSession[]> => {
    if (isMockMode()) return mockSessions();
-    return [];
+    const { data } = await api.get<UsageSession[]>("/usage/sessions", {
+      params: { limit },
+    });
+    return data;
  },
 };
@@ -1,3 +1,5 @@
 export { useUIStore } from "./ui-store";
 export { useNotificationStore } from "./notifications-store";
 export { useRateLimitStore } from "./rate-limit-store";
+export { useUsageStore } from "./usage-store";
+export type { UsageData } from "./usage-store";
@@ -0,0 +1,44 @@
+import { create } from "zustand";
+import type { ConnectionState } from "@/lib/websocket/connection";
+
+/**
+ * Live token/cost usage pushed over the /ws/system stream via USAGE_SNAPSHOT
+ * events. Mirrors the fields the UsageOverviewPanel renders, so the panel can
+ * swap polling for live data without reshaping anything.
+ */
+export interface UsageData {
+  /** Cumulative input tokens across currently-active agents. */
+  tokens_input: number;
+  /** Cumulative output tokens across currently-active agents. */
+  tokens_output: number;
+  /** Estimated USD cost for the snapshot. */
+  total_cost_usd: number;
+  /** Period label for the snapshot (e.g. "live"). */
+  period: string;
+  /** ISO timestamp of the snapshot. */
+  timestamp?: string;
+}
+
+interface UsageState {
+  /** Most-recent usage data received over WebSocket; null until first message arrives */
+  usageData: UsageData | null;
+  /** Current /ws/system WebSocket connection state, synced by useRateLimitWebSocket */
+  wsState: ConnectionState;
+
+  // Actions
+  /** Overwrite usageData with the latest payload from the WebSocket */
+  setUsageData: (data: UsageData) => void;
+  /** Clear usageData (e.g. on intentional disconnect or reset) */
+  clearUsageData: () => void;
+  /** Update the cached WebSocket connection state */
+  setWsState: (state: ConnectionState) => void;
+}
+
+export const useUsageStore = create<UsageState>((set) => ({
+  usageData: null,
+  wsState: "disconnected",
+
+  setUsageData: (data) => set({ usageData: data }),
+  clearUsageData: () => set({ usageData: null }),
+  setWsState: (state) => set({ wsState: state }),
+}));