From cf7b587f16a522600b465d136391f15725642fd4 Mon Sep 17 00:00:00 2001 From: core-be Date: Fri, 29 May 2026 03:25:12 -0700 Subject: [PATCH] feat(budget): multi-period per-workspace LLM budget (hourly/daily/weekly/monthly) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Extends the single monthly per-workspace budget to four independent ROLLING windows so a workspace can be capped per hour/day/week/month (#49 — gives the canvas Budget tab a real lever against runaway LLM spend, e.g. the reno-stars opus drain). SSOT design: - budget_periods.go = single source of truth: the period set + rolling windows, one FILTERed per-period spend query over the ledger, and the PURE parse/encode/exceededPeriods logic. Add a period = one line here. - migration: workspaces.budget_limits jsonb (canonical config, backfilled from the legacy monthly budget_limit) + workspace_spend_events ledger. - heartbeat (registry.go): derive the spend INCREMENT from the agent's existing cumulative report (delta vs prev; reset-aware) → ledger row. Server owns windowing; NO runtime change. - budget.go GET/PATCH: per-period limit/spend/remaining; accepts the new {budget_limits:{...}} shape AND the legacy {budget_limit} (→ monthly); legacy response fields still emitted + budget_limit kept synced (rollout back-compat). A limit of 0 = block-all (preserved); null/absent = no limit. - a2a_proxy.go checkWorkspaceBudget: 402 if ANY configured period's rolling window spend >= its limit; fail-open on DB error. - canvas BudgetSection: four period rows (USD limit input + spend/limit + bar). Tests: pure SSOT (parse/encode/exceededPeriods); GET/PATCH + multi-period + A2A enforcement (sqlmock, migrated to the new two-query flow); shared expectBudgetCheck helpers updated; canvas behavioral + per-period progress/aria. go build + vet + full handlers suite + migrations + canvas vitest all green. NOTE: the duplicate components/__tests__/BudgetSection.test.tsx (old single-limit UI) was repurposed to a focused per-period progress/aria suite — behavioral coverage now lives in tabs/__tests__/BudgetSection.test.tsx (one component, no parallel identical suites). Refs #49. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../__tests__/BudgetSection.test.tsx | 435 +++--------------- canvas/src/components/tabs/BudgetSection.tsx | 295 ++++++------ .../tabs/__tests__/BudgetSection.test.tsx | 275 ++++------- .../internal/handlers/a2a_proxy.go | 39 +- .../internal/handlers/a2a_queue_test.go | 8 +- workspace-server/internal/handlers/budget.go | 197 ++++---- .../internal/handlers/budget_periods.go | 160 +++++++ .../internal/handlers/budget_periods_test.go | 99 ++++ .../internal/handlers/budget_test.go | 280 +++++++---- .../internal/handlers/handlers_test.go | 8 +- .../internal/handlers/registry.go | 22 +- .../handlers/workspace_budget_test.go | 40 +- ...0000_workspace_multiperiod_budget.down.sql | 2 + ...000000_workspace_multiperiod_budget.up.sql | 30 ++ 14 files changed, 972 insertions(+), 918 deletions(-) create mode 100644 workspace-server/internal/handlers/budget_periods.go create mode 100644 workspace-server/internal/handlers/budget_periods_test.go create mode 100644 workspace-server/migrations/20260529000000_workspace_multiperiod_budget.down.sql create mode 100644 workspace-server/migrations/20260529000000_workspace_multiperiod_budget.up.sql diff --git a/canvas/src/components/__tests__/BudgetSection.test.tsx b/canvas/src/components/__tests__/BudgetSection.test.tsx index ed4778fa5..f259ca86d 100644 --- a/canvas/src/components/__tests__/BudgetSection.test.tsx +++ b/canvas/src/components/__tests__/BudgetSection.test.tsx @@ -1,411 +1,82 @@ // @vitest-environment jsdom /** - * Tests for BudgetSection (issue #541). + * Focused tests for BudgetSection's PER-PERIOD progress-bar math + aria (#49). * - * Covers: - * - Loading state - * - Stats row: used / limit, "Unlimited" when null - * - Progress bar: correct percentage, capped at 100%, absent when no limit - * - Budget remaining text - * - Input pre-fill (existing limit / blank when null) - * - Save: PATCH with number, PATCH with null (blank input) - * - 402 on GET → exceeded banner, no fetch-error text - * - 402 on PATCH → exceeded banner - * - Non-402 fetch error → error text - * - Non-402 save error → save error alert - * - Section header and subheading - * - Fetch error does not show stats + * Behavioral coverage (loading, save, 402 banners, USD formatting, legacy + * back-compat) lives in tabs/__tests__/BudgetSection.test.tsx — this file + * deliberately covers only the per-period progress percentage + aria-valuenow + * + the over-budget colouring, which that suite doesn't assert in detail. Kept + * separate to avoid duplicating the behavioral suite (one component, no + * parallel/identical suites). */ import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; -import { - render, - screen, - fireEvent, - waitFor, - cleanup, - act, -} from "@testing-library/react"; - -// ── Mock api ────────────────────────────────────────────────────────────────── +import { render, screen, waitFor, cleanup } from "@testing-library/react"; vi.mock("@/lib/api", () => ({ - api: { - get: vi.fn(), - patch: vi.fn(), - }, + api: { get: vi.fn(), patch: vi.fn() }, })); import { api } from "@/lib/api"; import { BudgetSection } from "../tabs/BudgetSection"; const mockGet = vi.mocked(api.get); -const mockPatch = vi.mocked(api.patch); -// ── Helpers ─────────────────────────────────────────────────────────────────── +type P = { limit: number | null; spend: number; remaining: number | null }; -function budgetResponse(overrides: Partial<{ - budget_limit: number | null; - budget_used: number; - budget_remaining: number | null; -}> = {}) { +// Build a periods response where the named period has the given limit/spend. +function withMonthly(limit: number | null, spend: number) { + const blank: P = { limit: null, spend: 0, remaining: null }; + const monthly: P = { limit, spend, remaining: limit == null ? null : limit - spend }; return { - budget_limit: 1000, - budget_used: 250, - budget_remaining: 750, - ...overrides, + periods: { hourly: blank, daily: blank, weekly: blank, monthly }, + budget_limit: limit, + monthly_spend: spend, + budget_remaining: monthly.remaining, }; } -function make402Error(): Error { - return new Error("API GET /workspaces/ws-1/budget: 402 Payment Required"); -} +beforeEach(() => vi.clearAllMocks()); +afterEach(() => cleanup()); -function make402PatchError(): Error { - return new Error("API PATCH /workspaces/ws-1/budget: 402 Payment Required"); -} - -function makeGenericError(msg = "network timeout"): Error { - return new Error(`API GET /workspaces/ws-1/budget: 500 ${msg}`); -} - -beforeEach(() => { - vi.clearAllMocks(); -}); - -afterEach(() => { - cleanup(); -}); - -// ── Rendering helpers ───────────────────────────────────────────────────────── - -async function renderLoaded(budgetData = budgetResponse()) { +async function renderLoaded(data: unknown) { // eslint-disable-next-line @typescript-eslint/no-explicit-any - mockGet.mockResolvedValueOnce(budgetData as any); + mockGet.mockResolvedValueOnce(data as any); render(); - // Wait for loading to finish await waitFor(() => expect(screen.queryByTestId("budget-loading")).toBeNull()); } -// ── Loading state ───────────────────────────────────────────────────────────── - -describe("BudgetSection — loading state", () => { - it("shows loading indicator while fetch is in flight", () => { - // Never resolve - mockGet.mockReturnValue(new Promise(() => {})); - render(); - expect(screen.getByTestId("budget-loading")).toBeTruthy(); - expect(screen.getByText("Loading…")).toBeTruthy(); +describe("BudgetSection — per-period progress bar", () => { + it("renders the bar for a limited period and omits it for an unlimited one", async () => { + await renderLoaded(withMonthly(1000, 250)); + expect(screen.getByTestId("budget-monthly-fill")).toBeTruthy(); + expect(screen.queryByTestId("budget-hourly-fill")).toBeNull(); // hourly unlimited }); - it("hides loading indicator after fetch resolves", async () => { - // eslint-disable-next-line @typescript-eslint/no-explicit-any - mockGet.mockResolvedValueOnce(budgetResponse() as any); - render(); - await waitFor(() => expect(screen.queryByTestId("budget-loading")).toBeNull()); - }); -}); - -// ── Section header ──────────────────────────────────────────────────────────── - -describe("BudgetSection — header and subheading", () => { - it("renders 'Budget' as the section heading", async () => { - await renderLoaded(); - expect(screen.getByText("Budget")).toBeTruthy(); - }); - - it("renders the subheading 'Limit total message credits for this workspace'", async () => { - await renderLoaded(); - expect( - screen.getByText("Limit total message credits for this workspace") - ).toBeTruthy(); - }); - - it("renders 'Budget limit (credits)' label for the input", async () => { - await renderLoaded(); - expect(screen.getByText("Budget limit (credits)")).toBeTruthy(); - }); -}); - -// ── Stats row ───────────────────────────────────────────────────────────────── - -describe("BudgetSection — stats row", () => { - it("shows budget_used in the stats row", async () => { - await renderLoaded(budgetResponse({ budget_used: 350, budget_limit: 1000 })); - expect(screen.getByTestId("budget-used-value").textContent).toBe("350"); - }); - - it("shows budget_limit in the stats row", async () => { - await renderLoaded(budgetResponse({ budget_used: 100, budget_limit: 500 })); - expect(screen.getByTestId("budget-limit-value").textContent).toBe("500"); - }); - - it("shows 'Unlimited' when budget_limit is null", async () => { - await renderLoaded(budgetResponse({ budget_limit: null, budget_remaining: null })); - expect(screen.getByTestId("budget-limit-value").textContent).toBe("Unlimited"); - }); - - it("shows budget_remaining when present", async () => { - await renderLoaded(budgetResponse({ budget_remaining: 750 })); - expect(screen.getByTestId("budget-remaining").textContent).toContain("750"); - expect(screen.getByTestId("budget-remaining").textContent).toContain("credits remaining"); - }); - - it("hides budget_remaining row when null", async () => { - await renderLoaded(budgetResponse({ budget_remaining: null })); - expect(screen.queryByTestId("budget-remaining")).toBeNull(); - }); - - it("does not crash when budget_used is missing from the response", async () => { - // Backend for a provisioning-stuck workspace may return a partial - // shape. Regression: previously this threw - // "Cannot read properties of undefined (reading 'toLocaleString')" - // and crashed the whole Details tab. - // eslint-disable-next-line @typescript-eslint/no-explicit-any - await renderLoaded({ budget_limit: 1000, budget_remaining: null } as any); - expect(screen.getByTestId("budget-used-value").textContent).toBe("0"); - }); -}); - -// ── Progress bar ────────────────────────────────────────────────────────────── - -describe("BudgetSection — progress bar", () => { - it("renders the progress bar when budget_limit is set", async () => { - await renderLoaded(budgetResponse({ budget_used: 250, budget_limit: 1000 })); - expect(screen.getByRole("progressbar")).toBeTruthy(); - }); - - it("does NOT render progress bar when budget_limit is null", async () => { - await renderLoaded(budgetResponse({ budget_limit: null, budget_remaining: null })); - expect(screen.queryByRole("progressbar")).toBeNull(); - }); - - it("fills to the correct percentage (25%)", async () => { - await renderLoaded(budgetResponse({ budget_used: 250, budget_limit: 1000 })); - const fill = screen.getByTestId("budget-progress-fill") as HTMLDivElement; - expect(fill.style.width).toBe("25%"); - }); - - it("fills to the correct percentage (50%)", async () => { - await renderLoaded(budgetResponse({ budget_used: 500, budget_limit: 1000 })); - const fill = screen.getByTestId("budget-progress-fill") as HTMLDivElement; - expect(fill.style.width).toBe("50%"); - }); - - it("caps fill at 100% when budget_used exceeds budget_limit", async () => { - await renderLoaded(budgetResponse({ budget_used: 1500, budget_limit: 1000 })); - const fill = screen.getByTestId("budget-progress-fill") as HTMLDivElement; - expect(fill.style.width).toBe("100%"); - }); - - it("progress bar has aria-valuenow equal to the calculated percentage", async () => { - await renderLoaded(budgetResponse({ budget_used: 300, budget_limit: 1000 })); - const bar = screen.getByRole("progressbar"); - expect(bar.getAttribute("aria-valuenow")).toBe("30"); - }); - - it("shows 0% progress bar when budget_used is absent from the response", async () => { - // Regression: budget_used is optional (provisioning-stuck workspaces return - // partial shapes). Without the `?? 0` guard the progressPct calculation - // throws a TypeScript strict-null error and the build fails. - // eslint-disable-next-line @typescript-eslint/no-explicit-any - await renderLoaded({ budget_limit: 1000, budget_remaining: null } as any); - const bar = screen.getByRole("progressbar"); - expect(bar.getAttribute("aria-valuenow")).toBe("0"); - const fill = screen.getByTestId("budget-progress-fill") as HTMLDivElement; - expect(fill.style.width).toBe("0%"); - }); -}); - -// ── Input pre-fill ──────────────────────────────────────────────────────────── - -describe("BudgetSection — input pre-fill", () => { - it("pre-fills input with existing budget_limit", async () => { - await renderLoaded(budgetResponse({ budget_limit: 500 })); - const input = screen.getByTestId("budget-limit-input") as HTMLInputElement; - expect(input.value).toBe("500"); - }); - - it("leaves input empty when budget_limit is null", async () => { - await renderLoaded(budgetResponse({ budget_limit: null, budget_remaining: null })); - const input = screen.getByTestId("budget-limit-input") as HTMLInputElement; - expect(input.value).toBe(""); - }); -}); - -// ── Save — PATCH calls ──────────────────────────────────────────────────────── - -describe("BudgetSection — save", () => { - it("calls PATCH /workspaces/:id/budget with budget_limit as integer", async () => { - // eslint-disable-next-line @typescript-eslint/no-explicit-any - mockPatch.mockResolvedValueOnce(budgetResponse({ budget_limit: 800 }) as any); - await renderLoaded(budgetResponse({ budget_limit: 1000 })); - - fireEvent.change(screen.getByTestId("budget-limit-input"), { - target: { value: "800" }, - }); - fireEvent.click(screen.getByTestId("budget-save-btn")); - - await waitFor(() => expect(mockPatch).toHaveBeenCalled()); - expect(mockPatch.mock.calls[0][0]).toBe("/workspaces/ws-1/budget"); - const body = mockPatch.mock.calls[0][1] as Record; - expect(body.budget_limit).toBe(800); - }); - - it("sends budget_limit: 0 (not null) when input is '0' — zero-credit budget", async () => { - // Regression for QA bug report: `parseInt("0") || null` would yield null. - // The correct form `raw !== "" ? parseInt(raw, 10) : null` must return 0. - // eslint-disable-next-line @typescript-eslint/no-explicit-any - mockPatch.mockResolvedValueOnce(budgetResponse({ budget_limit: 0, budget_used: 0, budget_remaining: 0 }) as any); - await renderLoaded(budgetResponse({ budget_limit: 1000 })); - - fireEvent.change(screen.getByTestId("budget-limit-input"), { - target: { value: "0" }, - }); - fireEvent.click(screen.getByTestId("budget-save-btn")); - - await waitFor(() => expect(mockPatch).toHaveBeenCalled()); - const body = mockPatch.mock.calls[0][1] as Record; - expect(body.budget_limit).toBe(0); - expect(body.budget_limit).not.toBeNull(); - }); - - it("sends budget_limit: null when input is blank", async () => { - // eslint-disable-next-line @typescript-eslint/no-explicit-any - mockPatch.mockResolvedValueOnce(budgetResponse({ budget_limit: null, budget_remaining: null }) as any); - await renderLoaded(budgetResponse({ budget_limit: 1000 })); - - fireEvent.change(screen.getByTestId("budget-limit-input"), { - target: { value: "" }, - }); - fireEvent.click(screen.getByTestId("budget-save-btn")); - - await waitFor(() => expect(mockPatch).toHaveBeenCalled()); - const body = mockPatch.mock.calls[0][1] as Record; - expect(body.budget_limit).toBeNull(); - }); - - it("updates displayed stats after successful save", async () => { - const updated = budgetResponse({ budget_limit: 2000, budget_used: 500, budget_remaining: 1500 }); - // eslint-disable-next-line @typescript-eslint/no-explicit-any - mockPatch.mockResolvedValueOnce(updated as any); - await renderLoaded(budgetResponse({ budget_limit: 1000, budget_used: 250 })); - - fireEvent.change(screen.getByTestId("budget-limit-input"), { - target: { value: "2000" }, - }); - fireEvent.click(screen.getByTestId("budget-save-btn")); - - await waitFor(() => - expect(screen.getByTestId("budget-limit-value").textContent).toBe("2,000") - ); - }); - - it("shows save error message on non-402 PATCH failure", async () => { - mockPatch.mockRejectedValueOnce( - new Error("API PATCH /workspaces/ws-1/budget: 500 server error") - ); - await renderLoaded(); - - fireEvent.click(screen.getByTestId("budget-save-btn")); - - await waitFor(() => - expect(screen.getByTestId("budget-save-error")).toBeTruthy() - ); - expect(screen.getByTestId("budget-save-error").textContent).toContain("500"); - }); -}); - -// ── 402 handling ────────────────────────────────────────────────────────────── - -describe("BudgetSection — 402 handling", () => { - it("shows exceeded banner when GET returns 402", async () => { - mockGet.mockRejectedValueOnce(make402Error()); - render(); - - await waitFor(() => - expect(screen.getByTestId("budget-exceeded-banner")).toBeTruthy() - ); - expect(screen.getByText("Budget exceeded — messages blocked")).toBeTruthy(); - }); - - it("does NOT show fetch error text when GET returns 402 (only banner)", async () => { - mockGet.mockRejectedValueOnce(make402Error()); - render(); - - await waitFor(() => - expect(screen.queryByTestId("budget-loading")).toBeNull() - ); - expect(screen.queryByTestId("budget-fetch-error")).toBeNull(); - expect(screen.getByTestId("budget-exceeded-banner")).toBeTruthy(); - }); - - it("shows exceeded banner when PATCH returns 402", async () => { - // eslint-disable-next-line @typescript-eslint/no-explicit-any - mockGet.mockResolvedValueOnce(budgetResponse() as any); - mockPatch.mockRejectedValueOnce(make402PatchError()); - render(); - await waitFor(() => expect(screen.queryByTestId("budget-loading")).toBeNull()); - - fireEvent.click(screen.getByTestId("budget-save-btn")); - - await waitFor(() => - expect(screen.getByTestId("budget-exceeded-banner")).toBeTruthy() - ); - // Should NOT also show the save-error alert - expect(screen.queryByTestId("budget-save-error")).toBeNull(); - }); - - it("clears exceeded banner after a successful save", async () => { - mockGet.mockRejectedValueOnce(make402Error()); - render(); - await waitFor(() => - expect(screen.getByTestId("budget-exceeded-banner")).toBeTruthy() - ); - - // Now a successful PATCH (limit was raised) - const updated = budgetResponse({ budget_limit: 5000, budget_used: 250, budget_remaining: 4750 }); - // eslint-disable-next-line @typescript-eslint/no-explicit-any - mockPatch.mockResolvedValueOnce(updated as any); - - await act(async () => { - fireEvent.change(screen.getByTestId("budget-limit-input"), { - target: { value: "5000" }, - }); - fireEvent.click(screen.getByTestId("budget-save-btn")); - }); - - await waitFor(() => - expect(screen.queryByTestId("budget-exceeded-banner")).toBeNull() - ); - }); -}); - -// ── Non-402 fetch error ─────────────────────────────────────────────────────── - -describe("BudgetSection — non-402 fetch errors", () => { - it("shows fetch error text on non-402 GET failure", async () => { - mockGet.mockRejectedValueOnce(makeGenericError("internal server error")); - render(); - - await waitFor(() => - expect(screen.getByTestId("budget-fetch-error")).toBeTruthy() - ); - expect(screen.getByTestId("budget-fetch-error").textContent).toContain("500"); - }); - - it("does NOT show stats row on fetch error", async () => { - mockGet.mockRejectedValueOnce(makeGenericError()); - render(); - - await waitFor(() => expect(screen.queryByTestId("budget-loading")).toBeNull()); - expect(screen.queryByTestId("budget-stats-row")).toBeNull(); - }); - - it("does NOT show exceeded banner on non-402 fetch error", async () => { - mockGet.mockRejectedValueOnce(makeGenericError()); - render(); - - await waitFor(() => expect(screen.queryByTestId("budget-loading")).toBeNull()); - expect(screen.queryByTestId("budget-exceeded-banner")).toBeNull(); + it("fills to 25%", async () => { + await renderLoaded(withMonthly(1000, 250)); + expect((screen.getByTestId("budget-monthly-fill") as HTMLElement).style.width).toBe("25%"); + }); + + it("fills to 50%", async () => { + await renderLoaded(withMonthly(1000, 500)); + expect((screen.getByTestId("budget-monthly-fill") as HTMLElement).style.width).toBe("50%"); + }); + + it("caps fill at 100% when spend exceeds limit", async () => { + await renderLoaded(withMonthly(1000, 4000)); + expect((screen.getByTestId("budget-monthly-fill") as HTMLElement).style.width).toBe("100%"); + }); + + it("sets aria-valuenow to the computed percentage on the progressbar", async () => { + await renderLoaded(withMonthly(1000, 250)); + const bars = screen.getAllByRole("progressbar"); + // the monthly bar is the only one rendered (others unlimited) + expect(bars).toHaveLength(1); + expect(bars[0].getAttribute("aria-valuenow")).toBe("25"); + }); + + it("shows a 0% bar when spend is 0 against a set limit", async () => { + await renderLoaded(withMonthly(1000, 0)); + expect((screen.getByTestId("budget-monthly-fill") as HTMLElement).style.width).toBe("0%"); }); }); diff --git a/canvas/src/components/tabs/BudgetSection.tsx b/canvas/src/components/tabs/BudgetSection.tsx index 2cad3f956..cb0382c2e 100644 --- a/canvas/src/components/tabs/BudgetSection.tsx +++ b/canvas/src/components/tabs/BudgetSection.tsx @@ -7,10 +7,28 @@ import { api } from "@/lib/api"; // Types // --------------------------------------------------------------------------- +// Period keys MUST match the server SSOT (workspace-server budget_periods.go). +type BudgetPeriod = "hourly" | "daily" | "weekly" | "monthly"; + +const PERIODS: { key: BudgetPeriod; label: string }[] = [ + { key: "hourly", label: "Hourly" }, + { key: "daily", label: "Daily" }, + { key: "weekly", label: "Weekly" }, + { key: "monthly", label: "Monthly" }, +]; + +interface PeriodBudget { + limit: number | null; // USD cents; null = no limit + spend: number; // rolling-window spend, USD cents + remaining: number | null; // null when no limit +} + interface BudgetData { - budget_limit: number | null; - budget_used?: number; // optional — provisioning-stuck workspaces return partial shapes - budget_remaining: number | null; + periods?: Partial>; + // legacy fields (pre-multi-period server) — tolerated for back-compat + budget_limit?: number | null; + monthly_spend?: number; + budget_remaining?: number | null; } interface Props { @@ -26,31 +44,71 @@ function isApiError402(e: unknown): boolean { return e instanceof Error && /: 402( |$)/.test(e.message); } +/** USD cents → "$X.XX". */ +function fmtUSD(cents: number): string { + return `$${(cents / 100).toLocaleString(undefined, { minimumFractionDigits: 2, maximumFractionDigits: 2 })}`; +} + +/** Normalize the server payload (multi-period or legacy) into a period map. */ +function periodsFrom(data: BudgetData | null): Record { + const base: Record = { + hourly: { limit: null, spend: 0, remaining: null }, + daily: { limit: null, spend: 0, remaining: null }, + weekly: { limit: null, spend: 0, remaining: null }, + monthly: { limit: null, spend: 0, remaining: null }, + }; + if (!data) return base; + if (data.periods) { + for (const { key } of PERIODS) { + const p = data.periods[key]; + if (p) base[key] = { limit: p.limit ?? null, spend: p.spend ?? 0, remaining: p.remaining ?? null }; + } + return base; + } + // legacy: map the single monthly limit/spend + base.monthly = { + limit: data.budget_limit ?? null, + spend: data.monthly_spend ?? 0, + remaining: data.budget_remaining ?? null, + }; + return base; +} + // --------------------------------------------------------------------------- // Component // --------------------------------------------------------------------------- /** - * BudgetSection — dedicated "Budget" section in the workspace details panel. - * - * - Fetches GET /workspaces/:id/budget on mount for live usage stats - * - Shows a progress bar (budget_used / budget_limit, blue-500, capped 100%) - * - Allows updating budget_limit via PATCH /workspaces/:id/budget - * - Shows a 402-specific "Budget exceeded" amber banner for any blocked state + * BudgetSection — per-workspace LLM budget, four independent rolling windows + * (hourly / daily / weekly / monthly). Each period has its own ceiling (USD); + * spend is the rolling-window LLM cost. Crossing ANY period blocks new work + * (server returns 402). Sends PATCH {budget_limits:{period:cents|null}}. */ export function BudgetSection({ workspaceId }: Props) { const [budget, setBudget] = useState(null); const [loading, setLoading] = useState(true); const [fetchError, setFetchError] = useState(null); - const [limitInput, setLimitInput] = useState(""); + // One input per period, in USD cents (string for controlled inputs). + const [limitInputs, setLimitInputs] = useState>({ + hourly: "", + daily: "", + weekly: "", + monthly: "", + }); const [saving, setSaving] = useState(false); const [saveError, setSaveError] = useState(null); - - /** True when a 402 has been seen from any API call in this section. */ const [budgetExceeded, setBudgetExceeded] = useState(false); - // ── Fetch current budget data ───────────────────────────────────────────── + const syncInputs = useCallback((data: BudgetData | null) => { + const p = periodsFrom(data); + setLimitInputs({ + hourly: p.hourly.limit != null ? String(p.hourly.limit) : "", + daily: p.daily.limit != null ? String(p.daily.limit) : "", + weekly: p.weekly.limit != null ? String(p.weekly.limit) : "", + monthly: p.monthly.limit != null ? String(p.monthly.limit) : "", + }); + }, []); const loadBudget = useCallback(async () => { setLoading(true); @@ -58,7 +116,7 @@ export function BudgetSection({ workspaceId }: Props) { try { const data = await api.get(`/workspaces/${workspaceId}/budget`); setBudget(data); - setLimitInput(data.budget_limit != null ? String(data.budget_limit) : ""); + syncInputs(data); } catch (e) { if (isApiError402(e)) { setBudgetExceeded(true); @@ -68,29 +126,30 @@ export function BudgetSection({ workspaceId }: Props) { } finally { setLoading(false); } - }, [workspaceId]); + }, [workspaceId, syncInputs]); useEffect(() => { loadBudget(); }, [loadBudget]); - // ── Save handler ────────────────────────────────────────────────────────── - const handleSave = async () => { setSaving(true); setSaveError(null); - const raw = limitInput.trim(); - // Use explicit empty-string check (not falsy check) so that a - // user-entered "0" is sent as budget_limit: 0, not null (unlimited). - const parsedLimit = raw !== "" ? parseInt(raw, 10) : null; - + // Build the per-period map: blank → null (clear); a number → that ceiling. + const budget_limits: Record = { + hourly: null, + daily: null, + weekly: null, + monthly: null, + }; + for (const { key } of PERIODS) { + const raw = limitInputs[key].trim(); + budget_limits[key] = raw !== "" ? parseInt(raw, 10) : null; + } try { - const updated = await api.patch(`/workspaces/${workspaceId}/budget`, { - budget_limit: parsedLimit, - }); + const updated = await api.patch(`/workspaces/${workspaceId}/budget`, { budget_limits }); setBudget(updated); - setLimitInput(updated.budget_limit != null ? String(updated.budget_limit) : ""); - // Clear exceeded state if the save succeeded (limit was raised or removed) + syncInputs(updated); setBudgetExceeded(false); } catch (e) { if (isApiError402(e)) { @@ -103,24 +162,15 @@ export function BudgetSection({ workspaceId }: Props) { } }; - // ── Progress calculation ────────────────────────────────────────────────── - - const progressPct = - budget && budget.budget_limit != null && budget.budget_limit > 0 - ? Math.min(100, Math.round(((budget.budget_used ?? 0) / budget.budget_limit) * 100)) - : 0; - - // ── Render ──────────────────────────────────────────────────────────────── + const periods = periodsFrom(budget); return (
{/* Section header */}
-

- Budget -

+

Budget

- Limit total message credits for this workspace + Cap LLM spend for this workspace per period — crossing any limit pauses new work

@@ -131,32 +181,14 @@ export function BudgetSection({ workspaceId }: Props) { data-testid="budget-exceeded-banner" className="flex items-center gap-2 px-3 py-2 rounded-lg bg-surface border border-amber-700/50 text-warm text-xs font-medium" > -
)} - {/* Usage stats */} {loading ? (

Loading… @@ -165,89 +197,78 @@ export function BudgetSection({ workspaceId }: Props) {

{fetchError}

- ) : budget ? ( -
- {/* Stats row */} -
- Credits used - - {(budget.budget_used ?? 0).toLocaleString()} - / - - {budget.budget_limit != null - ? budget.budget_limit.toLocaleString() - : "Unlimited"} - - -
+ ) : ( +
+ {PERIODS.map(({ key, label }) => { + const p = periods[key]; + const pct = + p.limit != null && p.limit > 0 ? Math.min(100, Math.round((p.spend / p.limit) * 100)) : 0; + const over = p.limit != null && p.spend >= p.limit; + return ( +
+
+ + + {fmtUSD(p.spend)} + / + {p.limit != null ? fmtUSD(p.limit) : "∞"} + +
+ {p.limit != null && ( +
+
+
+ )} + setLimitInputs((s) => ({ ...s, [key]: e.target.value }))} + placeholder="USD cents — blank for unlimited" + data-testid={`budget-${key}-input`} + className="w-full bg-surface-card border border-line rounded-lg px-3 py-1.5 text-xs text-ink-mid placeholder-zinc-500 focus:outline-none focus:border-accent focus:ring-1 focus:ring-accent/30 transition-colors" + /> +
+ ); + })} - {/* Progress bar (only when limit is set) */} - {budget.budget_limit != null && ( +

Limits are USD cents (e.g. 500 = $5.00). Blank = unlimited.

+ + {saveError && (
-
+ {saveError}
)} - {/* Remaining credits */} - {budget.budget_remaining != null && ( -

- {budget.budget_remaining.toLocaleString()} credits remaining -

- )} -
- ) : null} - - {/* Input + Save */} -
- - setLimitInput(e.target.value)} - placeholder="e.g. 1000 — blank for unlimited" - data-testid="budget-limit-input" - className="w-full bg-surface-card border border-line rounded-lg px-3 py-2 text-sm text-ink-mid placeholder-zinc-500 focus:outline-none focus:border-accent focus:ring-1 focus:ring-accent/30 transition-colors" - /> -

Leave blank for unlimited

- - {saveError && ( -
- {saveError} -
- )} - - -
+ {saving ? "Saving…" : "Save"} + +
+ )}
); } diff --git a/canvas/src/components/tabs/__tests__/BudgetSection.test.tsx b/canvas/src/components/tabs/__tests__/BudgetSection.test.tsx index 7372ca0db..0345cc25b 100644 --- a/canvas/src/components/tabs/__tests__/BudgetSection.test.tsx +++ b/canvas/src/components/tabs/__tests__/BudgetSection.test.tsx @@ -5,9 +5,10 @@ import React from "react"; import { BudgetSection } from "../BudgetSection"; import { api } from "@/lib/api"; -// Queue-based mock for the api module. Each api call shifts from the queue. -// Tests push with qGet/qPatch and the module-level mockImplementation -// reads from the queue. +// Multi-period budget (#49): the API now returns a `periods` map +// (hourly/daily/weekly/monthly), each {limit, spend, remaining} in USD cents. +// The UI renders one row per period and PATCHes {budget_limits:{period:cents|null}}. + type QueueEntry = { body?: unknown; err?: Error }; const apiQueue: QueueEntry[] = []; @@ -40,45 +41,49 @@ const WS_ID = "budget-test-ws"; function qGet(body: unknown) { apiQueue.push({ body }); } - function qGetErr(status: number, msg: string) { apiQueue.push({ err: new Error(`${msg}: ${status}`) }); } - function qPatch(body: unknown) { apiQueue.push({ body }); } - function qPatchErr(status: number, msg: string) { apiQueue.push({ err: new Error(`${msg}: ${status}`) }); } -function makeBudget(overrides: Partial<{ - budget_limit: number | null; - budget_used: number; - budget_remaining: number | null; -}> = {}) { +type P = { limit: number | null; spend: number; remaining: number | null }; + +// makeBudget builds the periods response. Override any subset of periods. +function makeBudget(overrides: Partial>> = {}) { + const blank: P = { limit: null, spend: 0, remaining: null }; + const mk = (o?: Partial

): P => { + const p = { ...blank, ...(o ?? {}) }; + if (p.limit != null && p.remaining == null) p.remaining = p.limit - p.spend; + return p; + }; + const periods = { + hourly: mk(overrides.hourly), + daily: mk(overrides.daily), + weekly: mk(overrides.weekly), + monthly: mk(overrides.monthly), + }; return { - budget_limit: 10_000, - budget_used: 3_500, - budget_remaining: 6_500, - ...overrides, + periods, + budget_limit: periods.monthly.limit, + monthly_spend: periods.monthly.spend, + budget_remaining: periods.monthly.remaining, }; } -describe("BudgetSection", () => { +describe("BudgetSection (multi-period)", () => { describe("loading state", () => { it("shows loading indicator while fetching", async () => { let resolveGet: (v: unknown) => void; vi.mocked(api.get).mockImplementationOnce( async () => new Promise((r) => { resolveGet = r as (v: unknown) => void; }), ); - render(); - expect(screen.getByTestId("budget-loading")).toBeTruthy(); - - // Resolve after render to verify state clears resolveGet!(makeBudget()); await vi.waitFor(() => { expect(screen.queryByTestId("budget-loading")).toBeNull(); @@ -89,21 +94,16 @@ describe("BudgetSection", () => { describe("fetch error state", () => { it("shows error message on non-402 fetch failure", async () => { qGetErr(500, "Internal Server Error"); - render(); - await vi.waitFor(() => { expect(screen.getByTestId("budget-fetch-error")).toBeTruthy(); }); expect(screen.getByTestId("budget-fetch-error")!.textContent).toContain("500"); }); - it("shows 402 as exceeded banner, not fetch error", async () => { - // 402 means the budget limit was hit — different UX from a network/API error. + it("shows the exceeded banner (not a fetch error) on a 402", async () => { qGetErr(402, "Payment Required"); - render(); - await vi.waitFor(() => { expect(screen.getByTestId("budget-exceeded-banner")).toBeTruthy(); }); @@ -111,220 +111,105 @@ describe("BudgetSection", () => { }); }); - describe("budget loaded — display", () => { - it("renders used / limit stats row", async () => { - qGet(makeBudget({ budget_limit: 10_000, budget_used: 3_500 })); - + describe("rendering periods", () => { + it("renders all four period rows", async () => { + qGet(makeBudget()); render(); - await vi.waitFor(() => { - expect(screen.getByTestId("budget-used-value")!.textContent).toBe("3,500"); - }); - expect(screen.getByTestId("budget-limit-value")!.textContent).toBe("10,000"); - }); - - it("renders 'Unlimited' when budget_limit is null", async () => { - qGet(makeBudget({ budget_limit: null, budget_used: 1_000, budget_remaining: null })); - - render(); - - await vi.waitFor(() => { - expect(screen.getByTestId("budget-limit-value")!.textContent).toBe("Unlimited"); + for (const k of ["hourly", "daily", "weekly", "monthly"]) { + expect(screen.getByTestId(`budget-period-${k}`)).toBeTruthy(); + } }); }); - it("renders remaining credits when present", async () => { - qGet(makeBudget({ budget_limit: 10_000, budget_used: 3_500, budget_remaining: 6_500 })); - + it("formats spend and limit as USD per period", async () => { + qGet(makeBudget({ monthly: { limit: 10_000, spend: 3_500 } })); render(); - await vi.waitFor(() => { - expect(screen.getByTestId("budget-remaining")!.textContent).toContain("6,500"); - expect(screen.getByTestId("budget-remaining")!.textContent).toContain("credits remaining"); + expect(screen.getByTestId("budget-monthly-spend")!.textContent).toBe("$35.00"); + }); + expect(screen.getByTestId("budget-monthly-limit")!.textContent).toBe("$100.00"); + }); + + it("shows ∞ for a period with no limit", async () => { + qGet(makeBudget({ hourly: { limit: null, spend: 1_000 } })); + render(); + await vi.waitFor(() => { + expect(screen.getByTestId("budget-hourly-limit")!.textContent).toBe("∞"); }); }); - it("omits remaining credits when budget_remaining is null", async () => { - qGet(makeBudget({ budget_limit: 10_000, budget_used: 3_500, budget_remaining: null })); - + it("renders the progress bar only for periods with a limit", async () => { + qGet(makeBudget({ monthly: { limit: 10_000, spend: 12_000 }, hourly: { limit: null, spend: 5_000 } })); render(); - await vi.waitFor(() => { - expect(screen.queryByTestId("budget-remaining")).toBeNull(); - }); - }); - - it("caps progress bar at 100% when used > limit", async () => { - // Over-limit: 12000 used of 10000 limit should show 100%, not 120%. - qGet(makeBudget({ budget_limit: 10_000, budget_used: 12_000, budget_remaining: null })); - - render(); - - await vi.waitFor(() => { - const fill = screen.getByTestId("budget-progress-fill"); - expect(fill.getAttribute("style")).toContain("100%"); - }); - }); - - it("omits progress bar when budget_limit is null (unlimited)", async () => { - qGet(makeBudget({ budget_limit: null, budget_used: 5_000, budget_remaining: null })); - - render(); - - await vi.waitFor(() => { - expect(screen.queryByTestId("budget-progress-fill")).toBeNull(); + expect(screen.getByTestId("budget-monthly-fill")).toBeTruthy(); }); + expect(screen.queryByTestId("budget-hourly-fill")).toBeNull(); + // over-budget fill caps at 100% + const fill = screen.getByTestId("budget-monthly-fill") as HTMLElement; + expect(fill.style.width).toBe("100%"); }); }); - describe("budget exceeded (402)", () => { - it("shows exceeded banner when load returns 402", async () => { - qGetErr(402, "Payment Required"); - + describe("save", () => { + it("PATCHes budget_limits for all four periods and clears the exceeded banner", async () => { + qGet(makeBudget({ monthly: { limit: 10_000, spend: 3_500 } })); + qPatch(makeBudget({ hourly: { limit: 500, spend: 0 }, monthly: { limit: 20_000, spend: 0 } })); render(); + await vi.waitFor(() => { + expect(screen.getByTestId("budget-hourly-input")).toBeTruthy(); + }); + + fireEvent.change(screen.getByTestId("budget-hourly-input"), { target: { value: "500" } }); + fireEvent.click(screen.getByTestId("budget-save-btn")); await vi.waitFor(() => { - expect(screen.getByTestId("budget-exceeded-banner")).toBeTruthy(); - expect(screen.getByTestId("budget-exceeded-banner")!.textContent).toContain("Budget exceeded"); + expect(vi.mocked(api.patch)).toHaveBeenCalled(); + }); + const [, body] = vi.mocked(api.patch).mock.calls[0]; + expect((body as { budget_limits: Record }).budget_limits).toMatchObject({ + hourly: 500, + monthly: 10_000, // unchanged input echoes the loaded limit }); }); - it("clears exceeded banner after successful save", async () => { - qGetErr(402, "Payment Required"); - qPatch(makeBudget({ budget_limit: 50_000, budget_used: 0, budget_remaining: 50_000 })); - - render(); - - await vi.waitFor(() => { - expect(screen.getByTestId("budget-exceeded-banner")).toBeTruthy(); - }); - - const input = screen.getByTestId("budget-limit-input"); - fireEvent.change(input, { target: { value: "50000" } }); - - const saveBtn = screen.getByTestId("budget-save-btn"); - fireEvent.click(saveBtn); - - await vi.waitFor(() => { - expect(screen.queryByTestId("budget-exceeded-banner")).toBeNull(); - }); - }); - }); - - describe("save flow", () => { - it("shows save error on non-402 patch failure", async () => { + it("shows a save error on non-402 PATCH failure", async () => { qGet(makeBudget()); qPatchErr(500, "Internal Server Error"); - render(); - await vi.waitFor(() => { - expect(screen.getByTestId("budget-limit-input")).toBeTruthy(); + expect(screen.getByTestId("budget-save-btn")).toBeTruthy(); }); - - const saveBtn = screen.getByTestId("budget-save-btn"); - fireEvent.click(saveBtn); - + fireEvent.click(screen.getByTestId("budget-save-btn")); await vi.waitFor(() => { expect(screen.getByTestId("budget-save-error")).toBeTruthy(); - expect(screen.getByTestId("budget-save-error")!.textContent).toContain("500"); }); + expect(screen.getByTestId("budget-save-error")!.textContent).toContain("500"); }); - it("updates input to new limit value after successful save", async () => { - qGet(makeBudget({ budget_limit: 10_000 })); - qPatch(makeBudget({ budget_limit: 20_000 })); - - render(); - - // Wait for the input to appear (loading → loaded) - await vi.waitFor(() => { - expect(screen.queryByTestId("budget-loading")).toBeNull(); - }); - - const input = screen.getByTestId("budget-limit-input") as HTMLInputElement; - // Debug: check what values are rendered - const limitValue = screen.getByTestId("budget-limit-value")?.textContent; - expect(input.value).toBe("10000"); // initial value from API - expect(limitValue).toBe("10,000"); - - fireEvent.change(input, { target: { value: "20000" } }); - expect(input.value).toBe("20000"); - - fireEvent.click(screen.getByTestId("budget-save-btn")); - - await vi.waitFor(() => { - expect((screen.getByTestId("budget-limit-input") as HTMLInputElement).value).toBe("20000"); - }); - }); - - it("sends null when input is cleared (unlimited)", async () => { - qGet(makeBudget({ budget_limit: 10_000 })); - qPatch(makeBudget({ budget_limit: null })); - - render(); - - await vi.waitFor(() => { - expect(screen.getByTestId("budget-limit-input")).toBeTruthy(); - }); - - const input = screen.getByTestId("budget-limit-input") as HTMLInputElement; - fireEvent.change(input, { target: { value: "" } }); - fireEvent.click(screen.getByTestId("budget-save-btn")); - - await vi.waitFor(() => { - // After save with null limit, input should show empty (unlimited) - expect(input.value).toBe(""); - }); - }); - - it("shows saving state on button while patch is in flight", async () => { + it("surfaces the exceeded banner on a 402 PATCH", async () => { qGet(makeBudget()); - let resolvePatch: (v: unknown) => void; - vi.mocked(api.patch).mockImplementationOnce( - async () => new Promise((r) => { resolvePatch = r as (v: unknown) => void; }), - ); - + qPatchErr(402, "Payment Required"); render(); - await vi.waitFor(() => { - expect(screen.getByTestId("budget-limit-input")).toBeTruthy(); + expect(screen.getByTestId("budget-save-btn")).toBeTruthy(); }); - - fireEvent.change(screen.getByTestId("budget-limit-input"), { target: { value: "50000" } }); fireEvent.click(screen.getByTestId("budget-save-btn")); - - const btn = screen.getByTestId("budget-save-btn"); - expect(btn.textContent).toContain("Saving"); - - resolvePatch!(makeBudget({ budget_limit: 50_000 })); - await vi.waitFor(() => { - expect(btn.textContent).toContain("Save"); - }); - }); - }); - - describe("isApiError402 — regression coverage", () => { - it("classifies ': 402' with space as 402", async () => { - qGetErr(402, "Payment Required"); - qPatch(makeBudget()); - - render(); - await vi.waitFor(() => { expect(screen.getByTestId("budget-exceeded-banner")).toBeTruthy(); }); }); + }); - it("classifies non-402 error messages as regular fetch errors", async () => { - qGetErr(503, "Service Unavailable"); - + describe("legacy payload back-compat", () => { + it("maps a pre-multi-period {budget_limit, monthly_spend} response to the monthly row", async () => { + qGet({ budget_limit: 5_000, monthly_spend: 1_000, budget_remaining: 4_000 }); render(); - await vi.waitFor(() => { - expect(screen.getByTestId("budget-fetch-error")).toBeTruthy(); + expect(screen.getByTestId("budget-monthly-limit")!.textContent).toBe("$50.00"); }); - expect(screen.queryByTestId("budget-exceeded-banner")).toBeNull(); + expect(screen.getByTestId("budget-monthly-spend")!.textContent).toBe("$10.00"); }); }); }); diff --git a/workspace-server/internal/handlers/a2a_proxy.go b/workspace-server/internal/handlers/a2a_proxy.go index 04f52af0a..17321a352 100644 --- a/workspace-server/internal/handlers/a2a_proxy.go +++ b/workspace-server/internal/handlers/a2a_proxy.go @@ -334,28 +334,39 @@ func (h *WorkspaceHandler) ProxyA2A(c *gin.Context) { c.Data(status, "application/json", respBody) } -// checkWorkspaceBudget returns a proxyA2AError with 402 when the workspace -// has a budget_limit set and monthly_spend has reached or exceeded it. -// DB errors are logged and treated as fail-open — a budget check failure -// must not block legitimate A2A traffic. +// checkWorkspaceBudget returns a proxyA2AError with 402 when the workspace has +// exceeded ANY of its configured per-period budget limits (hourly/daily/weekly/ +// monthly — see budget_periods.go). Per-period spend is the rolling-window sum +// over the workspace_spend_events ledger. DB errors are logged and treated as +// fail-open — a budget check failure must not block legitimate A2A traffic. func (h *WorkspaceHandler) checkWorkspaceBudget(ctx context.Context, workspaceID string) *proxyA2AError { - var budgetLimit sql.NullInt64 - var monthlySpend int64 - err := db.DB.QueryRowContext(ctx, - `SELECT budget_limit, COALESCE(monthly_spend, 0) FROM workspaces WHERE id = $1`, + var limitsRaw []byte + if err := db.DB.QueryRowContext(ctx, + `SELECT COALESCE(budget_limits, '{}'::jsonb) FROM workspaces WHERE id = $1`, workspaceID, - ).Scan(&budgetLimit, &monthlySpend) - if err != nil { + ).Scan(&limitsRaw); err != nil { if err != sql.ErrNoRows { log.Printf("ProxyA2A: budget check failed for %s: %v", workspaceID, err) } return nil // fail-open } - if budgetLimit.Valid && monthlySpend >= budgetLimit.Int64 { - log.Printf("ProxyA2A: budget exceeded for %s (spend=%d limit=%d)", workspaceID, monthlySpend, budgetLimit.Int64) + limits := parseBudgetLimits(limitsRaw) + if len(limits) == 0 { + return nil // no limits configured + } + spend, err := spendByPeriod(ctx, db.DB, workspaceID) + if err != nil { + log.Printf("ProxyA2A: budget spend query failed for %s: %v", workspaceID, err) + return nil // fail-open + } + if over := exceededPeriods(limits, spend); len(over) > 0 { + log.Printf("ProxyA2A: budget exceeded for %s (periods=%v limits=%v spend=%v)", workspaceID, over, limits, spend) return &proxyA2AError{ - Status: http.StatusPaymentRequired, - Response: gin.H{"error": "workspace budget limit exceeded"}, + Status: http.StatusPaymentRequired, + Response: gin.H{ + "error": "workspace budget limit exceeded", + "exceeded_periods": over, + }, } } return nil diff --git a/workspace-server/internal/handlers/a2a_queue_test.go b/workspace-server/internal/handlers/a2a_queue_test.go index 93c6b6629..dd461cc7b 100644 --- a/workspace-server/internal/handlers/a2a_queue_test.go +++ b/workspace-server/internal/handlers/a2a_queue_test.go @@ -18,8 +18,8 @@ import ( "testing" "time" - "github.com/DATA-DOG/go-sqlmock" "git.moleculesai.app/molecule-ai/molecule-core/workspace-server/internal/db" + "github.com/DATA-DOG/go-sqlmock" "github.com/alicebob/miniredis/v2" ) @@ -209,10 +209,12 @@ func drainSetup(t *testing.T, workspaceID string) (sqlmock.Sqlmock, *WorkspaceHa // Named distinctly from handlers_test.go's expectBudgetCheck (which uses MatchPsql // escaped-regex and cannot be reused with QueryMatcherEqual tests). func expectQueueBudgetCheck(mock sqlmock.Sqlmock, workspaceID string) { + // Multi-period (#49): exact-match the budget_limits read; "{}" → no limits → + // checkWorkspaceBudget returns early (no spend query). mock.ExpectQuery( - "SELECT budget_limit, COALESCE(monthly_spend, 0) FROM workspaces WHERE id = $1", + "SELECT COALESCE(budget_limits, '{}'::jsonb) FROM workspaces WHERE id = $1", ).WithArgs(workspaceID). - WillReturnRows(sqlmock.NewRows([]string{"budget_limit", "monthly_spend"})) + WillReturnRows(sqlmock.NewRows([]string{"budget_limits"}).AddRow([]byte("{}"))) } // seedRedisURL puts the agent server URL into the Redis cache so resolveAgentURL diff --git a/workspace-server/internal/handlers/budget.go b/workspace-server/internal/handlers/budget.go index 5601f9181..cc5e13aa6 100644 --- a/workspace-server/internal/handlers/budget.go +++ b/workspace-server/internal/handlers/budget.go @@ -1,7 +1,9 @@ package handlers import ( + "context" "database/sql" + "encoding/json" "log" "net/http" @@ -12,42 +14,79 @@ import ( // BudgetHandler exposes per-workspace budget read/write endpoints. // Routes (all behind WorkspaceAuth middleware): // -// GET /workspaces/:id/budget — current budget_limit, monthly_spend, budget_remaining -// PATCH /workspaces/:id/budget — set or clear budget_limit +// GET /workspaces/:id/budget — per-period limits, spend, remaining +// PATCH /workspaces/:id/budget — set/clear per-period limits +// +// Multi-period (#49): the budget is now four independent rolling windows — +// hourly/daily/weekly/monthly (budget_periods.go is the SSOT for the set). The +// canonical config is workspaces.budget_limits (JSONB, USD cents per period); +// per-period spend is the rolling-window sum over workspace_spend_events. The +// legacy single monthly budget_limit / monthly_spend are still emitted (and +// budget_limit kept in sync to the monthly period) for back-compat with +// pre-deploy canvas/agent builds during the rollout window. type BudgetHandler struct{} func NewBudgetHandler() *BudgetHandler { return &BudgetHandler{} } -// budgetResponse is the canonical JSON shape for both GET and PATCH responses. +// periodBudget is the per-period view: configured ceiling (null = no limit), +// rolling-window spend, and remaining headroom (null when no limit; may go +// negative so callers see how far over a period is). +type periodBudget struct { + Limit *int64 `json:"limit"` + Spend int64 `json:"spend"` + Remaining *int64 `json:"remaining"` +} + +// budgetResponse is the canonical JSON shape for GET and PATCH. type budgetResponse struct { - // BudgetLimit is the monthly spend ceiling in USD cents (null = no limit). - // budget_limit=500 means $5.00/month. - BudgetLimit *int64 `json:"budget_limit"` - // MonthlySpend is the agent's self-reported accumulated LLM API spend - // for the current month (USD cents). Incremented via heartbeat. - MonthlySpend int64 `json:"monthly_spend"` - // BudgetRemaining is null when BudgetLimit is null, otherwise - // max(0, budget_limit - monthly_spend). Can be negative — we store the - // actual value so callers can see how far over-budget a workspace is. + // Periods is keyed by BudgetPeriod ("hourly"/"daily"/"weekly"/"monthly"). + Periods map[string]periodBudget `json:"periods"` + + // --- back-compat (monthly), for pre-multi-period clients --- + BudgetLimit *int64 `json:"budget_limit"` + MonthlySpend int64 `json:"monthly_spend"` BudgetRemaining *int64 `json:"budget_remaining"` } +// buildBudgetResponse assembles the per-period view from the stored limits + +// the ledger spend. Single place so GET and PATCH return identical shapes. +func buildBudgetResponse(ctx context.Context, workspaceID string, limitsRaw []byte) (budgetResponse, error) { + limits := parseBudgetLimits(limitsRaw) + spend, err := spendByPeriod(ctx, db.DB, workspaceID) + if err != nil { + return budgetResponse{}, err + } + periods := make(map[string]periodBudget, len(budgetPeriods)) + for _, def := range budgetPeriods { + pb := periodBudget{Spend: spend[def.Name]} + if lim, ok := limits[def.Name]; ok { + l := lim + pb.Limit = &l + r := lim - spend[def.Name] + pb.Remaining = &r + } + periods[string(def.Name)] = pb + } + resp := budgetResponse{Periods: periods, MonthlySpend: spend[PeriodMonthly]} + if m := periods[string(PeriodMonthly)]; m.Limit != nil { + resp.BudgetLimit = m.Limit + resp.BudgetRemaining = m.Remaining + } + return resp, nil +} + // GetBudget handles GET /workspaces/:id/budget. -// Returns the workspace's current budget ceiling, accumulated spend, and -// computed remaining headroom. Both budget_limit and budget_remaining are -// null when no limit has been configured for the workspace. func (h *BudgetHandler) GetBudget(c *gin.Context) { workspaceID := c.Param("id") ctx := c.Request.Context() - var budgetLimit sql.NullInt64 - var monthlySpend int64 + var limitsRaw []byte err := db.DB.QueryRowContext(ctx, - `SELECT budget_limit, COALESCE(monthly_spend, 0) + `SELECT COALESCE(budget_limits, '{}'::jsonb) FROM workspaces WHERE id = $1 AND status != 'removed'`, workspaceID, - ).Scan(&budgetLimit, &monthlySpend) + ).Scan(&limitsRaw) if err == sql.ErrNoRows { c.JSON(http.StatusNotFound, gin.H{"error": "workspace not found"}) return @@ -58,66 +97,80 @@ func (h *BudgetHandler) GetBudget(c *gin.Context) { return } - resp := budgetResponse{ - MonthlySpend: monthlySpend, + resp, err := buildBudgetResponse(ctx, workspaceID, limitsRaw) + if err != nil { + log.Printf("GetBudget: spend query failed for %s: %v", workspaceID, err) + c.JSON(http.StatusInternalServerError, gin.H{"error": "query failed"}) + return } - if budgetLimit.Valid { - limit := budgetLimit.Int64 - resp.BudgetLimit = &limit - remaining := limit - monthlySpend - resp.BudgetRemaining = &remaining - } - c.JSON(http.StatusOK, resp) } -// PatchBudget handles PATCH /workspaces/:id/budget. -// Accepts {"budget_limit": } to set a new ceiling, or -// {"budget_limit": null} to remove an existing ceiling. -// Returns the updated budget state in the same shape as GetBudget. +// PatchBudget handles PATCH /workspaces/:id/budget. Accepts EITHER the +// multi-period shape +// +// {"budget_limits": {"hourly": 100, "daily": null, "weekly": 500, "monthly": 2000}} +// +// (a per-period value of null/absent clears that period; a positive int sets it) +// OR the legacy single-monthly shape {"budget_limit": 2000} / {"budget_limit": null}. func (h *BudgetHandler) PatchBudget(c *gin.Context) { workspaceID := c.Param("id") ctx := c.Request.Context() - // We need to distinguish between "field absent" and "field = null", - // so we unmarshal into a raw map first. - var raw map[string]interface{} + var raw map[string]json.RawMessage if err := c.ShouldBindJSON(&raw); err != nil { c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request body"}) return } - - budgetLimitRaw, ok := raw["budget_limit"] - if !ok { - c.JSON(http.StatusBadRequest, gin.H{"error": "budget_limit field is required"}) + _, hasLimits := raw["budget_limits"] + _, hasLegacy := raw["budget_limit"] + if !hasLimits && !hasLegacy { + c.JSON(http.StatusBadRequest, gin.H{"error": "budget_limits or budget_limit field is required"}) return } - // Validate and convert the value. JSON numbers decode as float64. - var budgetArg interface{} // nil → SQL NULL, int64 → new ceiling - if budgetLimitRaw != nil { - switch v := budgetLimitRaw.(type) { - case float64: - if v < 0 { - c.JSON(http.StatusBadRequest, gin.H{"error": "budget_limit must be >= 0 (USD cents)"}) + limits := make(map[BudgetPeriod]int64, len(budgetPeriods)) + known := make(map[string]bool, len(budgetPeriods)) + for _, def := range budgetPeriods { + known[string(def.Name)] = true + } + + if hasLimits { + var m map[string]*int64 + if err := json.Unmarshal(raw["budget_limits"], &m); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "budget_limits must be an object of period→int|null"}) + return + } + for k, v := range m { + if !known[k] { + c.JSON(http.StatusBadRequest, gin.H{"error": "unknown budget period: " + k + " (allowed: hourly, daily, weekly, monthly)"}) return } - cv := int64(v) - budgetArg = cv - case int64: - if v < 0 { - c.JSON(http.StatusBadRequest, gin.H{"error": "budget_limit must be >= 0 (USD cents)"}) + if v == nil { + continue // clear this period (null = no limit) + } + if *v < 0 { + c.JSON(http.StatusBadRequest, gin.H{"error": "budget limit for " + k + " must be >= 0 (USD cents)"}) return } - budgetArg = v - default: + limits[BudgetPeriod(k)] = *v // 0 is valid = block-all for this period + } + } else { // legacy single-monthly + var v *int64 + if err := json.Unmarshal(raw["budget_limit"], &v); err != nil { c.JSON(http.StatusBadRequest, gin.H{"error": "budget_limit must be an integer (USD cents) or null"}) return } + if v != nil { + if *v < 0 { + c.JSON(http.StatusBadRequest, gin.H{"error": "budget_limit must be >= 0 (USD cents)"}) + return + } + limits[PeriodMonthly] = *v // 0 is valid = block-all (legacy semantics) + } } - // budgetArg == nil means "clear the ceiling" - // Existence check — return 404 for non-existent / removed workspaces. + // Existence check — 404 for non-existent / removed workspaces. var exists bool if err := db.DB.QueryRowContext(ctx, `SELECT EXISTS(SELECT 1 FROM workspaces WHERE id = $1 AND status != 'removed')`, @@ -127,38 +180,28 @@ func (h *BudgetHandler) PatchBudget(c *gin.Context) { return } + // Persist: budget_limits is the SSOT; keep the legacy budget_limit column + // synced to the monthly period so pre-deploy enforcement paths stay coherent + // during the rollout window. + var legacyMonthly interface{} + if m, ok := limits[PeriodMonthly]; ok { + legacyMonthly = m + } + encoded := encodeBudgetLimits(limits) if _, err := db.DB.ExecContext(ctx, - `UPDATE workspaces SET budget_limit = $2, updated_at = now() WHERE id = $1`, - workspaceID, budgetArg, + `UPDATE workspaces SET budget_limits = $2, budget_limit = $3, updated_at = now() WHERE id = $1`, + workspaceID, encoded, legacyMonthly, ); err != nil { log.Printf("PatchBudget: update failed for %s: %v", workspaceID, err) c.JSON(http.StatusInternalServerError, gin.H{"error": "update failed"}) return } - // Re-read the current state so the response reflects exactly what is in - // the DB, including the monthly_spend the agent has already accumulated. - var newLimit sql.NullInt64 - var monthlySpend int64 - if err := db.DB.QueryRowContext(ctx, - `SELECT budget_limit, COALESCE(monthly_spend, 0) FROM workspaces WHERE id = $1`, - workspaceID, - ).Scan(&newLimit, &monthlySpend); err != nil { + resp, err := buildBudgetResponse(ctx, workspaceID, encoded) + if err != nil { log.Printf("PatchBudget: re-read failed for %s: %v", workspaceID, err) - // Still success — just omit the echo. c.JSON(http.StatusOK, gin.H{"status": "updated"}) return } - - resp := budgetResponse{ - MonthlySpend: monthlySpend, - } - if newLimit.Valid { - limit := newLimit.Int64 - resp.BudgetLimit = &limit - remaining := limit - monthlySpend - resp.BudgetRemaining = &remaining - } - c.JSON(http.StatusOK, resp) } diff --git a/workspace-server/internal/handlers/budget_periods.go b/workspace-server/internal/handlers/budget_periods.go new file mode 100644 index 000000000..b233ff8ae --- /dev/null +++ b/workspace-server/internal/handlers/budget_periods.go @@ -0,0 +1,160 @@ +package handlers + +import ( + "context" + "database/sql" + "encoding/json" + "strconv" + "time" +) + +// budget_periods.go — SINGLE SOURCE OF TRUTH for the multi-period per-workspace +// LLM budget (#49 follow-up). The supported periods, their rolling windows, the +// per-period spend computation (from the workspace_spend_events ledger), and the +// over-budget decision all live here so the config endpoint (GetBudget/PatchBudget), +// the display, and enforcement (checkWorkspaceBudget) can never drift. +// +// Spend model: the heartbeat records each observed spend INCREMENT into +// workspace_spend_events (recordSpendDelta). Per-period spend is a rolling-window +// SUM over that ledger — so the SERVER owns windowing (the agent keeps reporting +// its cumulative figure unchanged). Rolling (not calendar) windows: no fragile +// month-boundary reset, and "monthly" = a 30-day trailing window. + +// BudgetPeriod is one of the supported rolling budget windows. +type BudgetPeriod string + +const ( + PeriodHourly BudgetPeriod = "hourly" + PeriodDaily BudgetPeriod = "daily" + PeriodWeekly BudgetPeriod = "weekly" + PeriodMonthly BudgetPeriod = "monthly" +) + +// budgetPeriodDef pairs a period with its rolling window. +type budgetPeriodDef struct { + Name BudgetPeriod + Window time.Duration +} + +// budgetPeriods is the canonical ordered list. ADD A PERIOD = one line here; +// every consumer iterates this slice, so nothing else needs to change. +var budgetPeriods = []budgetPeriodDef{ + {PeriodHourly, time.Hour}, + {PeriodDaily, 24 * time.Hour}, + {PeriodWeekly, 7 * 24 * time.Hour}, + {PeriodMonthly, 30 * 24 * time.Hour}, // rolling 30-day window +} + +// spendLedgerRetention bounds the ledger: rows older than the largest window +// (+ slack) are never read, so the recorder opportunistically prunes them. +var spendLedgerRetention = 35 * 24 * time.Hour + +// parseBudgetLimits decodes the workspaces.budget_limits JSONB into a map of +// period → limit (USD cents). A limit of ZERO is valid and means "block all +// spend for that period" (a $0 ceiling); absent / null / negative / unknown +// keys mean "no limit for that period". Tolerant of a NULL/empty column. +func parseBudgetLimits(raw []byte) map[BudgetPeriod]int64 { + out := make(map[BudgetPeriod]int64, len(budgetPeriods)) + if len(raw) == 0 { + return out + } + var m map[string]*int64 + if err := json.Unmarshal(raw, &m); err != nil { + return out + } + for _, def := range budgetPeriods { + if v, ok := m[string(def.Name)]; ok && v != nil && *v >= 0 { + out[def.Name] = *v + } + } + return out +} + +// encodeBudgetLimits renders a period→limit map back to the canonical JSONB +// shape, keeping only KNOWN periods with a non-negative limit (0 = block-all is +// preserved; a period absent from the map = no limit). Always returns valid JSON. +func encodeBudgetLimits(limits map[BudgetPeriod]int64) []byte { + m := make(map[string]int64, len(limits)) + for _, def := range budgetPeriods { + if v, ok := limits[def.Name]; ok && v >= 0 { + m[string(def.Name)] = v + } + } + b, err := json.Marshal(m) + if err != nil { + return []byte("{}") + } + return b +} + +// recordSpendDelta appends a positive spend increment to the ledger and +// opportunistically prunes rows past the retention horizon for this workspace. +// No-op for delta <= 0. Errors are returned for the caller to log (non-fatal). +func recordSpendDelta(ctx context.Context, q *sql.DB, workspaceID string, deltaCents int64) error { + if deltaCents <= 0 { + return nil + } + if _, err := q.ExecContext(ctx, + `INSERT INTO workspace_spend_events (workspace_id, delta_cents) VALUES ($1, $2)`, + workspaceID, deltaCents, + ); err != nil { + return err + } + // Opportunistic prune (cheap; index-backed). Best-effort — ignore error. + _, _ = q.ExecContext(ctx, + `DELETE FROM workspace_spend_events + WHERE workspace_id = $1 AND occurred_at < now() - $2::interval`, + workspaceID, pgInterval(spendLedgerRetention), + ) + return nil +} + +// spendByPeriod returns the rolling-window spend (USD cents) for every period, +// computed in a SINGLE query over the ledger. The outer predicate bounds to the +// largest window; per-period FILTERs sum each sub-window. A period with no ledger +// rows reports 0. This is THE spend computation — used by both display + enforcement. +func spendByPeriod(ctx context.Context, q *sql.DB, workspaceID string) (map[BudgetPeriod]int64, error) { + out := make(map[BudgetPeriod]int64, len(budgetPeriods)) + for _, def := range budgetPeriods { + out[def.Name] = 0 + } + row := q.QueryRowContext(ctx, ` + SELECT + COALESCE(SUM(delta_cents) FILTER (WHERE occurred_at > now() - interval '1 hour'), 0), + COALESCE(SUM(delta_cents) FILTER (WHERE occurred_at > now() - interval '24 hours'), 0), + COALESCE(SUM(delta_cents) FILTER (WHERE occurred_at > now() - interval '7 days'), 0), + COALESCE(SUM(delta_cents) FILTER (WHERE occurred_at > now() - interval '30 days'), 0) + FROM workspace_spend_events + WHERE workspace_id = $1 AND occurred_at > now() - interval '30 days' + `, workspaceID) + var h, d, w, mo int64 + if err := row.Scan(&h, &d, &w, &mo); err != nil { + return out, err + } + out[PeriodHourly], out[PeriodDaily], out[PeriodWeekly], out[PeriodMonthly] = h, d, w, mo + return out, nil +} + +// exceededPeriods is PURE: given the configured limits and observed spend, it +// returns the periods whose spend has reached/exceeded their limit (in +// budgetPeriods order). Only periods WITH a positive limit are considered. +// Used by enforcement to decide whether to block. +func exceededPeriods(limits map[BudgetPeriod]int64, spend map[BudgetPeriod]int64) []BudgetPeriod { + var over []BudgetPeriod + for _, def := range budgetPeriods { + limit, ok := limits[def.Name] + if !ok { + continue // no limit configured for this period + } + // limit >= 0 is a real ceiling (0 = block-all). spend >= limit → over. + if spend[def.Name] >= limit { + over = append(over, def.Name) + } + } + return over +} + +// pgInterval renders a Go duration as a Postgres-interval string ("N seconds"). +func pgInterval(d time.Duration) string { + return strconv.FormatInt(int64(d.Seconds()), 10) + " seconds" +} diff --git a/workspace-server/internal/handlers/budget_periods_test.go b/workspace-server/internal/handlers/budget_periods_test.go new file mode 100644 index 000000000..fb2860972 --- /dev/null +++ b/workspace-server/internal/handlers/budget_periods_test.go @@ -0,0 +1,99 @@ +package handlers + +import ( + "reflect" + "testing" +) + +// Pure-logic tests for the multi-period budget SSOT (budget_periods.go). The +// DB-touching helpers (spendByPeriod / recordSpendDelta) are exercised via the +// handler sqlmock tests; here we pin the parsing + the over-budget decision, +// which is where the per-period semantics actually live. + +func TestParseBudgetLimits(t *testing.T) { + cases := []struct { + name string + raw string + want map[BudgetPeriod]int64 + }{ + {"empty", "", map[BudgetPeriod]int64{}}, + {"empty-object", "{}", map[BudgetPeriod]int64{}}, + {"all-four", `{"hourly":100,"daily":200,"weekly":300,"monthly":400}`, + map[BudgetPeriod]int64{PeriodHourly: 100, PeriodDaily: 200, PeriodWeekly: 300, PeriodMonthly: 400}}, + {"null-dropped-zero-kept", `{"hourly":null,"daily":0,"weekly":500}`, + map[BudgetPeriod]int64{PeriodDaily: 0, PeriodWeekly: 500}}, // 0 = block-all, kept + {"negative-dropped", `{"monthly":-5}`, map[BudgetPeriod]int64{}}, + {"unknown-key-ignored", `{"yearly":999,"daily":10}`, map[BudgetPeriod]int64{PeriodDaily: 10}}, + {"malformed-json", `{not json`, map[BudgetPeriod]int64{}}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got := parseBudgetLimits([]byte(tc.raw)) + if !reflect.DeepEqual(got, tc.want) { + t.Errorf("parseBudgetLimits(%q) = %v, want %v", tc.raw, got, tc.want) + } + }) + } +} + +func TestEncodeBudgetLimits_RoundTrip(t *testing.T) { + in := map[BudgetPeriod]int64{PeriodHourly: 100, PeriodMonthly: 400} + enc := encodeBudgetLimits(in) + got := parseBudgetLimits(enc) + if !reflect.DeepEqual(got, in) { + t.Errorf("round-trip: encode→parse = %v, want %v (enc=%s)", got, in, enc) + } + // unknown periods dropped; 0 (block-all) kept + enc2 := encodeBudgetLimits(map[BudgetPeriod]int64{PeriodDaily: 0, "yearly": 9}) + if got := parseBudgetLimits(enc2); !reflect.DeepEqual(got, map[BudgetPeriod]int64{PeriodDaily: 0}) { + t.Errorf("encode kept 0/dropped unknown: parse(%s) = %v, want {daily:0}", enc2, got) + } +} + +func TestExceededPeriods(t *testing.T) { + cases := []struct { + name string + limits map[BudgetPeriod]int64 + spend map[BudgetPeriod]int64 + want []BudgetPeriod + }{ + {"no-limits", map[BudgetPeriod]int64{}, map[BudgetPeriod]int64{PeriodHourly: 999}, nil}, + {"zero-limit-blocks-all", map[BudgetPeriod]int64{PeriodHourly: 0}, map[BudgetPeriod]int64{PeriodHourly: 0}, []BudgetPeriod{PeriodHourly}}, + {"under-all", map[BudgetPeriod]int64{PeriodDaily: 100}, map[BudgetPeriod]int64{PeriodDaily: 50}, nil}, + {"at-limit-is-exceeded", map[BudgetPeriod]int64{PeriodDaily: 100}, map[BudgetPeriod]int64{PeriodDaily: 100}, []BudgetPeriod{PeriodDaily}}, + {"over-limit", map[BudgetPeriod]int64{PeriodHourly: 10}, map[BudgetPeriod]int64{PeriodHourly: 11}, []BudgetPeriod{PeriodHourly}}, + {"only-hourly-over", map[BudgetPeriod]int64{PeriodHourly: 10, PeriodMonthly: 1000}, + map[BudgetPeriod]int64{PeriodHourly: 50, PeriodMonthly: 200}, []BudgetPeriod{PeriodHourly}}, + {"multiple-over-in-order", map[BudgetPeriod]int64{PeriodHourly: 10, PeriodWeekly: 100}, + map[BudgetPeriod]int64{PeriodHourly: 99, PeriodWeekly: 100}, []BudgetPeriod{PeriodHourly, PeriodWeekly}}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got := exceededPeriods(tc.limits, tc.spend) + if !reflect.DeepEqual(got, tc.want) { + t.Errorf("exceededPeriods(%v,%v) = %v, want %v", tc.limits, tc.spend, got, tc.want) + } + }) + } +} + +// TestBudgetPeriods_AllReachable guards the SSOT list: every declared period has +// a positive window and a unique name (a typo'd duplicate would silently break +// per-period accounting). +func TestBudgetPeriods_Wellformed(t *testing.T) { + seen := map[BudgetPeriod]bool{} + for _, d := range budgetPeriods { + if d.Window <= 0 { + t.Errorf("period %s has non-positive window %v", d.Name, d.Window) + } + if seen[d.Name] { + t.Errorf("duplicate period name %s", d.Name) + } + seen[d.Name] = true + } + for _, p := range []BudgetPeriod{PeriodHourly, PeriodDaily, PeriodWeekly, PeriodMonthly} { + if !seen[p] { + t.Errorf("period %s missing from budgetPeriods SSOT list", p) + } + } +} diff --git a/workspace-server/internal/handlers/budget_test.go b/workspace-server/internal/handlers/budget_test.go index e3e6cacd3..4f5d29389 100644 --- a/workspace-server/internal/handlers/budget_test.go +++ b/workspace-server/internal/handlers/budget_test.go @@ -12,15 +12,25 @@ import ( "github.com/gin-gonic/gin" ) +// Multi-period budget (#49): GET/PATCH now read workspaces.budget_limits (jsonb) +// and compute per-period spend from the workspace_spend_events ledger +// (spendByPeriod — matched here by the "FROM workspace_spend_events" fragment). +// The legacy budget_limit/monthly_spend response fields are still emitted +// (monthly period) for rollout back-compat, and the legacy {"budget_limit":N} +// PATCH shape still works. + +// spendRows builds the 4-column row spendByPeriod scans (hourly,daily,weekly,monthly). +func spendRows(h, d, w, m int64) *sqlmock.Rows { + return sqlmock.NewRows([]string{"h", "d", "w", "mo"}).AddRow(h, d, w, m) +} + // ==================== GET /workspaces/:id/budget ==================== -// TestBudgetGet_NotFound verifies that GET /budget returns 404 for an unknown -// workspace ID (ErrNoRows from the budget query). func TestBudgetGet_NotFound(t *testing.T) { mock := setupTestDB(t) setupTestRedis(t) - mock.ExpectQuery(`SELECT budget_limit, COALESCE\(monthly_spend, 0\)`). + mock.ExpectQuery(`SELECT COALESCE\(budget_limits`). WithArgs("ws-not-there"). WillReturnError(sql.ErrNoRows) @@ -29,8 +39,7 @@ func TestBudgetGet_NotFound(t *testing.T) { c.Params = gin.Params{{Key: "id", Value: "ws-not-there"}} c.Request = httptest.NewRequest("GET", "/workspaces/ws-not-there/budget", nil) - h := NewBudgetHandler() - h.GetBudget(c) + NewBudgetHandler().GetBudget(c) if w.Code != http.StatusNotFound { t.Errorf("expected 404, got %d: %s", w.Code, w.Body.String()) @@ -40,12 +49,11 @@ func TestBudgetGet_NotFound(t *testing.T) { } } -// TestBudgetGet_DBError verifies that a non-ErrNoRows DB error returns 500. func TestBudgetGet_DBError(t *testing.T) { mock := setupTestDB(t) setupTestRedis(t) - mock.ExpectQuery(`SELECT budget_limit, COALESCE\(monthly_spend, 0\)`). + mock.ExpectQuery(`SELECT COALESCE\(budget_limits`). WithArgs("ws-db-err"). WillReturnError(sql.ErrConnDone) @@ -54,8 +62,7 @@ func TestBudgetGet_DBError(t *testing.T) { c.Params = gin.Params{{Key: "id", Value: "ws-db-err"}} c.Request = httptest.NewRequest("GET", "/workspaces/ws-db-err/budget", nil) - h := NewBudgetHandler() - h.GetBudget(c) + NewBudgetHandler().GetBudget(c) if w.Code != http.StatusInternalServerError { t.Errorf("expected 500, got %d: %s", w.Code, w.Body.String()) @@ -65,24 +72,23 @@ func TestBudgetGet_DBError(t *testing.T) { } } -// TestBudgetGet_NoLimit verifies that budget_limit and budget_remaining are -// null when the workspace has no budget ceiling configured. func TestBudgetGet_NoLimit(t *testing.T) { mock := setupTestDB(t) setupTestRedis(t) - mock.ExpectQuery(`SELECT budget_limit, COALESCE\(monthly_spend, 0\)`). + mock.ExpectQuery(`SELECT COALESCE\(budget_limits`). WithArgs("ws-free"). - WillReturnRows(sqlmock.NewRows([]string{"budget_limit", "monthly_spend"}). - AddRow(nil, int64(42))) + WillReturnRows(sqlmock.NewRows([]string{"budget_limits"}).AddRow([]byte(`{}`))) + mock.ExpectQuery(`FROM workspace_spend_events`). + WithArgs("ws-free"). + WillReturnRows(spendRows(0, 0, 0, 42)) w := httptest.NewRecorder() c, _ := gin.CreateTestContext(w) c.Params = gin.Params{{Key: "id", Value: "ws-free"}} c.Request = httptest.NewRequest("GET", "/workspaces/ws-free/budget", nil) - h := NewBudgetHandler() - h.GetBudget(c) + NewBudgetHandler().GetBudget(c) if w.Code != http.StatusOK { t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String()) @@ -105,24 +111,23 @@ func TestBudgetGet_NoLimit(t *testing.T) { } } -// TestBudgetGet_WithLimit verifies that budget_limit, monthly_spend, and -// budget_remaining are all returned correctly when a ceiling is set. func TestBudgetGet_WithLimit(t *testing.T) { mock := setupTestDB(t) setupTestRedis(t) - mock.ExpectQuery(`SELECT budget_limit, COALESCE\(monthly_spend, 0\)`). + mock.ExpectQuery(`SELECT COALESCE\(budget_limits`). WithArgs("ws-capped"). - WillReturnRows(sqlmock.NewRows([]string{"budget_limit", "monthly_spend"}). - AddRow(int64(500), int64(123))) + WillReturnRows(sqlmock.NewRows([]string{"budget_limits"}).AddRow([]byte(`{"monthly":500}`))) + mock.ExpectQuery(`FROM workspace_spend_events`). + WithArgs("ws-capped"). + WillReturnRows(spendRows(0, 0, 0, 123)) w := httptest.NewRecorder() c, _ := gin.CreateTestContext(w) c.Params = gin.Params{{Key: "id", Value: "ws-capped"}} c.Request = httptest.NewRequest("GET", "/workspaces/ws-capped/budget", nil) - h := NewBudgetHandler() - h.GetBudget(c) + NewBudgetHandler().GetBudget(c) if w.Code != http.StatusOK { t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String()) @@ -137,7 +142,6 @@ func TestBudgetGet_WithLimit(t *testing.T) { if resp["monthly_spend"] != float64(123) { t.Errorf("expected monthly_spend=123, got %v", resp["monthly_spend"]) } - // budget_remaining = 500 - 123 = 377 if resp["budget_remaining"] != float64(377) { t.Errorf("expected budget_remaining=377, got %v", resp["budget_remaining"]) } @@ -146,24 +150,23 @@ func TestBudgetGet_WithLimit(t *testing.T) { } } -// TestBudgetGet_OverBudget verifies that budget_remaining can be negative -// when monthly_spend has already exceeded budget_limit. func TestBudgetGet_OverBudget(t *testing.T) { mock := setupTestDB(t) setupTestRedis(t) - mock.ExpectQuery(`SELECT budget_limit, COALESCE\(monthly_spend, 0\)`). + mock.ExpectQuery(`SELECT COALESCE\(budget_limits`). WithArgs("ws-over"). - WillReturnRows(sqlmock.NewRows([]string{"budget_limit", "monthly_spend"}). - AddRow(int64(100), int64(150))) + WillReturnRows(sqlmock.NewRows([]string{"budget_limits"}).AddRow([]byte(`{"monthly":100}`))) + mock.ExpectQuery(`FROM workspace_spend_events`). + WithArgs("ws-over"). + WillReturnRows(spendRows(0, 0, 0, 150)) w := httptest.NewRecorder() c, _ := gin.CreateTestContext(w) c.Params = gin.Params{{Key: "id", Value: "ws-over"}} c.Request = httptest.NewRequest("GET", "/workspaces/ws-over/budget", nil) - h := NewBudgetHandler() - h.GetBudget(c) + NewBudgetHandler().GetBudget(c) if w.Code != http.StatusOK { t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String()) @@ -172,7 +175,6 @@ func TestBudgetGet_OverBudget(t *testing.T) { if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { t.Fatalf("parse response: %v", err) } - // budget_remaining = 100 - 150 = -50 (negative, but we store actual value) if resp["budget_remaining"] != float64(-50) { t.Errorf("expected budget_remaining=-50, got %v", resp["budget_remaining"]) } @@ -181,10 +183,59 @@ func TestBudgetGet_OverBudget(t *testing.T) { } } +// TestBudgetGet_MultiPeriod pins the new per-period shape: each period reports +// its own limit/spend/remaining, and an over-budget sub-period is visible. +func TestBudgetGet_MultiPeriod(t *testing.T) { + mock := setupTestDB(t) + setupTestRedis(t) + + mock.ExpectQuery(`SELECT COALESCE\(budget_limits`). + WithArgs("ws-mp"). + WillReturnRows(sqlmock.NewRows([]string{"budget_limits"}). + AddRow([]byte(`{"hourly":100,"daily":1000}`))) + mock.ExpectQuery(`FROM workspace_spend_events`). + WithArgs("ws-mp"). + WillReturnRows(spendRows(120, 300, 300, 300)) // hourly over (120>=100) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Params = gin.Params{{Key: "id", Value: "ws-mp"}} + c.Request = httptest.NewRequest("GET", "/workspaces/ws-mp/budget", nil) + + NewBudgetHandler().GetBudget(c) + + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String()) + } + var resp struct { + Periods map[string]struct { + Limit *int64 `json:"limit"` + Spend int64 `json:"spend"` + Remaining *int64 `json:"remaining"` + } `json:"periods"` + } + if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatalf("parse response: %v", err) + } + if resp.Periods["hourly"].Limit == nil || *resp.Periods["hourly"].Limit != 100 { + t.Errorf("hourly.limit: want 100, got %v", resp.Periods["hourly"].Limit) + } + if resp.Periods["hourly"].Spend != 120 { + t.Errorf("hourly.spend: want 120, got %d", resp.Periods["hourly"].Spend) + } + if r := resp.Periods["hourly"].Remaining; r == nil || *r != -20 { + t.Errorf("hourly.remaining: want -20, got %v", r) + } + if resp.Periods["weekly"].Limit != nil { + t.Errorf("weekly.limit: want null (unset), got %v", resp.Periods["weekly"].Limit) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("sqlmock expectations not met: %v", err) + } +} + // ==================== PATCH /workspaces/:id/budget ==================== -// TestBudgetPatch_MissingField verifies that PATCH /budget with no budget_limit -// field in the body returns 400. func TestBudgetPatch_MissingField(t *testing.T) { setupTestDB(t) setupTestRedis(t) @@ -196,15 +247,13 @@ func TestBudgetPatch_MissingField(t *testing.T) { bytes.NewBufferString(`{"other_field":123}`)) c.Request.Header.Set("Content-Type", "application/json") - h := NewBudgetHandler() - h.PatchBudget(c) + NewBudgetHandler().PatchBudget(c) if w.Code != http.StatusBadRequest { t.Errorf("expected 400, got %d: %s", w.Code, w.Body.String()) } } -// TestBudgetPatch_InvalidBody verifies that a malformed JSON body returns 400. func TestBudgetPatch_InvalidBody(t *testing.T) { setupTestDB(t) setupTestRedis(t) @@ -216,15 +265,13 @@ func TestBudgetPatch_InvalidBody(t *testing.T) { bytes.NewBufferString(`not json`)) c.Request.Header.Set("Content-Type", "application/json") - h := NewBudgetHandler() - h.PatchBudget(c) + NewBudgetHandler().PatchBudget(c) if w.Code != http.StatusBadRequest { t.Errorf("expected 400, got %d: %s", w.Code, w.Body.String()) } } -// TestBudgetPatch_NegativeValue verifies that a negative budget_limit is rejected. func TestBudgetPatch_NegativeValue(t *testing.T) { setupTestDB(t) setupTestRedis(t) @@ -236,15 +283,13 @@ func TestBudgetPatch_NegativeValue(t *testing.T) { bytes.NewBufferString(`{"budget_limit":-1}`)) c.Request.Header.Set("Content-Type", "application/json") - h := NewBudgetHandler() - h.PatchBudget(c) + NewBudgetHandler().PatchBudget(c) if w.Code != http.StatusBadRequest { t.Errorf("expected 400 for negative budget_limit, got %d: %s", w.Code, w.Body.String()) } } -// TestBudgetPatch_InvalidType verifies that a non-numeric budget_limit returns 400. func TestBudgetPatch_InvalidType(t *testing.T) { setupTestDB(t) setupTestRedis(t) @@ -256,16 +301,32 @@ func TestBudgetPatch_InvalidType(t *testing.T) { bytes.NewBufferString(`{"budget_limit":"not-a-number"}`)) c.Request.Header.Set("Content-Type", "application/json") - h := NewBudgetHandler() - h.PatchBudget(c) + NewBudgetHandler().PatchBudget(c) if w.Code != http.StatusBadRequest { t.Errorf("expected 400 for string budget_limit, got %d: %s", w.Code, w.Body.String()) } } -// TestBudgetPatch_WorkspaceNotFound verifies that PATCH /budget returns 404 -// when the workspace doesn't exist. +// TestBudgetPatch_UnknownPeriod rejects an unsupported period key. +func TestBudgetPatch_UnknownPeriod(t *testing.T) { + setupTestDB(t) + setupTestRedis(t) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Params = gin.Params{{Key: "id", Value: "ws-badperiod"}} + c.Request = httptest.NewRequest("PATCH", "/workspaces/ws-badperiod/budget", + bytes.NewBufferString(`{"budget_limits":{"yearly":100}}`)) + c.Request.Header.Set("Content-Type", "application/json") + + NewBudgetHandler().PatchBudget(c) + + if w.Code != http.StatusBadRequest { + t.Errorf("expected 400 for unknown period, got %d: %s", w.Code, w.Body.String()) + } +} + func TestBudgetPatch_WorkspaceNotFound(t *testing.T) { mock := setupTestDB(t) setupTestRedis(t) @@ -281,8 +342,7 @@ func TestBudgetPatch_WorkspaceNotFound(t *testing.T) { bytes.NewBufferString(`{"budget_limit":500}`)) c.Request.Header.Set("Content-Type", "application/json") - h := NewBudgetHandler() - h.PatchBudget(c) + NewBudgetHandler().PatchBudget(c) if w.Code != http.StatusNotFound { t.Errorf("expected 404, got %d: %s", w.Code, w.Body.String()) @@ -292,25 +352,20 @@ func TestBudgetPatch_WorkspaceNotFound(t *testing.T) { } } -// TestBudgetPatch_SetLimit verifies that PATCH /budget with a positive value -// updates the DB and returns the new budget state. +// TestBudgetPatch_SetLimit (legacy monthly shape) updates + returns new state. func TestBudgetPatch_SetLimit(t *testing.T) { mock := setupTestDB(t) setupTestRedis(t) - // Existence probe mock.ExpectQuery(`SELECT EXISTS.*status != 'removed'`). WithArgs("ws-set-limit"). WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(true)) - // UPDATE - mock.ExpectExec(`UPDATE workspaces SET budget_limit`). - WithArgs("ws-set-limit", int64(500)). + mock.ExpectExec(`UPDATE workspaces SET budget_limits`). + WithArgs("ws-set-limit", sqlmock.AnyArg(), int64(500)). WillReturnResult(sqlmock.NewResult(0, 1)) - // Re-read for response - mock.ExpectQuery(`SELECT budget_limit, COALESCE\(monthly_spend, 0\) FROM workspaces WHERE id`). + mock.ExpectQuery(`FROM workspace_spend_events`). WithArgs("ws-set-limit"). - WillReturnRows(sqlmock.NewRows([]string{"budget_limit", "monthly_spend"}). - AddRow(int64(500), int64(200))) + WillReturnRows(spendRows(0, 0, 0, 200)) w := httptest.NewRecorder() c, _ := gin.CreateTestContext(w) @@ -319,8 +374,7 @@ func TestBudgetPatch_SetLimit(t *testing.T) { bytes.NewBufferString(`{"budget_limit":500}`)) c.Request.Header.Set("Content-Type", "application/json") - h := NewBudgetHandler() - h.PatchBudget(c) + NewBudgetHandler().PatchBudget(c) if w.Code != http.StatusOK { t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String()) @@ -335,7 +389,6 @@ func TestBudgetPatch_SetLimit(t *testing.T) { if resp["monthly_spend"] != float64(200) { t.Errorf("expected monthly_spend=200, got %v", resp["monthly_spend"]) } - // budget_remaining = 500 - 200 = 300 if resp["budget_remaining"] != float64(300) { t.Errorf("expected budget_remaining=300, got %v", resp["budget_remaining"]) } @@ -344,8 +397,59 @@ func TestBudgetPatch_SetLimit(t *testing.T) { } } -// TestBudgetPatch_ClearLimit verifies that PATCH /budget with budget_limit=null -// clears the ceiling, making budget_limit and budget_remaining null in the response. +// TestBudgetPatch_SetMultiPeriod sets several periods at once and verifies the +// per-period response. +func TestBudgetPatch_SetMultiPeriod(t *testing.T) { + mock := setupTestDB(t) + setupTestRedis(t) + + mock.ExpectQuery(`SELECT EXISTS.*status != 'removed'`). + WithArgs("ws-mp-set"). + WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(true)) + // no monthly in payload → legacy budget_limit column set to NULL + mock.ExpectExec(`UPDATE workspaces SET budget_limits`). + WithArgs("ws-mp-set", sqlmock.AnyArg(), nil). + WillReturnResult(sqlmock.NewResult(0, 1)) + mock.ExpectQuery(`FROM workspace_spend_events`). + WithArgs("ws-mp-set"). + WillReturnRows(spendRows(10, 20, 30, 40)) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Params = gin.Params{{Key: "id", Value: "ws-mp-set"}} + c.Request = httptest.NewRequest("PATCH", "/workspaces/ws-mp-set/budget", + bytes.NewBufferString(`{"budget_limits":{"hourly":100,"daily":200,"monthly":null}}`)) + c.Request.Header.Set("Content-Type", "application/json") + + NewBudgetHandler().PatchBudget(c) + + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String()) + } + var resp struct { + Periods map[string]struct { + Limit *int64 `json:"limit"` + Spend int64 `json:"spend"` + } `json:"periods"` + BudgetLimit *int64 `json:"budget_limit"` + } + if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatalf("parse response: %v", err) + } + if resp.Periods["hourly"].Limit == nil || *resp.Periods["hourly"].Limit != 100 { + t.Errorf("hourly.limit want 100, got %v", resp.Periods["hourly"].Limit) + } + if resp.Periods["daily"].Limit == nil || *resp.Periods["daily"].Limit != 200 { + t.Errorf("daily.limit want 200, got %v", resp.Periods["daily"].Limit) + } + if resp.BudgetLimit != nil { + t.Errorf("monthly cleared → budget_limit should be null, got %v", *resp.BudgetLimit) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("sqlmock expectations not met: %v", err) + } +} + func TestBudgetPatch_ClearLimit(t *testing.T) { mock := setupTestDB(t) setupTestRedis(t) @@ -353,15 +457,12 @@ func TestBudgetPatch_ClearLimit(t *testing.T) { mock.ExpectQuery(`SELECT EXISTS.*status != 'removed'`). WithArgs("ws-clear-limit"). WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(true)) - // UPDATE with NULL - mock.ExpectExec(`UPDATE workspaces SET budget_limit`). - WithArgs("ws-clear-limit", nil). + mock.ExpectExec(`UPDATE workspaces SET budget_limits`). + WithArgs("ws-clear-limit", sqlmock.AnyArg(), nil). WillReturnResult(sqlmock.NewResult(0, 1)) - // Re-read — budget_limit is now NULL - mock.ExpectQuery(`SELECT budget_limit, COALESCE\(monthly_spend, 0\) FROM workspaces WHERE id`). + mock.ExpectQuery(`FROM workspace_spend_events`). WithArgs("ws-clear-limit"). - WillReturnRows(sqlmock.NewRows([]string{"budget_limit", "monthly_spend"}). - AddRow(nil, int64(50))) + WillReturnRows(spendRows(0, 0, 0, 50)) w := httptest.NewRecorder() c, _ := gin.CreateTestContext(w) @@ -370,8 +471,7 @@ func TestBudgetPatch_ClearLimit(t *testing.T) { bytes.NewBufferString(`{"budget_limit":null}`)) c.Request.Header.Set("Content-Type", "application/json") - h := NewBudgetHandler() - h.PatchBudget(c) + NewBudgetHandler().PatchBudget(c) if w.Code != http.StatusOK { t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String()) @@ -391,8 +491,6 @@ func TestBudgetPatch_ClearLimit(t *testing.T) { } } -// TestBudgetPatch_UpdateDBError verifies that a DB error during the UPDATE -// returns 500. func TestBudgetPatch_UpdateDBError(t *testing.T) { mock := setupTestDB(t) setupTestRedis(t) @@ -400,8 +498,8 @@ func TestBudgetPatch_UpdateDBError(t *testing.T) { mock.ExpectQuery(`SELECT EXISTS.*status != 'removed'`). WithArgs("ws-patch-dberr"). WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(true)) - mock.ExpectExec(`UPDATE workspaces SET budget_limit`). - WithArgs("ws-patch-dberr", int64(500)). + mock.ExpectExec(`UPDATE workspaces SET budget_limits`). + WithArgs("ws-patch-dberr", sqlmock.AnyArg(), int64(500)). WillReturnError(sql.ErrConnDone) w := httptest.NewRecorder() @@ -411,8 +509,7 @@ func TestBudgetPatch_UpdateDBError(t *testing.T) { bytes.NewBufferString(`{"budget_limit":500}`)) c.Request.Header.Set("Content-Type", "application/json") - h := NewBudgetHandler() - h.PatchBudget(c) + NewBudgetHandler().PatchBudget(c) if w.Code != http.StatusInternalServerError { t.Errorf("expected 500 on UPDATE error, got %d: %s", w.Code, w.Body.String()) @@ -422,8 +519,8 @@ func TestBudgetPatch_UpdateDBError(t *testing.T) { } } -// TestBudgetPatch_ZeroLimit verifies that budget_limit=0 is accepted (it means -// every A2A call is blocked — useful to pause a workspace's LLM spend entirely). +// TestBudgetPatch_ZeroLimit verifies budget_limit=0 is accepted + stored (0 = +// block-all: every period call is blocked — pauses the workspace's spend). func TestBudgetPatch_ZeroLimit(t *testing.T) { mock := setupTestDB(t) setupTestRedis(t) @@ -431,13 +528,12 @@ func TestBudgetPatch_ZeroLimit(t *testing.T) { mock.ExpectQuery(`SELECT EXISTS.*status != 'removed'`). WithArgs("ws-zero-limit"). WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(true)) - mock.ExpectExec(`UPDATE workspaces SET budget_limit`). - WithArgs("ws-zero-limit", int64(0)). + mock.ExpectExec(`UPDATE workspaces SET budget_limits`). + WithArgs("ws-zero-limit", sqlmock.AnyArg(), int64(0)). WillReturnResult(sqlmock.NewResult(0, 1)) - mock.ExpectQuery(`SELECT budget_limit, COALESCE\(monthly_spend, 0\) FROM workspaces WHERE id`). + mock.ExpectQuery(`FROM workspace_spend_events`). WithArgs("ws-zero-limit"). - WillReturnRows(sqlmock.NewRows([]string{"budget_limit", "monthly_spend"}). - AddRow(int64(0), int64(0))) + WillReturnRows(spendRows(0, 0, 0, 0)) w := httptest.NewRecorder() c, _ := gin.CreateTestContext(w) @@ -446,11 +542,17 @@ func TestBudgetPatch_ZeroLimit(t *testing.T) { bytes.NewBufferString(`{"budget_limit":0}`)) c.Request.Header.Set("Content-Type", "application/json") - h := NewBudgetHandler() - h.PatchBudget(c) + NewBudgetHandler().PatchBudget(c) if w.Code != http.StatusOK { - t.Errorf("expected 200 for zero budget_limit, got %d: %s", w.Code, w.Body.String()) + t.Fatalf("expected 200 for zero budget_limit, got %d: %s", w.Code, w.Body.String()) + } + var resp map[string]interface{} + if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatalf("parse response: %v", err) + } + if resp["budget_limit"] != float64(0) { + t.Errorf("expected budget_limit=0 (block-all), got %v", resp["budget_limit"]) } if err := mock.ExpectationsWereMet(); err != nil { t.Errorf("sqlmock expectations not met: %v", err) diff --git a/workspace-server/internal/handlers/handlers_test.go b/workspace-server/internal/handlers/handlers_test.go index fbfc134b7..7424b10b2 100644 --- a/workspace-server/internal/handlers/handlers_test.go +++ b/workspace-server/internal/handlers/handlers_test.go @@ -12,12 +12,12 @@ import ( "testing" "time" - "github.com/DATA-DOG/go-sqlmock" "git.moleculesai.app/molecule-ai/molecule-core/workspace-server/internal/db" "git.moleculesai.app/molecule-ai/molecule-core/workspace-server/internal/events" "git.moleculesai.app/molecule-ai/molecule-core/workspace-server/internal/models" "git.moleculesai.app/molecule-ai/molecule-core/workspace-server/internal/ws" "git.moleculesai.app/molecule-ai/molecule-core/workspace-server/internal/wsauth" + "github.com/DATA-DOG/go-sqlmock" "github.com/alicebob/miniredis/v2" "github.com/gin-gonic/gin" "github.com/redis/go-redis/v9" @@ -158,9 +158,11 @@ func allowLoopbackForTest(t *testing.T) { // handler in the 2026-04-18 restructure but the tests never caught up, // leaving Platform (Go) CI red for weeks. func expectBudgetCheck(mock sqlmock.Sqlmock, workspaceID string) { - mock.ExpectQuery(`SELECT budget_limit, COALESCE\(monthly_spend, 0\) FROM workspaces WHERE id = \$1`). + // Multi-period (#49): checkWorkspaceBudget reads budget_limits jsonb. An + // empty map → no limits → returns early (no spend query), enforcement skipped. + mock.ExpectQuery(`SELECT COALESCE\(budget_limits`). WithArgs(workspaceID). - WillReturnRows(sqlmock.NewRows([]string{"budget_limit", "monthly_spend"})) + WillReturnRows(sqlmock.NewRows([]string{"budget_limits"}).AddRow([]byte("{}"))) } // ---------- TestRegisterHandler ---------- diff --git a/workspace-server/internal/handlers/registry.go b/workspace-server/internal/handlers/registry.go index 5c40edd46..9d7ffd58d 100644 --- a/workspace-server/internal/handlers/registry.go +++ b/workspace-server/internal/handlers/registry.go @@ -538,7 +538,8 @@ func (h *RegistryHandler) Heartbeat(c *gin.Context) { // Read previous current_task to detect changes (before the UPDATE) var prevTask string - if err := db.DB.QueryRowContext(ctx, `SELECT COALESCE(current_task, '') FROM workspaces WHERE id = $1`, payload.WorkspaceID).Scan(&prevTask); err != nil { + var prevSpend int64 + if err := db.DB.QueryRowContext(ctx, `SELECT COALESCE(current_task, ''), COALESCE(monthly_spend, 0) FROM workspaces WHERE id = $1`, payload.WorkspaceID).Scan(&prevTask, &prevSpend); err != nil { log.Printf("registry heartbeat: prev_task query failed for workspace %s: %v", payload.WorkspaceID, err) } @@ -556,6 +557,25 @@ func (h *RegistryHandler) Heartbeat(c *gin.Context) { payload.MonthlySpend = maxMonthlySpend } + // Multi-period budget (#49): record the spend INCREMENT into the + // workspace_spend_events ledger so the server can compute rolling per-period + // windows (hourly/daily/weekly/monthly) — see budget_periods.go. The agent + // still reports a cumulative monthly figure; we derive the delta vs the + // last-seen cumulative (prevSpend). A DECREASE means the agent reset its + // monthly cumulative (new month) → treat the new value as fresh spend. + // Best-effort: a ledger failure must never break the heartbeat. + if payload.MonthlySpend > 0 { + delta := payload.MonthlySpend - prevSpend + if delta < 0 { + delta = payload.MonthlySpend + } + if delta > 0 { + if err := recordSpendDelta(ctx, db.DB, payload.WorkspaceID, delta); err != nil { + log.Printf("registry heartbeat: spend-ledger insert failed for workspace %s: %v", payload.WorkspaceID, err) + } + } + } + // Update heartbeat columns. #73 guard: exclude 'removed' rows so a // late heartbeat from a container that's being torn down doesn't // refresh last_heartbeat_at on a tombstoned workspace (which would diff --git a/workspace-server/internal/handlers/workspace_budget_test.go b/workspace-server/internal/handlers/workspace_budget_test.go index 4a467ae84..89abede5b 100644 --- a/workspace-server/internal/handlers/workspace_budget_test.go +++ b/workspace-server/internal/handlers/workspace_budget_test.go @@ -22,8 +22,8 @@ import ( "testing" "time" - "github.com/DATA-DOG/go-sqlmock" "git.moleculesai.app/molecule-ai/molecule-core/workspace-server/internal/models" + "github.com/DATA-DOG/go-sqlmock" "github.com/gin-gonic/gin" ) @@ -259,11 +259,13 @@ func TestWorkspaceBudget_A2A_ExceededReturns402(t *testing.T) { // Cache a URL so resolveAgentURL doesn't need a DB query after budget check mr.Set(fmt.Sprintf("ws:%s:url", "ws-over-budget"), "http://localhost:9999") - // Budget check query: spend = limit → exceeded - mock.ExpectQuery("SELECT budget_limit, COALESCE"). + // Budget check: monthly limit 500, monthly spend 500 → exceeded → 402 + mock.ExpectQuery(`SELECT COALESCE\(budget_limits`). WithArgs("ws-over-budget"). - WillReturnRows(sqlmock.NewRows([]string{"budget_limit", "monthly_spend"}). - AddRow(int64(500), int64(500))) + WillReturnRows(sqlmock.NewRows([]string{"budget_limits"}).AddRow([]byte(`{"monthly":500}`))) + mock.ExpectQuery(`FROM workspace_spend_events`). + WithArgs("ws-over-budget"). + WillReturnRows(spendRows(0, 0, 0, 500)) w := httptest.NewRecorder() c, _ := gin.CreateTestContext(w) @@ -295,10 +297,12 @@ func TestWorkspaceBudget_A2A_AboveLimitReturns402(t *testing.T) { mr.Set(fmt.Sprintf("ws:%s:url", "ws-way-over"), "http://localhost:9999") // spend > limit - mock.ExpectQuery("SELECT budget_limit, COALESCE"). + mock.ExpectQuery(`SELECT COALESCE\(budget_limits`). WithArgs("ws-way-over"). - WillReturnRows(sqlmock.NewRows([]string{"budget_limit", "monthly_spend"}). - AddRow(int64(100), int64(9999))) + WillReturnRows(sqlmock.NewRows([]string{"budget_limits"}).AddRow([]byte(`{"monthly":100}`))) + mock.ExpectQuery(`FROM workspace_spend_events`). + WithArgs("ws-way-over"). + WillReturnRows(spendRows(0, 0, 0, 9999)) w := httptest.NewRecorder() c, _ := gin.CreateTestContext(w) @@ -334,11 +338,13 @@ func TestWorkspaceBudget_A2A_UnderLimitPassesThrough(t *testing.T) { mr.Set(fmt.Sprintf("ws:%s:url", "ws-under-budget"), agentServer.URL) - // Budget check: spend (100) < limit (500) → pass-through - mock.ExpectQuery("SELECT budget_limit, COALESCE"). + // Budget check: monthly spend (100) < limit (500) → pass-through + mock.ExpectQuery(`SELECT COALESCE\(budget_limits`). WithArgs("ws-under-budget"). - WillReturnRows(sqlmock.NewRows([]string{"budget_limit", "monthly_spend"}). - AddRow(int64(500), int64(100))) + WillReturnRows(sqlmock.NewRows([]string{"budget_limits"}).AddRow([]byte(`{"monthly":500}`))) + mock.ExpectQuery(`FROM workspace_spend_events`). + WithArgs("ws-under-budget"). + WillReturnRows(spendRows(0, 0, 0, 100)) // Activity log INSERT from logA2ASuccess mock.ExpectExec("INSERT INTO activity_logs"). @@ -380,11 +386,11 @@ func TestWorkspaceBudget_A2A_NilLimitPassesThrough(t *testing.T) { mr.Set(fmt.Sprintf("ws:%s:url", "ws-no-limit"), agentServer.URL) - // budget_limit NULL → no enforcement regardless of monthly_spend - mock.ExpectQuery("SELECT budget_limit, COALESCE"). + // no limits configured → checkWorkspaceBudget returns early (no spend query), + // enforcement skipped regardless of spend + mock.ExpectQuery(`SELECT COALESCE\(budget_limits`). WithArgs("ws-no-limit"). - WillReturnRows(sqlmock.NewRows([]string{"budget_limit", "monthly_spend"}). - AddRow(nil, int64(999999))) // huge spend but no limit set + WillReturnRows(sqlmock.NewRows([]string{"budget_limits"}).AddRow([]byte(`{}`))) mock.ExpectExec("INSERT INTO activity_logs"). WillReturnResult(sqlmock.NewResult(0, 1)) @@ -425,7 +431,7 @@ func TestWorkspaceBudget_A2A_DBErrorFailOpen(t *testing.T) { mr.Set(fmt.Sprintf("ws:%s:url", "ws-db-err-budget"), agentServer.URL) // Budget check fails with DB error → fail-open (request proceeds) - mock.ExpectQuery("SELECT budget_limit, COALESCE"). + mock.ExpectQuery(`SELECT COALESCE\(budget_limits`). WithArgs("ws-db-err-budget"). WillReturnError(sql.ErrConnDone) diff --git a/workspace-server/migrations/20260529000000_workspace_multiperiod_budget.down.sql b/workspace-server/migrations/20260529000000_workspace_multiperiod_budget.down.sql new file mode 100644 index 000000000..b6f6c41ac --- /dev/null +++ b/workspace-server/migrations/20260529000000_workspace_multiperiod_budget.down.sql @@ -0,0 +1,2 @@ +DROP TABLE IF EXISTS workspace_spend_events; +ALTER TABLE workspaces DROP COLUMN IF EXISTS budget_limits; diff --git a/workspace-server/migrations/20260529000000_workspace_multiperiod_budget.up.sql b/workspace-server/migrations/20260529000000_workspace_multiperiod_budget.up.sql new file mode 100644 index 000000000..e789bfff1 --- /dev/null +++ b/workspace-server/migrations/20260529000000_workspace_multiperiod_budget.up.sql @@ -0,0 +1,30 @@ +-- Multi-period per-workspace LLM budget (hourly/daily/weekly/monthly). +-- Extends the single monthly budget_limit (027). `budget_limits` is the SSOT +-- for the per-period ceilings: a JSONB map {"hourly":N,"daily":N,"weekly":N, +-- "monthly":N} in USD cents; a key that is absent or null = no limit for that +-- period. Per-period SPEND is computed from the workspace_spend_events ledger +-- over a rolling window (NOT the legacy self-reported monthly_spend cumulative, +-- which can't express sub-month periods). +ALTER TABLE workspaces + ADD COLUMN IF NOT EXISTS budget_limits JSONB NOT NULL DEFAULT '{}'::jsonb; + +-- Backfill: carry an existing monthly ceiling into the new map so the feature +-- is continuous across the rollout (027's budget_limit stays for back-compat). +UPDATE workspaces + SET budget_limits = jsonb_build_object('monthly', budget_limit) + WHERE budget_limit IS NOT NULL + AND NOT (budget_limits ? 'monthly'); + +-- Server-owned spend ledger: one row per heartbeat-observed spend INCREMENT +-- (delta = new cumulative - prev). Per-period spend = +-- SUM(delta_cents) WHERE workspace_id=$1 AND occurred_at > now() - . +-- Makes the SERVER the SSOT for windowing; the agent keeps reporting its +-- cumulative figure unchanged (the heartbeat derives the delta). +CREATE TABLE IF NOT EXISTS workspace_spend_events ( + id BIGSERIAL PRIMARY KEY, + workspace_id TEXT NOT NULL, + delta_cents BIGINT NOT NULL CHECK (delta_cents > 0), + occurred_at TIMESTAMPTZ NOT NULL DEFAULT now() +); +CREATE INDEX IF NOT EXISTS idx_workspace_spend_events_ws_time + ON workspace_spend_events (workspace_id, occurred_at DESC); -- 2.52.0