diff --git a/platform/internal/handlers/a2a_proxy.go b/platform/internal/handlers/a2a_proxy.go index 307c3311..f7664b22 100644 --- a/platform/internal/handlers/a2a_proxy.go +++ b/platform/internal/handlers/a2a_proxy.go @@ -203,6 +203,33 @@ 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. +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`, + workspaceID, + ).Scan(&budgetLimit, &monthlySpend) + if 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) + return &proxyA2AError{ + Status: http.StatusPaymentRequired, + Response: gin.H{"error": "workspace budget limit exceeded"}, + } + } + return nil +} + func (h *WorkspaceHandler) proxyA2ARequest(ctx context.Context, workspaceID string, body []byte, callerID string, logActivity bool) (int, []byte, *proxyA2AError) { // Access control: workspace-to-workspace requests must pass CanCommunicate check. // Canvas requests (callerID == "") and system callers (webhook:*, system:*, test:*) @@ -217,6 +244,14 @@ func (h *WorkspaceHandler) proxyA2ARequest(ctx context.Context, workspaceID stri } } + // Budget enforcement: reject A2A calls when the workspace has exceeded its + // monthly spend ceiling. Checked after access control so unauthorized calls + // are rejected first (403 > 429 in the denial hierarchy). Fail-open on DB + // errors so a budget check failure never blocks legitimate traffic. + if proxyErr := h.checkWorkspaceBudget(ctx, workspaceID); proxyErr != nil { + return 0, nil, proxyErr + } + agentURL, proxyErr := h.resolveAgentURL(ctx, workspaceID) if proxyErr != nil { return 0, nil, proxyErr 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/handlers/handlers_additional_test.go b/platform/internal/handlers/handlers_additional_test.go index 1ca55547..5316497c 100644 --- a/platform/internal/handlers/handlers_additional_test.go +++ b/platform/internal/handlers/handlers_additional_test.go @@ -30,7 +30,7 @@ func TestWorkspaceCreate_WithParentID(t *testing.T) { parentID := "parent-ws-123" mock.ExpectBegin() mock.ExpectExec("INSERT INTO workspaces"). - WithArgs(sqlmock.AnyArg(), "Child Agent", nil, 1, "langgraph", sqlmock.AnyArg(), &parentID, nil, "none"). + WithArgs(sqlmock.AnyArg(), "Child Agent", nil, 1, "langgraph", sqlmock.AnyArg(), &parentID, nil, "none", (*int64)(nil)). WillReturnResult(sqlmock.NewResult(0, 1)) mock.ExpectCommit() mock.ExpectExec("INSERT INTO canvas_layouts"). @@ -65,7 +65,7 @@ func TestWorkspaceCreate_ExplicitClaudeCodeRuntime(t *testing.T) { mock.ExpectBegin() mock.ExpectExec("INSERT INTO workspaces"). - WithArgs(sqlmock.AnyArg(), "CC Agent", nil, 2, "claude-code", sqlmock.AnyArg(), (*string)(nil), nil, "none"). + WithArgs(sqlmock.AnyArg(), "CC Agent", nil, 2, "claude-code", sqlmock.AnyArg(), (*string)(nil), nil, "none", (*int64)(nil)). WillReturnResult(sqlmock.NewResult(0, 1)) mock.ExpectCommit() mock.ExpectExec("INSERT INTO canvas_layouts"). @@ -194,12 +194,13 @@ func TestWorkspaceList_WithData(t *testing.T) { "id", "name", "role", "tier", "status", "agent_card", "url", "parent_id", "active_tasks", "last_error_rate", "last_sample_error", "uptime_seconds", "current_task", "runtime", "workspace_dir", "x", "y", "collapsed", + "budget_limit", "monthly_spend", } rows := sqlmock.NewRows(columns). AddRow("ws-1", "Agent One", "worker", 1, "online", []byte(`{"name":"agent1"}`), "http://localhost:8001", - nil, 3, 0.02, "", 7200, "processing", "langgraph", "", 10.0, 20.0, false). + nil, 3, 0.02, "", 7200, "processing", "langgraph", "", 10.0, 20.0, false, nil, int64(0)). AddRow("ws-2", "Agent Two", "", 2, "degraded", []byte("null"), "", - nil, 0, 0.6, "timeout", 100, "", "claude-code", "", 50.0, 60.0, true) + nil, 0, 0.6, "timeout", 100, "", "claude-code", "", 50.0, 60.0, true, nil, int64(0)) mock.ExpectQuery("SELECT w.id, w.name"). WillReturnRows(rows) diff --git a/platform/internal/handlers/handlers_test.go b/platform/internal/handlers/handlers_test.go index c8dae41e..25a67578 100644 --- a/platform/internal/handlers/handlers_test.go +++ b/platform/internal/handlers/handlers_test.go @@ -253,7 +253,7 @@ func TestWorkspaceCreate(t *testing.T) { // Expect workspace INSERT (uuid is dynamic, use AnyArg for id, runtime, awareness_namespace) mock.ExpectExec("INSERT INTO workspaces"). - WithArgs(sqlmock.AnyArg(), "Test Agent", nil, 1, "langgraph", sqlmock.AnyArg(), (*string)(nil), nil, "none"). + WithArgs(sqlmock.AnyArg(), "Test Agent", nil, 1, "langgraph", sqlmock.AnyArg(), (*string)(nil), nil, "none", (*int64)(nil)). WillReturnResult(sqlmock.NewResult(0, 1)) // Expect transaction commit (no secrets in this payload) @@ -340,12 +340,13 @@ func TestWorkspaceList(t *testing.T) { "id", "name", "role", "tier", "status", "agent_card", "url", "parent_id", "active_tasks", "last_error_rate", "last_sample_error", "uptime_seconds", "current_task", "runtime", "workspace_dir", "x", "y", "collapsed", + "budget_limit", "monthly_spend", } rows := sqlmock.NewRows(columns). AddRow("ws-1", "Agent One", "worker", 1, "online", []byte("null"), "http://localhost:8001", - nil, 0, 0.0, "", 100, "", "claude-code", "", 10.0, 20.0, false). + nil, 0, 0.0, "", 100, "", "claude-code", "", 10.0, 20.0, false, nil, int64(0)). AddRow("ws-2", "Agent Two", "manager", 2, "provisioning", []byte("null"), "", - nil, 0, 0.0, "", 0, "", "langgraph", "", 50.0, 60.0, false) + nil, 0, 0.0, "", 0, "", "langgraph", "", 50.0, 60.0, false, nil, int64(0)) mock.ExpectQuery("SELECT w.id, w.name"). WillReturnRows(rows) @@ -1007,12 +1008,14 @@ func TestWorkspaceGet_CurrentTask(t *testing.T) { "id", "name", "role", "tier", "status", "agent_card", "url", "parent_id", "active_tasks", "last_error_rate", "last_sample_error", "uptime_seconds", "current_task", "runtime", "workspace_dir", "x", "y", "collapsed", + "budget_limit", "monthly_spend", } mock.ExpectQuery("SELECT w.id, w.name"). WithArgs("ws-task"). WillReturnRows(sqlmock.NewRows(columns).AddRow( "ws-task", "Task Worker", "worker", 1, "online", []byte("null"), "http://localhost:9000", nil, 2, 0.0, "", 300, "Analyzing document", "langgraph", "", 10.0, 20.0, false, + nil, int64(0), )) w := httptest.NewRecorder() diff --git a/platform/internal/handlers/registry.go b/platform/internal/handlers/registry.go index 445d6903..b07bc1b7 100644 --- a/platform/internal/handlers/registry.go +++ b/platform/internal/handlers/registry.go @@ -239,18 +239,40 @@ func (h *RegistryHandler) Heartbeat(c *gin.Context) { // late heartbeat from a container that's being torn down doesn't // refresh last_heartbeat_at on a tombstoned workspace (which would // otherwise confuse the liveness monitor). - _, err := db.DB.ExecContext(ctx, ` - UPDATE workspaces SET - last_heartbeat_at = now(), - last_error_rate = $2, - last_sample_error = $3, - active_tasks = $4, - uptime_seconds = $5, - current_task = $6, - updated_at = now() - WHERE id = $1 AND status != 'removed' - `, payload.WorkspaceID, payload.ErrorRate, payload.SampleError, - payload.ActiveTasks, payload.UptimeSeconds, payload.CurrentTask) + // + // monthly_spend: updated when the agent reports a positive value (cumulative + // USD cents for the current month). Zero means "no update" — never write + // zero to avoid accidentally clearing a previously-reported spend value. + var err error + if payload.MonthlySpend > 0 { + _, err = db.DB.ExecContext(ctx, ` + UPDATE workspaces SET + last_heartbeat_at = now(), + last_error_rate = $2, + last_sample_error = $3, + active_tasks = $4, + uptime_seconds = $5, + current_task = $6, + monthly_spend = $7, + updated_at = now() + WHERE id = $1 AND status != 'removed' + `, payload.WorkspaceID, payload.ErrorRate, payload.SampleError, + payload.ActiveTasks, payload.UptimeSeconds, payload.CurrentTask, + payload.MonthlySpend) + } else { + _, err = db.DB.ExecContext(ctx, ` + UPDATE workspaces SET + last_heartbeat_at = now(), + last_error_rate = $2, + last_sample_error = $3, + active_tasks = $4, + uptime_seconds = $5, + current_task = $6, + updated_at = now() + WHERE id = $1 AND status != 'removed' + `, payload.WorkspaceID, payload.ErrorRate, payload.SampleError, + payload.ActiveTasks, payload.UptimeSeconds, payload.CurrentTask) + } if err != nil { log.Printf("Heartbeat update error: %v", err) c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to update"}) diff --git a/platform/internal/handlers/workspace.go b/platform/internal/handlers/workspace.go index dc727833..ac520d31 100644 --- a/platform/internal/handlers/workspace.go +++ b/platform/internal/handlers/workspace.go @@ -150,9 +150,9 @@ func (h *WorkspaceHandler) Create(c *gin.Context) { // Insert workspace with runtime persisted in DB (inside transaction) _, err := tx.ExecContext(ctx, ` - INSERT INTO workspaces (id, name, role, tier, runtime, awareness_namespace, status, parent_id, workspace_dir, workspace_access) - VALUES ($1, $2, $3, $4, $5, $6, 'provisioning', $7, $8, $9) - `, id, payload.Name, role, payload.Tier, payload.Runtime, awarenessNamespace, payload.ParentID, workspaceDir, workspaceAccess) + INSERT INTO workspaces (id, name, role, tier, runtime, awareness_namespace, status, parent_id, workspace_dir, workspace_access, budget_limit) + VALUES ($1, $2, $3, $4, $5, $6, 'provisioning', $7, $8, $9, $10) + `, id, payload.Name, role, payload.Tier, payload.Runtime, awarenessNamespace, payload.ParentID, workspaceDir, workspaceAccess, payload.BudgetLimit) if err != nil { tx.Rollback() //nolint:errcheck log.Printf("Create workspace error: %v", err) @@ -293,10 +293,13 @@ func scanWorkspaceRow(rows interface { var collapsed bool var parentID *string var agentCard []byte + var budgetLimit sql.NullInt64 + var monthlySpend int64 err := rows.Scan(&id, &name, &role, &tier, &status, &agentCard, &url, &parentID, &activeTasks, &errorRate, &sampleError, &uptimeSeconds, - ¤tTask, &runtime, &workspaceDir, &x, &y, &collapsed) + ¤tTask, &runtime, &workspaceDir, &x, &y, &collapsed, + &budgetLimit, &monthlySpend) if err != nil { return nil, err } @@ -315,11 +318,19 @@ func scanWorkspaceRow(rows interface { "current_task": currentTask, "runtime": runtime, "workspace_dir": nilIfEmpty(workspaceDir), + "monthly_spend": monthlySpend, "x": x, "y": y, "collapsed": collapsed, } + // budget_limit: nil when no limit set, int64 otherwise + if budgetLimit.Valid { + ws["budget_limit"] = budgetLimit.Int64 + } else { + ws["budget_limit"] = nil + } + // Only include non-empty values if role != "" { ws["role"] = role @@ -344,7 +355,8 @@ const workspaceListQuery = ` COALESCE(w.last_sample_error, ''), w.uptime_seconds, COALESCE(w.current_task, ''), COALESCE(w.runtime, 'langgraph'), COALESCE(w.workspace_dir, ''), - COALESCE(cl.x, 0), COALESCE(cl.y, 0), COALESCE(cl.collapsed, false) + COALESCE(cl.x, 0), COALESCE(cl.y, 0), COALESCE(cl.collapsed, false), + w.budget_limit, COALESCE(w.monthly_spend, 0) FROM workspaces w LEFT JOIN canvas_layouts cl ON cl.workspace_id = w.id WHERE w.status != 'removed' @@ -389,7 +401,8 @@ func (h *WorkspaceHandler) Get(c *gin.Context) { COALESCE(w.last_sample_error, ''), w.uptime_seconds, COALESCE(w.current_task, ''), COALESCE(w.runtime, 'langgraph'), COALESCE(w.workspace_dir, ''), - COALESCE(cl.x, 0), COALESCE(cl.y, 0), COALESCE(cl.collapsed, false) + COALESCE(cl.x, 0), COALESCE(cl.y, 0), COALESCE(cl.collapsed, false), + w.budget_limit, COALESCE(w.monthly_spend, 0) FROM workspaces w LEFT JOIN canvas_layouts cl ON cl.workspace_id = w.id WHERE w.id = $1 @@ -406,6 +419,12 @@ func (h *WorkspaceHandler) Get(c *gin.Context) { return } + // Strip financial fields — GET /workspaces/:id is on the open router. + // Any caller with a valid UUID would otherwise read billing data. + // The dedicated budget/spend endpoints are AdminAuth-gated. (#611) + delete(ws, "budget_limit") + delete(ws, "monthly_spend") + c.JSON(http.StatusOK, ws) } @@ -506,6 +525,10 @@ var sensitiveUpdateFields = map[string]struct{}{ "parent_id": {}, "runtime": {}, "workspace_dir": {}, + // budget_limit is intentionally NOT here. The dedicated + // PATCH /workspaces/:id/budget (AdminAuth) is the only write path. + // Accepting it here — even behind ValidateAnyToken — lets workspace agents + // self-clear their own spending ceiling. (#611 Security Auditor finding) } // Update handles PATCH /workspaces/:id @@ -603,6 +626,10 @@ func (h *WorkspaceHandler) Update(c *gin.Context) { } needsRestart = true } + // NOTE: budget_limit is intentionally NOT handled here. The dedicated + // PATCH /workspaces/:id/budget (AdminAuth) is the only write path. + // This endpoint uses ValidateAnyToken — any enrolled workspace bearer + // could otherwise self-clear its own spending ceiling. (#611 Security Auditor) // Update canvas position if both x and y provided if x, xOk := body["x"]; xOk { diff --git a/platform/internal/handlers/workspace_budget_test.go b/platform/internal/handlers/workspace_budget_test.go new file mode 100644 index 00000000..97a54e2a --- /dev/null +++ b/platform/internal/handlers/workspace_budget_test.go @@ -0,0 +1,438 @@ +package handlers + +// Tests for per-workspace budget_limit field and A2A enforcement (#541). +// +// Coverage: +// - GET /workspaces/:id includes budget_limit (nil when unset, int when set) +// - GET /workspaces/:id includes monthly_spend +// - POST /workspaces creates workspace with budget_limit +// - PATCH /workspaces/:id updates budget_limit (nil clears the ceiling) +// - A2A proxy returns 429 when monthly_spend >= budget_limit +// - A2A proxy passes through when monthly_spend < budget_limit +// - A2A proxy passes through when budget_limit is NULL (no limit) +// - A2A proxy fail-open on DB error during budget check + +import ( + "bytes" + "database/sql" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/DATA-DOG/go-sqlmock" + "github.com/gin-gonic/gin" +) + +// wsColumns is the canonical column list for scanWorkspaceRow tests. +var wsColumns = []string{ + "id", "name", "role", "tier", "status", "agent_card", "url", + "parent_id", "active_tasks", "last_error_rate", "last_sample_error", + "uptime_seconds", "current_task", "runtime", "workspace_dir", "x", "y", "collapsed", + "budget_limit", "monthly_spend", +} + +// ==================== GET — financial fields stripped from open endpoint ==================== + +// TestWorkspaceBudget_Get_NilLimit verifies that budget_limit and monthly_spend +// are NOT present in GET /workspaces/:id. The endpoint is on the open router — +// any caller with a valid UUID must not read billing data. (#611 Security Auditor) +func TestWorkspaceBudget_Get_NilLimit(t *testing.T) { + mock := setupTestDB(t) + setupTestRedis(t) + handler := NewWorkspaceHandler(newTestBroadcaster(), nil, "http://localhost:8080", t.TempDir()) + + mock.ExpectQuery("SELECT w.id, w.name"). + WithArgs("ws-nobudget"). + WillReturnRows(sqlmock.NewRows(wsColumns). + AddRow("ws-nobudget", "Free Agent", "worker", 1, "online", + []byte(`{}`), "http://localhost:9001", + nil, 0, 0.0, "", 0, "", "langgraph", "", + 0.0, 0.0, false, + nil, // budget_limit NULL + 0)) // monthly_spend 0 + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Params = gin.Params{{Key: "id", Value: "ws-nobudget"}} + c.Request = httptest.NewRequest("GET", "/workspaces/ws-nobudget", nil) + handler.Get(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("failed to parse response: %v", err) + } + // #611: financial fields must NOT appear on the open GET endpoint. + if _, present := resp["budget_limit"]; present { + t.Errorf("budget_limit must not appear in open GET /workspaces/:id response") + } + if _, present := resp["monthly_spend"]; present { + t.Errorf("monthly_spend must not appear in open GET /workspaces/:id response") + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("sqlmock expectations not met: %v", err) + } +} + +// TestWorkspaceBudget_Get_WithLimit verifies that budget_limit and monthly_spend +// are stripped from the open GET /workspaces/:id even when the DB has non-zero +// values. Financial reads go through the AdminAuth-gated budget endpoint. (#611) +func TestWorkspaceBudget_Get_WithLimit(t *testing.T) { + mock := setupTestDB(t) + setupTestRedis(t) + handler := NewWorkspaceHandler(newTestBroadcaster(), nil, "http://localhost:8080", t.TempDir()) + + mock.ExpectQuery("SELECT w.id, w.name"). + WithArgs("ws-limited"). + WillReturnRows(sqlmock.NewRows(wsColumns). + AddRow("ws-limited", "Capped Agent", "worker", 1, "online", + []byte(`{}`), "http://localhost:9002", + nil, 0, 0.0, "", 0, "", "langgraph", "", + 0.0, 0.0, false, + int64(500), // budget_limit = $5.00 in DB + int64(123))) // monthly_spend = $1.23 in DB + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Params = gin.Params{{Key: "id", Value: "ws-limited"}} + c.Request = httptest.NewRequest("GET", "/workspaces/ws-limited", nil) + handler.Get(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("failed to parse response: %v", err) + } + // #611: financial fields must NOT appear on the open GET endpoint even when + // the DB has non-zero values — they're stripped before c.JSON(). + if _, present := resp["budget_limit"]; present { + t.Errorf("budget_limit must not appear in open GET /workspaces/:id response (got %v)", resp["budget_limit"]) + } + if _, present := resp["monthly_spend"]; present { + t.Errorf("monthly_spend must not appear in open GET /workspaces/:id response (got %v)", resp["monthly_spend"]) + } + // Confirm non-financial fields are still present. + if resp["name"] != "Capped Agent" { + t.Errorf("expected name 'Capped Agent', got %v", resp["name"]) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("sqlmock expectations not met: %v", err) + } +} + +// ==================== POST — create with budget_limit ==================== + +// TestWorkspaceBudget_Create_WithLimit verifies that POST /workspaces with +// a budget_limit passes the value as the 10th INSERT parameter ($10). +func TestWorkspaceBudget_Create_WithLimit(t *testing.T) { + mock := setupTestDB(t) + setupTestRedis(t) + handler := NewWorkspaceHandler(newTestBroadcaster(), nil, "http://localhost:8080", t.TempDir()) + + budgetVal := int64(1000) // $10.00 + mock.ExpectBegin() + mock.ExpectExec("INSERT INTO workspaces"). + WithArgs( + sqlmock.AnyArg(), // id + "Budgeted Agent", // name + nil, // role + 1, // tier + "langgraph", // runtime + sqlmock.AnyArg(), // awareness_namespace + (*string)(nil), // parent_id + nil, // workspace_dir + "none", // workspace_access + &budgetVal, // budget_limit ($10) + ). + WillReturnResult(sqlmock.NewResult(0, 1)) + mock.ExpectCommit() + mock.ExpectExec("INSERT INTO canvas_layouts"). + WithArgs(sqlmock.AnyArg(), float64(0), float64(0)). + WillReturnResult(sqlmock.NewResult(0, 1)) + mock.ExpectExec("INSERT INTO structure_events"). + WillReturnResult(sqlmock.NewResult(0, 1)) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + body := `{"name":"Budgeted Agent","budget_limit":1000}` + c.Request = httptest.NewRequest("POST", "/workspaces", bytes.NewBufferString(body)) + c.Request.Header.Set("Content-Type", "application/json") + handler.Create(c) + + if w.Code != http.StatusCreated { + t.Errorf("expected 201, got %d: %s", w.Code, w.Body.String()) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("sqlmock expectations not met: %v", err) + } +} + +// ==================== PATCH — budget_limit silently ignored on general update ==================== + +// TestWorkspaceBudget_Update_SetLimit verifies that PATCH /workspaces/:id with +// budget_limit=500 does NOT issue any DB write for budget_limit. The only write +// path is the AdminAuth-gated PATCH /workspaces/:id/budget endpoint. (#611) +func TestWorkspaceBudget_Update_SetLimit(t *testing.T) { + mock := setupTestDB(t) + setupTestRedis(t) + handler := NewWorkspaceHandler(newTestBroadcaster(), nil, "http://localhost:8080", t.TempDir()) + + // Only the existence probe fires; no UPDATE for budget_limit. + mock.ExpectQuery("SELECT EXISTS.*workspaces WHERE id"). + WithArgs("ws-upd-budget"). + WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(true)) + // No ExpectExec for budget_limit — sqlmock will fail if one is issued. + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Params = gin.Params{{Key: "id", Value: "ws-upd-budget"}} + body := `{"budget_limit":500}` + c.Request = httptest.NewRequest("PATCH", "/workspaces/ws-upd-budget", bytes.NewBufferString(body)) + c.Request.Header.Set("Content-Type", "application/json") + handler.Update(c) + + if w.Code != http.StatusOK { + t.Errorf("expected 200, got %d: %s", w.Code, w.Body.String()) + } + // If a budget_limit UPDATE was issued, sqlmock would have an unexpected call. + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("unexpected DB activity — budget_limit must not be written via general Update: %v", err) + } +} + +// TestWorkspaceBudget_Update_ClearLimit verifies that PATCH /workspaces/:id +// with budget_limit=null does NOT issue any DB write for budget_limit. (#611) +func TestWorkspaceBudget_Update_ClearLimit(t *testing.T) { + mock := setupTestDB(t) + setupTestRedis(t) + handler := NewWorkspaceHandler(newTestBroadcaster(), nil, "http://localhost:8080", t.TempDir()) + + // Only the existence probe fires; no UPDATE for budget_limit. + mock.ExpectQuery("SELECT EXISTS.*workspaces WHERE id"). + WithArgs("ws-clear-budget"). + WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(true)) + // No ExpectExec — a budget_limit write here would re-open the vulnerability. + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Params = gin.Params{{Key: "id", Value: "ws-clear-budget"}} + body := `{"budget_limit":null}` + c.Request = httptest.NewRequest("PATCH", "/workspaces/ws-clear-budget", bytes.NewBufferString(body)) + c.Request.Header.Set("Content-Type", "application/json") + handler.Update(c) + + if w.Code != http.StatusOK { + t.Errorf("expected 200, got %d: %s", w.Code, w.Body.String()) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("unexpected DB activity — budget_limit must not be written via general Update: %v", err) + } +} + +// ==================== A2A enforcement ==================== + +// TestWorkspaceBudget_A2A_ExceededReturns402 verifies that the A2A proxy +// returns HTTP 402 {"error":"workspace budget limit exceeded"} when +// monthly_spend equals budget_limit. +func TestWorkspaceBudget_A2A_ExceededReturns402(t *testing.T) { + mock := setupTestDB(t) + mr := setupTestRedis(t) + handler := NewWorkspaceHandler(newTestBroadcaster(), nil, "http://localhost:8080", t.TempDir()) + + // 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"). + WithArgs("ws-over-budget"). + WillReturnRows(sqlmock.NewRows([]string{"budget_limit", "monthly_spend"}). + AddRow(int64(500), int64(500))) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Params = gin.Params{{Key: "id", Value: "ws-over-budget"}} + body := `{"message":{"role":"user","parts":[{"text":"hello"}]}}` + c.Request = httptest.NewRequest("POST", "/workspaces/ws-over-budget/a2a", bytes.NewBufferString(body)) + c.Request.Header.Set("Content-Type", "application/json") + handler.ProxyA2A(c) + + if w.Code != http.StatusPaymentRequired { + t.Errorf("expected 402 when budget exceeded, got %d: %s", w.Code, w.Body.String()) + } + var resp map[string]interface{} + json.Unmarshal(w.Body.Bytes(), &resp) + if resp["error"] != "workspace budget limit exceeded" { + t.Errorf("expected 'workspace budget limit exceeded', got %v", resp["error"]) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("sqlmock expectations not met: %v", err) + } +} + +// TestWorkspaceBudget_A2A_AboveLimitReturns402 verifies 402 when spend > limit. +func TestWorkspaceBudget_A2A_AboveLimitReturns402(t *testing.T) { + mock := setupTestDB(t) + mr := setupTestRedis(t) + handler := NewWorkspaceHandler(newTestBroadcaster(), nil, "http://localhost:8080", t.TempDir()) + + mr.Set(fmt.Sprintf("ws:%s:url", "ws-way-over"), "http://localhost:9999") + + // spend > limit + mock.ExpectQuery("SELECT budget_limit, COALESCE"). + WithArgs("ws-way-over"). + WillReturnRows(sqlmock.NewRows([]string{"budget_limit", "monthly_spend"}). + AddRow(int64(100), int64(9999))) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Params = gin.Params{{Key: "id", Value: "ws-way-over"}} + body := `{"message":{"role":"user","parts":[{"text":"test"}]}}` + c.Request = httptest.NewRequest("POST", "/workspaces/ws-way-over/a2a", bytes.NewBufferString(body)) + c.Request.Header.Set("Content-Type", "application/json") + handler.ProxyA2A(c) + + if w.Code != http.StatusPaymentRequired { + t.Errorf("expected 402 when spend > limit, got %d: %s", w.Code, w.Body.String()) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("sqlmock expectations not met: %v", err) + } +} + +// TestWorkspaceBudget_A2A_UnderLimitPassesThrough verifies that A2A calls +// succeed normally when monthly_spend is below budget_limit. +func TestWorkspaceBudget_A2A_UnderLimitPassesThrough(t *testing.T) { + mock := setupTestDB(t) + mr := setupTestRedis(t) + handler := NewWorkspaceHandler(newTestBroadcaster(), nil, "http://localhost:8080", t.TempDir()) + + // Stand up a minimal mock agent that returns a valid A2A response + agentServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + fmt.Fprint(w, `{"jsonrpc":"2.0","id":"1","result":{"status":"ok"}}`) + })) + defer agentServer.Close() + + 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"). + WithArgs("ws-under-budget"). + WillReturnRows(sqlmock.NewRows([]string{"budget_limit", "monthly_spend"}). + AddRow(int64(500), int64(100))) + + // Activity log INSERT from logA2ASuccess + mock.ExpectExec("INSERT INTO activity_logs"). + WillReturnResult(sqlmock.NewResult(0, 1)) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Params = gin.Params{{Key: "id", Value: "ws-under-budget"}} + body := `{"jsonrpc":"2.0","id":"1","method":"message/send","params":{"message":{"role":"user","parts":[{"text":"hello"}]}}}` + c.Request = httptest.NewRequest("POST", "/workspaces/ws-under-budget/a2a", bytes.NewBufferString(body)) + c.Request.Header.Set("Content-Type", "application/json") + handler.ProxyA2A(c) + + // Give the async logA2ASuccess goroutine a moment to fire + time.Sleep(50 * time.Millisecond) + + if w.Code != http.StatusOK { + t.Errorf("expected 200 when under budget, got %d: %s", w.Code, w.Body.String()) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("sqlmock expectations not met: %v", err) + } +} + +// TestWorkspaceBudget_A2A_NilLimitPassesThrough verifies that when +// budget_limit IS NULL (no ceiling set), A2A calls pass through unconditionally. +func TestWorkspaceBudget_A2A_NilLimitPassesThrough(t *testing.T) { + mock := setupTestDB(t) + mr := setupTestRedis(t) + handler := NewWorkspaceHandler(newTestBroadcaster(), nil, "http://localhost:8080", t.TempDir()) + + agentServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + fmt.Fprint(w, `{"jsonrpc":"2.0","id":"2","result":{"status":"ok"}}`) + })) + defer agentServer.Close() + + 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"). + WithArgs("ws-no-limit"). + WillReturnRows(sqlmock.NewRows([]string{"budget_limit", "monthly_spend"}). + AddRow(nil, int64(999999))) // huge spend but no limit set + + mock.ExpectExec("INSERT INTO activity_logs"). + WillReturnResult(sqlmock.NewResult(0, 1)) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Params = gin.Params{{Key: "id", Value: "ws-no-limit"}} + body := `{"jsonrpc":"2.0","id":"2","method":"message/send","params":{"message":{"role":"user","parts":[{"text":"hi"}]}}}` + c.Request = httptest.NewRequest("POST", "/workspaces/ws-no-limit/a2a", bytes.NewBufferString(body)) + c.Request.Header.Set("Content-Type", "application/json") + handler.ProxyA2A(c) + + time.Sleep(50 * time.Millisecond) + + if w.Code != http.StatusOK { + t.Errorf("expected 200 when no limit set, got %d: %s", w.Code, w.Body.String()) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("sqlmock expectations not met: %v", err) + } +} + +// TestWorkspaceBudget_A2A_DBErrorFailOpen verifies that a DB error during the +// budget check is fail-open — the request proceeds rather than being blocked. +func TestWorkspaceBudget_A2A_DBErrorFailOpen(t *testing.T) { + mock := setupTestDB(t) + mr := setupTestRedis(t) + handler := NewWorkspaceHandler(newTestBroadcaster(), nil, "http://localhost:8080", t.TempDir()) + + agentServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + fmt.Fprint(w, `{"jsonrpc":"2.0","id":"3","result":{"status":"ok"}}`) + })) + defer agentServer.Close() + + 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"). + WithArgs("ws-db-err-budget"). + WillReturnError(sql.ErrConnDone) + + mock.ExpectExec("INSERT INTO activity_logs"). + WillReturnResult(sqlmock.NewResult(0, 1)) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Params = gin.Params{{Key: "id", Value: "ws-db-err-budget"}} + body := `{"jsonrpc":"2.0","id":"3","method":"message/send","params":{"message":{"role":"user","parts":[{"text":"fail-open test"}]}}}` + c.Request = httptest.NewRequest("POST", "/workspaces/ws-db-err-budget/a2a", bytes.NewBufferString(body)) + c.Request.Header.Set("Content-Type", "application/json") + handler.ProxyA2A(c) + + time.Sleep(50 * time.Millisecond) + + if w.Code != http.StatusOK { + t.Errorf("expected 200 on DB error (fail-open), 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/handlers/workspace_test.go b/platform/internal/handlers/workspace_test.go index e36665d0..b524d412 100644 --- a/platform/internal/handlers/workspace_test.go +++ b/platform/internal/handlers/workspace_test.go @@ -24,13 +24,15 @@ func TestWorkspaceGet_Success(t *testing.T) { "id", "name", "role", "tier", "status", "agent_card", "url", "parent_id", "active_tasks", "last_error_rate", "last_sample_error", "uptime_seconds", "current_task", "runtime", "workspace_dir", "x", "y", "collapsed", + "budget_limit", "monthly_spend", } mock.ExpectQuery("SELECT w.id, w.name"). WithArgs("ws-get-1"). WillReturnRows(sqlmock.NewRows(columns). AddRow("ws-get-1", "My Agent", "worker", 1, "online", []byte(`{"name":"test"}`), "http://localhost:8001", nil, 2, 0.05, "", 3600, "working", "langgraph", - "", 10.0, 20.0, false)) + "", 10.0, 20.0, false, + nil, 0)) w := httptest.NewRecorder() c, _ := gin.CreateTestContext(w) @@ -149,7 +151,7 @@ func TestWorkspaceCreate_DBInsertError(t *testing.T) { // Transaction begins, workspace INSERT fails, transaction is rolled back. mock.ExpectBegin() mock.ExpectExec("INSERT INTO workspaces"). - WithArgs(sqlmock.AnyArg(), "Failing Agent", nil, 1, "langgraph", sqlmock.AnyArg(), (*string)(nil), nil, "none"). + WithArgs(sqlmock.AnyArg(), "Failing Agent", nil, 1, "langgraph", sqlmock.AnyArg(), (*string)(nil), nil, "none", (*int64)(nil)). WillReturnError(sql.ErrConnDone) mock.ExpectRollback() @@ -181,7 +183,7 @@ func TestWorkspaceCreate_DefaultsApplied(t *testing.T) { mock.ExpectBegin() // Expect workspace INSERT with defaulted tier=1, runtime="langgraph" mock.ExpectExec("INSERT INTO workspaces"). - WithArgs(sqlmock.AnyArg(), "Default Agent", nil, 1, "langgraph", sqlmock.AnyArg(), (*string)(nil), nil, "none"). + WithArgs(sqlmock.AnyArg(), "Default Agent", nil, 1, "langgraph", sqlmock.AnyArg(), (*string)(nil), nil, "none", (*int64)(nil)). WillReturnResult(sqlmock.NewResult(0, 1)) mock.ExpectCommit() @@ -232,7 +234,7 @@ func TestWorkspaceCreate_WithSecrets_Persists(t *testing.T) { mock.ExpectBegin() mock.ExpectExec("INSERT INTO workspaces"). - WithArgs(sqlmock.AnyArg(), "Hermes Agent", nil, 1, "hermes", sqlmock.AnyArg(), (*string)(nil), nil, "none"). + WithArgs(sqlmock.AnyArg(), "Hermes Agent", nil, 1, "hermes", sqlmock.AnyArg(), (*string)(nil), nil, "none", (*int64)(nil)). WillReturnResult(sqlmock.NewResult(0, 1)) // Secret inserted inside the same transaction. mock.ExpectExec("INSERT INTO workspace_secrets"). @@ -344,6 +346,7 @@ func TestWorkspaceList_Empty(t *testing.T) { "id", "name", "role", "tier", "status", "agent_card", "url", "parent_id", "active_tasks", "last_error_rate", "last_sample_error", "uptime_seconds", "current_task", "runtime", "workspace_dir", "x", "y", "collapsed", + "budget_limit", "monthly_spend", })) w := httptest.NewRecorder() @@ -857,3 +860,132 @@ func TestWorkspaceUpdate_SensitiveField_NoTokensYet_FailOpen(t *testing.T) { t.Errorf("bootstrap fail-open: got %d, want 200 (%s)", w.Code, w.Body.String()) } } + +// ==================== #611 Security Auditor regressions ==================== + +// TestWorkspaceGet_FinancialFieldsStripped verifies that GET /workspaces/:id +// does NOT expose budget_limit or monthly_spend. The endpoint is on the open +// router — any caller with a UUID would otherwise read billing data. (#611 Fix 2) +func TestWorkspaceGet_FinancialFieldsStripped(t *testing.T) { + mock := setupTestDB(t) + setupTestRedis(t) + broadcaster := newTestBroadcaster() + handler := NewWorkspaceHandler(broadcaster, nil, "http://localhost:8080", t.TempDir()) + + columns := []string{ + "id", "name", "role", "tier", "status", "agent_card", "url", + "parent_id", "active_tasks", "last_error_rate", "last_sample_error", + "uptime_seconds", "current_task", "runtime", "workspace_dir", "x", "y", "collapsed", + "budget_limit", "monthly_spend", + } + // Populate with non-zero financial values to confirm they are stripped. + mock.ExpectQuery("SELECT w.id, w.name"). + WithArgs("ws-fin-1"). + WillReturnRows(sqlmock.NewRows(columns). + AddRow("ws-fin-1", "Finance Test", "worker", 1, "online", []byte(`{}`), + "http://localhost:9001", nil, 0, 0.0, "", 0, "", "langgraph", + "", 0.0, 0.0, false, + int64(50000), int64(12500))) // budget_limit=500 USD, spend=125 USD + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Params = gin.Params{{Key: "id", Value: "ws-fin-1"}} + c.Request = httptest.NewRequest("GET", "/workspaces/ws-fin-1", nil) + + handler.Get(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("failed to parse response: %v", err) + } + if _, present := resp["budget_limit"]; present { + t.Errorf("budget_limit must not appear in GET /workspaces/:id response (got %v)", resp["budget_limit"]) + } + if _, present := resp["monthly_spend"]; present { + t.Errorf("monthly_spend must not appear in GET /workspaces/:id response (got %v)", resp["monthly_spend"]) + } + // Sanity-check that normal fields are still present. + if resp["name"] != "Finance Test" { + t.Errorf("expected name 'Finance Test', got %v", resp["name"]) + } + + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("unmet sqlmock expectations: %v", err) + } +} + +// TestWorkspaceUpdate_BudgetLimitIgnored verifies that including budget_limit +// in a PATCH /workspaces/:id body does NOT trigger a DB write. The only write +// path for budget_limit is PATCH /workspaces/:id/budget (AdminAuth-gated). +// Any workspace bearer must not be able to self-clear its spending ceiling. +// (#611 Fix 1) +func TestWorkspaceUpdate_BudgetLimitIgnored(t *testing.T) { + mock := setupTestDB(t) + setupTestRedis(t) + broadcaster := newTestBroadcaster() + handler := NewWorkspaceHandler(broadcaster, nil, "http://localhost:8080", t.TempDir()) + + // Only the existence probe fires — no UPDATE for budget_limit. + mock.ExpectQuery("SELECT EXISTS.*workspaces WHERE id"). + WithArgs("ws-budget-test"). + WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(true)) + // name update is the only expected write + mock.ExpectExec("UPDATE workspaces SET name"). + WithArgs("ws-budget-test", "Safe Name"). + WillReturnResult(sqlmock.NewResult(0, 1)) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Params = gin.Params{{Key: "id", Value: "ws-budget-test"}} + // Send budget_limit alongside an innocuous field. + body := `{"name":"Safe Name","budget_limit":null}` + c.Request = httptest.NewRequest("PATCH", "/workspaces/ws-budget-test", + bytes.NewBufferString(body)) + c.Request.Header.Set("Content-Type", "application/json") + + handler.Update(c) + + if w.Code != http.StatusOK { + t.Errorf("expected 200, got %d: %s", w.Code, w.Body.String()) + } + + // sqlmock will fail if any unexpected DB call was made (e.g. for budget_limit). + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("unexpected DB call — budget_limit must not be written via Update: %v", err) + } +} + +// TestWorkspaceUpdate_BudgetLimitOnly_Ignored verifies that a body containing +// ONLY budget_limit results in no DB writes at all (besides the existence probe). +func TestWorkspaceUpdate_BudgetLimitOnly_Ignored(t *testing.T) { + mock := setupTestDB(t) + setupTestRedis(t) + broadcaster := newTestBroadcaster() + handler := NewWorkspaceHandler(broadcaster, nil, "http://localhost:8080", t.TempDir()) + + mock.ExpectQuery("SELECT EXISTS.*workspaces WHERE id"). + WithArgs("ws-budget-only"). + WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(true)) + // No UPDATE expected — budget_limit must be silently skipped. + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Params = gin.Params{{Key: "id", Value: "ws-budget-only"}} + c.Request = httptest.NewRequest("PATCH", "/workspaces/ws-budget-only", + bytes.NewBufferString(`{"budget_limit":999999}`)) + c.Request.Header.Set("Content-Type", "application/json") + + handler.Update(c) + + if w.Code != http.StatusOK { + t.Errorf("expected 200, got %d: %s", w.Code, w.Body.String()) + } + + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("unexpected DB call for budget_limit: %v", err) + } +} diff --git a/platform/internal/models/workspace.go b/platform/internal/models/workspace.go index 4bf9ed9a..e7c642f2 100644 --- a/platform/internal/models/workspace.go +++ b/platform/internal/models/workspace.go @@ -44,6 +44,11 @@ type HeartbeatPayload struct { ActiveTasks int `json:"active_tasks"` UptimeSeconds int `json:"uptime_seconds"` CurrentTask string `json:"current_task"` + // MonthlySpend is the agent's self-reported accumulated LLM API spend for + // the current month, in USD cents. Zero means "no update" — the platform + // only writes to monthly_spend when this field is > 0. Agents should + // report their cumulative spend each heartbeat (not the delta). + MonthlySpend int64 `json:"monthly_spend"` } type UpdateCardPayload struct { @@ -63,6 +68,9 @@ type CreateWorkspacePayload struct { WorkspaceDir string `json:"workspace_dir"` // host path to mount as /workspace (empty = isolated volume) WorkspaceAccess string `json:"workspace_access"` // "none" (default), "read_only", or "read_write" — see #65 ParentID *string `json:"parent_id"` + // BudgetLimit is the optional monthly spend ceiling in USD cents. + // NULL (omitted) means no limit. budget_limit=500 means $5.00/month. + BudgetLimit *int64 `json:"budget_limit"` // Secrets is an optional map of key→plaintext-value pairs to persist as // workspace secrets at creation time. Stored encrypted (same path as // POST /workspaces/:id/secrets). Nil/empty map is a no-op. diff --git a/platform/internal/router/router.go b/platform/internal/router/router.go index d41b653a..8e735e45 100644 --- a/platform/internal/router/router.go +++ b/platform/internal/router/router.go @@ -256,6 +256,14 @@ 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 stays on wsAuth — a workspace agent reading its own budget is legitimate. + // PATCH is admin-only — workspace agents must not be able to self-clear their + // spending ceiling (that would defeat the entire budget enforcement feature). + budgeth := handlers.NewBudgetHandler() + wsAuth.GET("/budget", budgeth.GetBudget) + r.PATCH("/workspaces/:id/budget", middleware.AdminAuth(db.DB), budgeth.PatchBudget) + // Token management (user-facing create/list/revoke) tokh := handlers.NewTokenHandler() wsAuth.GET("/tokens", tokh.List) diff --git a/platform/migrations/027_workspace_budget.down.sql b/platform/migrations/027_workspace_budget.down.sql new file mode 100644 index 00000000..c7cd48e7 --- /dev/null +++ b/platform/migrations/027_workspace_budget.down.sql @@ -0,0 +1,3 @@ +ALTER TABLE workspaces + DROP COLUMN IF EXISTS budget_limit, + DROP COLUMN IF EXISTS monthly_spend; diff --git a/platform/migrations/027_workspace_budget.up.sql b/platform/migrations/027_workspace_budget.up.sql new file mode 100644 index 00000000..afe61da5 --- /dev/null +++ b/platform/migrations/027_workspace_budget.up.sql @@ -0,0 +1,11 @@ +-- Per-workspace monthly budget limit (#541). +-- NULL means no limit. When monthly_spend reaches budget_limit, the A2A +-- proxy returns 402 {"error":"workspace budget limit exceeded"} and rejects +-- further A2A calls until budget_limit is raised or monthly_spend is reset. +-- +-- Units: USD cents (integer). budget_limit=500 means $5.00/month. +-- monthly_spend is updated by the workspace via the heartbeat endpoint; +-- agents report their accumulated LLM API cost each heartbeat cycle. +ALTER TABLE workspaces + ADD COLUMN IF NOT EXISTS budget_limit BIGINT DEFAULT NULL, + ADD COLUMN IF NOT EXISTS monthly_spend BIGINT NOT NULL DEFAULT 0;