Claude Subscription Usage Gate Implementation Plan¶

func TestKindBlocked(t *testing.T) {
    future := time.Now().Add(time.Hour)
    sub := Subscription{FiveHourPercent: 42, FiveHourReset: future, WeeklyPercent: 10, WeeklyReset: future}
    base := Config{Enabled: true, Mode: ModeClaudeSubscription,
        SpawnCeilingByKind: map[string]int{"brainstorm": 40, "incident": 98}}
    cases := []struct {
        name    string
        cfg     Config
        kind    string
        blocked bool
    }{
        {"brainstorm over ceiling", base, "brainstorm", true},   // 42 >= 40
        {"incident under ceiling", base, "incident", false},     // 42 < 98
        {"kind without ceiling not blocked", base, "implement", false},
        {"disabled never blocks", Config{Mode: ModeClaudeSubscription, SpawnCeilingByKind: base.SpawnCeilingByKind}, "brainstorm", false},
        {"customWindow mode never per-kind blocks", Config{Enabled: true, Mode: ModeCustomWindow, SpawnCeilingByKind: base.SpawnCeilingByKind}, "brainstorm", false},
    }
    for _, tc := range cases {
        t.Run(tc.name, func(t *testing.T) {
            if got := KindBlocked(tc.cfg, sub, tc.kind, time.Now()); got != tc.blocked {
                t.Fatalf("KindBlocked=%v want %v", got, tc.blocked)
            }
        })
    }
}

func TestKindBlockedIgnoresExpiredWindow(t *testing.T) {
    past := time.Now().Add(-time.Hour)
    sub := Subscription{FiveHourPercent: 99, FiveHourReset: past} // expired -> ignored
    cfg := Config{Enabled: true, Mode: ModeClaudeSubscription, SpawnCeilingByKind: map[string]int{"brainstorm": 40}}
    if KindBlocked(cfg, sub, "brainstorm", time.Now()) {
        t.Fatal("expired window must not block")
    }
}

    // SpawnCeilingByKind gates each Task kind independently in claudeSubscription
    // mode: work of kind K is held once account usage reaches SpawnCeilingByKind[K]
    // percent. Kinds absent from the map are not per-kind gated (they fall through
    // to the pool-class proactive/emergency thresholds). Ignored in customWindow mode.
    SpawnCeilingByKind map[string]int

    OpusPercent    float64
    OpusReset      time.Time
    SonnetPercent  float64
    SonnetReset    time.Time
    OverageEnabled bool
    OveragePercent float64

// KindBlocked reports whether work of the given kind must be held, given the
// account subscription usage. It applies only in claudeSubscription mode with a
// configured per-kind ceiling; every other case returns false so the caller's
// pool-class Decision remains authoritative.
func KindBlocked(cfg Config, sub Subscription, kind string, now time.Time) bool {
    if !cfg.Enabled || cfg.Mode != ModeClaudeSubscription {
        return false
    }
    ceiling, ok := cfg.SpawnCeilingByKind[kind]
    if !ok || ceiling <= 0 {
        return false
    }
    return subscriptionUsedPercent(sub, now) >= float64(ceiling)
}

package accountusage

import (
    "testing"
    "time"
)

func TestStoreGetSetIsConcurrencySafeAndCopies(t *testing.T) {
    s := &Store{}
    if got := s.Get(); got.Healthy || !got.UpdatedAt.IsZero() {
        t.Fatal("zero store must be unhealthy/empty")
    }
    reset := time.Now().Add(time.Hour)
    s.Set(Snapshot{FiveHour: Window{Percent: 55, Reset: reset}, Healthy: true, UpdatedAt: time.Now()})
    got := s.Get()
    if got.FiveHour.Percent != 55 || !got.Healthy {
        t.Fatalf("Get mismatch: %+v", got)
    }
}

func TestSnapshotSubscriptionProjection(t *testing.T) {
    reset := time.Now().Add(time.Hour)
    snap := Snapshot{
        FiveHour: Window{Percent: 42, Reset: reset},
        Weekly:   Window{Percent: 71, Reset: reset},
        Opus:     Window{Percent: 80, Reset: reset},
    }
    sub := snap.Subscription()
    if sub.FiveHourPercent != 42 || sub.WeeklyPercent != 71 || sub.OpusPercent != 80 {
        t.Fatalf("projection mismatch: %+v", sub)
    }
    if !sub.FiveHourReset.Equal(reset) {
        t.Fatal("reset not projected")
    }
}

package accountusage

import (
    "sync"
    "time"

    "github.com/szymonrychu/tatara-operator/internal/budget"
)

// Window is one usage window: percent used (0..100) and when it resets.
type Window struct {
    Percent float64
    Reset   time.Time
}

// Overage is the read-only pay-as-you-go overage pool ("monthly").
type Overage struct {
    Enabled bool
    Percent float64
    Used    float64
    Limit   float64
}

// Snapshot is the fleet-wide account usage at a point in time.
type Snapshot struct {
    FiveHour  Window
    Weekly    Window
    Opus      Window // per-model weekly; zero when N/A on the plan
    Sonnet    Window
    Overage   Overage
    Healthy   bool
    UpdatedAt time.Time
}

// Subscription projects the snapshot into the budget gate's input type. The gate
// uses 5h + weekly; per-model + overage ride along for metrics only.
func (s Snapshot) Subscription() budget.Subscription {
    return budget.Subscription{
        FiveHourPercent: s.FiveHour.Percent,
        FiveHourReset:   s.FiveHour.Reset,
        WeeklyPercent:   s.Weekly.Percent,
        WeeklyReset:     s.Weekly.Reset,
        OpusPercent:     s.Opus.Percent,
        OpusReset:       s.Opus.Reset,
        SonnetPercent:   s.Sonnet.Percent,
        SonnetReset:     s.Sonnet.Reset,
        OverageEnabled:  s.Overage.Enabled,
        OveragePercent:  s.Overage.Percent,
    }
}

// Store holds the single fleet-wide snapshot, safe for concurrent reads by every
// project's admission reconcile and writes by the poller.
type Store struct {
    mu   sync.RWMutex
    snap Snapshot
}

func (s *Store) Get() Snapshot {
    s.mu.RLock()
    defer s.mu.RUnlock()
    return s.snap
}

func (s *Store) Set(snap Snapshot) {
    s.mu.Lock()
    defer s.mu.Unlock()
    s.snap = snap
}

package accountusage

import (
    "context"
    "net/http"
    "net/http/httptest"
    "testing"
)

const sampleUsage = `{
  "five_hour": {"utilization": 42.5, "resets_at": "2999-01-01T00:00:00Z"},
  "seven_day": {"utilization": 71.0, "resets_at": "2999-01-02T00:00:00Z"},
  "seven_day_opus": {"utilization": 80.0, "resets_at": "2999-01-02T00:00:00Z"},
  "seven_day_sonnet": null,
  "extra_usage": {"is_enabled": true, "monthly_limit": 100, "used_credits": 25, "utilization": 25.0}
}`

func TestFetchParsesWindowsAndHeaders(t *testing.T) {
    var gotUA, gotAuth, gotBeta string
    srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        gotUA = r.Header.Get("User-Agent")
        gotAuth = r.Header.Get("Authorization")
        gotBeta = r.Header.Get("anthropic-beta")
        w.Write([]byte(sampleUsage))
    }))
    defer srv.Close()

    c := NewClient(ClientConfig{
        BaseURL:     srv.URL,
        TokenSource: func() (string, error) { return "tok123", nil },
        UserAgent:   "claude-code/9.9.9",
        AuthMode:    "bearer",
    })
    snap, err := c.Fetch(context.Background())
    if err != nil {
        t.Fatalf("Fetch: %v", err)
    }
    if snap.FiveHour.Percent != 42.5 || snap.Weekly.Percent != 71.0 || snap.Opus.Percent != 80.0 {
        t.Fatalf("windows: %+v", snap)
    }
    if snap.Sonnet.Percent != 0 { // null -> zero window
        t.Fatalf("null sonnet must be zero, got %+v", snap.Sonnet)
    }
    if !snap.Overage.Enabled || snap.Overage.Percent != 25.0 {
        t.Fatalf("overage: %+v", snap.Overage)
    }
    if !snap.Healthy {
        t.Fatal("successful fetch must be Healthy")
    }
    if gotUA != "claude-code/9.9.9" || gotAuth != "Bearer tok123" || gotBeta != "oauth-2025-04-20" {
        t.Fatalf("headers UA=%q auth=%q beta=%q", gotUA, gotAuth, gotBeta)
    }
}

func TestFetchXAPIKeyMode(t *testing.T) {
    var gotKey, gotAuth string
    srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        gotKey = r.Header.Get("x-api-key")
        gotAuth = r.Header.Get("Authorization")
        w.Write([]byte(sampleUsage))
    }))
    defer srv.Close()
    c := NewClient(ClientConfig{BaseURL: srv.URL, TokenSource: func() (string, error) { return "k", nil }, UserAgent: "claude-cli/1", AuthMode: "x-api-key"})
    if _, err := c.Fetch(context.Background()); err != nil {
        t.Fatal(err)
    }
    if gotKey != "k" || gotAuth != "" {
        t.Fatalf("x-api-key mode: key=%q auth=%q", gotKey, gotAuth)
    }
}

func TestFetchNon200IsError(t *testing.T) {
    srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(429) }))
    defer srv.Close()
    c := NewClient(ClientConfig{BaseURL: srv.URL, TokenSource: func() (string, error) { return "k", nil }, UserAgent: "x", AuthMode: "bearer"})
    if _, err := c.Fetch(context.Background()); err == nil {
        t.Fatal("429 must error")
    }
}

func TestNormalizeUtilizationFraction(t *testing.T) {
    // utilization may be reported 0..1 (fraction) instead of 0..100; normalize.
    if got := normalizeUtil(0.42); got != 42.0 {
        t.Fatalf("fraction normalize: %v", got)
    }
    if got := normalizeUtil(42.0); got != 42.0 {
        t.Fatalf("percent passthrough: %v", got)
    }
}

package accountusage

import (
    "context"
    "encoding/json"
    "fmt"
    "io"
    "net/http"
    "time"
)

const (
    defaultBaseURL  = "https://api.anthropic.com"
    usagePath       = "/api/oauth/usage"
    anthropicBeta   = "oauth-2025-04-20"
    requestTimeout  = 15 * time.Second
)

type ClientConfig struct {
    BaseURL     string
    TokenSource func() (string, error)
    UserAgent   string
    AuthMode    string // "bearer" (default) or "x-api-key"
    HTTP        *http.Client
}

type Client struct {
    base  string
    token func() (string, error)
    ua    string
    auth  string
    http  *http.Client
}

func NewClient(cfg ClientConfig) *Client {
    base := cfg.BaseURL
    if base == "" {
        base = defaultBaseURL
    }
    h := cfg.HTTP
    if h == nil {
        h = &http.Client{Timeout: requestTimeout}
    }
    auth := cfg.AuthMode
    if auth == "" {
        auth = "bearer"
    }
    return &Client{base: base, token: cfg.TokenSource, ua: cfg.UserAgent, auth: auth, http: h}
}

type usageWindow struct {
    Utilization float64 `json:"utilization"`
    ResetsAt    string  `json:"resets_at"`
}

type usageResponse struct {
    FiveHour *usageWindow `json:"five_hour"`
    SevenDay *usageWindow `json:"seven_day"`
    Opus     *usageWindow `json:"seven_day_opus"`
    Sonnet   *usageWindow `json:"seven_day_sonnet"`
    Extra    *struct {
        Enabled     bool    `json:"is_enabled"`
        MonthlyLimit float64 `json:"monthly_limit"`
        UsedCredits  float64 `json:"used_credits"`
        Utilization  float64 `json:"utilization"`
    } `json:"extra_usage"`
}

func (c *Client) Fetch(ctx context.Context) (Snapshot, error) {
    tok, err := c.token()
    if err != nil {
        return Snapshot{}, fmt.Errorf("accountusage: token: %w", err)
    }
    req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.base+usagePath, nil)
    if err != nil {
        return Snapshot{}, err
    }
    req.Header.Set("User-Agent", c.ua)
    req.Header.Set("Content-Type", "application/json")
    if c.auth == "x-api-key" {
        req.Header.Set("x-api-key", tok)
        req.Header.Set("anthropic-beta", "claude-code-20250219,"+anthropicBeta)
    } else {
        req.Header.Set("Authorization", "Bearer "+tok)
        req.Header.Set("anthropic-beta", anthropicBeta)
    }
    resp, err := c.http.Do(req)
    if err != nil {
        return Snapshot{}, fmt.Errorf("accountusage: request: %w", err)
    }
    defer resp.Body.Close()
    body, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<16))
    if resp.StatusCode != http.StatusOK {
        return Snapshot{}, fmt.Errorf("accountusage: status %d: %s", resp.StatusCode, string(body))
    }
    var ur usageResponse
    if err := json.Unmarshal(body, &ur); err != nil {
        return Snapshot{}, fmt.Errorf("accountusage: decode: %w", err)
    }
    if ur.FiveHour == nil && ur.SevenDay == nil {
        return Snapshot{}, fmt.Errorf("accountusage: schema drift: no five_hour/seven_day fields")
    }
    snap := Snapshot{Healthy: true, UpdatedAt: time.Now()}
    snap.FiveHour = toWindow(ur.FiveHour)
    snap.Weekly = toWindow(ur.SevenDay)
    snap.Opus = toWindow(ur.Opus)
    snap.Sonnet = toWindow(ur.Sonnet)
    if ur.Extra != nil {
        snap.Overage = Overage{Enabled: ur.Extra.Enabled, Percent: normalizeUtil(ur.Extra.Utilization), Used: ur.Extra.UsedCredits, Limit: ur.Extra.MonthlyLimit}
    }
    return snap, nil
}

func toWindow(w *usageWindow) Window {
    if w == nil {
        return Window{}
    }
    out := Window{Percent: normalizeUtil(w.Utilization)}
    if w.ResetsAt != "" {
        if t, err := time.Parse(time.RFC3339, w.ResetsAt); err == nil {
            out.Reset = t
        }
    }
    return out
}

// normalizeUtil coerces a utilization reported as a 0..1 fraction into 0..100.
// Confirm the actual scale in the Task 0 spike; this handles either.
func normalizeUtil(v float64) float64 {
    if v > 0 && v <= 1 {
        return v * 100
    }
    return v
}

package accountusage

import (
    "context"
    "errors"
    "testing"
    "time"
)

type fakeFetcher struct {
    snaps []Snapshot
    errs  []error
    i     int
}

func (f *fakeFetcher) Fetch(context.Context) (Snapshot, error) {
    defer func() { f.i++ }()
    if f.i < len(f.errs) && f.errs[f.i] != nil {
        return Snapshot{}, f.errs[f.i]
    }
    if f.i < len(f.snaps) {
        return f.snaps[f.i], nil
    }
    return Snapshot{}, errors.New("exhausted")
}

func TestPollOnceSuccessSetsHealthy(t *testing.T) {
    f := &fakeFetcher{snaps: []Snapshot{{FiveHour: Window{Percent: 33}, Healthy: true}}}
    st := &Store{}
    p := &Poller{Fetcher: f, Store: st, FailureThreshold: 3, Now: time.Now}
    p.pollOnce(context.Background())
    if got := st.Get(); !got.Healthy || got.FiveHour.Percent != 33 {
        t.Fatalf("store after success: %+v", got)
    }
}

func TestConsecutiveFailuresMarkStaleAfterThreshold(t *testing.T) {
    f := &fakeFetcher{
        snaps: []Snapshot{{FiveHour: Window{Percent: 50}, Healthy: true}},
        errs:  []error{nil, errors.New("x"), errors.New("x"), errors.New("x")},
    }
    st := &Store{}
    p := &Poller{Fetcher: f, Store: st, FailureThreshold: 3, Now: time.Now}
    p.pollOnce(context.Background()) // success -> healthy, keep last-known percent
    p.pollOnce(context.Background()) // fail 1 -> still healthy (last-known)
    if !st.Get().Healthy {
        t.Fatal("single failure must not mark stale")
    }
    p.pollOnce(context.Background()) // fail 2
    p.pollOnce(context.Background()) // fail 3 -> stale
    got := st.Get()
    if got.Healthy {
        t.Fatal("threshold failures must mark stale")
    }
    if got.FiveHour.Percent != 50 {
        t.Fatal("stale snapshot must retain last-known windows")
    }
}

package accountusage

import (
    "context"
    "log/slog"
    "time"
)

type fetcher interface {
    Fetch(context.Context) (Snapshot, error)
}

type Poller struct {
    Fetcher          fetcher
    Store            *Store
    Interval         time.Duration
    FailureThreshold int
    Now              func() time.Time
    // onUpdate, when set, is called after each successful poll (ConfigMap mirror +
    // metrics wired in Task A9). Kept as a hook so the ticker stays testable.
    onUpdate func(Snapshot)

    failures int
}

func (p *Poller) NeedLeaderElection() bool { return true }

func (p *Poller) Start(ctx context.Context) error {
    if p.Interval < 180*time.Second {
        p.Interval = 180 * time.Second
    }
    if p.Now == nil {
        p.Now = time.Now
    }
    p.pollOnce(ctx) // immediate first poll
    t := time.NewTicker(p.Interval)
    defer t.Stop()
    for {
        select {
        case <-ctx.Done():
            return nil
        case <-t.C:
            p.pollOnce(ctx)
        }
    }
}

func (p *Poller) pollOnce(ctx context.Context) {
    snap, err := p.Fetcher.Fetch(ctx)
    if err != nil {
        p.failures++
        slog.Warn("accountusage poll failed", "failures", p.failures, "error", err)
        if p.failures >= p.FailureThreshold {
            cur := p.Store.Get()
            cur.Healthy = false // keep last-known windows, mark stale
            p.Store.Set(cur)
        }
        return
    }
    p.failures = 0
    snap.Healthy = true
    snap.UpdatedAt = p.Now()
    p.Store.Set(snap)
    slog.Info("accountusage poll ok", "five_hour_pct", snap.FiveHour.Percent, "weekly_pct", snap.Weekly.Percent)
    if p.onUpdate != nil {
        p.onUpdate(snap)
    }
}

func TestBudgetConfigCopiesSpawnCeilingByKind(t *testing.T) {
    p := &Project{Spec: ProjectSpec{TokenBudget: &TokenBudgetSpec{
        Enabled: true, Mode: "claudeSubscription",
        SpawnCeilingByKind: map[string]int32{"brainstorm": 40, "incident": 98},
    }}}
    cfg := p.BudgetConfig(budget.Config{})
    if cfg.SpawnCeilingByKind["brainstorm"] != 40 || cfg.SpawnCeilingByKind["incident"] != 98 {
        t.Fatalf("ceilings not copied: %+v", cfg.SpawnCeilingByKind)
    }
}

    // SpawnCeilingByKind gates each Task kind independently in claudeSubscription
    // mode: work of kind K is held once account usage reaches the given percent.
    // Keys are Task kinds; kinds absent here fall through to proactive/emergency.
    // +kubebuilder:validation:XValidation:rule="self.all(k, self[k] >= 0 && self[k] <= 100)",message="spawnCeilingByKind values must be 0..100"
    // +kubebuilder:validation:XValidation:rule="self.all(k, k in ['implement','review','selfImprove','triageIssue','brainstorm','issueLifecycle','incident','healthCheck','refine'])",message="spawnCeilingByKind keys must be valid Task kinds"
    // +optional
    SpawnCeilingByKind map[string]int32 `json:"spawnCeilingByKind,omitempty"`
    // PollIntervalSeconds is how often the operator polls Claude account usage
    // (claudeSubscription mode). Floor 180 (enforced operator-side too).
    // +kubebuilder:validation:Minimum=180
    // +optional
    PollIntervalSeconds *int32 `json:"pollIntervalSeconds,omitempty"`
    // MonitorOverage surfaces the pay-as-you-go overage pool on dashboards. It is
    // read-only and never gates spawning.
    // +optional
    MonitorOverage *bool `json:"monitorOverage,omitempty"`

    if len(s.SpawnCeilingByKind) > 0 {
        cfg.SpawnCeilingByKind = make(map[string]int, len(s.SpawnCeilingByKind))
        for k, v := range s.SpawnCeilingByKind {
            cfg.SpawnCeilingByKind[k] = int(v)
        }
    }

func TestDispatcherBlockKindUsesStoreInSubscriptionMode(t *testing.T) {
    future := time.Now().Add(time.Hour)
    store := &accountusage.Store{}
    store.Set(accountusage.Snapshot{FiveHour: accountusage.Window{Percent: 60, Reset: future}, Healthy: true})
    r := &DispatcherReconciler{Usage: store}
    proj := &tatarav1alpha1.Project{Spec: tatarav1alpha1.ProjectSpec{TokenBudget: &tatarav1alpha1.TokenBudgetSpec{
        Enabled: true, Mode: "claudeSubscription",
        SpawnCeilingByKind: map[string]int32{"brainstorm": 40, "incident": 98},
    }}}
    cfg := proj.BudgetConfig(r.BudgetDefaults)
    block := r.blockKindFunc(proj, cfg, time.Now())
    if !block("brainstorm") { // 60 >= 40
        t.Fatal("brainstorm must be held at 60%")
    }
    if block("incident") { // 60 < 98
        t.Fatal("incident must not be held at 60%")
    }
}

    // Usage is the fleet-wide Claude account usage snapshot (claudeSubscription
    // mode). Nil-safe: a nil store yields an empty snapshot (nothing per-kind held).
    Usage *accountusage.Store

// blockKindFunc returns a predicate that reports whether a queued event of a
// given kind must be held on the account usage gate. Subscription mode reads the
// fleet-wide store; other modes never per-kind block.
func (r *DispatcherReconciler) blockKindFunc(proj *tatarav1alpha1.Project, cfg budget.Config, now time.Time) func(string) bool {
    var sub budget.Subscription
    if r.Usage != nil {
        sub = r.Usage.Get().Subscription()
    }
    return func(kind string) bool {
        return budget.KindBlocked(cfg, sub, kind, now)
    }
}

func (r *DispatcherReconciler) admit(ctx context.Context, proj *tatarav1alpha1.Project,
    qes []tatarav1alpha1.QueuedEvent, tasks []tatarav1alpha1.Task, d budget.Decision, blockKind func(string) bool) (requeue bool, err error) {

        if blockKind != nil && blockKind(kindOf(qe)) {
            continue // held on its per-kind account-usage ceiling; leave Queued
        }

    sub := proj.BudgetSubscription()
    if budgetCfg.Mode == budget.ModeClaudeSubscription && r.Usage != nil {
        sub = r.Usage.Get().Subscription()
    }
    decision := budget.Evaluate(budgetCfg, proj.BudgetWindowState(), sub, time.Now())
    ...
    requeue, err := r.admit(ctx, proj, qes, tasks, decision, r.blockKindFunc(proj, budgetCfg, time.Now()))