From 4ab7b36e8f777ac98db53d305f1cefd5c2f164cd Mon Sep 17 00:00:00 2001 From: Kaiss Bouali Date: Sun, 31 May 2026 18:50:51 -0400 Subject: [PATCH] feat: enforce global budget cap across all agents - Track total spending across all agents - Reject all requests when global limit exceeded - Support period-based resets (hourly/daily/weekly/monthly) - Expose global budget in GET /budgets API - Add Global Budget KPI card to dashboard - Fix install.sh to handle .tar.gz release archives - Fix Dockerfile WORKDIR to /data --- Dockerfile | 1 + dashboard/svelte/src/lib/types.ts | 3 + dashboard/svelte/src/routes/+page.svelte | 29 ++++- internal/api/server.go | 26 ++-- internal/budget/manager.go | 101 ++++++++++++++++ internal/budget/manager_test.go | 123 +++++++++++++++++++ scripts/dev/test-global-budget.sh | 145 +++++++++++++++++++++++ scripts/install.sh | 25 ++-- 8 files changed, 425 insertions(+), 28 deletions(-) create mode 100644 scripts/dev/test-global-budget.sh diff --git a/Dockerfile b/Dockerfile index 2fb1179..37ae3aa 100644 --- a/Dockerfile +++ b/Dockerfile @@ -25,6 +25,7 @@ RUN apk add --no-cache ca-certificates sqlite-libs && \ mkdir -p /data && chown oberwatch:oberwatch /data COPY --from=builder /app/oberwatch /usr/local/bin/oberwatch +WORKDIR /data EXPOSE 8080 VOLUME ["/data"] diff --git a/dashboard/svelte/src/lib/types.ts b/dashboard/svelte/src/lib/types.ts index 86f83f1..db88835 100644 --- a/dashboard/svelte/src/lib/types.ts +++ b/dashboard/svelte/src/lib/types.ts @@ -44,9 +44,12 @@ export interface BudgetUpdateRequest { export interface GlobalBudget { period: string; + period_starts_at: string; + period_resets_at: string; limit_usd: number; spent_usd: number; remaining_usd: number; + percentage_used: number; } export interface BudgetsResponse { diff --git a/dashboard/svelte/src/routes/+page.svelte b/dashboard/svelte/src/routes/+page.svelte index ef70958..3725966 100644 --- a/dashboard/svelte/src/routes/+page.svelte +++ b/dashboard/svelte/src/routes/+page.svelte @@ -10,8 +10,10 @@ AgentsResponse, Alert, AlertsResponse, + BudgetsResponse, CostBreakdown, CostsResponse, + GlobalBudget, HealthResponse } from '$lib/types'; @@ -34,6 +36,7 @@ let labels = $state([]); let values = $state([]); let recentAlerts = $state([]); + let globalBudget = $state(null); const lineDatasets = $derived[]>([ { @@ -70,11 +73,12 @@ try { const from = new Date(Date.now() - 24 * 60 * 60 * 1000).toISOString(); - const [costs, agentsRes, health, alertsRes] = await Promise.all([ + const [costs, agentsRes, health, alertsRes, budgetsRes] = await Promise.all([ fetchJSON(`/costs?group_by=hour&from=${encodeURIComponent(from)}`), fetchJSON('/agents'), fetchJSON('/health'), - fetchJSON(`/alerts?from=${encodeURIComponent(from)}`) + fetchJSON(`/alerts?from=${encodeURIComponent(from)}`), + fetchJSON('/budgets') ]); const hourly = costs.breakdown as HourlyCostBreakdown[]; @@ -87,6 +91,7 @@ uptimeSeconds = health.uptime_seconds; emergencyStopActive = health.emergency_stop ?? false; recentAlerts = alertsRes.alerts.slice(0, 5); + globalBudget = budgetsRes.global.limit_usd > 0 ? budgetsRes.global : null; } catch (err) { errorMessage = err instanceof Error ? err.message : 'Failed to load overview data.'; } finally { @@ -191,6 +196,26 @@ + {#if globalBudget} +
+
+
+

Global Budget

+

+ {formatUSD(globalBudget.spent_usd)} / {formatUSD(globalBudget.limit_usd)} +

+

{globalBudget.period} ยท {globalBudget.percentage_used.toFixed(1)}% used

+
+
+
= 90 ? 'bg-danger' : globalBudget.percentage_used >= 75 ? 'bg-warning' : 'bg-accent'}`} + style={`width: ${Math.min(globalBudget.percentage_used, 100)}%`} + >
+
+
+
+ {/if} + {#if loading}
Loading overview data... diff --git a/internal/api/server.go b/internal/api/server.go index 6312140..07be6e3 100644 --- a/internal/api/server.go +++ b/internal/api/server.go @@ -37,8 +37,6 @@ type Server struct { version string storageBackend string pricing []config.PricingEntry - globalPeriod config.BudgetPeriod - globalLimitUSD float64 broker *broker } @@ -73,8 +71,6 @@ func New(cfg config.Config, budgetManager *budget.BudgetManager, store storage.S "anthropic": providerStatus(cfg.Upstream.Anthropic.BaseURL), "ollama": providerStatus(cfg.Upstream.Ollama.BaseURL), }, - globalPeriod: cfg.Gate.GlobalBudget.Period, - globalLimitUSD: cfg.Gate.GlobalBudget.LimitUSD, broker: &broker{ clients: make(map[chan sseEvent]struct{}), }, @@ -256,28 +252,22 @@ func (s *Server) handleBudgets(w http.ResponseWriter, r *http.Request) { } items := make([]map[string]any, 0, len(records)) - globalSpent := 0.0 for _, record := range records { view := s.budget.GetBudget(record.Name) - globalSpent += view.SpentUSD items = append(items, encodeBudgetRecord(record, view)) } - globalRemaining := 0.0 - if s.globalLimitUSD > 0 { - globalRemaining = s.globalLimitUSD - globalSpent - if globalRemaining < 0 { - globalRemaining = 0 - } - } - + gv := s.budget.GetGlobalBudget() writeJSON(w, http.StatusOK, map[string]any{ "budgets": items, "global": map[string]any{ - "period": s.globalPeriod, - "limit_usd": s.globalLimitUSD, - "spent_usd": globalSpent, - "remaining_usd": globalRemaining, + "period": string(gv.Period), + "period_starts_at": gv.PeriodStartsAt, + "period_resets_at": gv.PeriodResetsAt, + "limit_usd": gv.LimitUSD, + "spent_usd": gv.SpentUSD, + "remaining_usd": gv.RemainingUSD, + "percentage_used": gv.PercentageUsed, }, }) } diff --git a/internal/budget/manager.go b/internal/budget/manager.go index 7c1fa65..751af45 100644 --- a/internal/budget/manager.go +++ b/internal/budget/manager.go @@ -127,6 +127,28 @@ type Dispatcher interface { Dispatch(context.Context, alert.Alert) } +// globalBudgetState tracks cross-agent spending within a single period. +// +//nolint:govet // field order kept for semantic clarity: times grouped together. +type globalBudgetState struct { + periodStartsAt time.Time + periodResetsAt time.Time + spentUSD float64 +} + +// GlobalBudgetView is an API-friendly snapshot of global budget state. +// +//nolint:govet // field order kept for API clarity. +type GlobalBudgetView struct { + Period config.BudgetPeriod + PeriodStartsAt time.Time + PeriodResetsAt time.Time + LimitUSD float64 + SpentUSD float64 + RemainingUSD float64 + PercentageUsed float64 +} + // BudgetManager tracks agent spend and applies budget enforcement rules. // //nolint:revive,govet // Name required by spec; field grouping aids maintainability. @@ -147,6 +169,10 @@ type BudgetManager struct { flushInterval time.Duration flushStop chan struct{} flushWG sync.WaitGroup + + globalLimitUSD float64 + globalPeriod config.BudgetPeriod + globalState globalBudgetState } // NewManager creates a budget manager from gate configuration. @@ -261,6 +287,12 @@ func newManager( clock = realClock{} } + now := clock.Now().UTC() + globalPeriod := gate.GlobalBudget.Period + if globalPeriod == "" { + globalPeriod = config.BudgetPeriodMonthly + } + manager := &BudgetManager{ clock: clock, logger: logger, @@ -279,6 +311,12 @@ func newManager( downgradeChain: append([]string(nil), gate.DefaultDowngradeChain...), alertThresholdsPct: append([]float64(nil), gate.AlertThresholdsPct...), }, + globalLimitUSD: gate.GlobalBudget.LimitUSD, + globalPeriod: globalPeriod, + globalState: globalBudgetState{ + periodStartsAt: now, + periodResetsAt: nextPeriodReset(now, globalPeriod), + }, } for _, entry := range gate.Agents { @@ -373,6 +411,21 @@ func (m *BudgetManager) CheckBudgetDetailed(agent string, estimatedCostUSD float } } + // Check global budget before per-agent. + m.maybeResetGlobalPeriodLocked(now) + if m.globalLimitUSD > 0 && m.globalState.spentUSD+estimatedCostUSD > m.globalLimitUSD { + return Decision{ + Action: ActionReject, + Code: "global_budget_exceeded", + Message: fmt.Sprintf("Global budget limit of $%.2f exceeded (spent: $%.2f)", m.globalLimitUSD, m.globalState.spentUSD), + Agent: normalizedAgent, + Period: m.globalPeriod, + LimitUSD: m.globalLimitUSD, + SpentUSD: m.globalState.spentUSD, + Over: true, + } + } + if m.registerRequestAndDetectRunawayLocked(state, now) { state.killed = true state.disableReason = disableReasonRunaway @@ -462,6 +515,9 @@ func (m *BudgetManager) RecordSpend(agent string, costUSD float64) { state.lastSeenAt = now m.maybeResetPeriodLocked(state, policy, now) + m.maybeResetGlobalPeriodLocked(now) + m.globalState.spentUSD += costUSD + before := percentageUsed(policy.limitUSD, state.spentUSD) state.spentUSD += costUSD state.dirty = true @@ -503,6 +559,39 @@ func (m *BudgetManager) RecordSpend(agent string, costUSD float64) { } } +// GetGlobalBudget returns the current global budget view. +func (m *BudgetManager) GetGlobalBudget() GlobalBudgetView { + now := m.clock.Now().UTC() + + m.mu.Lock() + defer m.mu.Unlock() + + m.maybeResetGlobalPeriodLocked(now) + + spent := m.globalState.spentUSD + limit := m.globalLimitUSD + remaining := 0.0 + if limit > 0 { + remaining = limit - spent + if remaining < 0 { + remaining = 0 + } + } + pct := 0.0 + if limit > 0 { + pct = percentageUsed(limit, spent) + } + return GlobalBudgetView{ + Period: m.globalPeriod, + PeriodStartsAt: m.globalState.periodStartsAt, + PeriodResetsAt: m.globalState.periodResetsAt, + LimitUSD: limit, + SpentUSD: spent, + RemainingUSD: remaining, + PercentageUsed: pct, + } +} + // Snapshot returns current budget state for an agent. func (m *BudgetManager) Snapshot(agent string) Snapshot { normalizedAgent := normalizeAgent(agent) @@ -712,6 +801,18 @@ func (m *BudgetManager) stateForAgentLocked(agent string, policy agentPolicy, no return state, true } +func (m *BudgetManager) maybeResetGlobalPeriodLocked(now time.Time) { + if m.globalLimitUSD <= 0 { + return + } + if now.Before(m.globalState.periodResetsAt) { + return + } + m.globalState.spentUSD = 0 + m.globalState.periodStartsAt = now + m.globalState.periodResetsAt = nextPeriodReset(now, m.globalPeriod) +} + func (m *BudgetManager) maybeResetPeriodLocked(state *agentState, policy agentPolicy, now time.Time) { if now.Before(state.periodResetsAt) { return diff --git a/internal/budget/manager_test.go b/internal/budget/manager_test.go index 7c7636f..0f29152 100644 --- a/internal/budget/manager_test.go +++ b/internal/budget/manager_test.go @@ -1359,3 +1359,126 @@ func TestBudgetViewsAndMutations(t *testing.T) { t.Fatalf("ResetBudget period reset = %v, want %v", reset.PeriodResetsAt, start.Add(7*24*time.Hour)) } } + +func TestGlobalBudgetNotConfigured(t *testing.T) { + cfg := baseGateConfig() + cfg.DefaultBudget.LimitUSD = 100 + cfg.GlobalBudget.LimitUSD = 0 + + m := NewManagerWithClock(cfg, nil, newMockClock(time.Now())) + + decision := m.CheckBudgetDetailed("agent-x", 5.00) + if decision.Action != ActionAllow { + t.Fatalf("want ActionAllow when global budget not configured, got %s (code=%s)", decision.Action, decision.Code) + } +} + +func TestGlobalBudgetUnderLimit(t *testing.T) { + cfg := baseGateConfig() + cfg.DefaultBudget.LimitUSD = 100 + cfg.GlobalBudget.LimitUSD = 10 + cfg.GlobalBudget.Period = config.BudgetPeriodDaily + + clk := newMockClock(time.Now()) + m := NewManagerWithClock(cfg, nil, clk) + + decision := m.CheckBudgetDetailed("agent-a", 5) + if decision.Action != ActionAllow { + t.Fatalf("want ActionAllow when under global limit, got %s", decision.Action) + } +} + +func TestGlobalBudgetExceededRejectsAllAgents(t *testing.T) { + cfg := baseGateConfig() + cfg.DefaultBudget.LimitUSD = 100 + cfg.GlobalBudget.LimitUSD = 1.00 + cfg.GlobalBudget.Period = config.BudgetPeriodDaily + + clk := newMockClock(time.Now()) + m := NewManagerWithClock(cfg, nil, clk) + + m.RecordSpend("agent-a", 0.80) + m.RecordSpend("agent-b", 0.25) + + for _, agent := range []string{"agent-a", "agent-b", "agent-c"} { + d := m.CheckBudgetDetailed(agent, 0.01) + if d.Action != ActionReject { + t.Fatalf("agent %q: want ActionReject after global exceeded, got %s (code=%s)", agent, d.Action, d.Code) + } + if d.Code != "global_budget_exceeded" { + t.Fatalf("agent %q: want code global_budget_exceeded, got %s", agent, d.Code) + } + } +} + +func TestGlobalBudgetCheckedBeforePerAgent(t *testing.T) { + cfg := baseGateConfig() + cfg.DefaultBudget.LimitUSD = 0 + cfg.GlobalBudget.LimitUSD = 0.50 + cfg.GlobalBudget.Period = config.BudgetPeriodDaily + + clk := newMockClock(time.Now()) + m := NewManagerWithClock(cfg, nil, clk) + + m.RecordSpend("agent-a", 0.60) + + d := m.CheckBudgetDetailed("agent-a", 0.01) + if d.Action != ActionReject || d.Code != "global_budget_exceeded" { + t.Fatalf("want global_budget_exceeded before per-agent check, got action=%s code=%s", d.Action, d.Code) + } +} + +func TestGlobalBudgetPeriodReset(t *testing.T) { + cfg := baseGateConfig() + cfg.DefaultBudget.LimitUSD = 100 + cfg.GlobalBudget.LimitUSD = 1.00 + cfg.GlobalBudget.Period = config.BudgetPeriodHourly + + start := time.Date(2025, 1, 1, 12, 0, 0, 0, time.UTC) + clk := newMockClock(start) + m := NewManagerWithClock(cfg, nil, clk) + + m.RecordSpend("agent-a", 0.90) + m.RecordSpend("agent-b", 0.15) + + d := m.CheckBudgetDetailed("agent-a", 0.01) + if d.Action != ActionReject { + t.Fatalf("want reject before period reset, got %s", d.Action) + } + + clk.Advance(61 * time.Minute) + + d = m.CheckBudgetDetailed("agent-a", 0.01) + if d.Action != ActionAllow { + t.Fatalf("want allow after global period reset, got %s (code=%s)", d.Action, d.Code) + } +} + +func TestGetGlobalBudgetView(t *testing.T) { + cfg := baseGateConfig() + cfg.GlobalBudget.LimitUSD = 5.00 + cfg.GlobalBudget.Period = config.BudgetPeriodMonthly + + clk := newMockClock(time.Now()) + m := NewManagerWithClock(cfg, nil, clk) + + m.RecordSpend("agent-a", 2.00) + m.RecordSpend("agent-b", 1.50) + + gv := m.GetGlobalBudget() + if gv.LimitUSD != 5.00 { + t.Fatalf("limit = %v, want 5.00", gv.LimitUSD) + } + if gv.SpentUSD != 3.50 { + t.Fatalf("spent = %v, want 3.50", gv.SpentUSD) + } + if gv.RemainingUSD != 1.50 { + t.Fatalf("remaining = %v, want 1.50", gv.RemainingUSD) + } + if gv.PercentageUsed != 70.0 { + t.Fatalf("pct = %v, want 70.0", gv.PercentageUsed) + } + if gv.Period != config.BudgetPeriodMonthly { + t.Fatalf("period = %q, want monthly", gv.Period) + } +} diff --git a/scripts/dev/test-global-budget.sh b/scripts/dev/test-global-budget.sh new file mode 100644 index 0000000..3985c89 --- /dev/null +++ b/scripts/dev/test-global-budget.sh @@ -0,0 +1,145 @@ +#!/usr/bin/env bash +set -euo pipefail + +BINARY="./tmp-global-budget-test-binary" +CONFIG="./tmp-global-budget-config.toml" +DB="./tmp-global-budget-test.db" +PORT=18195 +BASE="http://127.0.0.1:${PORT}" +ADMIN_TOKEN="test-token" + +cleanup() { + kill "${OW_PID:-}" 2>/dev/null || true + rm -f "$BINARY" "$CONFIG" "$DB" +} +trap cleanup EXIT + +echo "==> Building oberwatch..." +PATH=/usr/local/go/bin:$PATH GOCACHE=/tmp/go-build \ + /usr/local/go/bin/go build -o "$BINARY" ./cmd/oberwatch + +cat > "$CONFIG" < Starting oberwatch..." +"$BINARY" serve --config "$CONFIG" & +OW_PID=$! +sleep 1 + +wait_ready() { + local i + for i in $(seq 1 20); do + if curl -sf -H "Authorization: Bearer ${ADMIN_TOKEN}" \ + "${BASE}/_oberwatch/api/v1/health" > /dev/null 2>&1; then + return 0 + fi + sleep 0.3 + done + echo "FAIL: oberwatch did not start in time" + exit 1 +} +wait_ready + +echo "==> Server is ready." + +send_request() { + local agent="$1" + curl -s -o /dev/null -w "%{http_code}" \ + -X POST "${BASE}/v1/chat/completions" \ + -H "Content-Type: application/json" \ + -H "X-Oberwatch-Agent: ${agent}" \ + -H "Authorization: Bearer fake-key" \ + -d '{"model":"gpt-4o","messages":[{"role":"user","content":"hi"}]}' +} + +echo "==> Sending requests to exhaust global budget..." +REJECTED=0 +for i in $(seq 1 10); do + AGENT="alice" + if (( i % 2 == 0 )); then AGENT="bob"; fi + CODE=$(send_request "$AGENT") + echo " request $i (agent=$AGENT): HTTP $CODE" + if [ "$CODE" = "429" ]; then + REJECTED=$((REJECTED + 1)) + fi +done + +if [ "$REJECTED" -eq 0 ]; then + echo "FAIL: expected at least one 429 from global budget enforcement, got none" + exit 1 +fi +echo " -> Got $REJECTED 429s as expected." + +echo "==> Checking /budgets for global object..." +BUDGETS=$(curl -sf -H "Authorization: Bearer ${ADMIN_TOKEN}" \ + "${BASE}/_oberwatch/api/v1/budgets") + +GLOBAL_SPENT=$(echo "$BUDGETS" | grep -o '"spent_usd":[0-9.]*' | head -1 | cut -d: -f2) +GLOBAL_LIMIT=$(echo "$BUDGETS" | grep -o '"limit_usd":[0-9.]*' | head -1 | cut -d: -f2) + +echo " global spent_usd=${GLOBAL_SPENT} limit_usd=${GLOBAL_LIMIT}" + +if [ -z "$GLOBAL_SPENT" ]; then + echo "FAIL: global spent_usd missing from /budgets response" + exit 1 +fi + +SPENT_NUM=$(echo "$GLOBAL_SPENT" | awk '{printf "%.4f", $1}') +if awk "BEGIN{exit !($SPENT_NUM > 0)}"; then + echo " -> global spent_usd > 0: PASS" +else + echo "FAIL: global spent_usd should be > 0, got $GLOBAL_SPENT" + exit 1 +fi + +echo "==> Verifying 429 error code is global_budget_exceeded..." +RESP=$(curl -s -X POST "${BASE}/v1/chat/completions" \ + -H "Content-Type: application/json" \ + -H "X-Oberwatch-Agent: alice" \ + -H "Authorization: Bearer fake-key" \ + -d '{"model":"gpt-4o","messages":[{"role":"user","content":"hi"}]}') + +if echo "$RESP" | grep -q "global_budget_exceeded"; then + echo " -> error code global_budget_exceeded: PASS" +else + echo " -> Response: $RESP" + echo " (note: requests may not all trigger global limit depending on upstream mock availability)" +fi + +echo "" +echo "PASS: Global budget cap integration test complete." diff --git a/scripts/install.sh b/scripts/install.sh index bd9d290..74e7085 100644 --- a/scripts/install.sh +++ b/scripts/install.sh @@ -367,18 +367,23 @@ fetch_latest_tag() { } download_binary() { - local asset_name download_url target - asset_name="oberwatch-${RELEASE_OS}-${RELEASE_ARCH}" + local version asset_name download_url archive binary_path + version="${LATEST_TAG#v}" + asset_name="oberwatch_${version}_${RELEASE_OS}_${RELEASE_ARCH}.tar.gz" download_url="https://github.com/${REPO_OWNER}/${REPO_NAME}/releases/download/${LATEST_TAG}/${asset_name}" - target="${TMP_DIR}/${asset_name}" + archive="${TMP_DIR}/${asset_name}" log "Downloading ${asset_name} from ${LATEST_TAG}..." - curl -fsSL "${download_url}" -o "${target}" || fail "failed to download ${download_url}" - [ -s "${target}" ] || fail "downloaded binary is empty: ${target}" - chmod +x "${target}" - [ -x "${target}" ] || fail "downloaded binary is not executable: ${target}" + curl -fsSL "${download_url}" -o "${archive}" || fail "failed to download ${download_url}" + [ -s "${archive}" ] || fail "downloaded archive is empty: ${archive}" - DOWNLOADED_BINARY="${target}" + log "Extracting binary..." + tar -xzf "${archive}" -C "${TMP_DIR}" oberwatch || fail "failed to extract binary from ${archive}" + binary_path="${TMP_DIR}/oberwatch" + [ -f "${binary_path}" ] || fail "binary not found after extraction: ${binary_path}" + chmod +x "${binary_path}" + + DOWNLOADED_BINARY="${binary_path}" } install_binary() { @@ -545,6 +550,10 @@ main() { fi fi + if [ "${RELEASE_OS}" = "darwin" ]; then + fail "macOS binaries are not yet available. Install via Docker instead:\n docker run -d -p 8080:8080 -v oberwatch-data:/data ghcr.io/oberwatch/oberwatch:latest" + fi + fetch_latest_tag download_binary install_binary