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