From 1e18dfc8a3d2bcbf8d3122941c9cb315f36388a4 Mon Sep 17 00:00:00 2001 From: Evanfeenstra Date: Mon, 14 Sep 2026 11:59:20 -0700 Subject: [PATCH] gateway: phase-6 cap walk (PIPELINE 2) behind enforce_budgets MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CheckCaps reads the accumulators PostLLMHook writes and rejects with 402 when a verified macaroon is over any spend cap: - run_cost_exceeded / run_step_exceeded: every chain layer, leaf first (a parent's exhausted cap stops its children) - realm_budget_exceeded (phase 11) / ua_budget_exceeded: cost:ua - agent_cost_exceeded: agent_budgets bucket; cap_usd 0 blocks outright - budget_check_unavailable on Redis errors (fail closed) Gated by a new enforce_budgets flag (config key or BIFROST_PLUGIN_ENFORCE_BUDGETS), effective only alongside enforce_macaroons. Off (default), over-cap calls log 'auth: budget shadow …' and pass through with claims stamped so accounting keeps running. The accumulator now writes cost:ua when the UA carries a realm cap for this swarm's realm, not only an org-wide max_total_usd — the cap walk reads the same counter for both. --- gateway/internal/auth/accumulator.go | 18 +- gateway/internal/auth/accumulator_test.go | 33 ++ gateway/internal/auth/capwalk.go | 264 +++++++++++++ gateway/internal/auth/capwalk_test.go | 346 ++++++++++++++++++ gateway/internal/auth/config.go | 59 ++- gateway/internal/auth/config_test.go | 44 +++ gateway/internal/auth/doc.go | 37 +- gateway/internal/auth/enforcement.go | 44 ++- gateway/internal/env/env.go | 26 +- gateway/main.go | 4 +- .../phases/phase-6-plugin-enforcement.md | 24 +- 11 files changed, 860 insertions(+), 39 deletions(-) create mode 100644 gateway/internal/auth/capwalk.go create mode 100644 gateway/internal/auth/capwalk_test.go diff --git a/gateway/internal/auth/accumulator.go b/gateway/internal/auth/accumulator.go index 33c7bf90e..671652e31 100644 --- a/gateway/internal/auth/accumulator.go +++ b/gateway/internal/auth/accumulator.go @@ -108,10 +108,20 @@ func accumulate( pipe.Expire(ctx, stepsKey, ttl) } - // UA cumulative envelope — only when the org actually set one. - // No bucket ⇒ no enforcement; per-invocation caps are checked at - // signature time and need no Redis state. - if claims.UABudget != nil && claims.UABudget.MaxTotalUSD > 0 && claims.UANonce != "" { + // UA cumulative envelope — only when the org actually set one, + // either as an org-wide max_total_usd or (phase 11) a cap for + // this swarm's realm; the cap walk reads the same counter for + // both. No bucket ⇒ no enforcement; per-invocation caps are + // checked at signature time and need no Redis state. + var uaCap float64 + if claims.UABudget != nil { + uaCap = claims.UABudget.MaxTotalUSD + } + var realmID string + if reg := getRegistry(); reg != nil { + realmID = reg.RealmID() + } + if claims.UANonce != "" && (uaCap > 0 || realmCapFor(claims, realmID) > 0) { uaKey := redisclient.Key(costUAPrefix + claims.UANonce) uaTTL := runKeyTTL(parseRFC3339(claims.UAExp), now) pipe.HIncrByFloat(ctx, uaKey, "total", costUSD) diff --git a/gateway/internal/auth/accumulator_test.go b/gateway/internal/auth/accumulator_test.go index 35609bd57..b65ef2f6a 100644 --- a/gateway/internal/auth/accumulator_test.go +++ b/gateway/internal/auth/accumulator_test.go @@ -222,3 +222,36 @@ func TestApplyToLLMPost_NilClaimsOrNoRedis_NoOp(t *testing.T) { t.Fatalf("nil claims wrote keys: %v", keys) } } + +func TestAccumulate_UAEnvelope_WrittenForRealmCap(t *testing.T) { + mr := newMiniRedis(t) + reg := newTestRegistry(t) + if _, err := reg.SetRealmID("w1"); err != nil { + t.Fatal(err) + } + SetTrustRegistry(reg) + t.Cleanup(func() { SetTrustRegistry(nil) }) + + // No org-wide max_total_usd, but a cap for this swarm's realm: + // the cap walk compares cost:ua against it, so it must be written. + claims := chainClaims([]string{"r_1"}, []string{"2026-05-14T11:00:00Z"}) + claims.EffectiveCaveats.Budget = &macaroon.Budget{ + RealmBudgets: map[string]macaroon.RealmBudget{"w1": {MaxTotalUSD: 3}}, + } + if err := accumulate(context.Background(), claims, 0.25, nil, testNow()); err != nil { + t.Fatal(err) + } + if got := mr.HGet("bifrost:cost:ua:"+claims.UANonce, "total"); got != "0.25" { + t.Errorf("cost:ua total = %q, want 0.25", got) + } + + // A realm cap for some other swarm doesn't count here. + mr.FlushAll() + claims.EffectiveCaveats.Budget.RealmBudgets = map[string]macaroon.RealmBudget{"w2": {MaxTotalUSD: 3}} + if err := accumulate(context.Background(), claims, 0.25, nil, testNow()); err != nil { + t.Fatal(err) + } + if mr.Exists("bifrost:cost:ua:" + claims.UANonce) { + t.Fatal("cost:ua written for a realm cap that isn't this swarm's") + } +} diff --git a/gateway/internal/auth/capwalk.go b/gateway/internal/auth/capwalk.go new file mode 100644 index 000000000..724b14fa6 --- /dev/null +++ b/gateway/internal/auth/capwalk.go @@ -0,0 +1,264 @@ +package auth + +import ( + "context" + "errors" + "fmt" + "time" + + "github.com/redis/go-redis/v9" + + macaroon "github.com/stakwork/stakgraph/gateway/auth/go" + "github.com/stakwork/stakgraph/gateway/internal/duration" + "github.com/stakwork/stakgraph/gateway/internal/redisclient" +) + +// The cap walk — phase-6 "Hot path" PIPELINE 2. Reads the +// accumulators PostLLMHook writes (accumulator.go) and compares them +// against every cap the verified macaroon and the plugin config +// declare. One pipelined Redis round-trip, then a fixed-order +// evaluation so the rejection reason is deterministic: +// +// 1. cost:run: >= layer.MaxCostUSD for every chain layer → run_cost_exceeded +// 2. cost:ua: >= realm cap (phase 11) → realm_budget_exceeded +// 3. cost:ua: >= ua.MaxTotalUSD → ua_budget_exceeded +// 4. steps:run: >= layer.MaxSteps for every chain layer → run_step_exceeded +// 5. cost:agent:: >= agent_budgets cap → agent_cost_exceeded +// +// The walk covers every layer, leaf first, then each ancestor: a +// child whose own $2 cap is untouched is still rejected when its +// parent's $5 is gone, because the parent's counter includes every +// descendant's spend. Steps are walked the same way (the plan only +// names the leaf; walking ancestors is strictly stronger and matches +// what the accumulator increments). +// +// Comparisons use *past* spend — the call's own price isn't known +// before it runs — so a run can overshoot by up to one call's cost +// (times concurrency). Phase 6 "Pipelining and atomicity" accepts +// that. `>=` rather than `>` so a run that lands exactly on its cap +// is done. +// +// A cap of 0 means "no cap" for macaroon caveats (the field is +// optional at every layer). For agent_budgets it means the opposite: +// an entry with cap_usd 0 blocks the agent outright — that is how +// phase 6 says to stop an agent permanently (a kill:agent key only +// lasts 24h). +// +// Fail-closed on Redis errors (402 budget_check_unavailable); no-op +// in observability mode (Redis unconfigured), like CheckRevocations. + +// capLayer is one budgeted run in the walk order. +type capLayer struct { + runID string + maxCostUSD float64 + maxSteps int +} + +// CheckCaps runs the cap walk. realmID is the swarm's own realm (from +// the trust registry; "" for single-swarm deployments), used to pick +// the per-realm cap out of the macaroon's realm_budgets. Returns nil +// on "within every cap". +func CheckCaps(ctx context.Context, claims *macaroon.Claims, realmID string, now time.Time) *AdapterError { + rdb := redisclient.Client() + if rdb == nil { + return nil + } + if claims == nil { + return &AdapterError{ + Code: "verify_internal_error", + HTTPStatus: 401, + Message: "nil claims passed to cap walk", + } + } + + layers := capLayers(claims) + realmCap := realmCapFor(claims, realmID) + var uaCap float64 + if claims.UABudget != nil { + uaCap = claims.UABudget.MaxTotalUSD + } + agentBudget, agentConfigured := GetConfig().AgentBudgets[claims.AgentName] + + pctx, cancel := context.WithTimeout(ctx, pipelineTimeout) + defer cancel() + pipe := rdb.Pipeline() + + costCmds := make([]*redis.StringCmd, len(layers)) + stepCmds := make([]*redis.StringCmd, len(layers)) + for i, l := range layers { + if l.maxCostUSD > 0 { + costCmds[i] = pipe.HGet(pctx, redisclient.Key(costRunPrefix+l.runID), "total") + } + if l.maxSteps > 0 { + stepCmds[i] = pipe.HGet(pctx, redisclient.Key(stepsRunPrefix+l.runID), "total") + } + } + var uaCmd *redis.StringCmd + if claims.UANonce != "" && (uaCap > 0 || realmCap > 0) { + uaCmd = pipe.HGet(pctx, redisclient.Key(costUAPrefix+claims.UANonce), "total") + } + var agentCmd *redis.StringCmd + var agentBucket string + if agentConfigured && agentBudget.CapUSD > 0 && agentBudget.Window != "" { + if w, err := duration.Parse(agentBudget.Window); err == nil { + agentBucket = w.BucketKey(now) + agentCmd = pipe.HGet(pctx, redisclient.Key(costAgentPrefix+claims.AgentName+":"+agentBucket), "total") + } + // An unparseable window is logged by the accumulator on + // every write; nothing to compare here, so no cap applies. + } + + if _, err := pipe.Exec(pctx); err != nil && !errors.Is(err, redis.Nil) { + return budgetUnavailable(fmt.Sprintf("redis pipeline: %v", err)) + } + + // 1. Per-run cost, leaf first. + for i, l := range layers { + if costCmds[i] == nil { + continue + } + spent, err := hashFloat(costCmds[i]) + if err != nil { + return budgetUnavailable(fmt.Sprintf("cost:run:%s: %v", l.runID, err)) + } + if spent >= l.maxCostUSD { + return &AdapterError{ + Code: "run_cost_exceeded", + HTTPStatus: 402, + Message: fmt.Sprintf("run %s spent $%.4f of its $%.2f cap", l.runID, spent, l.maxCostUSD), + } + } + } + + // 2-3. UA envelope: realm cap (phase 11) then the org-wide total. + if uaCmd != nil { + spent, err := hashFloat(uaCmd) + if err != nil { + return budgetUnavailable(fmt.Sprintf("cost:ua: %v", err)) + } + if realmCap > 0 && spent >= realmCap { + return &AdapterError{ + Code: "realm_budget_exceeded", + HTTPStatus: 402, + Message: fmt.Sprintf("user %s spent $%.4f of the $%.2f cap for realm %s", + claims.UserID, spent, realmCap, realmID), + } + } + if uaCap > 0 && spent >= uaCap { + return &AdapterError{ + Code: "ua_budget_exceeded", + HTTPStatus: 402, + Message: fmt.Sprintf("user %s spent $%.4f of the $%.2f authorization envelope", + claims.UserID, spent, uaCap), + } + } + } + + // 4. Per-run steps, same walk. + for i, l := range layers { + if stepCmds[i] == nil { + continue + } + steps, err := hashInt(stepCmds[i]) + if err != nil { + return budgetUnavailable(fmt.Sprintf("steps:run:%s: %v", l.runID, err)) + } + if steps >= int64(l.maxSteps) { + return &AdapterError{ + Code: "run_step_exceeded", + HTTPStatus: 402, + Message: fmt.Sprintf("run %s used %d of its %d steps", l.runID, steps, l.maxSteps), + } + } + } + + // 5. Per-agent windowed budget. + if agentConfigured && agentBudget.CapUSD <= 0 { + return &AdapterError{ + Code: "agent_cost_exceeded", + HTTPStatus: 402, + Message: fmt.Sprintf("agent %s has a $0 budget (blocked by operator config)", claims.AgentName), + } + } + if agentCmd != nil { + spent, err := hashFloat(agentCmd) + if err != nil { + return budgetUnavailable(fmt.Sprintf("cost:agent:%s: %v", claims.AgentName, err)) + } + if spent >= agentBudget.CapUSD { + return &AdapterError{ + Code: "agent_cost_exceeded", + HTTPStatus: 402, + Message: fmt.Sprintf("agent %s spent $%.4f of its $%.2f/%s cap (bucket %s)", + claims.AgentName, spent, agentBudget.CapUSD, agentBudget.Window, agentBucket), + } + } + } + + return nil +} + +// capLayers returns the chain's budgeted runs in walk order: leaf +// first, then each ancestor outward. Claims.Chain is outermost-first; +// when it's empty (older callers, tests) the leaf is synthesized from +// Claims.RunID + EffectiveCaveats. +func capLayers(claims *macaroon.Claims) []capLayer { + if len(claims.Chain) == 0 { + if claims.RunID == "" { + return nil + } + return []capLayer{{ + runID: claims.RunID, + maxCostUSD: claims.EffectiveCaveats.MaxCostUSD, + maxSteps: claims.EffectiveCaveats.MaxSteps, + }} + } + out := make([]capLayer, 0, len(claims.Chain)) + seen := make(map[string]bool, len(claims.Chain)) + for i := len(claims.Chain) - 1; i >= 0; i-- { + l := claims.Chain[i] + if l.RunID == "" || seen[l.RunID] { + continue + } + seen[l.RunID] = true + out = append(out, capLayer{runID: l.RunID, maxCostUSD: l.MaxCostUSD, maxSteps: l.MaxSteps}) + } + return out +} + +// realmCapFor returns the narrowed per-realm cap for this swarm's +// realm, or 0 when the macaroon carries no realm_budgets, the swarm +// has no realm_id, or the realm isn't listed (membership is checked +// separately by CheckRealmMembership; here absent just means no cap). +func realmCapFor(claims *macaroon.Claims, realmID string) float64 { + if realmID == "" || claims.EffectiveCaveats.Budget == nil { + return 0 + } + return claims.EffectiveCaveats.Budget.RealmBudgets[realmID].MaxTotalUSD +} + +func budgetUnavailable(detail string) *AdapterError { + return &AdapterError{ + Code: "budget_check_unavailable", + HTTPStatus: 402, + Message: detail, + } +} + +// hashFloat / hashInt read an HGET result, treating a missing key or +// field (redis.Nil) as zero — a run that has never spent is at $0. +func hashFloat(cmd *redis.StringCmd) (float64, error) { + v, err := cmd.Float64() + if errors.Is(err, redis.Nil) { + return 0, nil + } + return v, err +} + +func hashInt(cmd *redis.StringCmd) (int64, error) { + v, err := cmd.Int64() + if errors.Is(err, redis.Nil) { + return 0, nil + } + return v, err +} diff --git a/gateway/internal/auth/capwalk_test.go b/gateway/internal/auth/capwalk_test.go new file mode 100644 index 000000000..523867eff --- /dev/null +++ b/gateway/internal/auth/capwalk_test.go @@ -0,0 +1,346 @@ +package auth + +import ( + "context" + "testing" + "time" + + macaroon "github.com/stakwork/stakgraph/gateway/auth/go" + "github.com/stakwork/stakgraph/gateway/internal/pluginctx" +) + +// capClaims: a two-layer chain (parent $5/100 steps → child $2/40 +// steps) under a UA with a $10 envelope, agent "coder". +func capClaims() *macaroon.Claims { + return &macaroon.Claims{ + OrgID: testOrgID, + UserID: testUserID, + AgentName: "coder", + RunID: "r_child", + UANonce: "aaaa000000000000000000000000aaaa", + UABudget: &macaroon.Budget{MaxTotalUSD: 10}, + Nonces: []string{"aaaa000000000000000000000000aaaa", "bbbb000000000000000000000000bbbb"}, + IAT: time.Now().UTC().Format(time.RFC3339), + EffectiveCaveats: macaroon.EffectiveCaveats{ + MaxCostUSD: 2, + MaxSteps: 40, + }, + Chain: []macaroon.ChainLayer{ + {RunID: "r_parent", MaxCostUSD: 5, MaxSteps: 100}, + {RunID: "r_child", MaxCostUSD: 2, MaxSteps: 40}, + }, + } +} + +func check(t *testing.T, claims *macaroon.Claims, realmID string) *AdapterError { + t.Helper() + return CheckCaps(context.Background(), claims, realmID, time.Now().UTC()) +} + +func wantCode(t *testing.T, err *AdapterError, code string) { + t.Helper() + if err == nil { + t.Fatalf("want %s, got nil", code) + } + if err.Code != code { + t.Fatalf("want %s, got %s (%s)", code, err.Code, err.Message) + } + if err.HTTPStatus != 402 { + t.Fatalf("want 402, got %d", err.HTTPStatus) + } +} + +func TestCheckCaps_ObservabilityMode_NoOp(t *testing.T) { + if err := check(t, capClaims(), ""); err != nil { + t.Fatalf("no redis ⇒ no-op, got %+v", err) + } +} + +func TestCheckCaps_NothingSpent_Passes(t *testing.T) { + _ = newMiniRedis(t) + if err := check(t, capClaims(), ""); err != nil { + t.Fatalf("fresh run must pass, got %+v", err) + } +} + +func TestCheckCaps_UnderEveryCap_Passes(t *testing.T) { + mr := newMiniRedis(t) + mr.HSet("bifrost:cost:run:r_child", "total", "1.99") + mr.HSet("bifrost:cost:run:r_parent", "total", "4.99") + mr.HSet("bifrost:steps:run:r_child", "total", "39") + mr.HSet("bifrost:cost:ua:aaaa000000000000000000000000aaaa", "total", "9.99") + if err := check(t, capClaims(), ""); err != nil { + t.Fatalf("under cap must pass, got %+v", err) + } +} + +func TestCheckCaps_LeafCostExceeded(t *testing.T) { + mr := newMiniRedis(t) + mr.HSet("bifrost:cost:run:r_child", "total", "2.00") // exactly at cap ⇒ done + wantCode(t, check(t, capClaims(), ""), "run_cost_exceeded") +} + +func TestCheckCaps_AncestorCostExceeded_KillsChild(t *testing.T) { + mr := newMiniRedis(t) + mr.HSet("bifrost:cost:run:r_child", "total", "0.50") // child fine + mr.HSet("bifrost:cost:run:r_parent", "total", "5.10") // parent's tree is over + err := check(t, capClaims(), "") + wantCode(t, err, "run_cost_exceeded") + if want := "run r_parent"; !contains(err.Message, want) { + t.Fatalf("message should name the parent: %q", err.Message) + } +} + +func TestCheckCaps_StepsExceeded(t *testing.T) { + mr := newMiniRedis(t) + mr.HSet("bifrost:steps:run:r_child", "total", "40") + wantCode(t, check(t, capClaims(), ""), "run_step_exceeded") +} + +func TestCheckCaps_AncestorStepsExceeded(t *testing.T) { + mr := newMiniRedis(t) + mr.HSet("bifrost:steps:run:r_parent", "total", "100") + wantCode(t, check(t, capClaims(), ""), "run_step_exceeded") +} + +func TestCheckCaps_UAEnvelopeExceeded(t *testing.T) { + mr := newMiniRedis(t) + mr.HSet("bifrost:cost:ua:aaaa000000000000000000000000aaaa", "total", "10") + wantCode(t, check(t, capClaims(), ""), "ua_budget_exceeded") +} + +func TestCheckCaps_NoUABudget_NoUARead(t *testing.T) { + mr := newMiniRedis(t) + mr.HSet("bifrost:cost:ua:aaaa000000000000000000000000aaaa", "total", "999") + c := capClaims() + c.UABudget = nil + if err := check(t, c, ""); err != nil { + t.Fatalf("no UA budget ⇒ no UA cap, got %+v", err) + } +} + +func TestCheckCaps_RealmBudget(t *testing.T) { + mr := newMiniRedis(t) + mr.HSet("bifrost:cost:ua:aaaa000000000000000000000000aaaa", "total", "3") + c := capClaims() + c.EffectiveCaveats.Budget = &macaroon.Budget{ + RealmBudgets: map[string]macaroon.RealmBudget{ + "w1": {MaxTotalUSD: 3}, + "w2": {MaxTotalUSD: 50}, + }, + } + // This swarm is w1: $3 of $3 ⇒ realm cap first (before the $10 UA cap). + wantCode(t, check(t, c, "w1"), "realm_budget_exceeded") + // Swarm w2 has a $50 realm cap; $3 is fine. + if err := check(t, c, "w2"); err != nil { + t.Fatalf("w2 under realm cap, got %+v", err) + } + // Single-swarm deployment: no realm cap applies. + if err := check(t, c, ""); err != nil { + t.Fatalf("no realm_id ⇒ no realm cap, got %+v", err) + } +} + +func TestCheckCaps_RealmCapOnly_ReadsUACounter(t *testing.T) { + mr := newMiniRedis(t) + mr.HSet("bifrost:cost:ua:aaaa000000000000000000000000aaaa", "total", "3") + c := capClaims() + c.UABudget = nil // no org-wide total; only a realm cap + c.EffectiveCaveats.Budget = &macaroon.Budget{ + RealmBudgets: map[string]macaroon.RealmBudget{"w1": {MaxTotalUSD: 3}}, + } + wantCode(t, check(t, c, "w1"), "realm_budget_exceeded") +} + +func TestCheckCaps_AgentBudget(t *testing.T) { + mr := newMiniRedis(t) + SetConfigForTest(Config{AgentBudgets: map[string]AgentBudget{ + "coder": {CapUSD: 1, Window: "1d"}, + }}) + t.Cleanup(func() { SetConfigForTest(Config{}) }) + today := time.Now().UTC().Format("2006-01-02") + + mr.HSet("bifrost:cost:agent:coder:"+today, "total", "0.99") + if err := check(t, capClaims(), ""); err != nil { + t.Fatalf("under agent cap, got %+v", err) + } + mr.HSet("bifrost:cost:agent:coder:"+today, "total", "1.00") + wantCode(t, check(t, capClaims(), ""), "agent_cost_exceeded") +} + +func TestCheckCaps_AgentBudgetZero_Blocks(t *testing.T) { + _ = newMiniRedis(t) + SetConfigForTest(Config{AgentBudgets: map[string]AgentBudget{ + "coder": {CapUSD: 0, Window: "1d"}, + }}) + t.Cleanup(func() { SetConfigForTest(Config{}) }) + wantCode(t, check(t, capClaims(), ""), "agent_cost_exceeded") + + // An unrelated agent is unaffected. + c := capClaims() + c.AgentName = "web-search" + if err := check(t, c, ""); err != nil { + t.Fatalf("other agent must pass, got %+v", err) + } +} + +func TestCheckCaps_ZeroCaveatCaps_MeanNoCap(t *testing.T) { + mr := newMiniRedis(t) + mr.HSet("bifrost:cost:run:r_solo", "total", "1000") + mr.HSet("bifrost:steps:run:r_solo", "total", "1000") + c := &macaroon.Claims{UserID: testUserID, RunID: "r_solo"} // no Chain, no caps + if err := check(t, c, ""); err != nil { + t.Fatalf("zero caps ⇒ uncapped, got %+v", err) + } +} + +func TestCheckCaps_NoChain_FallsBackToEffectiveCaveats(t *testing.T) { + mr := newMiniRedis(t) + mr.HSet("bifrost:cost:run:r_solo", "total", "2") + c := &macaroon.Claims{ + UserID: testUserID, + RunID: "r_solo", + EffectiveCaveats: macaroon.EffectiveCaveats{MaxCostUSD: 2}, + } + wantCode(t, check(t, c, ""), "run_cost_exceeded") +} + +func TestCheckCaps_Order_CostBeforeSteps(t *testing.T) { + mr := newMiniRedis(t) + mr.HSet("bifrost:cost:run:r_child", "total", "9") + mr.HSet("bifrost:steps:run:r_child", "total", "99") + mr.HSet("bifrost:cost:ua:aaaa000000000000000000000000aaaa", "total", "99") + wantCode(t, check(t, capClaims(), ""), "run_cost_exceeded") +} + +func TestCheckCaps_RedisDown_FailsClosed(t *testing.T) { + mr := newMiniRedis(t) + mr.Close() + err := check(t, capClaims(), "") + wantCode(t, err, "budget_check_unavailable") +} + +func TestCheckCaps_KeyNamespace(t *testing.T) { + mr := newMiniRedis(t) + mr.HSet("cost:run:r_child", "total", "999") // missing bifrost: prefix + if err := check(t, capClaims(), ""); err != nil { + t.Fatalf("un-prefixed key must not match, got %+v", err) + } +} + +// ─── through Evaluate / ApplyToLLMPre ──────────────────────────────── + +func overCapMacaroon(t *testing.T, mr interface{ HSet(string, ...string) }) string { + t.Helper() + opts := defaultMacaroonOptions(time.Now()) // $5 cap + mr.HSet("bifrost:cost:run:"+opts.runID, "total", "5") + return buildMacaroon(t, opts) +} + +func TestEvaluate_OverCap_ClaimsAndCapErr(t *testing.T) { + reg := newTestRegistry(t) + SetTrustRegistry(reg) + t.Cleanup(func() { SetTrustRegistry(nil) }) + mr := newMiniRedis(t) + + d := Evaluate(context.Background(), overCapMacaroon(t, mr)) + if d.Err != nil { + t.Fatalf("auth must pass: %+v", d.Err) + } + if d.Claims == nil { + t.Fatal("claims must be set even when over cap") + } + if d.CapErr == nil || d.CapErr.Code != "run_cost_exceeded" { + t.Fatalf("want run_cost_exceeded, got %+v", d.CapErr) + } +} + +func TestApplyToLLMPre_BudgetsEnforced_OverCap402(t *testing.T) { + reg := newTestRegistry(t) + SetTrustRegistry(reg) + t.Cleanup(func() { SetTrustRegistry(nil) }) + SetConfigForTest(Config{EnforceMacaroons: true, EnforceBudgets: true}) + t.Cleanup(func() { SetConfigForTest(Config{}) }) + mr := newMiniRedis(t) + + bctx := newBifrostCtx() + pluginctx.SetRawMacaroon(bctx, overCapMacaroon(t, mr)) + sc := ApplyToLLMPre(bctx) + if sc == nil || sc.Error == nil || sc.Error.StatusCode == nil || *sc.Error.StatusCode != 402 { + t.Fatalf("want 402 short-circuit, got %+v", sc) + } + if *sc.Error.Error.Code != "run_cost_exceeded" { + t.Fatalf("want run_cost_exceeded, got %s", *sc.Error.Error.Code) + } + if pluginctx.VerifiedClaims(bctx) != nil { + t.Fatal("rejected request must not stamp claims") + } +} + +func TestApplyToLLMPre_BudgetShadow_OverCapPassesAndStamps(t *testing.T) { + reg := newTestRegistry(t) + SetTrustRegistry(reg) + t.Cleanup(func() { SetTrustRegistry(nil) }) + SetConfigForTest(Config{EnforceMacaroons: true, EnforceBudgets: false}) + t.Cleanup(func() { SetConfigForTest(Config{}) }) + mr := newMiniRedis(t) + + bctx := newBifrostCtx() + pluginctx.SetRawMacaroon(bctx, overCapMacaroon(t, mr)) + if sc := ApplyToLLMPre(bctx); sc != nil { + t.Fatalf("budget shadow must pass through, got %+v", sc) + } + if pluginctx.VerifiedClaims(bctx) == nil { + t.Fatal("shadow pass-through must stamp claims so the accumulator keeps counting") + } +} + +func TestApplyToLLMPre_BudgetsWithoutMacaroonEnforce_Shadow(t *testing.T) { + reg := newTestRegistry(t) + SetTrustRegistry(reg) + t.Cleanup(func() { SetTrustRegistry(nil) }) + // enforce_budgets alone is inert — budgets would be bypassable + // by dropping the macaroon. + SetConfigForTest(Config{EnforceMacaroons: false, EnforceBudgets: true}) + t.Cleanup(func() { SetConfigForTest(Config{}) }) + mr := newMiniRedis(t) + + bctx := newBifrostCtx() + pluginctx.SetRawMacaroon(bctx, overCapMacaroon(t, mr)) + if sc := ApplyToLLMPre(bctx); sc != nil { + t.Fatalf("enforce_budgets without enforce_macaroons must not reject, got %+v", sc) + } +} + +func TestApplyToLLMPre_BudgetsEnforced_UnderCapPasses(t *testing.T) { + reg := newTestRegistry(t) + SetTrustRegistry(reg) + t.Cleanup(func() { SetTrustRegistry(nil) }) + SetConfigForTest(Config{EnforceMacaroons: true, EnforceBudgets: true}) + t.Cleanup(func() { SetConfigForTest(Config{}) }) + mr := newMiniRedis(t) + + opts := defaultMacaroonOptions(time.Now()) + mr.HSet("bifrost:cost:run:"+opts.runID, "total", "4.99") + bctx := newBifrostCtx() + pluginctx.SetRawMacaroon(bctx, buildMacaroon(t, opts)) + if sc := ApplyToLLMPre(bctx); sc != nil { + t.Fatalf("under cap must pass, got %+v", sc) + } + if pluginctx.VerifiedClaims(bctx) == nil { + t.Fatal("claims not stamped") + } +} + +func contains(s, sub string) bool { + return len(sub) == 0 || (len(s) >= len(sub) && indexOf(s, sub) >= 0) +} + +func indexOf(s, sub string) int { + for i := 0; i+len(sub) <= len(s); i++ { + if s[i:i+len(sub)] == sub { + return i + } + } + return -1 +} diff --git a/gateway/internal/auth/config.go b/gateway/internal/auth/config.go index b44c1bfb6..335e5de98 100644 --- a/gateway/internal/auth/config.go +++ b/gateway/internal/auth/config.go @@ -51,6 +51,22 @@ type Config struct { // persisted — grep the boot line for it. EnforceMacaroonsSource string `json:"-"` + // EnforceBudgets gates the phase-6 cap walk (capwalk.go). When + // false (default) a request over any spend cap is logged as + // "budget shadow: would reject" and continues; when true it + // short-circuits with 402. Separate from EnforceMacaroons on + // purpose: the cap walk is only as good as the accumulators it + // reads, and a swarm that already enforces macaroons must be + // able to watch its budget decisions before they bite. + // + // Only effective together with EnforceMacaroons — see + // BudgetsEnforced. Env override: BIFROST_PLUGIN_ENFORCE_BUDGETS. + EnforceBudgets bool `json:"enforce_budgets"` + + // EnforceBudgetsSource mirrors EnforceMacaroonsSource for the + // budgets flag. + EnforceBudgetsSource string `json:"-"` + // AgentBudgets is the per-agent windowed spend cap declared // in plugin.yaml. Phase 6's PreLLMHook reads these to gate // inference; phase 8's dashboard reads them to render the @@ -98,6 +114,7 @@ type ModelPrice struct { // "key present and false" (source=config) stay distinguishable. type pluginConfigEnvelope struct { EnforceMacaroons *bool `json:"enforce_macaroons"` + EnforceBudgets *bool `json:"enforce_budgets"` AgentBudgets map[string]AgentBudget `json:"agent_budgets"` ModelPricing map[string]ModelPrice `json:"model_pricing"` } @@ -155,8 +172,16 @@ func SetConfigForTest(c Config) { cfgMu.Unlock() } +// BudgetsEnforced reports whether the cap walk should reject. Both +// flags must be on: with macaroons in shadow mode a caller can skip +// the macaroon entirely, so a budget 402 would be trivially +// bypassable and only punish the callers doing the right thing. +func (c Config) BudgetsEnforced() bool { + return c.EnforceMacaroons && c.EnforceBudgets +} + func parseConfig(raw any) (Config, error) { - cfg := Config{EnforceMacaroonsSource: "default"} + cfg := Config{EnforceMacaroonsSource: "default", EnforceBudgetsSource: "default"} if raw != nil { // Round-trip through JSON so we accept whatever map shape // Bifrost decoded the block into. No reflection on field names. @@ -174,18 +199,34 @@ func parseConfig(raw any) (Config, error) { cfg.EnforceMacaroons = *envelope.EnforceMacaroons cfg.EnforceMacaroonsSource = "config" } + if envelope.EnforceBudgets != nil { + cfg.EnforceBudgets = *envelope.EnforceBudgets + cfg.EnforceBudgetsSource = "config" + } } - // Env override wins over the config block. A value we can't + // Env overrides win over the config block. A value we can't // parse keeps the plugin alive on the config value (see Init for // why fatal would be worse) but must be impossible to miss: an // ERROR line here and "source=env-invalid" on the boot line. - if v, set, err := env.EnforceMacaroonsValue(); err != nil { - pluginlog.Errf("auth: %v — ignoring the override; enforce_macaroons=%t from %s stands", - err, cfg.EnforceMacaroons, cfg.EnforceMacaroonsSource) - cfg.EnforceMacaroonsSource = "env-invalid" - } else if set { - cfg.EnforceMacaroons = v - cfg.EnforceMacaroonsSource = "env" + applyEnvOverride("enforce_macaroons", env.EnforceMacaroonsValue, &cfg.EnforceMacaroons, &cfg.EnforceMacaroonsSource) + applyEnvOverride("enforce_budgets", env.EnforceBudgetsValue, &cfg.EnforceBudgets, &cfg.EnforceBudgetsSource) + + if cfg.EnforceBudgets && !cfg.EnforceMacaroons { + pluginlog.Warnf("auth: enforce_budgets=true has no effect while enforce_macaroons=false — cap walk stays in shadow") } return cfg, nil } + +// applyEnvOverride layers one strict-bool env var over a config +// value, recording where the final value came from. +func applyEnvOverride(name string, read func() (bool, bool, error), dst *bool, source *string) { + v, set, err := read() + switch { + case err != nil: + pluginlog.Errf("auth: %v — ignoring the override; %s=%t from %s stands", err, name, *dst, *source) + *source = "env-invalid" + case set: + *dst = v + *source = "env" + } +} diff --git a/gateway/internal/auth/config_test.go b/gateway/internal/auth/config_test.go index 1c9f59139..0141d8eb1 100644 --- a/gateway/internal/auth/config_test.go +++ b/gateway/internal/auth/config_test.go @@ -161,3 +161,47 @@ func TestInit_EnvOverride_GarbageFallsBackToConfig(t *testing.T) { t.Fatalf("nil config + garbage env: got enforce=%v source=%q, want false/env-invalid", got.EnforceMacaroons, got.EnforceMacaroonsSource) } } + +// --- enforce_budgets + BIFROST_PLUGIN_ENFORCE_BUDGETS ------------------- + +func TestInit_EnforceBudgets_ConfigAndEnv(t *testing.T) { + t.Cleanup(func() { SetConfigForTest(Config{}) }) + + if err := Init(map[string]any{"enforce_budgets": true}); err != nil { + t.Fatalf("Init: %v", err) + } + got := GetConfig() + if !got.EnforceBudgets || got.EnforceBudgetsSource != "config" { + t.Fatalf("config: got %v/%q, want true/config", got.EnforceBudgets, got.EnforceBudgetsSource) + } + if got.BudgetsEnforced() { + t.Fatal("enforce_budgets alone must not enforce (macaroons in shadow)") + } + + t.Setenv(env.EnforceBudgets, "off") + if err := Init(map[string]any{"enforce_macaroons": true, "enforce_budgets": true}); err != nil { + t.Fatalf("Init: %v", err) + } + got = GetConfig() + if got.EnforceBudgets || got.EnforceBudgetsSource != "env" { + t.Fatalf("env off over config true: got %v/%q, want false/env", got.EnforceBudgets, got.EnforceBudgetsSource) + } + + t.Setenv(env.EnforceBudgets, "1") + if err := Init(map[string]any{"enforce_macaroons": true}); err != nil { + t.Fatalf("Init: %v", err) + } + got = GetConfig() + if !got.BudgetsEnforced() || got.EnforceBudgetsSource != "env" { + t.Fatalf("env on + macaroons on: got effective=%v source=%q", got.BudgetsEnforced(), got.EnforceBudgetsSource) + } + + t.Setenv(env.EnforceBudgets, "yep") + if err := Init(map[string]any{"enforce_budgets": true}); err != nil { + t.Fatalf("Init must not fail on a bad override: %v", err) + } + got = GetConfig() + if !got.EnforceBudgets || got.EnforceBudgetsSource != "env-invalid" { + t.Fatalf("garbage env: got %v/%q, want true/env-invalid", got.EnforceBudgets, got.EnforceBudgetsSource) + } +} diff --git a/gateway/internal/auth/doc.go b/gateway/internal/auth/doc.go index ebb84da6f..1f0dadcf3 100644 --- a/gateway/internal/auth/doc.go +++ b/gateway/internal/auth/doc.go @@ -8,14 +8,18 @@ // // What's in scope here // -------------------- -// - config.go enforce_macaroons flag (shadow → enforce rollout), -// agent_budgets, model_pricing +// - config.go enforce_macaroons + enforce_budgets flags (shadow → +// enforce rollout, each with an env override), agent_budgets, +// model_pricing // - verifier.go Verify() — header extraction + trust lookup + pure verify // - revocation.go CheckRevocations() — phase-6 PIPELINE 1: // bifrost:revoke:* / revoke_user_before:* (401) and the // kill: (every chain layer) / kill:agent: switches (402) // - kill.go KillRun/KillAgent + Unkill*, GetRunState/GetAgentState — // the admin primitives behind /_plugin/{runs,agents}/:id/{kill,state} +// - capwalk.go CheckCaps() — phase-6 PIPELINE 2: per-run cost/steps +// for every chain layer, UA envelope, realm cap, agent bucket (402s; +// gated by enforce_budgets) // - ttl.go clamp(exp-now+1h, 1h, 7d) shared by revocation + accumulators // - enforcement.go Evaluate() + ApplyToLLMPre() — hook glue // - accumulator.go ApplyToLLMPost() — phase-6 PostLLMHook pipeline: @@ -24,20 +28,19 @@ // - pricing.go PriceCall() — model_pricing table → dollars // - admin.go revoke primitives behind /_plugin/revoke/* // -// What's still out of scope (phase 6 read side) -// --------------------------------------------- -// - Per-run cost/step cap walk in PreLLMHook (PIPELINE 2 — reads -// the accumulators this package now writes) -// - ua_budget / realm_budget / agent budget 402 rejections +// What's still out of scope (phase 6) +// ----------------------------------- // - Tool-loop detection (reads tools:run) -// - hard_ceiling / tool_loop config + the /_plugin/config/* overrides +// - hard_ceiling defense-in-depth + the user_id == customer_id cross-check +// - tool_loop config + the /_plugin/config/* overrides // -// The write side landing first is deliberate: accumulators are -// shadow-safe (they reject nothing), they light up the phase-8 -// budget endpoint that currently falls back to logs.db, and the cap -// walk needs weeks of real accumulated state to validate against -// before it starts rejecting. Kill switches don't depend on that -// state, so they shipped ahead of the cap walk. +// Rollout order was deliberate: accumulators first (shadow-safe, +// they reject nothing), then kill switches (no dependency on +// accumulated state), then the cap walk behind its own +// enforce_budgets flag — it needs real accumulated spend to validate +// against before it starts rejecting, and a swarm that already +// enforces macaroons must be able to watch "budget shadow: would +// reject" lines before flipping it. // // Operational posture // ------------------- @@ -51,8 +54,10 @@ // // With enforce_macaroons=true the failure path becomes 401 (bad, // missing or revoked macaroon) or 402 (valid macaroon, but the run -// or agent was killed) with a stable AdapterError.Code. Operators -// flip the flag per-swarm once +// or agent was killed) with a stable AdapterError.Code. Spend caps +// are a second knob: enforce_budgets=true (only effective alongside +// enforce_macaroons) turns the cap walk's "budget shadow" log lines +// into 402s. Operators flip the flags per-swarm once // the shadow-mode logs show no false positives — either in the // config.json plugin block or, without rebuilding the image, via the // BIFROST_PLUGIN_ENFORCE_MACAROONS env var (which wins when set; an diff --git a/gateway/internal/auth/enforcement.go b/gateway/internal/auth/enforcement.go index 6e6f379ab..269faccad 100644 --- a/gateway/internal/auth/enforcement.go +++ b/gateway/internal/auth/enforcement.go @@ -50,6 +50,14 @@ type Decision struct { // hooks. nil when the pure verifier itself failed (no claims // to enrich with) or when the request had no header at all. PostVerifyClaims *macaroon.Claims + + // CapErr is the cap-walk verdict (capwalk.go), set only when + // Claims is non-nil: auth passed, but the run / user / realm / + // agent is over a spend cap. Kept apart from Err because it is + // gated by its own flag — ApplyToLLMPre rejects on it only when + // Config.BudgetsEnforced(), and otherwise logs "would reject" + // and lets the request through with claims stamped. + CapErr *AdapterError } // Evaluate runs the full verify pipeline (signature + revocation + @@ -101,6 +109,15 @@ func Evaluate(ctx context.Context, rawMacaroon string) Decision { } d.Claims = claims + + // Phase-6 PIPELINE 2: the cap walk. Auth has passed by now, so + // the verdict goes in CapErr rather than Err — the caller + // decides whether budgets are enforced or shadowed. + var realmID string + if registry != nil { + realmID = registry.RealmID() + } + d.CapErr = CheckCaps(ctx, claims, realmID, now) return d } @@ -197,12 +214,18 @@ func ApplyToLLMPre(bctx *schemas.BifrostContext) *schemas.LLMPluginShortCircuit decision := Evaluate(bctx, rawMacaroon) cfg := GetConfig() - logDecision(decision, cfg.EnforceMacaroons) + logDecision(decision, cfg.EnforceMacaroons, cfg.BudgetsEnforced()) switch { + case decision.Claims != nil && decision.CapErr != nil && cfg.BudgetsEnforced(): + // Verified, but over a spend cap and budgets are live. + return shortCircuitFromError(decision.CapErr) + case decision.Claims != nil: - // Happy path. Always stamp claims, in both modes — shadow - // mode wants downstream hooks to see the verified shape. + // Happy path (or budget shadow). Always stamp claims, in + // both modes — shadow mode wants downstream hooks to see + // the verified shape, and the accumulator must keep + // counting a run that is over cap in shadow. StampClaims(bctx, decision.Claims) return nil @@ -258,12 +281,25 @@ func shortCircuitFromError(e *AdapterError) *schemas.LLMPluginShortCircuit { } } -func logDecision(d Decision, enforce bool) { +func logDecision(d Decision, enforce, enforceBudgets bool) { mode := "shadow" if enforce { mode = "enforce" } switch { + case d.Claims != nil && d.CapErr != nil: + // Verified but over cap. In budget shadow this line is the + // whole point of the rollout — it's what operators grep + // for before flipping enforce_budgets. + budgetMode := "shadow" + if enforceBudgets { + budgetMode = "enforce" + } + pluginlog.Warnf( + "auth: budget %s code=%s status=%d org=%s user=%s agent=%s run_id=%s detail=%q", + budgetMode, d.CapErr.Code, d.CapErr.HTTPStatus, + d.Claims.OrgID, d.Claims.UserID, d.Claims.AgentName, d.Claims.RunID, d.CapErr.Message, + ) case d.Claims != nil: pluginlog.Logf( "auth: verify ok mode=%s org=%s user=%s agent=%s run_id=%s permitted_realms=%v", diff --git a/gateway/internal/env/env.go b/gateway/internal/env/env.go index fa9b411c9..0066d4480 100644 --- a/gateway/internal/env/env.go +++ b/gateway/internal/env/env.go @@ -87,6 +87,15 @@ const ( // failing the plugin. EnforceMacaroons = "BIFROST_PLUGIN_ENFORCE_MACAROONS" + // EnforceBudgets overrides the plugin config block's + // `enforce_budgets` flag: whether the phase-6 cap walk (per-run / + // UA / realm / agent spend caps) rejects with 402 or only logs + // "would reject". Same grammar and error handling as + // EnforceMacaroons. Only effective when macaroons are enforced + // too — without that, a caller bypasses budgets by omitting the + // macaroon. + EnforceBudgets = "BIFROST_PLUGIN_ENFORCE_BUDGETS" + // RedisURL is the connection string for the macaroon-enforcement // Redis. In sphinx-swarm this points at the shared redis.sphinx // instance; in docker-compose it points at the sidecar `redis` @@ -255,7 +264,20 @@ func RedisURLValue() (string, bool) { // recognised truthy/falsy sets — the caller decides what to do with // a typo; this package never guesses which way it was meant. func EnforceMacaroonsValue() (value bool, set bool, err error) { - raw := strings.TrimSpace(os.Getenv(EnforceMacaroons)) + return strictBool(EnforceMacaroons) +} + +// EnforceBudgetsValue parses BIFROST_PLUGIN_ENFORCE_BUDGETS with the +// same contract as EnforceMacaroonsValue. +func EnforceBudgetsValue() (value bool, set bool, err error) { + return strictBool(EnforceBudgets) +} + +// strictBool reads a boolean env var that must be spelled one of the +// recognised ways — unlike IsProduction, a typo is surfaced rather +// than silently read as false, because these flip enforcement. +func strictBool(name string) (value bool, set bool, err error) { + raw := strings.TrimSpace(os.Getenv(name)) if raw == "" { return false, false, nil } @@ -265,7 +287,7 @@ func EnforceMacaroonsValue() (value bool, set bool, err error) { case "0", "false", "no", "off": return false, true, nil } - return false, true, fmt.Errorf("%s=%q: want one of 1/true/yes/on or 0/false/no/off", EnforceMacaroons, raw) + return false, true, fmt.Errorf("%s=%q: want one of 1/true/yes/on or 0/false/no/off", name, raw) } // IsProduction reports whether the plugin is running in a diff --git a/gateway/main.go b/gateway/main.go index 505d458e5..510620510 100644 --- a/gateway/main.go +++ b/gateway/main.go @@ -90,7 +90,9 @@ func Init(config any) error { } auth.SetTrustRegistry(reg) authCfg := auth.GetConfig() - pluginlog.Logf("auth: macaroon adapter wired enforce=%t source=%s", authCfg.EnforceMacaroons, authCfg.EnforceMacaroonsSource) + pluginlog.Logf("auth: macaroon adapter wired enforce=%t source=%s enforce_budgets=%t source=%s (effective=%t)", + authCfg.EnforceMacaroons, authCfg.EnforceMacaroonsSource, + authCfg.EnforceBudgets, authCfg.EnforceBudgetsSource, authCfg.BudgetsEnforced()) // Model-price catalog for the phase-6 accumulator: loads the // persisted datasheet, then fetches bifrost's published sheet in diff --git a/gateway/plans/phases/phase-6-plugin-enforcement.md b/gateway/plans/phases/phase-6-plugin-enforcement.md index 6bd33caf3..c1efd940b 100644 --- a/gateway/plans/phases/phase-6-plugin-enforcement.md +++ b/gateway/plans/phases/phase-6-plugin-enforcement.md @@ -40,9 +40,27 @@ > hotstate.go`, `revoke.go`): `/_plugin/runs/:id/{kill,state}`, > `/_plugin/agents/:name/{kill,state}` (cookie-or-bearer, CSRF on > cookie mutations) and `/_plugin/revoke/{nonce,user}/:id` -> (bearer-only). Still open: the PreLLMHook cost/step cap walk -> (PIPELINE 2) and its 402s, tool-loop detection, and the -> `/_plugin/config/*` override layer. +> (bearer-only). +> +> **Status (cap walk landed):** PIPELINE 2 is implemented in +> `gateway/internal/auth/capwalk.go` — `CheckCaps` reads +> `cost:run` / `steps:run` for every chain layer (leaf first, then +> ancestors; steps are walked like cost, one step stronger than the +> leaf-only check below), `cost:ua` when the UA carries a +> `max_total_usd` **or** a `realm_budgets` cap for this swarm's +> realm (phase 11 — the accumulator now writes `cost:ua` in that +> case too), and `cost:agent::` for configured +> agents. Rejection codes are as listed under "Hot path" step 3 +> plus `realm_budget_exceeded`; Redis errors ⇒ 402 +> `budget_check_unavailable`. An `agent_budgets` entry with +> `cap_usd: 0` blocks the agent outright (the permanent form of +> `kill:agent`). The walk is gated by a **separate** flag, +> `enforce_budgets` (config key or `BIFROST_PLUGIN_ENFORCE_BUDGETS`), +> effective only alongside `enforce_macaroons`; off, it logs +> `auth: budget shadow code=… detail=…` and lets the call through +> with claims stamped so accounting continues. Still open: +> tool-loop detection, `hard_ceiling`, the `user_id == customer_id` +> cross-check, and the `/_plugin/config/*` override layer. > > **Status (phase 11 cutover):** Redis bucket keys and hot-path > flow are unchanged from the description below. Phase 11