feat: token budget enforcement per virtual key (Story 3.3)

docs/swagger/docs.go

@@ -1332,6 +1332,9 @@ const docTemplate = `{
                 "target": {
                     "type": "string"
                 },
+                "token_budget": {
+                    "type": "integer"
+                },
                 "used": {
                     "type": "boolean"
                 }

docs/swagger/swagger.json

@@ -1326,6 +1326,9 @@
                 "target": {
                     "type": "string"
                 },
+                "token_budget": {
+                    "type": "integer"
+                },
                 "used": {
                     "type": "boolean"
                 }

docs/swagger/swagger.yaml

@@ -16,6 +16,8 @@ definitions:
         type: string
       target:
         type: string
+      token_budget:
+        type: integer
       used:
         type: boolean
     type: object

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;

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.

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()
 	}

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())
+	}
+}