From 22af070ef36ee31caea2efef7015eef48099958e Mon Sep 17 00:00:00 2001 From: Molecule AI Backend Engineer Date: Fri, 17 Apr 2026 02:02:22 +0000 Subject: [PATCH] feat(#541): add dedicated GET/PATCH /workspaces/:id/budget endpoints - New BudgetHandler with GetBudget and PatchBudget methods - GET returns budget_limit (null or int64 USD cents), monthly_spend, and computed budget_remaining (null when no limit, can be negative when over-budget so callers can see the magnitude of the overage) - PATCH accepts {budget_limit: int64|null}; null clears the ceiling; validates non-negative values; re-reads DB to echo final state - Both handlers are wired in router.go under the WorkspaceAuth group - 14 unit tests covering happy paths, 404, 400 validation, DB errors, over-budget state, zero limit, and clear-limit round-trip - All 20 packages pass go test ./... and go build ./... is clean Co-Authored-By: Claude Sonnet 4.6 --- platform/internal/handlers/budget.go | 171 ++++++++ platform/internal/handlers/budget_test.go | 458 ++++++++++++++++++++++ platform/internal/router/router.go | 6 + 3 files changed, 635 insertions(+) create mode 100644 platform/internal/handlers/budget.go create mode 100644 platform/internal/handlers/budget_test.go diff --git a/platform/internal/handlers/budget.go b/platform/internal/handlers/budget.go new file mode 100644 index 00000000..0af2ee8e --- /dev/null +++ b/platform/internal/handlers/budget.go @@ -0,0 +1,171 @@ +package handlers + +import ( + "database/sql" + "log" + "net/http" + + "github.com/Molecule-AI/molecule-monorepo/platform/internal/db" + "github.com/gin-gonic/gin" +) + +// 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 +type BudgetHandler struct{} + +func NewBudgetHandler() *BudgetHandler { return &BudgetHandler{} } + +// budgetResponse is the canonical JSON shape for both GET and PATCH responses. +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. + BudgetRemaining *int64 `json:"budget_remaining"` +} + +// 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 + err := db.DB.QueryRowContext(ctx, + `SELECT budget_limit, COALESCE(monthly_spend, 0) + FROM workspaces + WHERE id = $1 AND status != 'removed'`, + workspaceID, + ).Scan(&budgetLimit, &monthlySpend) + if err == sql.ErrNoRows { + c.JSON(http.StatusNotFound, gin.H{"error": "workspace not found"}) + return + } + if err != nil { + log.Printf("GetBudget: query failed for %s: %v", workspaceID, err) + c.JSON(http.StatusInternalServerError, gin.H{"error": "query failed"}) + return + } + + resp := budgetResponse{ + MonthlySpend: monthlySpend, + } + if budgetLimit.Valid { + limit := budgetLimit.Int64 + resp.BudgetLimit = &limit + remaining := limit - monthlySpend + resp.BudgetRemaining = &remaining + } + + c.JSON(http.StatusOK, resp) +} + +// patchBudgetRequest is the expected JSON body for PATCH /workspaces/:id/budget. +// budget_limit=null removes the ceiling; a positive integer sets it (USD cents). +type patchBudgetRequest struct { + // BudgetLimit pointer so JSON null → nil, absent → parse error (required field). + BudgetLimit *int64 `json:"budget_limit"` +} + +// 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. +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{} + if err := c.ShouldBindJSON(&raw); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + budgetLimitRaw, ok := raw["budget_limit"] + if !ok { + c.JSON(http.StatusBadRequest, gin.H{"error": "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)"}) + 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)"}) + return + } + budgetArg = v + default: + c.JSON(http.StatusBadRequest, gin.H{"error": "budget_limit must be an integer (USD cents) or null"}) + return + } + } + // budgetArg == nil means "clear the ceiling" + + // Existence check — return 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')`, + workspaceID, + ).Scan(&exists); err != nil || !exists { + c.JSON(http.StatusNotFound, gin.H{"error": "workspace not found"}) + return + } + + if _, err := db.DB.ExecContext(ctx, + `UPDATE workspaces SET budget_limit = $2, updated_at = now() WHERE id = $1`, + workspaceID, budgetArg, + ); 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 { + 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/platform/internal/handlers/budget_test.go b/platform/internal/handlers/budget_test.go new file mode 100644 index 00000000..e3e6cacd --- /dev/null +++ b/platform/internal/handlers/budget_test.go @@ -0,0 +1,458 @@ +package handlers + +import ( + "bytes" + "database/sql" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/DATA-DOG/go-sqlmock" + "github.com/gin-gonic/gin" +) + +// ==================== 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\)`). + WithArgs("ws-not-there"). + WillReturnError(sql.ErrNoRows) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + 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) + + if w.Code != http.StatusNotFound { + t.Errorf("expected 404, got %d: %s", w.Code, w.Body.String()) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("sqlmock expectations not met: %v", err) + } +} + +// 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\)`). + WithArgs("ws-db-err"). + WillReturnError(sql.ErrConnDone) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + 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) + + if w.Code != http.StatusInternalServerError { + t.Errorf("expected 500, got %d: %s", w.Code, w.Body.String()) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("sqlmock expectations not met: %v", err) + } +} + +// 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\)`). + WithArgs("ws-free"). + WillReturnRows(sqlmock.NewRows([]string{"budget_limit", "monthly_spend"}). + AddRow(nil, int64(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) + + if w.Code != http.StatusOK { + t.Fatalf("expected 200, 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"] != nil { + t.Errorf("expected budget_limit=null, got %v", resp["budget_limit"]) + } + if resp["budget_remaining"] != nil { + t.Errorf("expected budget_remaining=null, got %v", resp["budget_remaining"]) + } + if resp["monthly_spend"] != float64(42) { + t.Errorf("expected monthly_spend=42, got %v", resp["monthly_spend"]) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("sqlmock expectations not met: %v", err) + } +} + +// 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\)`). + WithArgs("ws-capped"). + WillReturnRows(sqlmock.NewRows([]string{"budget_limit", "monthly_spend"}). + AddRow(int64(500), int64(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) + + if w.Code != http.StatusOK { + t.Fatalf("expected 200, 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(500) { + t.Errorf("expected budget_limit=500, got %v", resp["budget_limit"]) + } + 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"]) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("sqlmock expectations not met: %v", err) + } +} + +// 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\)`). + WithArgs("ws-over"). + WillReturnRows(sqlmock.NewRows([]string{"budget_limit", "monthly_spend"}). + AddRow(int64(100), int64(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) + + if w.Code != http.StatusOK { + t.Fatalf("expected 200, 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) + } + // 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"]) + } + 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) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Params = gin.Params{{Key: "id", Value: "ws-patch-missing"}} + c.Request = httptest.NewRequest("PATCH", "/workspaces/ws-patch-missing/budget", + bytes.NewBufferString(`{"other_field":123}`)) + c.Request.Header.Set("Content-Type", "application/json") + + h := NewBudgetHandler() + h.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) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Params = gin.Params{{Key: "id", Value: "ws-patch-bad"}} + c.Request = httptest.NewRequest("PATCH", "/workspaces/ws-patch-bad/budget", + bytes.NewBufferString(`not json`)) + c.Request.Header.Set("Content-Type", "application/json") + + h := NewBudgetHandler() + h.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) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Params = gin.Params{{Key: "id", Value: "ws-negative"}} + c.Request = httptest.NewRequest("PATCH", "/workspaces/ws-negative/budget", + bytes.NewBufferString(`{"budget_limit":-1}`)) + c.Request.Header.Set("Content-Type", "application/json") + + h := NewBudgetHandler() + h.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) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Params = gin.Params{{Key: "id", Value: "ws-badtype"}} + c.Request = httptest.NewRequest("PATCH", "/workspaces/ws-badtype/budget", + bytes.NewBufferString(`{"budget_limit":"not-a-number"}`)) + c.Request.Header.Set("Content-Type", "application/json") + + h := NewBudgetHandler() + h.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. +func TestBudgetPatch_WorkspaceNotFound(t *testing.T) { + mock := setupTestDB(t) + setupTestRedis(t) + + mock.ExpectQuery(`SELECT EXISTS.*status != 'removed'`). + WithArgs("ws-no-exist"). + WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(false)) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Params = gin.Params{{Key: "id", Value: "ws-no-exist"}} + c.Request = httptest.NewRequest("PATCH", "/workspaces/ws-no-exist/budget", + bytes.NewBufferString(`{"budget_limit":500}`)) + c.Request.Header.Set("Content-Type", "application/json") + + h := NewBudgetHandler() + h.PatchBudget(c) + + if w.Code != http.StatusNotFound { + t.Errorf("expected 404, got %d: %s", w.Code, w.Body.String()) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("sqlmock expectations not met: %v", err) + } +} + +// TestBudgetPatch_SetLimit verifies that PATCH /budget with a positive value +// updates the DB and returns the new budget 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)). + WillReturnResult(sqlmock.NewResult(0, 1)) + // Re-read for response + mock.ExpectQuery(`SELECT budget_limit, COALESCE\(monthly_spend, 0\) FROM workspaces WHERE id`). + WithArgs("ws-set-limit"). + WillReturnRows(sqlmock.NewRows([]string{"budget_limit", "monthly_spend"}). + AddRow(int64(500), int64(200))) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Params = gin.Params{{Key: "id", Value: "ws-set-limit"}} + c.Request = httptest.NewRequest("PATCH", "/workspaces/ws-set-limit/budget", + bytes.NewBufferString(`{"budget_limit":500}`)) + c.Request.Header.Set("Content-Type", "application/json") + + h := NewBudgetHandler() + h.PatchBudget(c) + + if w.Code != http.StatusOK { + t.Fatalf("expected 200, 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(500) { + t.Errorf("expected budget_limit=500, got %v", resp["budget_limit"]) + } + 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"]) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("sqlmock expectations not met: %v", err) + } +} + +// TestBudgetPatch_ClearLimit verifies that PATCH /budget with budget_limit=null +// clears the ceiling, making budget_limit and budget_remaining null in the response. +func TestBudgetPatch_ClearLimit(t *testing.T) { + mock := setupTestDB(t) + setupTestRedis(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). + 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`). + WithArgs("ws-clear-limit"). + WillReturnRows(sqlmock.NewRows([]string{"budget_limit", "monthly_spend"}). + AddRow(nil, int64(50))) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Params = gin.Params{{Key: "id", Value: "ws-clear-limit"}} + c.Request = httptest.NewRequest("PATCH", "/workspaces/ws-clear-limit/budget", + bytes.NewBufferString(`{"budget_limit":null}`)) + c.Request.Header.Set("Content-Type", "application/json") + + h := NewBudgetHandler() + h.PatchBudget(c) + + if w.Code != http.StatusOK { + t.Fatalf("expected 200, 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"] != nil { + t.Errorf("expected budget_limit=null after clear, got %v", resp["budget_limit"]) + } + if resp["budget_remaining"] != nil { + t.Errorf("expected budget_remaining=null after clear, got %v", resp["budget_remaining"]) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("sqlmock expectations not met: %v", err) + } +} + +// TestBudgetPatch_UpdateDBError verifies that a DB error during the UPDATE +// returns 500. +func TestBudgetPatch_UpdateDBError(t *testing.T) { + mock := setupTestDB(t) + setupTestRedis(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)). + WillReturnError(sql.ErrConnDone) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Params = gin.Params{{Key: "id", Value: "ws-patch-dberr"}} + c.Request = httptest.NewRequest("PATCH", "/workspaces/ws-patch-dberr/budget", + bytes.NewBufferString(`{"budget_limit":500}`)) + c.Request.Header.Set("Content-Type", "application/json") + + h := NewBudgetHandler() + h.PatchBudget(c) + + if w.Code != http.StatusInternalServerError { + t.Errorf("expected 500 on UPDATE error, got %d: %s", w.Code, w.Body.String()) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("sqlmock expectations not met: %v", err) + } +} + +// 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). +func TestBudgetPatch_ZeroLimit(t *testing.T) { + mock := setupTestDB(t) + setupTestRedis(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)). + WillReturnResult(sqlmock.NewResult(0, 1)) + mock.ExpectQuery(`SELECT budget_limit, COALESCE\(monthly_spend, 0\) FROM workspaces WHERE id`). + WithArgs("ws-zero-limit"). + WillReturnRows(sqlmock.NewRows([]string{"budget_limit", "monthly_spend"}). + AddRow(int64(0), int64(0))) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Params = gin.Params{{Key: "id", Value: "ws-zero-limit"}} + c.Request = httptest.NewRequest("PATCH", "/workspaces/ws-zero-limit/budget", + bytes.NewBufferString(`{"budget_limit":0}`)) + c.Request.Header.Set("Content-Type", "application/json") + + h := NewBudgetHandler() + h.PatchBudget(c) + + if w.Code != http.StatusOK { + t.Errorf("expected 200 for zero budget_limit, got %d: %s", w.Code, w.Body.String()) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("sqlmock expectations not met: %v", err) + } +} diff --git a/platform/internal/router/router.go b/platform/internal/router/router.go index d41b653a..b6669059 100644 --- a/platform/internal/router/router.go +++ b/platform/internal/router/router.go @@ -256,6 +256,12 @@ func Setup(hub *ws.Hub, broadcaster *events.Broadcaster, prov *provisioner.Provi // (mirrors the /workspaces/:id/a2a pattern). Issue #249. r.GET("/workspaces/:id/schedules/health", schedh.Health) + // Budget — per-workspace spend ceiling and current usage (#541). + // GET returns the current state; PATCH sets or clears the ceiling. + budgeth := handlers.NewBudgetHandler() + wsAuth.GET("/budget", budgeth.GetBudget) + wsAuth.PATCH("/budget", budgeth.PatchBudget) + // Token management (user-facing create/list/revoke) tokh := handlers.NewTokenHandler() wsAuth.GET("/tokens", tokh.List)