External/roboco
1
0
Fork 0
mirror of https://github.com/rennf93/roboco.git synced 2026-08-03 07:23:24 +02:00
Code Issues Packages Projects Releases Wiki Activity

[4865ff8b] Add WebSocket support to the usage dashboard (#115)

Browse Source 
* [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
2026-06-11 23:19:50 +02:00
committed by GitHub
co-authored by Frontend Developer 1 Backend Developer 1 Renn F
parent 1d1ec88aad
commit 547fe444f2
22 changed files with 1300 additions and 187 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CHANGELOG.md
+4 -2
View File
@@ -23,8 +23,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
and daily rollups, with provider-aware pricing (Anthropic models priced;
local/Ollama models intentionally $0). Visible on the usage dashboard.
- **`/ws/system` operator WebSocket stream** with a `websocket_bridge` that
forwards system events (the rate-limit lifecycle) from the event bus to panel
clients in real time.
forwards system events from the event bus to panel clients in real time — the
rate-limit lifecycle and live token/cost usage (`USAGE_UPDATE` /
`USAGE_SNAPSHOT`), so the dashboard's "Token Usage & Cost" panel updates over
the socket and falls back to HTTP polling when it drops.
### Fixed
CLAUDE.md
+5 -2
View File
@@ -472,7 +472,7 @@ The orchestrator exposes WebSocket endpoints under `/ws` (router in
| Endpoint | Purpose |
|----------|---------|
| `/ws/channels/{id}`, `/ws/agents/{id}`, `/ws/sessions/{id}`, `/ws/notifications/{id}` | Per-resource live streams |
| `/ws/system` | Operator/system-wide stream (no per-agent keying) — currently the rate-limit lifecycle (`RATE_LIMIT_HIT` / `RATE_LIMIT_LIFTED`) |
| `/ws/system` | Operator/system-wide stream (no per-agent keying) — the rate-limit lifecycle (`RATE_LIMIT_HIT` / `RATE_LIMIT_LIFTED`) and live usage (`USAGE_UPDATE` / `USAGE_SNAPSHOT`, pushed to the usage dashboard) |
Server-side events reach these sockets through `roboco/api/websocket_bridge.py`,
which subscribes to the `StreamEventBus` and forwards each event to the matching
@@ -492,7 +492,10 @@ stand up a parallel endpoint or client stack.
via the SDK server's `/usage/sync` (hook → orchestrator finalize →
`agent_spawn_sessions` → `daily_usage_rollups` → dashboard). Cost uses
provider-aware pricing in `roboco/billing/pricing.py` (Anthropic priced;
local/Ollama intentionally `$0`).
local/Ollama intentionally `$0`). The token sweep also publishes
`USAGE_UPDATE`/`USAGE_SNAPSHOT` to `/ws/system`, so the dashboard's
"Token Usage & Cost" panel updates live and falls back to HTTP polling when
the stream is down.
### Startup Sequence
docs/rag/tools/api-endpoints.md
+1 -1
View File
@@ -90,7 +90,7 @@ Base URL: `http://{host}:{port}/api/v1`
| Method | Endpoint | Description |
|--------|----------|-------------|
| GET | `/api/system/rate-limits` | Active per-provider rate-limit state (`{ entries: [...] }`) |
| WS | `/ws/system` | Operator stream — rate-limit lifecycle (`RATE_LIMIT_HIT` / `RATE_LIMIT_LIFTED`) |
| WS | `/ws/system` | Operator stream — rate-limit lifecycle (`RATE_LIMIT_HIT` / `RATE_LIMIT_LIFTED`) and live usage (`USAGE_UPDATE` / `USAGE_SNAPSHOT`) pushed to the usage dashboard |
| WS | `/ws/agents/{id}`, `/ws/channels/{id}`, `/ws/sessions/{id}`, `/ws/notifications/{id}` | Per-resource live streams |
## Documentation
panel/src/components/dashboard/key-metrics-panel.tsx
+13 -16
View File
@@ -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}`,
},
];
panel/src/components/dashboard/usage-overview-panel.tsx
+82 -13
View File
@@ -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>
)}
panel/src/hooks/use-rate-limit-websocket.ts
+40 -5
View File
@@ -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]);
panel/src/lib/api/usage.ts
+7 -8
View File
@@ -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;
},
};
panel/src/store/index.ts
+2
View File
@@ -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";
panel/src/store/usage-store.ts
+44
View File
@@ -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 }),
}));
roboco/agent_sdk/server.py
+3 -31
View File
@@ -41,6 +41,9 @@ from roboco.agent_sdk.models import (
VerbAttemptRequest,
VerbCircuitStatus,
)
from roboco.agent_sdk.transcript_usage import (
sum_transcript_usage as _sum_transcript_usage,
)
from roboco.foundation.policy.agent_loop import DEFAULT_BUDGET as _BUDGET
from roboco.foundation.policy.agent_loop import retry_limit_for
from roboco.services.gateway.envelope import Envelope
@@ -741,37 +744,6 @@ def _token_usage_snapshot() -> TokenUsageStatus:
)
def _sum_transcript_usage(path: Path) -> tuple[int, int, int, int]:
"""Sum per-message token usage across a Claude Code JSONL transcript.
Each assistant entry carries a ``message.usage`` block with the token
counts for that API response; summing them yields the session total.
Returns ``(input, output, cache_read, cache_write)``. Malformed lines are
skipped — a single bad line must never lose the whole count.
"""
tin = tout = tcr = tcw = 0
with path.open("r", encoding="utf-8", errors="ignore") as fh:
for raw in fh:
stripped = raw.strip()
if not stripped:
continue
try:
entry = json.loads(stripped)
except (ValueError, TypeError):
continue
message = entry.get("message")
if not isinstance(message, dict):
continue
usage = message.get("usage")
if not isinstance(usage, dict):
continue
tin += int(usage.get("input_tokens", 0) or 0)
tout += int(usage.get("output_tokens", 0) or 0)
tcr += int(usage.get("cache_read_input_tokens", 0) or 0)
tcw += int(usage.get("cache_creation_input_tokens", 0) or 0)
return tin, tout, tcr, tcw
@app.post("/usage/sync", response_model=TokenUsageStatus)
async def usage_sync(req: TranscriptSyncRequest) -> TokenUsageStatus:
"""Parse the Claude Code transcript and *set* cumulative token totals.
roboco/agent_sdk/transcript_usage.py
+84
View File
@@ -0,0 +1,84 @@
"""Claude Code transcript token-usage parsing.
A dependency-light helper (only ``json`` + ``pathlib``) so callers that need
durable token counts — notably the orchestrator's session-finalization path —
can read them without importing the agent SDK server, which pulls in the
FastAPI / RAG (piragi / openai) stack.
"""
from __future__ import annotations
import json
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from pathlib import Path
def _coerce_int(value: Any) -> int:
"""Coerce a transcript usage value to int, treating null/garbage as zero."""
try:
return int(value)
except (TypeError, ValueError):
return 0
def _line_usage(line: str) -> tuple[str | None, tuple[int, int, int, int]] | None:
"""Parse one transcript line into ``(message_id, token deltas)``.
Returns ``None`` for blank lines, malformed JSON, or entries without a
``message.usage`` block. ``message_id`` lets the caller de-duplicate:
Claude Code logs a single assistant message as several lines (one per
content block — thinking, text, tool_use), each repeating the *same*
``usage``, so counting every line would multiply the totals.
"""
stripped = line.strip()
if not stripped:
return None
try:
entry = json.loads(stripped)
except (ValueError, TypeError):
return None
message = entry.get("message")
if not isinstance(message, dict):
return None
usage = message.get("usage")
if not isinstance(usage, dict):
return None
deltas = (
_coerce_int(usage.get("input_tokens")),
_coerce_int(usage.get("output_tokens")),
_coerce_int(usage.get("cache_read_input_tokens")),
_coerce_int(usage.get("cache_creation_input_tokens")),
)
return message.get("id"), deltas
def sum_transcript_usage(path: Path) -> tuple[int, int, int, int]:
"""Sum per-message token usage across a Claude Code JSONL transcript.
Each assistant message carries a ``message.usage`` block with the token
counts for that API response; summing them yields the session total.
Messages that span several lines (same ``message.id``) are counted once —
Claude Code emits one line per content block, each repeating the message's
usage, so naive summing roughly doubles the totals. Returns
``(input, output, cache_read, cache_write)``. Malformed lines are skipped —
a single bad line must never lose the whole count.
"""
tin = tout = tcr = tcw = 0
seen: set[str] = set()
with path.open("r", encoding="utf-8", errors="ignore") as fh:
for raw in fh:
parsed = _line_usage(raw)
if parsed is None:
continue
message_id, (line_in, line_out, line_cr, line_cw) = parsed
if message_id is not None:
if message_id in seen:
continue
seen.add(message_id)
tin += line_in
tout += line_out
tcr += line_cr
tcw += line_cw
return tin, tout, tcr, tcw
roboco/api/routes/usage.py
+21
View File
@@ -145,3 +145,24 @@ async def get_cache_efficiency(
"""
svc = get_usage_service(db)
return await svc.get_cache_efficiency(period)
# =============================================================================
# SESSIONS
# =============================================================================
@router.get("/sessions")
async def get_usage_sessions(
db: DbSession,
limit: Annotated[
int, Query(ge=1, le=200, description="Max sessions to return")
] = 50,
) -> list[dict[str, Any]]:
"""Return the most recent agent spawn sessions, newest first.
Each row carries per-session token totals (input / output / cache) and the
estimated cost — the raw rows behind the aggregate usage panels.
"""
svc = get_usage_service(db)
return await svc.get_recent_sessions(limit)
roboco/api/websocket_bridge.py
+28
View File
@@ -20,6 +20,11 @@ _RATE_LIMIT_WS_TYPES = {
EventType.RATE_LIMIT_LIFTED: "RATE_LIMIT_LIFTED",
}
_USAGE_WS_TYPES = {
EventType.USAGE_UPDATE: "USAGE_UPDATE",
EventType.USAGE_SNAPSHOT: "USAGE_SNAPSHOT",
}
# Handler for notification events
async def _handle_notification_sent(event: Event) -> None:
@@ -144,6 +149,25 @@ async def _handle_rate_limit_event(event: Event) -> None:
await manager.broadcast_system({"type": ws_type, **event.data})
async def _handle_usage_event(event: Event) -> None:
"""Forward USAGE_UPDATE/SNAPSHOT events to operator system WS clients.
Both event types carry all the fields the panel needs directly in
``event.data``; we tag them with the discriminating ``type`` string the
panel switches on (the same UPPER_SNAKE mapping the rate-limit handler
uses), so the panel can distinguish per-agent updates from aggregate
snapshots.
"""
ws_type = _USAGE_WS_TYPES.get(event.type)
if ws_type is None:
return
await manager.broadcast_system({"type": ws_type, **event.data})
logger.debug(
"Usage event forwarded to system WebSocket",
event_type=ws_type,
)
def register_websocket_bridge_handlers() -> None:
"""
Register event handlers that forward events to WebSocket clients.
@@ -172,6 +196,10 @@ def register_websocket_bridge_handlers() -> None:
bus.subscribe(EventType.RATE_LIMIT_HIT, _handle_rate_limit_event)
bus.subscribe(EventType.RATE_LIMIT_LIFTED, _handle_rate_limit_event)
# Usage events -> system WebSocket (panel dashboard)
bus.subscribe(EventType.USAGE_UPDATE, _handle_usage_event)
bus.subscribe(EventType.USAGE_SNAPSHOT, _handle_usage_event)
logger.info("WebSocket bridge handlers registered")
roboco/models/events.py
+4
View File
@@ -69,6 +69,10 @@ class EventType(StrEnum):
RATE_LIMIT_HIT = "rate_limit.hit"
RATE_LIMIT_LIFTED = "rate_limit.lifted"
# Usage events
USAGE_UPDATE = "usage.update"
USAGE_SNAPSHOT = "usage.snapshot"
# Question events
QUESTION_ASKED = "question.asked"
QUESTION_ANSWERED = "question.answered"
roboco/runtime/orchestrator.py
+239 -106
View File
@@ -3218,7 +3218,7 @@ class AgentOrchestrator:
fetch, which misses whenever the agent container is short-lived or
already torn down. Returns zeros when no transcript is found.
"""
from roboco.agent_sdk.server import _sum_transcript_usage
from roboco.agent_sdk.transcript_usage import sum_transcript_usage
projects = Path.home() / ".claude" / "projects"
try:
@@ -3228,12 +3228,49 @@ class AgentOrchestrator:
if d.is_dir()
for f in d.glob("*.jsonl")
]
if not jsonl:
return (0, 0, 0, 0)
newest = max(jsonl, key=lambda f: f.stat().st_mtime)
return sum_transcript_usage(newest)
except OSError:
return (0, 0, 0, 0)
if not jsonl:
return (0, 0, 0, 0)
newest = max(jsonl, key=lambda f: f.stat().st_mtime)
return _sum_transcript_usage(newest)
async def _resolve_final_token_usage(
self, agent_id: str
) -> tuple[int, int, int, int]:
"""Resolve final token counts for a stopping agent.
Tries the live SDK ``/usage/status`` first; if that misses — the SDK's
in-memory counts race container teardown for short-lived agents — it
falls back to the agent's Claude Code transcript, which is durable and
mounted into this container. Returns
``(input, output, cache_read, cache_write)``.
"""
tokens = (0, 0, 0, 0)
sdk_url = f"http://roboco-agent-{agent_id}:{SDK_PORT}/usage/status"
try:
async with httpx.AsyncClient(timeout=3.0) as client:
resp = await client.get(sdk_url)
if resp.status_code == http_status.HTTP_200_OK:
data = resp.json()
tokens = (
data.get("tokens_input", 0),
data.get("tokens_output", 0),
data.get("tokens_cache_read", 0),
data.get("tokens_cache_write", 0),
)
except Exception as sdk_exc:
logger.debug(
"Could not fetch final token counts from SDK",
agent_id=agent_id,
error=str(sdk_exc),
)
if not tokens[0] and not tokens[1]:
tin, tout, cr, cw = self._usage_from_transcript(agent_id)
if tin or tout:
tokens = (tin, tout, cr, cw)
return tokens
async def _finalize_spawn_session(
self,
@@ -3242,9 +3279,9 @@ class AgentOrchestrator:
) -> None:
"""Close the open agent_spawn_sessions row for this agent.
Fetches final token counts from the agent SDK's /usage/status endpoint,
calculates cost via pricing module, then updates the DB row with
ended_at, token totals, exit_reason, and estimated_cost_usd.
Resolves final token counts (live SDK, with a durable transcript
fallback), calculates cost via the pricing module, then updates the DB
row with ended_at, token totals, exit_reason, and estimated_cost_usd.
Errors are caught and logged — finalization must never block stop_agent.
"""
try:
@@ -3252,44 +3289,16 @@ class AgentOrchestrator:
from roboco.db.base import get_session_factory
from roboco.db.tables import AgentSpawnSessionTable
# Fetch final token counts from the agent's SDK
sdk_url = f"http://roboco-agent-{agent_id}:{SDK_PORT}/usage/status"
tokens_input = 0
tokens_output = 0
tokens_cache_read = 0
tokens_cache_write = 0
model = "unknown"
try:
async with httpx.AsyncClient(timeout=3.0) as client:
resp = await client.get(sdk_url)
if resp.status_code == http_status.HTTP_200_OK:
data = resp.json()
tokens_input = data.get("tokens_input", 0)
tokens_output = data.get("tokens_output", 0)
tokens_cache_read = data.get("tokens_cache_read", 0)
tokens_cache_write = data.get("tokens_cache_write", 0)
except Exception as sdk_exc:
logger.debug(
"Could not fetch final token counts from SDK",
agent_id=agent_id,
error=str(sdk_exc),
)
# The live SDK fetch above races the container teardown and misses
# for short-lived agents (counts live in the SDK server's memory,
# which dies with the container). Fall back to the durable source of
# truth: the agent's Claude Code transcript, mounted into this
# container — so usage is captured regardless of container timing.
if not tokens_input and not tokens_output:
tin, tout, cr, cw = self._usage_from_transcript(agent_id)
if tin or tout:
tokens_input = tin
tokens_output = tout
tokens_cache_read = cr
tokens_cache_write = cw
# Resolve final token counts (live SDK, with transcript fallback).
(
tokens_input,
tokens_output,
tokens_cache_read,
tokens_cache_write,
) = await self._resolve_final_token_usage(agent_id)
# Look up the model and usage_session_id from the running instance config.
model = "unknown"
instance = self._instances.get(agent_id)
if instance and instance.config:
model = instance.config.model or "unknown"
@@ -3357,6 +3366,113 @@ class AgentOrchestrator:
error=str(exc),
)
@staticmethod
async def _fetch_agent_tokens(
client: httpx.AsyncClient, agent_id: str
) -> tuple[int, int, int, int] | None:
"""Fetch cumulative token counts from an agent's SDK usage endpoint.
Returns ``(input, output, cache_read, cache_write)`` or ``None`` when the
agent returns a non-200 status or has not accrued any tokens yet.
"""
sdk_url = f"http://roboco-agent-{agent_id}:{SDK_PORT}/usage/status"
resp = await client.get(sdk_url)
if resp.status_code != http_status.HTTP_200_OK:
return None
data = resp.json()
tokens = (
data.get("tokens_input", 0),
data.get("tokens_output", 0),
data.get("tokens_cache_read", 0),
data.get("tokens_cache_write", 0),
)
if sum(tokens) == 0:
return None
return tokens
async def _resolve_active_tokens(
self, client: httpx.AsyncClient, agent_id: str
) -> tuple[int, int, int, int] | None:
"""Resolve live token counts for an active agent.
Tries the agent SDK's ``/usage/status`` first; on a zero/miss falls
back to the durable transcript (the SDK can report zero mid-run, the
same race the finalize path handles). Returns ``None`` when neither
source has any usage yet.
"""
tokens = await self._fetch_agent_tokens(client, agent_id)
if tokens is not None:
return tokens
transcript = self._usage_from_transcript(agent_id)
return transcript if any(transcript) else None
@staticmethod
async def _persist_token_snapshot(
session_factory: Any,
agent_id: str,
instance: AgentInstance,
tokens: tuple[int, int, int, int],
) -> bool:
"""Insert a token_usage_snapshots row and refresh the open session totals.
Returns True when a snapshot was written; False when the agent has no
open spawn-session row to attach it to.
"""
from uuid import uuid4
from sqlalchemy import select, update
from roboco.db.tables import AgentSpawnSessionTable, TokenUsageSnapshotTable
tokens_input, tokens_output, tokens_cache_read, tokens_cache_write = tokens
async with session_factory() as db:
# Prefer a direct lookup by the session UUID captured at spawn time;
# fall back to the agent_slug heuristic for instances that pre-date
# the usage_session_id field.
if instance.usage_session_id is not None:
result = await db.execute(
select(AgentSpawnSessionTable).where(
AgentSpawnSessionTable.id == instance.usage_session_id
)
)
else:
result = await db.execute(
select(AgentSpawnSessionTable)
.where(
AgentSpawnSessionTable.agent_slug == agent_id,
AgentSpawnSessionTable.ended_at.is_(None),
)
.order_by(AgentSpawnSessionTable.started_at.desc())
.limit(1)
)
session_row = result.scalar_one_or_none()
if session_row is None:
return False
db.add(
TokenUsageSnapshotTable(
id=uuid4(),
agent_spawn_session_id=session_row.id,
snapshotted_at=datetime.now(UTC),
tokens_input=tokens_input,
tokens_output=tokens_output,
tokens_cache_read=tokens_cache_read,
tokens_cache_write=tokens_cache_write,
)
)
await db.execute(
update(AgentSpawnSessionTable)
.where(AgentSpawnSessionTable.id == session_row.id)
.values(
tokens_input=tokens_input,
tokens_output=tokens_output,
tokens_cache_read=tokens_cache_read,
tokens_cache_write=tokens_cache_write,
)
)
await db.commit()
return True
async def _sweep_token_snapshots(self) -> None:
"""Write a token_usage_snapshots row for each active agent with non-zero tokens.
@@ -3364,18 +3480,26 @@ class AgentOrchestrator:
token counts on the open agent_spawn_sessions row so the DB reflects
current progress without waiting for session close.
Errors per-agent are caught so one bad agent doesn't abort the whole sweep.
Additionally publishes USAGE_UPDATE events per agent (throttled to at most
one per 5-second window) and a USAGE_SNAPSHOT aggregate after the loop.
"""
if not self._instances:
return
try:
from roboco.db.base import get_session_factory
from roboco.db.tables import AgentSpawnSessionTable, TokenUsageSnapshotTable
except ImportError:
return
session_factory = get_session_factory()
# Accumulators for the post-loop USAGE_SNAPSHOT event.
_usage_by_agent: list[dict[str, Any]] = []
_usage_total_input = 0
_usage_total_output = 0
_usage_total_cost = 0.0
async with httpx.AsyncClient(timeout=3.0) as client:
for agent_id, instance in list(self._instances.items()):
if instance.state not in (
@@ -3384,80 +3508,60 @@ class AgentOrchestrator:
):
continue
sdk_url = f"http://roboco-agent-{agent_id}:{SDK_PORT}/usage/status"
try:
resp = await client.get(sdk_url)
if resp.status_code != http_status.HTTP_200_OK:
tokens = await self._resolve_active_tokens(client, agent_id)
if tokens is None:
continue
data = resp.json()
tokens_input = data.get("tokens_input", 0)
tokens_output = data.get("tokens_output", 0)
tokens_cache_read = data.get("tokens_cache_read", 0)
tokens_cache_write = data.get("tokens_cache_write", 0)
# Skip agents with no token usage yet
total = (
tokens_input
+ tokens_output
+ tokens_cache_read
+ tokens_cache_write
persisted = await self._persist_token_snapshot(
session_factory, agent_id, instance, tokens
)
if total == 0:
if not persisted:
continue
async with session_factory() as db:
from sqlalchemy import select, update
tokens_input, tokens_output = tokens[0], tokens[1]
model = instance.config.model if instance.config else "unknown"
# Prefer a direct lookup by the session UUID captured at
# spawn time; fall back to the agent_slug heuristic for
# instances that pre-date the usage_session_id field.
if instance.usage_session_id is not None:
result = await db.execute(
select(AgentSpawnSessionTable).where(
AgentSpawnSessionTable.id
== instance.usage_session_id
)
)
else:
result = await db.execute(
select(AgentSpawnSessionTable)
.where(
AgentSpawnSessionTable.agent_slug == agent_id,
AgentSpawnSessionTable.ended_at.is_(None),
)
.order_by(AgentSpawnSessionTable.started_at.desc())
.limit(1)
)
session_row = result.scalar_one_or_none()
if session_row is None:
continue
# Publish USAGE_UPDATE event for this agent (throttled).
with contextlib.suppress(Exception):
from roboco.events import get_event_bus
from roboco.services.usage_events import (
UsageUpdate,
publish_usage_update,
)
# Insert snapshot
from uuid import uuid4 as _uuid4
await publish_usage_update(
get_event_bus(),
UsageUpdate(
agent_id=agent_id,
task_id=instance.current_task_id,
input_tokens=tokens_input,
output_tokens=tokens_output,
model=model,
),
)
snapshot = TokenUsageSnapshotTable(
id=_uuid4(),
agent_spawn_session_id=session_row.id,
snapshotted_at=datetime.now(UTC),
# Accumulate per-agent data for the aggregate snapshot.
with contextlib.suppress(Exception):
from roboco.billing.pricing import calculate_cost
agent_cost = calculate_cost(
model=model,
tokens_input=tokens_input,
tokens_output=tokens_output,
tokens_cache_read=tokens_cache_read,
tokens_cache_write=tokens_cache_write,
)
db.add(snapshot)
# Update cumulative totals on the session row
await db.execute(
update(AgentSpawnSessionTable)
.where(AgentSpawnSessionTable.id == session_row.id)
.values(
tokens_input=tokens_input,
tokens_output=tokens_output,
tokens_cache_read=tokens_cache_read,
tokens_cache_write=tokens_cache_write,
)
_usage_by_agent.append(
{
"agent_id": agent_id,
"input_tokens": tokens_input,
"output_tokens": tokens_output,
"model": model,
"cost_estimate": agent_cost,
}
)
await db.commit()
_usage_total_input += tokens_input
_usage_total_output += tokens_output
_usage_total_cost += agent_cost
except Exception as agent_exc:
logger.debug(
@@ -3466,6 +3570,28 @@ class AgentOrchestrator:
error=str(agent_exc),
)
# Publish a USAGE_SNAPSHOT aggregate if any active agents had token data.
if _usage_by_agent:
with contextlib.suppress(Exception):
from roboco.events import get_event_bus
from roboco.services.usage_events import (
UsageSnapshot,
publish_usage_snapshot,
)
await publish_usage_snapshot(
get_event_bus(),
UsageSnapshot(
period="live",
totals={
"input_tokens": _usage_total_input,
"output_tokens": _usage_total_output,
},
cost_estimate=_usage_total_cost,
by_agent=_usage_by_agent,
),
)
async def _sweep_daily_rollup(self) -> None:
"""Upsert daily_usage_rollups from closed agent_spawn_sessions.
@@ -3910,6 +4036,13 @@ Start by:
container_id=cid,
exit_code=exit_code,
)
# The agent self-exited (a graceful i_am_idle shutdown, or a crash), so
# stop_agent() — which normally finalizes — was never called. Finalize
# here to capture token usage from the transcript; otherwise the
# spawn-session row is left open with zero tokens.
await self._finalize_spawn_session(
agent_id, exit_reason="completed" if graceful else "crashed"
)
instance.state = AgentState.OFFLINE
instance.container_id = None
if graceful:
roboco/services/usage.py
+30
View File
@@ -35,6 +35,23 @@ def _row_tokens(row: Any) -> tuple[int, int, int, int]:
)
def _session_row(row: Any) -> dict[str, Any]:
"""Shape one spawn-session row for the dashboard's sessions table."""
tin, tout, tcr, tcw = _row_tokens(row)
return {
"id": str(row.id),
"agent_slug": row.agent_slug,
"model": row.model,
"started_at": row.started_at.isoformat(),
"ended_at": row.ended_at.isoformat() if row.ended_at else None,
"tokens_input": tin,
"tokens_output": tout,
"tokens_cache": tcr + tcw,
"total_tokens": tin + tout + tcr + tcw,
"cost": float(row.estimated_cost_usd or 0.0),
}
def _parse_period(period: str) -> tuple[datetime, int]:
"""Parse period string into (start_dt, hours).
@@ -441,6 +458,19 @@ class UsageService(BaseService):
"cost_today_usd": round(float(row.total_cost_usd or 0.0), 6),
}
async def get_recent_sessions(self, limit: int = 50) -> list[dict[str, Any]]:
"""Return the most recent spawn sessions, newest first.
These are the raw per-session rows behind the aggregate panels — the
dashboard's "Recent Sessions" table.
"""
result = await self.session.execute(
select(AgentSpawnSessionTable)
.order_by(AgentSpawnSessionTable.started_at.desc())
.limit(limit)
)
return [_session_row(row) for row in result.scalars().all()]
def get_usage_service(db: AsyncSession) -> UsageService:
"""Factory function matching the pattern used by other services."""
roboco/services/usage_events.py
+129
View File
@@ -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()))
tests/unit/agent_sdk/test_usage_sync.py
+48
View File
@@ -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"],
)
tests/unit/api/test_websocket_bridge.py
+71 -2
View File
@@ -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
tests/unit/runtime/test_orchestrator_write_hooks.py
+108 -1
View File
@@ -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()
tests/unit/services/test_usage.py
+75
View File
@@ -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
tests/unit/services/test_usage_events.py
+262
View File
@@ -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