From 895c256b5639474e30e82ddc4d7b549ac996d713 Mon Sep 17 00:00:00 2001 From: Victor Faro Date: Sun, 5 Apr 2026 20:01:01 -0300 Subject: [PATCH] feat: token budget enforcement per virtual key (Story 3.3) - Add token_budget field to VirtualKey (0 = unlimited) - Proxy checks UsageStore.TotalTokens before forwarding; returns 429 when budget reached - Migration 010_add_virtual_key_token_budget.sql - Tests: under budget allowed, exceeded blocked, zero means unlimited Co-Authored-By: Claude Sonnet 4.6 --- docs/swagger/docs.go | 3 + docs/swagger/swagger.json | 3 + docs/swagger/swagger.yaml | 2 + .../010_add_virtual_key_token_budget.sql | 1 + pkg/keys/virtualkey.go | 7 +- routes/v1/proxy.go | 7 ++ tests/token_budget_test.go | 116 ++++++++++++++++++ 7 files changed, 136 insertions(+), 3 deletions(-) create mode 100644 migrations/010_add_virtual_key_token_budget.sql create mode 100644 tests/token_budget_test.go diff --git a/docs/swagger/docs.go b/docs/swagger/docs.go index 8df9ef7..1232a4f 100644 --- a/docs/swagger/docs.go +++ b/docs/swagger/docs.go @@ -1332,6 +1332,9 @@ const docTemplate = `{ "target": { "type": "string" }, + "token_budget": { + "type": "integer" + }, "used": { "type": "boolean" } diff --git a/docs/swagger/swagger.json b/docs/swagger/swagger.json index 1ae4930..454e092 100644 --- a/docs/swagger/swagger.json +++ b/docs/swagger/swagger.json @@ -1326,6 +1326,9 @@ "target": { "type": "string" }, + "token_budget": { + "type": "integer" + }, "used": { "type": "boolean" } diff --git a/docs/swagger/swagger.yaml b/docs/swagger/swagger.yaml index 1510fd2..ec1345e 100644 --- a/docs/swagger/swagger.yaml +++ b/docs/swagger/swagger.yaml @@ -16,6 +16,8 @@ definitions: type: string target: type: string + token_budget: + type: integer used: type: boolean type: object diff --git a/migrations/010_add_virtual_key_token_budget.sql b/migrations/010_add_virtual_key_token_budget.sql new file mode 100644 index 0000000..166dfc3 --- /dev/null +++ b/migrations/010_add_virtual_key_token_budget.sql @@ -0,0 +1 @@ +ALTER TABLE virtual_keys ADD COLUMN IF NOT EXISTS token_budget INTEGER NOT NULL DEFAULT 0; diff --git a/pkg/keys/virtualkey.go b/pkg/keys/virtualkey.go index 9a54d6e..dd47ddb 100644 --- a/pkg/keys/virtualkey.go +++ b/pkg/keys/virtualkey.go @@ -10,9 +10,10 @@ type VirtualKey struct { ExpiresAt time.Time `json:"expires_at" gorm:"not null"` Target string `json:"target" gorm:"not null"` RateLimit int `json:"rate_limit" gorm:"not null"` - Source string `json:"source,omitempty" gorm:"size:16;default:''"` - OneShot bool `json:"one_shot,omitempty" gorm:"default:false"` - Used bool `json:"used,omitempty" gorm:"default:false"` + Source string `json:"source,omitempty" gorm:"size:16;default:''"` + OneShot bool `json:"one_shot,omitempty" gorm:"default:false"` + Used bool `json:"used,omitempty" gorm:"default:false"` + TokenBudget int `json:"token_budget,omitempty" gorm:"default:0"` } // SourceMCP is the source label for keys issued via the MCP tool. diff --git a/routes/v1/proxy.go b/routes/v1/proxy.go index 2a9c02d..614da05 100644 --- a/routes/v1/proxy.go +++ b/routes/v1/proxy.go @@ -85,6 +85,13 @@ func (h *Handler) Proxy(w http.ResponseWriter, r *http.Request) { return } + if k.TokenBudget > 0 && h.UsageStore != nil { + if h.UsageStore.TotalTokens(k.ID) >= k.TokenBudget { + writeError(w, "token budget exceeded", http.StatusTooManyRequests) + return + } + } + if config.MetricsEnabled() { metrics.KeyUsageTotal.WithLabelValues(k.ID).Inc() } diff --git a/tests/token_budget_test.go b/tests/token_budget_test.go new file mode 100644 index 0000000..4d43c5b --- /dev/null +++ b/tests/token_budget_test.go @@ -0,0 +1,116 @@ +package tests + +import ( + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/farovictor/bifrost/pkg/keys" + "github.com/farovictor/bifrost/pkg/rootkeys" + routes "github.com/farovictor/bifrost/routes" + "github.com/farovictor/bifrost/pkg/services" + "github.com/farovictor/bifrost/pkg/usage" +) + +func setupBudgetEnv(t *testing.T, budget int) (s *routes.Server, router http.Handler, keyID string, backend *httptest.Server) { + t.Helper() + s = newTestServer(t) + + backend = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{"id":"resp"}`)) + })) + t.Cleanup(backend.Close) + + rk := rootkeys.RootKey{ID: "rk-budget", APIKey: "real"} + s.RootKeyStore.Create(rk) + svc := services.Service{ID: "svc-budget", Endpoint: backend.URL, RootKeyID: rk.ID} + s.ServiceStore.Create(svc) + + keyID = "vk-budget" + k := keys.VirtualKey{ + ID: keyID, + Target: svc.ID, + Scope: keys.ScopeWrite, + RateLimit: 100, + ExpiresAt: time.Now().Add(time.Hour), + TokenBudget: budget, + } + s.KeyStore.Create(k) + + router = setupRouter(s) + return +} + +func TestTokenBudget_AllowedWhenUnderBudget(t *testing.T) { + s, router, keyID, _ := setupBudgetEnv(t, 1000) + + // Pre-seed some usage below the budget + s.UsageStore.Record(usage.Event{ + KeyID: keyID, + Timestamp: time.Now(), + StatusCode: 200, + Service: "svc-budget", + LatencyMS: 10, + TotalTokens: 500, + }) + + req := httptest.NewRequest(http.MethodPost, "/v1/proxy/chat", nil) + req.Header.Set("X-Virtual-Key", keyID) + rr := httptest.NewRecorder() + router.ServeHTTP(rr, req) + + if rr.Code != http.StatusOK { + t.Fatalf("expected 200 under budget, got %d: %s", rr.Code, rr.Body.String()) + } +} + +func TestTokenBudget_BlockedWhenExceeded(t *testing.T) { + s, router, keyID, _ := setupBudgetEnv(t, 100) + + // Pre-seed usage at or above the budget + s.UsageStore.Record(usage.Event{ + KeyID: keyID, + Timestamp: time.Now(), + StatusCode: 200, + Service: "svc-budget", + LatencyMS: 10, + TotalTokens: 100, + }) + + req := httptest.NewRequest(http.MethodPost, "/v1/proxy/chat", nil) + req.Header.Set("X-Virtual-Key", keyID) + rr := httptest.NewRecorder() + router.ServeHTTP(rr, req) + + if rr.Code != http.StatusTooManyRequests { + t.Fatalf("expected 429 when budget exceeded, got %d: %s", rr.Code, rr.Body.String()) + } + if msg := errorBody(t, rr); msg != "token budget exceeded" { + t.Errorf("unexpected error message: %s", msg) + } +} + +func TestTokenBudget_ZeroMeansUnlimited(t *testing.T) { + s, router, keyID, _ := setupBudgetEnv(t, 0) + + // Seed a huge amount — should still be allowed (budget=0 → unlimited) + s.UsageStore.Record(usage.Event{ + KeyID: keyID, + Timestamp: time.Now(), + StatusCode: 200, + Service: "svc-budget", + LatencyMS: 10, + TotalTokens: 999_999_999, + }) + + req := httptest.NewRequest(http.MethodPost, "/v1/proxy/chat", nil) + req.Header.Set("X-Virtual-Key", keyID) + rr := httptest.NewRecorder() + router.ServeHTTP(rr, req) + + if rr.Code != http.StatusOK { + t.Fatalf("expected 200 when budget=0 (unlimited), got %d: %s", rr.Code, rr.Body.String()) + } +}