molecule-ai/molecule-core
52
0
Fork 2
Code Issues 183 Pull Requests 51 Actions 5 Packages Projects Releases Wiki Activity

feat(budget): multi-period per-workspace LLM budget (hourly/daily/weekly/monthly) #2009

Merged
hongming merged 1 commits from feat/mc-multiperiod-workspace-budget into main 2026-05-29 11:45:32 +00:00
Conversation 4 Commits 1 Files Changed 14
+972 -918
14 changed files with 972 additions and 918 deletions
Download Patch File Download Diff File Expand all files Collapse all files
canvas/src/components/__tests__/BudgetSection.test.tsx
+53 -382
View File
@@ -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(<BudgetSection workspaceId="ws-1" />);
// 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(<BudgetSection workspaceId="ws-1" />);
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(<BudgetSection workspaceId="ws-1" />);
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<string, unknown>;
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<string, unknown>;
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<string, unknown>;
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(<BudgetSection workspaceId="ws-1" />);
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(<BudgetSection workspaceId="ws-1" />);
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(<BudgetSection workspaceId="ws-1" />);
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(<BudgetSection workspaceId="ws-1" />);
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(<BudgetSection workspaceId="ws-1" />);
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(<BudgetSection workspaceId="ws-1" />);
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(<BudgetSection workspaceId="ws-1" />);
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%");
});
});
canvas/src/components/tabs/BudgetSection.tsx
+158 -137
View File
@@ -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<Record<BudgetPeriod, PeriodBudget>>;
// 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<BudgetPeriod, PeriodBudget> {
const base: Record<BudgetPeriod, PeriodBudget> = {
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<BudgetData | null>(null);
const [loading, setLoading] = useState(true);
const [fetchError, setFetchError] = useState<string | null>(null);
const [limitInput, setLimitInput] = useState("");
// One input per period, in USD cents (string for controlled inputs).
const [limitInputs, setLimitInputs] = useState<Record<BudgetPeriod, string>>({
hourly: "",
daily: "",
weekly: "",
monthly: "",
});
const [saving, setSaving] = useState(false);
const [saveError, setSaveError] = useState<string | null>(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<BudgetData>(`/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<BudgetPeriod, number | null> = {
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<BudgetData>(`/workspaces/${workspaceId}/budget`, {
budget_limit: parsedLimit,
});
const updated = await api.patch<BudgetData>(`/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 (
<div className="space-y-3" data-testid="budget-section">
{/* Section header */}
<div>
<h3 className="text-xs font-semibold text-ink-mid uppercase tracking-wider">
Budget
</h3>
<h3 className="text-xs font-semibold text-ink-mid uppercase tracking-wider">Budget</h3>
<p className="text-[11px] text-ink-mid mt-0.5">
Limit total message credits for this workspace
Cap LLM spend for this workspace per period — crossing any limit pauses new work
</p>
</div>
@@ -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"
>
<svg
width="13"
height="13"
viewBox="0 0 13 13"
fill="none"
aria-hidden="true"
className="shrink-0"
>
<path
d="M6.5 1.5L11.5 10.5H1.5L6.5 1.5Z"
stroke="currentColor"
strokeWidth="1.4"
strokeLinejoin="round"
/>
<path
d="M6.5 5.5V7.5M6.5 9.5h.01"
stroke="currentColor"
strokeWidth="1.4"
strokeLinecap="round"
/>
<svg width="13" height="13" viewBox="0 0 13 13" fill="none" aria-hidden="true" className="shrink-0">
<path d="M6.5 1.5L11.5 10.5H1.5L6.5 1.5Z" stroke="currentColor" strokeWidth="1.4" strokeLinejoin="round" />
<path d="M6.5 5.5V7.5M6.5 9.5h.01" stroke="currentColor" strokeWidth="1.4" strokeLinecap="round" />
</svg>
Budget exceeded — messages blocked
Budget exceeded — new work paused
</div>
)}
{/* Usage stats */}
{loading ? (
<p className="text-xs text-ink-mid" data-testid="budget-loading">
Loading…
@@ -165,89 +197,78 @@ export function BudgetSection({ workspaceId }: Props) {
<p className="text-xs text-bad" data-testid="budget-fetch-error">
{fetchError}
</p>
) : budget ? (
<div className="space-y-2">
{/* Stats row */}
<div className="flex items-baseline justify-between" data-testid="budget-stats-row">
<span className="text-xs text-ink-mid">Credits used</span>
<span className="text-xs font-mono text-ink-mid">
<span data-testid="budget-used-value">{(budget.budget_used ?? 0).toLocaleString()}</span>
<span className="text-ink-mid mx-1">/</span>
<span data-testid="budget-limit-value">
{budget.budget_limit != null
? budget.budget_limit.toLocaleString()
: "Unlimited"}
</span>
</span>
</div>
) : (
<div className="space-y-3">
{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 (
<div key={key} className="space-y-1" data-testid={`budget-period-${key}`}>
<div className="flex items-baseline justify-between">
<label htmlFor={`budget-${key}-${workspaceId}`} className="text-xs text-ink-mid">
{label}
</label>
<span className="text-[11px] font-mono text-ink-mid">
<span data-testid={`budget-${key}-spend`}>{fmtUSD(p.spend)}</span>
<span className="mx-1">/</span>
<span data-testid={`budget-${key}-limit`}>{p.limit != null ? fmtUSD(p.limit) : "∞"}</span>
</span>
</div>
{p.limit != null && (
<div
role="progressbar"
aria-label={`${label} budget usage`}
aria-valuenow={pct}
aria-valuemin={0}
aria-valuemax={100}
className="h-1.5 w-full rounded-full bg-surface-card overflow-hidden"
>
<div
data-testid={`budget-${key}-fill`}
className={`h-full rounded-full transition-all duration-300 ${over ? "bg-bad" : "bg-accent"}`}
style={{ width: `${pct}%` }}
/>
</div>
)}
<input
id={`budget-${key}-${workspaceId}`}
type="number"
min="0"
step="1"
value={limitInputs[key]}
onChange={(e) => 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"
/>
</div>
);
})}
{/* Progress bar (only when limit is set) */}
{budget.budget_limit != null && (
<p className="text-[11px] text-ink-mid">Limits are USD cents (e.g. 500 = $5.00). Blank = unlimited.</p>
{saveError && (
<div
role="progressbar"
aria-label="Budget usage"
aria-valuenow={progressPct}
aria-valuemin={0}
aria-valuemax={100}
className="h-1.5 w-full rounded-full bg-surface-card overflow-hidden"
role="alert"
data-testid="budget-save-error"
className="px-3 py-1.5 rounded-lg bg-red-950/40 border border-red-800/50 text-xs text-bad"
>
<div
data-testid="budget-progress-fill"
className="h-full rounded-full bg-accent transition-all duration-300"
style={{ width: `${progressPct}%` }}
/>
{saveError}
</div>
)}
{/* Remaining credits */}
{budget.budget_remaining != null && (
<p className="text-[11px] text-ink-mid" data-testid="budget-remaining">
{budget.budget_remaining.toLocaleString()} credits remaining
</p>
)}
</div>
) : null}
{/* Input + Save */}
<div className="space-y-1.5 pt-1">
<label
htmlFor={`budget-limit-input-${workspaceId}`}
className="text-[11px] text-ink-mid block"
>
Budget limit (credits)
</label>
<input
id={`budget-limit-input-${workspaceId}`}
type="number"
min="0"
step="1"
value={limitInput}
onChange={(e) => 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"
/>
<p className="text-xs text-ink-mid">Leave blank for unlimited</p>
{saveError && (
<div
role="alert"
data-testid="budget-save-error"
className="px-3 py-1.5 rounded-lg bg-red-950/40 border border-red-800/50 text-xs text-bad"
<button
onClick={handleSave}
disabled={saving}
data-testid="budget-save-btn"
className="px-4 py-1.5 bg-accent-strong hover:bg-accent active:bg-accent-strong rounded-lg text-xs font-medium text-white disabled:opacity-50 transition-colors focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-accent focus-visible:ring-offset-1 focus-visible:ring-offset-zinc-900"
>
{saveError}
</div>
)}
<button
onClick={handleSave}
disabled={saving}
data-testid="budget-save-btn"
className="px-4 py-1.5 bg-accent-strong hover:bg-accent active:bg-accent-strong rounded-lg text-xs font-medium text-white disabled:opacity-50 transition-colors focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-accent focus-visible:ring-offset-1 focus-visible:ring-offset-zinc-900"
>
{saving ? "Saving…" : "Save"}
</button>
</div>
{saving ? "Saving…" : "Save"}
</button>
</div>
)}
</div>
);
}
canvas/src/components/tabs/__tests__/BudgetSection.test.tsx
+80 -195
View File
@@ -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<Record<"hourly" | "daily" | "weekly" | "monthly", Partial<P>>> = {}) {
const blank: P = { limit: null, spend: 0, remaining: null };
const mk = (o?: Partial<P>): 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(<BudgetSection workspaceId={WS_ID} />);
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(<BudgetSection workspaceId={WS_ID} />);
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(<BudgetSection workspaceId={WS_ID} />);
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(<BudgetSection workspaceId={WS_ID} />);
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(<BudgetSection workspaceId={WS_ID} />);
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(<BudgetSection workspaceId={WS_ID} />);
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(<BudgetSection workspaceId={WS_ID} />);
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(<BudgetSection workspaceId={WS_ID} />);
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(<BudgetSection workspaceId={WS_ID} />);
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(<BudgetSection workspaceId={WS_ID} />);
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(<BudgetSection workspaceId={WS_ID} />);
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<string, number | null> }).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(<BudgetSection workspaceId={WS_ID} />);
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(<BudgetSection workspaceId={WS_ID} />);
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(<BudgetSection workspaceId={WS_ID} />);
// 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(<BudgetSection workspaceId={WS_ID} />);
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(<BudgetSection workspaceId={WS_ID} />);
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(<BudgetSection workspaceId={WS_ID} />);
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(<BudgetSection workspaceId={WS_ID} />);
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");
});
});
});
workspace-server/internal/handlers/a2a_proxy.go
+25 -14
View File
@@ -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
workspace-server/internal/handlers/a2a_queue_test.go
+5 -3
View File
@@ -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
workspace-server/internal/handlers/budget.go
+120 -77
View File
@@ -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": <int64>} 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)
}
workspace-server/internal/handlers/budget_periods.go
+160
View File
@@ -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"
}
workspace-server/internal/handlers/budget_periods_test.go
+99
View File
@@ -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)
}
}
}
workspace-server/internal/handlers/budget_test.go
+191 -89
View File
@@ -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)
workspace-server/internal/handlers/handlers_test.go
+5 -3
View File
@@ -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 ----------
workspace-server/internal/handlers/registry.go
+21 -1
View File
@@ -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
workspace-server/internal/handlers/workspace_budget_test.go
+23 -17
View File
@@ -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)
workspace-server/migrations/20260529000000_workspace_multiperiod_budget.down.sql
+2
View File
@@ -0,0 +1,2 @@
DROP TABLE IF EXISTS workspace_spend_events;
ALTER TABLE workspaces DROP COLUMN IF EXISTS budget_limits;
workspace-server/migrations/20260529000000_workspace_multiperiod_budget.up.sql
+30
View File
@@ -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() - <window>.
-- 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);