diff --git a/authbridge/authlib/go.mod b/authbridge/authlib/go.mod index a2ca3e191..2e8c24047 100644 --- a/authbridge/authlib/go.mod +++ b/authbridge/authlib/go.mod @@ -13,6 +13,7 @@ require ( github.com/rossoctl/context-guru v0.0.0-20260720181432-8fc7c7b36563 github.com/spiffe/go-spiffe/v2 v2.8.1 golang.org/x/net v0.57.0 + golang.org/x/sync v0.22.0 golang.org/x/sys v0.47.0 google.golang.org/genproto/googleapis/rpc v0.0.0-20260526163538-3dc84a4a5aaa google.golang.org/grpc v1.82.0 @@ -117,7 +118,6 @@ require ( go.yaml.in/yaml/v3 v3.0.4 // indirect golang.org/x/arch v0.23.0 // indirect golang.org/x/crypto v0.54.0 // indirect - golang.org/x/sync v0.22.0 // indirect golang.org/x/text v0.40.0 // indirect golang.org/x/time v0.15.0 // indirect google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af // indirect diff --git a/authbridge/authlib/listener/forwardproxy/server.go b/authbridge/authlib/listener/forwardproxy/server.go index e3d33cc79..66cc19c74 100644 --- a/authbridge/authlib/listener/forwardproxy/server.go +++ b/authbridge/authlib/listener/forwardproxy/server.go @@ -609,7 +609,7 @@ func (s *Server) handleStreamingResponse(w http.ResponseWriter, r *http.Request, defer func() { // Use a detached context for finalization: the client may have // cancelled the request context after reading the full stream, - // but aggregating plugins (inference-parser, token-budget) still + // but aggregating plugins (inference-parser, session-budget) still // need their last=true dispatch to finalize state. finalCtx := context.WithoutCancel(r.Context()) finalAction := s.OutboundPipeline.RunResponseFrame(finalCtx, pctx, nil, true) diff --git a/authbridge/authlib/plugins/sessionbudget/e2e_test.go b/authbridge/authlib/plugins/sessionbudget/e2e_test.go new file mode 100644 index 000000000..daf2cf342 --- /dev/null +++ b/authbridge/authlib/plugins/sessionbudget/e2e_test.go @@ -0,0 +1,459 @@ +package sessionbudget + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "net/url" + "sync" + "testing" + "time" + + "github.com/rossoctl/cortex/authbridge/authlib/listener/forwardproxy" + "github.com/rossoctl/cortex/authbridge/authlib/pipeline" + "github.com/rossoctl/cortex/authbridge/authlib/session" +) + +func newE2EPlugin(t *testing.T, maxTokens int64, store *memStore) *SessionBudget { + t.Helper() + p := New() + cfg, _ := json.Marshal(config{ + RedisURL: "mem://test", + MaxTokens: maxTokens, + OnExceed: "deny", + RefreshInterval: "30ms", + RedisUnavailable: "fail_open", + }) + if err := p.Configure(cfg); err != nil { + t.Fatalf("Configure: %v", err) + } + p.store = store + go p.refreshLoop(30 * time.Millisecond) + t.Cleanup(func() { close(p.stopCh); <-p.stopped }) + return p +} + +// newE2EPluginPause builds a plugin in on_exceed=pause mode so cold-cache +// hydrate runs on the OnRequest path (see plugin.go). The caller supplies a +// webhook URL — pass a deny-returning stub if you want breaches to reject. +func newE2EPluginPause(t *testing.T, maxTokens int64, store *memStore, webhookURL string) *SessionBudget { + t.Helper() + p := New() + cfg, _ := json.Marshal(config{ + RedisURL: "mem://test", + MaxTokens: maxTokens, + OnExceed: "pause", + PauseWebhook: webhookURL, + PauseTimeout: "2s", + PauseTimeoutAction: "deny", + PauseGracePeriod: "0s", + RefreshInterval: "30ms", + RedisUnavailable: "fail_open", + }) + if err := p.Configure(cfg); err != nil { + t.Fatalf("Configure: %v", err) + } + p.store = store + p.httpClient = &http.Client{Timeout: 0} + go p.refreshLoop(30 * time.Millisecond) + t.Cleanup(func() { close(p.stopCh); <-p.stopped }) + return p +} + +func respond(p *SessionBudget, sessionID string, tokens int) { + p.OnResponseFrame(context.Background(), makePctx(sessionID, tokens), nil, true) +} + +func request(p *SessionBudget, sessionID string) pipeline.Action { + pctx := &pipeline.Context{ + Direction: pipeline.Outbound, + Headers: http.Header{}, + Session: &pipeline.SessionView{ID: sessionID}, + } + return p.OnRequest(context.Background(), pctx) +} + +// TestE2E_HTTPRoundTrip wires session-budget into a real forward proxy. +// Under-budget requests reach the backend; the proxy is functional. +func TestE2E_HTTPRoundTrip(t *testing.T) { + store := newMemStore() + p := newE2EPlugin(t, 1000, store) + + pipe, err := pipeline.New([]pipeline.Plugin{p}) + if err != nil { + t.Fatal(err) + } + sessions := session.New(5*time.Minute, 100, 0) + defer sessions.Close() + + srv, err := forwardproxy.NewServer(pipeline.NewHolder(pipe), sessions, nil) + if err != nil { + t.Fatal(err) + } + proxy := httptest.NewServer(srv.Handler()) + defer proxy.Close() + + backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("ok")) + })) + defer backend.Close() + + proxyURL, _ := url.Parse(proxy.URL) + client := &http.Client{Transport: &http.Transport{Proxy: http.ProxyURL(proxyURL)}} + + req, _ := http.NewRequest(http.MethodGet, backend.URL+"/v1/chat/completions", nil) + resp, err := client.Do(req) + if err != nil { + t.Fatalf("request through proxy: %v", err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + body, _ := io.ReadAll(resp.Body) + t.Fatalf("expected 200, got %d: %s", resp.StatusCode, body) + } +} + +// TestE2E_AccumulateAndDeny verifies the full lifecycle: accumulate +// tokens via OnResponseFrame, then OnRequest denies with a 403. +func TestE2E_AccumulateAndDeny(t *testing.T) { + p := newE2EPlugin(t, 150, newMemStore()) + + for i := 0; i < 3; i++ { + respond(p, "sess", 60) + } + + action := request(p, "sess") + if action.Type != pipeline.Reject { + t.Fatalf("expected Reject, got %v", action.Type) + } + status, _, body := action.Violation.Render() + if status != http.StatusForbidden { + t.Errorf("status = %d, want 403", status) + } + var parsed map[string]any + if err := json.Unmarshal(body, &parsed); err != nil { + t.Fatal(err) + } + if parsed["error"] != "budget.exceeded" { + t.Errorf("error = %v, want budget.exceeded", parsed["error"]) + } +} + +// TestE2E_MultiSession verifies independent session budgets. +func TestE2E_MultiSession(t *testing.T) { + p := newE2EPlugin(t, 100, newMemStore()) + + for i := 0; i < 3; i++ { + respond(p, "A", 40) // 120 > 100 + } + respond(p, "B", 20) // 20 < 100 + + if a := request(p, "A"); a.Type != pipeline.Reject { + t.Fatalf("session A: expected Reject, got %v", a.Type) + } + if a := request(p, "B"); a.Type != pipeline.Continue { + t.Fatalf("session B: expected Continue, got %v", a.Type) + } +} + +// TestE2E_LocalCacheEnforcesDuringOutage confirms that a populated +// cache enforces even when the backing store is unreachable. +func TestE2E_LocalCacheEnforcesDuringOutage(t *testing.T) { + // Build without starting refreshLoop so we can swap store safely. + p := New() + cfg, _ := json.Marshal(config{ + RedisURL: "mem://test", + MaxTokens: 100, + OnExceed: "deny", + RefreshInterval: "30ms", + RedisUnavailable: "fail_open", + }) + if err := p.Configure(cfg); err != nil { + t.Fatalf("Configure: %v", err) + } + p.store = &failingStore{} + go p.refreshLoop(30 * time.Millisecond) + t.Cleanup(func() { close(p.stopCh); <-p.stopped }) + + p.mu.Lock() + p.cache["s"] = &counters{tokens: 110, calls: 5, startedAt: time.Now()} + p.mu.Unlock() + + if a := request(p, "s"); a.Type != pipeline.Reject { + t.Fatalf("expected Reject from cache with store down, got %v", a.Type) + } +} + +// TestE2E_RefreshRecovery confirms that refreshCache picks up +// authoritative store values after an outage resolves. +func TestE2E_RefreshRecovery(t *testing.T) { + inner := newMemStore() + cs := &controllableStore{inner: inner} + // Build without a background refreshLoop so refreshCache is invoked + // deterministically and store swaps are race-free. + p := New() + cfg, _ := json.Marshal(config{ + RedisURL: "mem://test", + MaxTokens: 200, + OnExceed: "deny", + RefreshInterval: "30ms", + RedisUnavailable: "fail_open", + }) + if err := p.Configure(cfg); err != nil { + t.Fatalf("Configure: %v", err) + } + p.store = cs + + ctx := context.Background() + inner.HashIncr(ctx, "session-budget:s", "tokens", 180) + inner.HashIncr(ctx, "session-budget:s", "calls", 7) + inner.HashSetNX(ctx, "session-budget:s", "started_at", "1700000000") + + p.mu.Lock() + p.cache["s"] = &counters{tokens: 50} + p.mu.Unlock() + + cs.setFailing(true) + p.refreshCache() + p.mu.RLock() + gotDuring := p.cache["s"].tokens + p.mu.RUnlock() + if gotDuring != 50 { + t.Fatalf("during outage: tokens = %d, want 50", gotDuring) + } + + cs.setFailing(false) + p.refreshCache() + p.mu.RLock() + gotAfter := p.cache["s"].tokens + p.mu.RUnlock() + if gotAfter != 180 { + t.Errorf("after recovery: tokens = %d, want 180", gotAfter) + } +} + +// TestE2E_PodRestart verifies that a fresh plugin with an empty cache +// hydrates from Redis on the first request in pause mode — no cold-cache +// overshoot for sessions already over-budget on Redis. Only pause mode +// hydrates on the request path (see plugin.go OnRequest); deny and observe +// intentionally skip on cold cache and let OnResponseFrame + the refresh +// loop populate counters, at the cost of a one-request-per-pod overshoot +// for pre-existing sessions. +func TestE2E_PodRestart(t *testing.T) { + store := newMemStore() + webhook := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, _ = w.Write([]byte(`{"action":"deny"}`)) + })) + defer webhook.Close() + p := newE2EPluginPause(t, 200, store, webhook.URL) + + ctx := context.Background() + // Pre-seed Redis above the limit (210 > 200). + store.HashIncr(ctx, "session-budget:s", "tokens", 210) + store.HashIncr(ctx, "session-budget:s", "calls", 8) + store.HashSetNX(ctx, "session-budget:s", "started_at", "1700000000") + + // Cold cache — first request hydrates from Redis, fires the webhook, + // gets a deny, and rejects. + if a := request(p, "s"); a.Type != pipeline.Reject { + t.Fatalf("cold cache with over-budget Redis: expected Reject, got %v", a.Type) + } +} + +// TestE2E_HydrateSingleflight verifies concurrent cold-cache requests +// for the same session share one Redis lookup instead of stampeding. +// Uses pause mode because that's the only mode where OnRequest hydrates +// (deny/observe intentionally skip cold-cache to keep Redis off the hot path). +func TestE2E_HydrateSingleflight(t *testing.T) { + inner := newMemStore() + cs := &controllableStore{inner: inner, hashGetDelay: 50 * time.Millisecond} + webhook := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, _ = w.Write([]byte(`{"action":"deny"}`)) + })) + defer webhook.Close() + // Build without starting refreshLoop so background HashGet calls don't + // pollute the singleflight counter we're asserting on. + p := New() + cfg, _ := json.Marshal(config{ + RedisURL: "mem://test", + MaxTokens: 200, + OnExceed: "pause", + PauseWebhook: webhook.URL, + PauseTimeout: "2s", + PauseTimeoutAction: "deny", + PauseGracePeriod: "0s", + RefreshInterval: "30ms", + RedisUnavailable: "fail_open", + }) + if err := p.Configure(cfg); err != nil { + t.Fatalf("Configure: %v", err) + } + p.store = cs + p.httpClient = &http.Client{Timeout: 0} + + ctx := context.Background() + inner.HashIncr(ctx, "session-budget:s", "tokens", 210) + inner.HashIncr(ctx, "session-budget:s", "calls", 8) + inner.HashSetNX(ctx, "session-budget:s", "started_at", "1700000000") + + // This test asserts singleflight dedup on hydrate — it counts HashGet + // calls, not per-request outcomes. Pause-mode concurrent breaches + // piggyback on pendingApproval (one leader gets Reject, others continue + // via pause_pending_approval), so we don't assert per-request Reject. + const N = 20 + var wg sync.WaitGroup + wg.Add(N) + for i := 0; i < N; i++ { + go func() { + defer wg.Done() + _ = request(p, "s") + }() + } + wg.Wait() + + // Without singleflight all N calls would hit HashGet. With it, only the + // first flight does — later arrivals see the populated cache and skip hydrate. + // Allow a small slack for goroutines that raced past the cache-check before + // the first flight populated it. + got := cs.hashGetCalls() + if got > 3 { + t.Errorf("HashGet called %d times for %d concurrent cold-cache requests; expected ≤3 (singleflight dedup)", got, N) + } +} + +// controllableStore delegates to inner memStore but can be toggled to fail, +// counts HashGet calls, and optionally injects latency into HashGet. +type controllableStore struct { + inner *memStore + failing bool + hashGetDelay time.Duration + mu sync.Mutex + hashGets int +} + +func (c *controllableStore) setFailing(v bool) { c.mu.Lock(); c.failing = v; c.mu.Unlock() } +func (c *controllableStore) isFailing() bool { c.mu.Lock(); defer c.mu.Unlock(); return c.failing } +func (c *controllableStore) hashGetCalls() int { c.mu.Lock(); defer c.mu.Unlock(); return c.hashGets } +func (c *controllableStore) err() error { return context.DeadlineExceeded } + +func (c *controllableStore) Get(ctx context.Context, key string) (string, error) { + if c.isFailing() { + return "", c.err() + } + return c.inner.Get(ctx, key) +} +func (c *controllableStore) Set(ctx context.Context, key, value string, ttl time.Duration) error { + if c.isFailing() { + return c.err() + } + return c.inner.Set(ctx, key, value, ttl) +} +func (c *controllableStore) Incr(ctx context.Context, key string, delta int64) (int64, error) { + if c.isFailing() { + return 0, c.err() + } + return c.inner.Incr(ctx, key, delta) +} +func (c *controllableStore) HashIncr(ctx context.Context, key, field string, delta int64) (int64, error) { + if c.isFailing() { + return 0, c.err() + } + return c.inner.HashIncr(ctx, key, field, delta) +} +func (c *controllableStore) HashGet(ctx context.Context, key string) (map[string]string, error) { + c.mu.Lock() + c.hashGets++ + delay := c.hashGetDelay + failing := c.failing + c.mu.Unlock() + if delay > 0 { + time.Sleep(delay) + } + if failing { + return nil, c.err() + } + return c.inner.HashGet(ctx, key) +} +func (c *controllableStore) HashSetNX(ctx context.Context, key, field, value string) (bool, error) { + if c.isFailing() { + return false, c.err() + } + return c.inner.HashSetNX(ctx, key, field, value) +} +func (c *controllableStore) Expire(ctx context.Context, key string, ttl time.Duration) error { + if c.isFailing() { + return c.err() + } + return c.inner.Expire(ctx, key, ttl) +} +func (c *controllableStore) Close() error { return nil } + +// TestE2E_PauseMode covers the full lifecycle for both webhook outcomes. +// The 'approve' row also verifies the request body carries session_id; +// the 'deny' row also verifies the 403 response schema. +func TestE2E_PauseMode(t *testing.T) { + tests := []struct { + name string + response string + want pipeline.ActionType + }{ + {"approve", `{"action":"approve"}`, pipeline.Continue}, + {"deny", `{"action":"deny"}`, pipeline.Reject}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var req map[string]any + json.NewDecoder(r.Body).Decode(&req) + if req["session_id"] != "sess" { + t.Errorf("webhook got session_id=%v, want sess", req["session_id"]) + } + w.WriteHeader(http.StatusOK) + w.Write([]byte(tt.response)) + })) + defer srv.Close() + + p := New() + cfg, _ := json.Marshal(config{ + RedisURL: "mem://test", + MaxCalls: 3, + OnExceed: "pause", + PauseWebhook: srv.URL, + PauseTimeout: "5s", + PauseTimeoutAction: "deny", + RefreshInterval: "30ms", + RedisUnavailable: "fail_open", + }) + if err := p.Configure(cfg); err != nil { + t.Fatalf("Configure: %v", err) + } + p.store = newMemStore() + p.httpClient = &http.Client{} + go p.refreshLoop(30 * time.Millisecond) + t.Cleanup(func() { close(p.stopCh); <-p.stopped }) + + p.mu.Lock() + p.cache["sess"] = &counters{tokens: 100, calls: 3, startedAt: time.Now()} + p.mu.Unlock() + + action := request(p, "sess") + if action.Type != tt.want { + t.Fatalf("action = %v, want %v", action.Type, tt.want) + } + if tt.want == pipeline.Reject { + status, _, body := action.Violation.Render() + if status != http.StatusForbidden { + t.Errorf("status = %d, want 403", status) + } + var parsed map[string]any + json.Unmarshal(body, &parsed) + if parsed["error"] != "budget.exceeded" { + t.Errorf("error = %v, want budget.exceeded", parsed["error"]) + } + } + }) + } +} diff --git a/authbridge/authlib/plugins/tokenbudget/lifecycle_test.go b/authbridge/authlib/plugins/sessionbudget/lifecycle_test.go similarity index 88% rename from authbridge/authlib/plugins/tokenbudget/lifecycle_test.go rename to authbridge/authlib/plugins/sessionbudget/lifecycle_test.go index ee0a129e4..2543474d4 100644 --- a/authbridge/authlib/plugins/tokenbudget/lifecycle_test.go +++ b/authbridge/authlib/plugins/sessionbudget/lifecycle_test.go @@ -1,4 +1,4 @@ -package tokenbudget +package sessionbudget import ( "context" @@ -20,9 +20,9 @@ func TestFullLifecycle_RefreshFromRedis(t *testing.T) { p.store = store ctx := context.Background() - store.HashIncr(ctx, "token-budget:remote-sess", "tokens", 450) - store.HashIncr(ctx, "token-budget:remote-sess", "calls", 10) - store.HashSetNX(ctx, "token-budget:remote-sess", "started_at", "1700000000") + store.HashIncr(ctx, "session-budget:remote-sess", "tokens", 450) + store.HashIncr(ctx, "session-budget:remote-sess", "calls", 10) + store.HashSetNX(ctx, "session-budget:remote-sess", "started_at", "1700000000") // Seed cache so refreshCache picks it up. p.mu.Lock() diff --git a/authbridge/authlib/plugins/sessionbudget/plugin.go b/authbridge/authlib/plugins/sessionbudget/plugin.go new file mode 100644 index 000000000..e9e9a3d7c --- /dev/null +++ b/authbridge/authlib/plugins/sessionbudget/plugin.go @@ -0,0 +1,638 @@ +// Package sessionbudget enforces per-session lifetime budgets on tokens, +// inference calls, and wall-clock duration. Must run before inference-parser +// in the declared plugin order (response path is reverse: inference-parser +// finalizes counts first, then this plugin reads them). +package sessionbudget + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "log/slog" + "net/http" + "strconv" + "sync" + "time" + + "github.com/rossoctl/cortex/authbridge/authlib/pipeline" + "github.com/rossoctl/cortex/authbridge/authlib/plugins" + "github.com/rossoctl/cortex/authbridge/authlib/storage" + "golang.org/x/sync/singleflight" +) + +type config struct { + RedisURL string `json:"redis_url" required:"true" description:"Redis/Valkey connection URL."` + MaxTokens int64 `json:"max_tokens" description:"Cumulative token ceiling per session. 0 = no limit."` + MaxCalls int64 `json:"max_calls" description:"Max LLM/inference calls per session (counted from inference-parser output; MCP tool calls and other outbound traffic do not count). 0 = no limit."` + MaxDurationSeconds int64 `json:"max_duration_seconds" description:"Wall-clock session lifetime in seconds. 0 = no limit."` + OnExceed string `json:"on_exceed" description:"Action on breach: deny, observe (shadow), or pause (HITL webhook approval)." default:"deny" enum:"deny,observe,pause"` + PauseWebhook string `json:"pause_webhook" description:"URL to POST for approval when on_exceed=pause. Required when on_exceed=pause."` + PauseTimeout string `json:"pause_timeout" description:"How long to wait for webhook response." default:"30s"` + PauseTimeoutAction string `json:"pause_timeout_action" description:"Action on webhook timeout/error: deny or allow." default:"deny" enum:"deny,allow"` + PauseGracePeriod string `json:"pause_grace_period" description:"After approval, suppress further webhooks for this duration." default:"5m"` + SessionTTLSeconds int `json:"session_ttl_seconds" description:"Redis key TTL; should be >= max_duration_seconds." default:"7200"` + RefreshInterval string `json:"refresh_interval" description:"How often to sync local cache from Redis." default:"5s"` + RedisUnavailable string `json:"redis_unavailable" description:"Behavior when Redis is unreachable. Only fail_open is supported; fail_closed is reserved." default:"fail_open"` +} + +// approvalFlight carries the outcome of one webhook call. The leader writes +// approved before closing done; followers read it only after receiving from +// done, so the happens-before edge is safe without further synchronization. +// Attaching the result to the flight (instead of the cache entry) makes each +// waiter observe the outcome of the flight it actually waited on — a new +// leader that starts a second webhook after this one closes cannot clobber +// this flight's approved, and a refreshCache that deletes the cache entry +// mid-flight cannot make followers read a stale zero value. +type approvalFlight struct { + done chan struct{} + approved bool +} + +type counters struct { + tokens int64 + calls int64 + startedAt time.Time + lastApprovedAt time.Time + // pendingApproval is non-nil while a webhook call for this session is in + // flight. Concurrent breaches wait on flight.done; the leader publishes + // flight.approved before closing done, then clears this field. + pendingApproval *approvalFlight + // pendingWrites counts in-flight accumulate goroutines whose Redis writes + // haven't landed yet. refreshCache leaves the local entry alone (rather + // than deleting a Redis-missing session) while this is > 0, so a race + // between accumulate and a refresh tick can't erase local counters. + pendingWrites int +} + +// SessionBudget is the plugin state. Redis provides cross-pod durability; +// the local cache provides zero-I/O enforcement on the request path. +type SessionBudget struct { + cfg config + store storage.Store + log *slog.Logger + httpClient *http.Client + gracePeriod time.Duration + pauseTimeout time.Duration + + mu sync.RWMutex + cache map[string]*counters + hydrateG singleflight.Group + stopCh chan struct{} + stopped chan struct{} +} + +func New() *SessionBudget { + return &SessionBudget{ + cache: make(map[string]*counters), + stopCh: make(chan struct{}), + stopped: make(chan struct{}), + log: slog.Default().With("plugin", "session-budget"), + } +} + +func init() { + plugins.RegisterPlugin("session-budget", func() pipeline.Plugin { return New() }) +} + +func (p *SessionBudget) Name() string { return "session-budget" } + +func (p *SessionBudget) Capabilities() pipeline.PluginCapabilities { + return pipeline.PluginCapabilities{ + Description: "Enforce per-session token, call, and duration budgets via Redis.", + } +} + +func (p *SessionBudget) Configure(raw json.RawMessage) error { + p.cfg = config{ + OnExceed: "deny", + SessionTTLSeconds: 7200, + RefreshInterval: "5s", + RedisUnavailable: "fail_open", + } + if err := json.Unmarshal(raw, &p.cfg); err != nil { + return fmt.Errorf("session-budget config: %w", err) + } + if p.cfg.RedisURL == "" { + return fmt.Errorf("session-budget: redis_url is required") + } + if p.cfg.MaxTokens <= 0 && p.cfg.MaxCalls <= 0 && p.cfg.MaxDurationSeconds <= 0 { + return fmt.Errorf("session-budget: at least one limit (max_tokens, max_calls, max_duration_seconds) must be > 0") + } + switch p.cfg.OnExceed { + case "deny", "observe", "pause": + default: + return fmt.Errorf("session-budget: on_exceed must be \"deny\", \"observe\", or \"pause\" (got %q)", p.cfg.OnExceed) + } + if p.cfg.OnExceed == "pause" { + if p.cfg.PauseWebhook == "" { + return fmt.Errorf("session-budget: pause_webhook is required when on_exceed=\"pause\"") + } + if p.cfg.PauseTimeout == "" { + p.cfg.PauseTimeout = "30s" + } + if d, err := time.ParseDuration(p.cfg.PauseTimeout); err != nil { + return fmt.Errorf("session-budget: invalid pause_timeout %q: %w", p.cfg.PauseTimeout, err) + } else if d <= 0 { + return fmt.Errorf("session-budget: pause_timeout must be > 0 (got %q)", p.cfg.PauseTimeout) + } else { + p.pauseTimeout = d + } + if p.cfg.PauseTimeoutAction == "" { + p.cfg.PauseTimeoutAction = "deny" + } + if p.cfg.PauseTimeoutAction != "deny" && p.cfg.PauseTimeoutAction != "allow" { + return fmt.Errorf("session-budget: pause_timeout_action must be \"deny\" or \"allow\" (got %q)", p.cfg.PauseTimeoutAction) + } + if p.cfg.PauseGracePeriod == "" { + p.cfg.PauseGracePeriod = "5m" + } + if d, err := time.ParseDuration(p.cfg.PauseGracePeriod); err != nil { + return fmt.Errorf("session-budget: invalid pause_grace_period %q: %w", p.cfg.PauseGracePeriod, err) + } else if d < 0 { + return fmt.Errorf("session-budget: pause_grace_period must be >= 0 (got %q); use \"0s\" to fire the webhook on every breach", p.cfg.PauseGracePeriod) + } else { + p.gracePeriod = d + } + } + if d, err := time.ParseDuration(p.cfg.RefreshInterval); err != nil { + return fmt.Errorf("session-budget: invalid refresh_interval %q: %w", p.cfg.RefreshInterval, err) + } else if d <= 0 { + return fmt.Errorf("session-budget: refresh_interval must be > 0 (got %q)", p.cfg.RefreshInterval) + } + if p.cfg.RedisUnavailable == "fail_closed" { + return fmt.Errorf("session-budget: redis_unavailable=fail_closed is not yet implemented; use fail_open") + } + return nil +} + +func (p *SessionBudget) Init(_ context.Context) error { + // "redis" driver handles both Redis and Valkey (wire-compatible); URL must use redis:// scheme. + store, err := storage.Open("redis", p.cfg.RedisURL) + if err != nil { + return fmt.Errorf("session-budget: redis connect: %w", err) + } + p.store = store + + if p.cfg.OnExceed == "pause" && p.httpClient == nil { + // Timeout: 0 — per-request deadline is set via context in callPauseWebhook. + p.httpClient = &http.Client{Timeout: 0} + } + + interval, _ := time.ParseDuration(p.cfg.RefreshInterval) + go p.refreshLoop(interval) + return nil +} + +// In-flight accumulate goroutines get ErrClosed after store.Close — bounded by their 2s ctx. +func (p *SessionBudget) Shutdown(ctx context.Context) error { + close(p.stopCh) + select { + case <-p.stopped: + case <-ctx.Done(): + // refreshLoop is still running and calls p.store.HashGet; closing the + // store here would race that call. Leave the store open and let the + // process exit reclaim it — refreshLoop terminates when p.stopCh + // closes, and the shutdown-deadline is already exceeded. + p.log.Warn("shutdown timed out waiting for refresh loop; leaving store open") + return ctx.Err() + } + if p.store != nil { + return p.store.Close() + } + return nil +} + +// OnRequest evaluates cached counters against limits. On cold cache the first +// miss hydrates from Redis so pre-existing sessions enforce immediately. +func (p *SessionBudget) OnRequest(ctx context.Context, pctx *pipeline.Context) pipeline.Action { + sessionID := p.sessionID(pctx) + if sessionID == "" { + pctx.Skip("no_session_id") + return pipeline.Action{Type: pipeline.Continue} + } + + p.mu.Lock() + c, ok := p.cache[sessionID] + if !ok { + p.mu.Unlock() + // Cold cache handling is mode-dependent: + // - pause: synchronously hydrate from Redis so pre-existing sessions + // (seeded by another pod) fire the webhook on request #1. Pause is + // the only mode where a one-request-per-pod overshoot would defeat + // the point — HITL only works if we ask before continuing. + // - deny / observe: skip with cold_cache. The local counters populate + // via OnResponseFrame + the background refresh loop; a single pod + // may under-enforce by up to one request for a pre-existing session, + // which is the same tradeoff these modes have always had. This + // avoids putting Redis on the request path for the common modes. + if p.cfg.OnExceed == "pause" { + if !p.hydrateCache(ctx, sessionID) { + pctx.Skip("cold_cache") + return pipeline.Action{Type: pipeline.Continue} + } + p.mu.Lock() + c, ok = p.cache[sessionID] + if !ok { + p.mu.Unlock() + pctx.Skip("cold_cache") + return pipeline.Action{Type: pipeline.Continue} + } + } else { + pctx.Skip("cold_cache") + return pipeline.Action{Type: pipeline.Continue} + } + } + snap := *c + if reason := p.evaluate(&snap); reason != "" { + switch p.cfg.OnExceed { + case "observe": + p.mu.Unlock() + pctx.Observe("shadow_budget_exceeded") + p.log.Warn("budget exceeded (shadow mode)", + "session", sessionID, + "reason", reason, + "tokens", snap.tokens, + "calls", snap.calls) + return pipeline.Action{Type: pipeline.Continue} + + case "pause": + // Grace window: skip webhook if recently approved. + if p.gracePeriod > 0 && !c.lastApprovedAt.IsZero() && time.Since(c.lastApprovedAt) < p.gracePeriod { + p.mu.Unlock() + pctx.Allow("pause_grace_window") + return pipeline.Action{Type: pipeline.Continue} + } + if c.pendingApproval != nil { + // Another goroutine is already calling the webhook — wait for + // its outcome so followers honor a deny instead of racing past. + flight := c.pendingApproval + p.mu.Unlock() + select { + case <-flight.done: + case <-ctx.Done(): + pctx.Record(pipeline.Invocation{Action: pipeline.ActionDeny, Reason: "pause_wait_canceled"}) + return pipeline.DenyWithDetails("budget.exceeded", reason+" (client canceled during pause)", p.buildDetails(&snap)) + } + if flight.approved { + pctx.Allow("pause_follower_approved") + return pipeline.Action{Type: pipeline.Continue} + } + pctx.Record(pipeline.Invocation{Action: pipeline.ActionDeny, Reason: "pause_follower_denied"}) + return pipeline.DenyWithDetails("budget.exceeded", reason+" (approval denied)", p.buildDetails(&snap)) + } + flight := &approvalFlight{done: make(chan struct{})} + c.pendingApproval = flight + p.mu.Unlock() + p.log.Info("budget exceeded, requesting approval", + "session", sessionID, + "reason", reason) + // Deferred cleanup so a panic (or runtime.Goexit) in + // callPauseWebhook can't wedge the session: without this, + // pendingApproval would stay non-nil forever and every + // future request for this session would block on the dead + // flight.done. Order matters: publish outcome to the flight + // object first, then close done (channel-close is the + // happens-before edge for followers), then clear + // pendingApproval under the lock. + approved := false + defer func() { + flight.approved = approved + close(flight.done) + p.mu.Lock() + if cc, ok := p.cache[sessionID]; ok { + cc.pendingApproval = nil + if approved { + cc.lastApprovedAt = time.Now() + } + } + p.mu.Unlock() + }() + approved = p.callPauseWebhook(ctx, sessionID, reason, &snap) + if approved { + pctx.Allow("pause_approved") + return pipeline.Action{Type: pipeline.Continue} + } + details := p.buildDetails(&snap) + pctx.Record(pipeline.Invocation{Action: pipeline.ActionDeny, Reason: "pause_denied"}) + return pipeline.DenyWithDetails("budget.exceeded", reason+" (approval denied)", details) + + default: // "deny" + p.mu.Unlock() + details := p.buildDetails(&snap) + pctx.Record(pipeline.Invocation{Action: pipeline.ActionDeny, Reason: "budget_exceeded"}) + return pipeline.DenyWithDetails("budget.exceeded", reason, details) + } + } + p.mu.Unlock() + // Counts are incremented in OnResponseFrame when inference lands, so + // max_calls only counts LLM/inference calls (see plugin doc). + pctx.Allow("under_budget") + return pipeline.Action{Type: pipeline.Continue} +} + +// OnResponse is a no-op; see OnResponseFrame. +func (p *SessionBudget) OnResponse(_ context.Context, _ *pipeline.Context) pipeline.Action { + return pipeline.Action{Type: pipeline.Continue} +} + +// OnResponseFrame accumulates token counts on finalization (last=true). +func (p *SessionBudget) OnResponseFrame(_ context.Context, pctx *pipeline.Context, _ []byte, last bool) pipeline.Action { + if !last { + return pipeline.Action{Type: pipeline.Continue} + } + + sessionID := p.sessionID(pctx) + if sessionID == "" { + return pipeline.Action{Type: pipeline.Continue} + } + + inf := pctx.Extensions.Inference + if inf == nil { + return pipeline.Action{Type: pipeline.Continue} + } + + tokens := int64(inf.TotalTokens) + + p.mu.Lock() + c, ok := p.cache[sessionID] + if !ok { + c = &counters{startedAt: time.Now()} + p.cache[sessionID] = c + } + c.tokens += tokens + c.calls++ + c.pendingWrites++ + p.mu.Unlock() + + go p.accumulate(sessionID, tokens) + + return pipeline.Action{Type: pipeline.Continue} +} + +func (p *SessionBudget) buildDetails(snap *counters) map[string]any { + details := map[string]any{ + "spent_tokens": snap.tokens, + "spent_calls": snap.calls, + "token_limit": p.cfg.MaxTokens, + "call_limit": p.cfg.MaxCalls, + } + if p.cfg.MaxDurationSeconds > 0 && !snap.startedAt.IsZero() { + details["duration_seconds"] = int64(time.Since(snap.startedAt).Seconds()) + details["duration_limit"] = p.cfg.MaxDurationSeconds + } + return details +} + +type pauseRequest struct { + SessionID string `json:"session_id"` + Reason string `json:"reason"` + SpentTokens int64 `json:"spent_tokens"` + SpentCalls int64 `json:"spent_calls"` + TokenLimit int64 `json:"token_limit"` + CallLimit int64 `json:"call_limit"` + DurationSeconds int64 `json:"duration_seconds,omitempty"` + DurationLimit int64 `json:"duration_limit,omitempty"` +} + +type pauseResponse struct { + Action string `json:"action"` +} + +func (p *SessionBudget) callPauseWebhook(ctx context.Context, sessionID, reason string, snap *counters) bool { + ctx, cancel := context.WithTimeout(ctx, p.pauseTimeout) + defer cancel() + + body := pauseRequest{ + SessionID: sessionID, + Reason: reason, + SpentTokens: snap.tokens, + SpentCalls: snap.calls, + TokenLimit: p.cfg.MaxTokens, + CallLimit: p.cfg.MaxCalls, + } + if p.cfg.MaxDurationSeconds > 0 && !snap.startedAt.IsZero() { + body.DurationSeconds = int64(time.Since(snap.startedAt).Seconds()) + body.DurationLimit = p.cfg.MaxDurationSeconds + } + + payload, _ := json.Marshal(body) + req, err := http.NewRequestWithContext(ctx, http.MethodPost, p.cfg.PauseWebhook, bytes.NewReader(payload)) + if err != nil { + p.log.Warn("pause webhook request build failed", "session", sessionID, "err", err) + return p.cfg.PauseTimeoutAction == "allow" + } + req.Header.Set("Content-Type", "application/json") + + resp, err := p.httpClient.Do(req) + if err != nil { + p.log.Warn("pause webhook call failed", "session", sessionID, "err", err) + return p.cfg.PauseTimeoutAction == "allow" + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + responseBody, _ := io.ReadAll(io.LimitReader(resp.Body, 512)) + p.log.Warn("pause webhook non-200", "session", sessionID, "status", resp.StatusCode, "response_bytes", len(responseBody)) + return p.cfg.PauseTimeoutAction == "allow" + } + + var result pauseResponse + if err := json.NewDecoder(io.LimitReader(resp.Body, 4096)).Decode(&result); err != nil { + p.log.Warn("pause webhook response decode failed", "session", sessionID, "err", err) + return p.cfg.PauseTimeoutAction == "allow" + } + switch result.Action { + case "approve": + return true + case "deny": + return false + default: + p.log.Warn("pause webhook unknown action; treating as deny", "session", sessionID, "action", result.Action) + return false + } +} + +func (p *SessionBudget) evaluate(c *counters) string { + if p.cfg.MaxTokens > 0 && c.tokens >= p.cfg.MaxTokens { + return fmt.Sprintf("token limit reached: %d/%d", c.tokens, p.cfg.MaxTokens) + } + if p.cfg.MaxCalls > 0 && c.calls >= p.cfg.MaxCalls { + return fmt.Sprintf("call limit reached: %d/%d", c.calls, p.cfg.MaxCalls) + } + if p.cfg.MaxDurationSeconds > 0 && !c.startedAt.IsZero() { + elapsed := time.Since(c.startedAt).Seconds() + if int64(elapsed) >= p.cfg.MaxDurationSeconds { + return fmt.Sprintf("duration limit reached: %ds/%ds", int64(elapsed), p.cfg.MaxDurationSeconds) + } + } + return "" +} + +// accumulate writes counters to Redis. On failure, writes are dropped (fail-open). +func (p *SessionBudget) accumulate(sessionID string, tokens int64) { + defer func() { + p.mu.Lock() + if cc, ok := p.cache[sessionID]; ok && cc.pendingWrites > 0 { + cc.pendingWrites-- + } + p.mu.Unlock() + }() + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + + key := p.redisKey(sessionID) + ttl := time.Duration(p.cfg.SessionTTLSeconds) * time.Second + + if tokens > 0 { + if _, err := p.store.HashIncr(ctx, key, "tokens", tokens); err != nil { + p.log.Warn("redis HashIncr tokens failed", "session", sessionID, "err", err) + } + } + + if _, err := p.store.HashIncr(ctx, key, "calls", 1); err != nil { + p.log.Warn("redis HashIncr calls failed", "session", sessionID, "err", err) + } + + set, _ := p.store.HashSetNX(ctx, key, "started_at", strconv.FormatInt(time.Now().Unix(), 10)) + if set { + _ = p.store.Expire(ctx, key, ttl) + } +} + +func (p *SessionBudget) refreshLoop(interval time.Duration) { + defer close(p.stopped) + ticker := time.NewTicker(interval) + defer ticker.Stop() + + for { + select { + case <-p.stopCh: + return + case <-ticker.C: + p.refreshCache() + } + } +} + +// hydrateCache pulls one session's counters from Redis on cold-cache miss. +// Concurrent callers for the same session share one Redis lookup via singleflight. +func (p *SessionBudget) hydrateCache(ctx context.Context, sessionID string) bool { + v, _, _ := p.hydrateG.Do(sessionID, func() (any, error) { + lookupCtx, cancel := context.WithTimeout(ctx, 200*time.Millisecond) + defer cancel() + fields, err := p.store.HashGet(lookupCtx, p.redisKey(sessionID)) + if err != nil { + p.log.Debug("hydrate: redis lookup failed", "session", sessionID, "err", err) + return false, nil + } + if len(fields) == 0 { + return false, nil + } + tokens, _ := strconv.ParseInt(fields["tokens"], 10, 64) + calls, _ := strconv.ParseInt(fields["calls"], 10, 64) + var startedAt time.Time + if ts, err := strconv.ParseInt(fields["started_at"], 10, 64); err == nil { + startedAt = time.Unix(ts, 0) + } + p.mu.Lock() + if _, exists := p.cache[sessionID]; !exists { + p.cache[sessionID] = &counters{tokens: tokens, calls: calls, startedAt: startedAt} + } + p.mu.Unlock() + return true, nil + }) + return v.(bool) +} + +// refreshCache replaces local counters with authoritative Redis values. +func (p *SessionBudget) refreshCache() { + p.mu.RLock() + keys := make([]string, 0, len(p.cache)) + for k := range p.cache { + keys = append(keys, k) + } + p.mu.RUnlock() + + for _, sessionID := range keys { + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + fields, err := p.store.HashGet(ctx, p.redisKey(sessionID)) + cancel() + + if err != nil { + p.log.Warn("redis refresh failed", "session", sessionID, "err", err) + continue + } + + if len(fields) == 0 { + p.mu.Lock() + // Only delete when there are no in-flight accumulate goroutines — + // otherwise a slow Redis write racing with a refresh tick would + // wipe the local entry before its counters land in Redis. + if existing, ok := p.cache[sessionID]; !ok || existing.pendingWrites == 0 { + delete(p.cache, sessionID) + } + p.mu.Unlock() + continue + } + + tokens, _ := strconv.ParseInt(fields["tokens"], 10, 64) + calls, _ := strconv.ParseInt(fields["calls"], 10, 64) + var startedAt time.Time + if ts, err := strconv.ParseInt(fields["started_at"], 10, 64); err == nil { + startedAt = time.Unix(ts, 0) + } + + p.mu.Lock() + var lastApprovedAt time.Time + var pendingApproval *approvalFlight + var pendingWrites int + if existing, ok := p.cache[sessionID]; ok { + // Take the max of local and Redis to avoid regressing counters when + // in-flight accumulate goroutines haven't committed to Redis yet. + if tokens < existing.tokens { + tokens = existing.tokens + } + if calls < existing.calls { + calls = existing.calls + } + if startedAt.IsZero() && !existing.startedAt.IsZero() { + startedAt = existing.startedAt + } + lastApprovedAt = existing.lastApprovedAt + // Preserve mid-webhook: dropping would let a concurrent breach fire a duplicate. + pendingApproval = existing.pendingApproval + pendingWrites = existing.pendingWrites + } + p.cache[sessionID] = &counters{ + tokens: tokens, + calls: calls, + startedAt: startedAt, + lastApprovedAt: lastApprovedAt, + pendingApproval: pendingApproval, + pendingWrites: pendingWrites, + } + p.mu.Unlock() + } +} + +func (p *SessionBudget) sessionID(pctx *pipeline.Context) string { + if pctx.Session != nil && pctx.Session.ID != "" { + return pctx.Session.ID + } + return "" +} + +func (p *SessionBudget) redisKey(sessionID string) string { + return "session-budget:" + sessionID +} + +var ( + _ pipeline.Plugin = (*SessionBudget)(nil) + _ pipeline.Configurable = (*SessionBudget)(nil) + _ pipeline.Initializer = (*SessionBudget)(nil) + _ pipeline.Shutdowner = (*SessionBudget)(nil) + _ pipeline.StreamingResponder = (*SessionBudget)(nil) +) diff --git a/authbridge/authlib/plugins/sessionbudget/plugin_test.go b/authbridge/authlib/plugins/sessionbudget/plugin_test.go new file mode 100644 index 000000000..b193b31e4 --- /dev/null +++ b/authbridge/authlib/plugins/sessionbudget/plugin_test.go @@ -0,0 +1,931 @@ +package sessionbudget + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/rossoctl/cortex/authbridge/authlib/pipeline" + "github.com/rossoctl/cortex/authbridge/authlib/storage" +) + +// memStore is a minimal in-memory storage.Store for testing. +type memStore struct { + mu sync.Mutex + hashes map[string]map[string]string + kvs map[string]string + ttls map[string]time.Duration +} + +func newMemStore() *memStore { + return &memStore{ + hashes: make(map[string]map[string]string), + kvs: make(map[string]string), + ttls: make(map[string]time.Duration), + } +} + +func (m *memStore) Get(_ context.Context, key string) (string, error) { + m.mu.Lock() + defer m.mu.Unlock() + return m.kvs[key], nil +} + +func (m *memStore) Set(_ context.Context, key, value string, ttl time.Duration) error { + m.mu.Lock() + defer m.mu.Unlock() + m.kvs[key] = value + if ttl > 0 { + m.ttls[key] = ttl + } + return nil +} + +func (m *memStore) Incr(_ context.Context, key string, delta int64) (int64, error) { + m.mu.Lock() + defer m.mu.Unlock() + var cur int64 + if v, ok := m.kvs[key]; ok { + fmt.Sscanf(v, "%d", &cur) + } + cur += delta + m.kvs[key] = fmt.Sprintf("%d", cur) + return cur, nil +} + +func (m *memStore) HashIncr(_ context.Context, key, field string, delta int64) (int64, error) { + m.mu.Lock() + defer m.mu.Unlock() + if m.hashes[key] == nil { + m.hashes[key] = make(map[string]string) + } + var cur int64 + if v, ok := m.hashes[key][field]; ok { + fmt.Sscanf(v, "%d", &cur) + } + cur += delta + m.hashes[key][field] = fmt.Sprintf("%d", cur) + return cur, nil +} + +func (m *memStore) HashGet(_ context.Context, key string) (map[string]string, error) { + m.mu.Lock() + defer m.mu.Unlock() + h := m.hashes[key] + if h == nil { + return map[string]string{}, nil + } + out := make(map[string]string, len(h)) + for k, v := range h { + out[k] = v + } + return out, nil +} + +func (m *memStore) HashSetNX(_ context.Context, key, field, value string) (bool, error) { + m.mu.Lock() + defer m.mu.Unlock() + if m.hashes[key] == nil { + m.hashes[key] = make(map[string]string) + } + if _, exists := m.hashes[key][field]; exists { + return false, nil + } + m.hashes[key][field] = value + return true, nil +} + +func (m *memStore) Expire(_ context.Context, key string, ttl time.Duration) error { + m.mu.Lock() + defer m.mu.Unlock() + m.ttls[key] = ttl + return nil +} + +func (m *memStore) Close() error { return nil } + +var _ storage.Store = (*memStore)(nil) + +// failingStore always returns errors (simulates total store unavailability). +type failingStore struct{} + +func (failingStore) Get(context.Context, string) (string, error) { return "", context.DeadlineExceeded } +func (failingStore) Set(context.Context, string, string, time.Duration) error { + return context.DeadlineExceeded +} +func (failingStore) Incr(context.Context, string, int64) (int64, error) { + return 0, context.DeadlineExceeded +} +func (failingStore) HashIncr(context.Context, string, string, int64) (int64, error) { + return 0, context.DeadlineExceeded +} +func (failingStore) HashGet(context.Context, string) (map[string]string, error) { + return nil, context.DeadlineExceeded +} +func (failingStore) HashSetNX(context.Context, string, string, string) (bool, error) { + return false, context.DeadlineExceeded +} +func (failingStore) Expire(context.Context, string, time.Duration) error { + return context.DeadlineExceeded +} +func (failingStore) Close() error { return nil } + +func init() { + storage.Register("mem", func(_ string) (storage.Store, error) { + return newMemStore(), nil + }) +} + +func newTestPlugin(maxTokens, maxCalls, maxDuration int64) *SessionBudget { + p := New() + cfg := fmt.Sprintf(`{ + "redis_url": "mem://test", + "max_tokens": %d, + "max_calls": %d, + "max_duration_seconds": %d, + "refresh_interval": "100ms" + }`, maxTokens, maxCalls, maxDuration) + if err := p.Configure(json.RawMessage(cfg)); err != nil { + panic(err) + } + store := newMemStore() + p.store = store + return p +} + +func makePctx(sessionID string, totalTokens int) *pipeline.Context { + pctx := &pipeline.Context{ + Direction: pipeline.Outbound, + Headers: http.Header{}, + Session: &pipeline.SessionView{ID: sessionID}, + Extensions: pipeline.Extensions{ + Inference: &pipeline.InferenceExtension{ + TotalTokens: totalTokens, + }, + }, + } + return pctx +} + +func TestOnRequest_UnderLimit(t *testing.T) { + p := newTestPlugin(1000, 0, 0) + pctx := makePctx("sess-1", 0) + + action := p.OnRequest(context.Background(), pctx) + if action.Type != pipeline.Continue { + t.Fatalf("expected Continue, got %v", action.Type) + } +} + +func TestOnResponseFrame_Accumulates(t *testing.T) { + p := newTestPlugin(1000, 0, 0) + pctx := makePctx("sess-1", 42) + + action := p.OnResponseFrame(context.Background(), pctx, nil, true) + if action.Type != pipeline.Continue { + t.Fatalf("expected Continue, got %v", action.Type) + } + + // Check in-memory cache was updated. + p.mu.RLock() + c := p.cache["sess-1"] + p.mu.RUnlock() + + if c == nil { + t.Fatal("expected cache entry") + } + if c.tokens != 42 { + t.Errorf("tokens = %d, want 42", c.tokens) + } + if c.calls != 1 { + t.Errorf("calls = %d, want 1", c.calls) + } +} + +func TestOnResponseFrame_SkipsNonLast(t *testing.T) { + p := newTestPlugin(1000, 0, 0) + pctx := makePctx("sess-1", 42) + + action := p.OnResponseFrame(context.Background(), pctx, []byte("data"), false) + if action.Type != pipeline.Continue { + t.Fatalf("expected Continue, got %v", action.Type) + } + + p.mu.RLock() + _, ok := p.cache["sess-1"] + p.mu.RUnlock() + if ok { + t.Error("expected no cache entry on non-last frame") + } +} + +func TestOnResponseFrame_NoInference(t *testing.T) { + p := newTestPlugin(1000, 0, 0) + pctx := &pipeline.Context{ + Direction: pipeline.Outbound, + Headers: http.Header{}, + Session: &pipeline.SessionView{ID: "sess-1"}, + } + + action := p.OnResponseFrame(context.Background(), pctx, nil, true) + if action.Type != pipeline.Continue { + t.Fatalf("expected Continue, got %v", action.Type) + } + + p.mu.RLock() + _, ok := p.cache["sess-1"] + p.mu.RUnlock() + if ok { + t.Error("expected no cache entry when no inference data") + } +} + +func TestOnRequest_NoSession(t *testing.T) { + p := newTestPlugin(100, 0, 0) + pctx := &pipeline.Context{ + Direction: pipeline.Outbound, + Headers: http.Header{}, + } + + action := p.OnRequest(context.Background(), pctx) + if action.Type != pipeline.Continue { + t.Fatalf("expected Continue for nil session, got %v", action.Type) + } +} + +func TestAccumulate_WritesToStore(t *testing.T) { + p := newTestPlugin(1000, 0, 0) + store := newMemStore() + p.store = store + + p.accumulate("sess-1", 100) + + fields, _ := store.HashGet(context.Background(), "session-budget:sess-1") + if fields["tokens"] != "100" { + t.Errorf("tokens in store = %q, want 100", fields["tokens"]) + } + if fields["calls"] != "1" { + t.Errorf("calls in store = %q, want 1", fields["calls"]) + } + if fields["started_at"] == "" { + t.Error("started_at not set in store") + } +} + +func TestAccumulate_ZeroTokens(t *testing.T) { + p := newTestPlugin(1000, 10, 0) + store := newMemStore() + p.store = store + + p.accumulate("sess-1", 0) + + fields, _ := store.HashGet(context.Background(), "session-budget:sess-1") + if fields["tokens"] != "" { + t.Errorf("tokens in store = %q, want empty (no HINCRBY for 0)", fields["tokens"]) + } + if fields["calls"] != "1" { + t.Errorf("calls in store = %q, want 1", fields["calls"]) + } + if fields["started_at"] == "" { + t.Error("started_at not set in store") + } +} + +func TestOnResponseFrame_ZeroTokensCountsCalls(t *testing.T) { + p := newTestPlugin(1000, 5, 0) + pctx := makePctx("sess-1", 0) + + action := p.OnResponseFrame(context.Background(), pctx, nil, true) + if action.Type != pipeline.Continue { + t.Fatalf("expected Continue, got %v", action.Type) + } + + p.mu.RLock() + c := p.cache["sess-1"] + p.mu.RUnlock() + if c == nil { + t.Fatal("expected cache entry for zero-token response") + } + if c.calls != 1 { + t.Errorf("calls = %d, want 1", c.calls) + } + if c.tokens != 0 { + t.Errorf("tokens = %d, want 0", c.tokens) + } +} + +func TestConfigure_Validation(t *testing.T) { + tests := []struct { + name string + cfg string + wantErr bool + }{ + {"valid", `{"redis_url":"redis://localhost","max_tokens":100}`, false}, + {"missing redis_url", `{"max_tokens":100}`, true}, + {"no limits", `{"redis_url":"redis://localhost"}`, true}, + {"invalid json", `{broken}`, true}, + {"zero refresh_interval", `{"redis_url":"redis://localhost","max_tokens":100,"refresh_interval":"0s"}`, true}, + {"negative refresh_interval", `{"redis_url":"redis://localhost","max_tokens":100,"refresh_interval":"-1s"}`, true}, + {"unparseable refresh_interval", `{"redis_url":"redis://localhost","max_tokens":100,"refresh_interval":"abc"}`, true}, + {"fail_closed rejected", `{"redis_url":"redis://localhost","max_tokens":100,"redis_unavailable":"fail_closed"}`, true}, + {"invalid on_exceed", `{"redis_url":"redis://localhost","max_tokens":100,"on_exceed":"block"}`, true}, + {"pause valid", `{"redis_url":"redis://localhost","max_calls":10,"on_exceed":"pause","pause_webhook":"http://localhost:9999/approve"}`, false}, + {"pause missing webhook", `{"redis_url":"redis://localhost","max_calls":10,"on_exceed":"pause"}`, true}, + {"pause invalid timeout", `{"redis_url":"redis://localhost","max_calls":10,"on_exceed":"pause","pause_webhook":"http://x","pause_timeout":"nope"}`, true}, + {"pause invalid timeout_action", `{"redis_url":"redis://localhost","max_calls":10,"on_exceed":"pause","pause_webhook":"http://x","pause_timeout_action":"maybe"}`, true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + p := New() + err := p.Configure(json.RawMessage(tt.cfg)) + if (err != nil) != tt.wantErr { + t.Errorf("Configure() error = %v, wantErr = %v", err, tt.wantErr) + } + }) + } +} + +func TestOnRequest_ShadowMode(t *testing.T) { + p := New() + cfg := `{ + "redis_url": "mem://test", + "max_tokens": 100, + "on_exceed": "observe", + "refresh_interval": "100ms" + }` + if err := p.Configure(json.RawMessage(cfg)); err != nil { + t.Fatalf("Configure: %v", err) + } + p.store = newMemStore() + + p.OnResponseFrame(context.Background(), makePctx("sess-1", 60), nil, true) + p.OnResponseFrame(context.Background(), makePctx("sess-1", 60), nil, true) + + p.mu.RLock() + c := p.cache["sess-1"] + p.mu.RUnlock() + if c.tokens != 120 { + t.Errorf("tokens = %d, want 120", c.tokens) + } + + action := p.OnRequest(context.Background(), makePctx("sess-1", 0)) + if action.Type != pipeline.Continue { + t.Fatalf("shadow mode: expected Continue past limit, got %v", action.Type) + } + + // Calls are counted in OnResponseFrame (2 responses -> 2 calls). + // OnRequest in observe mode does not increment. + p.mu.RLock() + calls := p.cache["sess-1"].calls + p.mu.RUnlock() + if calls != 2 { + t.Errorf("calls after shadow OnRequest = %d, want 2 (from 2 OnResponseFrame calls)", calls) + } + + // Accumulation continues past the limit — observe never blocks writes. + p.OnResponseFrame(context.Background(), makePctx("sess-1", 60), nil, true) + p.mu.RLock() + tokens := p.cache["sess-1"].tokens + p.mu.RUnlock() + if tokens != 180 { + t.Errorf("tokens after post-limit response = %d, want 180", tokens) + } + if a := p.OnRequest(context.Background(), makePctx("sess-1", 0)); a.Type != pipeline.Continue { + t.Fatalf("shadow mode (2nd request past limit): expected Continue, got %v", a.Type) + } +} + +// TestOnRequest_RejectsAtCallLimit verifies that once cache reflects +// calls >= max_calls (populated by prior OnResponseFrame or hydrate), +// further OnRequests reject. Calls are counted on inference response, +// not on request, so this test seeds the cache directly. +func TestOnRequest_RejectsAtCallLimit(t *testing.T) { + p := newTestPlugin(1000, 10, 0) + + p.mu.Lock() + p.cache["sess-1"] = &counters{tokens: 50, calls: 10, startedAt: time.Now()} + p.mu.Unlock() + + action := p.OnRequest(context.Background(), makePctx("sess-1", 0)) + if action.Type != pipeline.Reject { + t.Fatalf("at limit: expected Reject, got %v", action.Type) + } +} + +func newPausePlugin(t *testing.T, maxCalls int64, webhookURL, timeoutAction string) *SessionBudget { + t.Helper() + p := New() + cfg := fmt.Sprintf(`{ + "redis_url": "mem://test", + "max_calls": %d, + "on_exceed": "pause", + "pause_webhook": %q, + "pause_timeout": "200ms", + "pause_timeout_action": %q, + "refresh_interval": "100ms" + }`, maxCalls, webhookURL, timeoutAction) + if err := p.Configure(json.RawMessage(cfg)); err != nil { + t.Fatalf("Configure: %v", err) + } + p.store = newMemStore() + p.httpClient = &http.Client{Timeout: 0} + return p +} + +// Approve / deny happy paths are covered by TestE2E_PauseMode in e2e_test.go +// (table-driven, asserts webhook request body + 403 response schema). + +// TestOnRequest_PauseWebhookFailureFallback covers every "webhook doesn't +// return a valid approve" path: timeout, non-200, malformed body. All fall +// back to pause_timeout_action. The 'allow' row also proves the allow +// branch (only place it's exercised). +func TestOnRequest_PauseWebhookFailureFallback(t *testing.T) { + hang := func(done <-chan struct{}) http.HandlerFunc { + return func(w http.ResponseWriter, _ *http.Request) { <-done } + } + tests := []struct { + name string + handler http.HandlerFunc + action string + want pipeline.ActionType + }{ + {"timeout_deny", nil, "deny", pipeline.Reject}, + {"timeout_allow", nil, "allow", pipeline.Continue}, + {"non200_deny", func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusInternalServerError) }, "deny", pipeline.Reject}, + {"badjson_deny", func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + w.Write([]byte(`not json`)) + }, "deny", pipeline.Reject}, + // Well-formed JSON with an action the plugin doesn't recognize + // must deny (protocol-safe default) regardless of pause_timeout_action — + // unlike the transport/decode failures above, this is not a webhook + // outage, so pause_timeout_action does not apply. + {"unknown_action_deny", func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{"action":"maybe"}`)) + }, "allow", pipeline.Reject}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + done := make(chan struct{}) + h := tt.handler + if h == nil { + h = hang(done) + } + srv := httptest.NewServer(h) + defer func() { close(done); srv.Close() }() + + p := newPausePlugin(t, 3, srv.URL, tt.action) + p.mu.Lock() + p.cache["sess-1"] = &counters{tokens: 0, calls: 3, startedAt: time.Now()} + p.mu.Unlock() + + got := p.OnRequest(context.Background(), makePctx("sess-1", 0)) + if got.Type != tt.want { + t.Fatalf("action = %v, want %v", got.Type, tt.want) + } + }) + } +} + +func TestOnRequest_PauseWebhookUnreachable(t *testing.T) { + p := newPausePlugin(t, 3, "http://127.0.0.1:1", "deny") + p.mu.Lock() + p.cache["sess-1"] = &counters{tokens: 0, calls: 3, startedAt: time.Now()} + p.mu.Unlock() + + action := p.OnRequest(context.Background(), makePctx("sess-1", 0)) + if action.Type != pipeline.Reject { + t.Fatalf("expected Reject on unreachable webhook (pause_timeout_action=deny), got %v", action.Type) + } +} + +func TestRefreshCache_PreservesLastApprovedAt(t *testing.T) { + var webhookCalls atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + webhookCalls.Add(1) + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{"action":"approve"}`)) + })) + defer srv.Close() + + p := New() + cfg := fmt.Sprintf(`{ + "redis_url": "mem://test", + "max_calls": 3, + "on_exceed": "pause", + "pause_webhook": %q, + "pause_timeout": "2s", + "pause_timeout_action": "deny", + "pause_grace_period": "10m", + "refresh_interval": "100ms" + }`, srv.URL) + if err := p.Configure(json.RawMessage(cfg)); err != nil { + t.Fatalf("Configure: %v", err) + } + store := newMemStore() + p.store = store + p.httpClient = &http.Client{Timeout: 0} + + // Seed cache at limit and fire first request to get approval. + p.mu.Lock() + p.cache["sess"] = &counters{tokens: 0, calls: 3, startedAt: time.Now()} + p.mu.Unlock() + + action := p.OnRequest(context.Background(), makePctx("sess", 0)) + if action.Type != pipeline.Continue { + t.Fatalf("expected Continue after approval, got %v", action.Type) + } + if c := webhookCalls.Load(); c != 1 { + t.Fatalf("expected 1 webhook call, got %d", c) + } + + // Simulate Redis having authoritative counters. + ctx := context.Background() + store.HashIncr(ctx, "session-budget:sess", "tokens", 100) + store.HashIncr(ctx, "session-budget:sess", "calls", 5) + store.HashSetNX(ctx, "session-budget:sess", "started_at", "1700000000") + + // Refresh replaces counters from Redis — must preserve lastApprovedAt. + p.refreshCache() + + // Second request should still be within grace (no new webhook call). + action = p.OnRequest(context.Background(), makePctx("sess", 0)) + if action.Type != pipeline.Continue { + t.Fatalf("expected Continue within grace after refresh, got %v", action.Type) + } + if c := webhookCalls.Load(); c != 1 { + t.Fatalf("expected still 1 webhook call after refresh, got %d", c) + } +} + +func TestOnRequest_PauseGraceWindow(t *testing.T) { + var webhookCalls atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + webhookCalls.Add(1) + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{"action":"approve"}`)) + })) + defer srv.Close() + + p := New() + cfg := fmt.Sprintf(`{ + "redis_url": "mem://test", + "max_calls": 3, + "on_exceed": "pause", + "pause_webhook": %q, + "pause_timeout": "2s", + "pause_timeout_action": "deny", + "pause_grace_period": "10m", + "refresh_interval": "100ms" + }`, srv.URL) + if err := p.Configure(json.RawMessage(cfg)); err != nil { + t.Fatalf("Configure: %v", err) + } + p.store = newMemStore() + p.httpClient = &http.Client{Timeout: 0} + + p.mu.Lock() + p.cache["sess"] = &counters{tokens: 0, calls: 3, startedAt: time.Now()} + p.mu.Unlock() + + // First request fires the webhook. + action := p.OnRequest(context.Background(), makePctx("sess", 0)) + if action.Type != pipeline.Continue { + t.Fatalf("first request: expected Continue, got %v", action.Type) + } + if c := webhookCalls.Load(); c != 1 { + t.Fatalf("expected 1 webhook call, got %d", c) + } + + // Second request within grace window skips the webhook. + action = p.OnRequest(context.Background(), makePctx("sess", 0)) + if action.Type != pipeline.Continue { + t.Fatalf("second request (grace): expected Continue, got %v", action.Type) + } + if c := webhookCalls.Load(); c != 1 { + t.Fatalf("expected still 1 webhook call after grace, got %d", c) + } +} + +func TestOnRequest_PauseGraceExpired(t *testing.T) { + var webhookCalls atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + webhookCalls.Add(1) + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{"action":"approve"}`)) + })) + defer srv.Close() + + p := New() + cfg := fmt.Sprintf(`{ + "redis_url": "mem://test", + "max_calls": 3, + "on_exceed": "pause", + "pause_webhook": %q, + "pause_timeout": "2s", + "pause_timeout_action": "deny", + "pause_grace_period": "10m", + "refresh_interval": "100ms" + }`, srv.URL) + if err := p.Configure(json.RawMessage(cfg)); err != nil { + t.Fatalf("Configure: %v", err) + } + p.store = newMemStore() + p.httpClient = &http.Client{Timeout: 0} + + // Seed cache with lastApprovedAt already expired (11 minutes ago > 10m grace). + p.mu.Lock() + p.cache["sess"] = &counters{ + tokens: 0, + calls: 3, + startedAt: time.Now(), + lastApprovedAt: time.Now().Add(-11 * time.Minute), + } + p.mu.Unlock() + + // Request after grace expired fires webhook and continues on approve. + action := p.OnRequest(context.Background(), makePctx("sess", 0)) + if action.Type != pipeline.Continue { + t.Fatalf("expected Continue after grace expired + approve, got %v", action.Type) + } + if c := webhookCalls.Load(); c != 1 { + t.Fatalf("expected 1 webhook call after grace expired, got %d", c) + } +} + +// TestOnRequest_PausePendingApprovalSentinel verifies that concurrent +// breaches share one webhook call AND that the follower waits for the +// leader's outcome (rather than racing past optimistically). +func TestOnRequest_PausePendingApprovalSentinel(t *testing.T) { + var webhookCalls atomic.Int32 + started := make(chan struct{}) + // OnceFunc: a regression that lets a second webhook fire produces a + // readable failure instead of close-of-closed-channel panic. + signalStart := sync.OnceFunc(func() { close(started) }) + proceed := make(chan struct{}) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + webhookCalls.Add(1) + signalStart() + <-proceed + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{"action":"approve"}`)) + })) + defer srv.Close() + + p := newPausePlugin(t, 3, srv.URL, "deny") + p.mu.Lock() + p.cache["sess"] = &counters{tokens: 0, calls: 3, startedAt: time.Now()} + p.mu.Unlock() + + // Two concurrent requests: leader fires the webhook, follower must wait. + var wg sync.WaitGroup + results := make([]pipeline.ActionType, 2) + for i := 0; i < 2; i++ { + wg.Add(1) + go func(idx int) { + defer wg.Done() + a := p.OnRequest(context.Background(), makePctx("sess", 0)) + results[idx] = a.Type + }(i) + } + <-started // leader's webhook is in-flight; follower is now waiting on the channel + + // Unblock the webhook — both goroutines should complete after leader returns. + close(proceed) + wg.Wait() + + if c := webhookCalls.Load(); c != 1 { + t.Errorf("webhook called %d times, want exactly 1 (sentinel prevents thundering herd)", c) + } + for i, got := range results { + if got != pipeline.Continue { + t.Errorf("request %d: got %v, want Continue (webhook approved)", i, got) + } + } +} + +// TestOnRequest_PauseFollowerHonorsDeny verifies that a follower waiting +// on the leader's webhook call gets Reject when the leader is denied, +// instead of racing past optimistically. +func TestOnRequest_PauseFollowerHonorsDeny(t *testing.T) { + started := make(chan struct{}) + signalStart := sync.OnceFunc(func() { close(started) }) + proceed := make(chan struct{}) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + signalStart() + <-proceed + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{"action":"deny"}`)) + })) + defer srv.Close() + + p := newPausePlugin(t, 3, srv.URL, "deny") + p.mu.Lock() + p.cache["sess"] = &counters{tokens: 0, calls: 3, startedAt: time.Now()} + p.mu.Unlock() + + var wg sync.WaitGroup + results := make([]pipeline.ActionType, 2) + for i := 0; i < 2; i++ { + wg.Add(1) + go func(idx int) { + defer wg.Done() + a := p.OnRequest(context.Background(), makePctx("sess", 0)) + results[idx] = a.Type + }(i) + } + <-started + close(proceed) + wg.Wait() + + for i, got := range results { + if got != pipeline.Reject { + t.Errorf("request %d: got %v, want Reject (webhook denied)", i, got) + } + } +} + +// TestOnRequest_ColdCacheModeGating pins the documented contract: cold-cache +// hydrates from Redis on the OnRequest path only in pause mode. deny and +// observe skip cold-cache and let counters populate via OnResponseFrame + +// refresh loop — accepting one-request-per-pod overshoot for pre-existing +// over-budget sessions. See docs/session-budget-plugin.md "Cold-cache +// behavior". +func TestOnRequest_ColdCacheModeGating(t *testing.T) { + for _, mode := range []string{"deny", "observe"} { + t.Run(mode, func(t *testing.T) { + p := New() + cfg, _ := json.Marshal(config{ + RedisURL: "mem://test", + MaxTokens: 100, + OnExceed: mode, + RefreshInterval: "1s", + RedisUnavailable: "fail_open", + }) + if err := p.Configure(cfg); err != nil { + t.Fatalf("Configure: %v", err) + } + // Pre-seed Redis above budget. If cold-cache hydrated, evaluate() + // would fire. deny/observe must NOT hydrate — request continues. + store := newMemStore() + ctx := context.Background() + store.HashIncr(ctx, "session-budget:s", "tokens", 500) + store.HashIncr(ctx, "session-budget:s", "calls", 9) + p.store = store + + a := p.OnRequest(ctx, makePctx("s", 0)) + if a.Type != pipeline.Continue { + t.Fatalf("%s cold cache: expected Continue (no hydrate on request path), got %v", mode, a.Type) + } + }) + } +} + +func TestEvaluate_MultipleLimits(t *testing.T) { + p := newTestPlugin(100, 10, 60) + + tests := []struct { + name string + c *counters + wantDeny bool + }{ + {"all under", &counters{tokens: 50, calls: 5, startedAt: time.Now()}, false}, + {"tokens over", &counters{tokens: 100, calls: 5, startedAt: time.Now()}, true}, + {"calls over", &counters{tokens: 50, calls: 10, startedAt: time.Now()}, true}, + {"duration over", &counters{tokens: 50, calls: 5, startedAt: time.Now().Add(-90 * time.Second)}, true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + reason := p.evaluate(tt.c) + if tt.wantDeny && reason == "" { + t.Error("expected denial reason, got empty") + } + if !tt.wantDeny && reason != "" { + t.Errorf("expected no denial, got %q", reason) + } + }) + } +} + +// Sequential flights with opposite outcomes must not interfere. Regression +// for the pendingResult-on-cache-entry bug: flight 2's deny cannot flip +// flight 1's already-returned approve because outcome rides on the flight. +func TestOnRequest_PauseSequentialFlightsIndependent(t *testing.T) { + var callIdx atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + n := callIdx.Add(1) + w.WriteHeader(http.StatusOK) + if n == 1 { + w.Write([]byte(`{"action":"approve"}`)) + } else { + w.Write([]byte(`{"action":"deny"}`)) + } + })) + defer srv.Close() + + p := newPausePlugin(t, 3, srv.URL, "deny") + p.gracePeriod = 0 // Configure default is 5m; otherwise flight 2 is short-circuited. + + p.mu.Lock() + p.cache["sess"] = &counters{tokens: 0, calls: 3, startedAt: time.Now()} + p.mu.Unlock() + + if a := p.OnRequest(context.Background(), makePctx("sess", 0)); a.Type != pipeline.Continue { + t.Fatalf("flight 1: got %v, want Continue", a.Type) + } + + p.mu.Lock() + p.cache["sess"] = &counters{tokens: 0, calls: 3, startedAt: time.Now()} + p.mu.Unlock() + + if a := p.OnRequest(context.Background(), makePctx("sess", 0)); a.Type != pipeline.Reject { + t.Fatalf("flight 2: got %v, want Reject", a.Type) + } + if n := callIdx.Load(); n != 2 { + t.Fatalf("webhook calls = %d, want 2", n) + } +} + +// TestShutdown_TimeoutDoesNotCloseStore is the regression test for the +// Shutdown store-close race. On the ctx.Done() path, the old code called +// p.store.Close() while refreshLoop was still running and could be inside +// a HashGet call — concurrent use of a closed client. The fix is to leave +// the store open on the timeout path and let process exit reclaim it. +// This test uses a store that records Close() calls and forces the +// timeout branch by passing an already-expired context. +// Shutdown's ctx.Done() branch must NOT call store.Close(). No refreshLoop +// runs, so p.stopped never closes and ctx.Done is the only ready select case. +func TestShutdown_TimeoutDoesNotCloseStore(t *testing.T) { + p := newTestPlugin(0, 3, 0) + rec := &closeRecordingStore{Store: p.store} + p.store = rec + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + if err := p.Shutdown(ctx); err == nil { + t.Fatal("expected non-nil error from Shutdown on canceled ctx") + } + if n := rec.closes.Load(); n != 0 { + t.Errorf("store.Close() called %d times on timeout path, want 0", n) + } +} + +// closeRecordingStore wraps a Store and counts Close() calls. +type closeRecordingStore struct { + storage.Store + closes atomic.Int32 +} + +func (s *closeRecordingStore) Close() error { + s.closes.Add(1) + return s.Store.Close() +} + +// panicTransport panics from RoundTrip so the panic surfaces inside +// p.httpClient.Do(req) — i.e. inside callPauseWebhook after +// pendingApproval is published. Simulates a runtime crash mid-webhook. +type panicTransport struct{} + +func (panicTransport) RoundTrip(*http.Request) (*http.Response, error) { + panic("simulated webhook client panic") +} + +// TestOnRequest_PauseWebhookPanicClearsFlight proves the deferred cleanup +// in OnRequest's pause branch: if callPauseWebhook panics, the defer must +// clear pendingApproval so a follow-up request re-enters as a new leader +// instead of following a dead flight forever. No sleeps, no timers — a +// follow-up call that returned proves the defer ran. +func TestOnRequest_PauseWebhookPanicClearsFlight(t *testing.T) { + p := newPausePlugin(t, 3, "http://unused.invalid", "deny") + p.httpClient = &http.Client{Transport: panicTransport{}} + p.mu.Lock() + p.cache["sess"] = &counters{tokens: 0, calls: 3, startedAt: time.Now()} + p.mu.Unlock() + + call := func() { + defer func() { _ = recover() }() + _ = p.OnRequest(context.Background(), makePctx("sess", 0)) + } + call() + + p.mu.RLock() + pending := p.cache["sess"].pendingApproval + p.mu.RUnlock() + if pending != nil { + t.Fatal("pendingApproval not cleared after panic — session would wedge") + } + + // If the defer had NOT closed flight.done, this second call would + // follower-wait on the dead channel and the test would deadlock + // (caught by go test -timeout). Returning at all proves the fix. + call() +} diff --git a/authbridge/authlib/plugins/tokenbudget/e2e_test.go b/authbridge/authlib/plugins/tokenbudget/e2e_test.go deleted file mode 100644 index 3f4993788..000000000 --- a/authbridge/authlib/plugins/tokenbudget/e2e_test.go +++ /dev/null @@ -1,296 +0,0 @@ -package tokenbudget - -import ( - "context" - "encoding/json" - "io" - "net/http" - "net/http/httptest" - "net/url" - "sync" - "testing" - "time" - - "github.com/rossoctl/cortex/authbridge/authlib/listener/forwardproxy" - "github.com/rossoctl/cortex/authbridge/authlib/pipeline" - "github.com/rossoctl/cortex/authbridge/authlib/session" -) - -func newE2EPlugin(t *testing.T, maxTokens int64, store *memStore) *TokenBudget { - t.Helper() - p := New() - cfg, _ := json.Marshal(config{ - RedisURL: "mem://test", - MaxTokens: maxTokens, - OnExceed: "deny", - RefreshInterval: "30ms", - RedisUnavailable: "fail_open", - }) - if err := p.Configure(cfg); err != nil { - t.Fatalf("Configure: %v", err) - } - p.store = store - go p.refreshLoop(30 * time.Millisecond) - t.Cleanup(func() { close(p.stopCh); <-p.stopped }) - return p -} - -func respond(p *TokenBudget, sessionID string, tokens int) { - p.OnResponseFrame(context.Background(), makePctx(sessionID, tokens), nil, true) -} - -func request(p *TokenBudget, sessionID string) pipeline.Action { - pctx := &pipeline.Context{ - Direction: pipeline.Outbound, - Headers: http.Header{}, - Session: &pipeline.SessionView{ID: sessionID}, - } - return p.OnRequest(context.Background(), pctx) -} - -// TestE2E_HTTPRoundTrip wires token-budget into a real forward proxy. -// Under-budget requests reach the backend; the proxy is functional. -func TestE2E_HTTPRoundTrip(t *testing.T) { - store := newMemStore() - p := newE2EPlugin(t, 1000, store) - - pipe, err := pipeline.New([]pipeline.Plugin{p}) - if err != nil { - t.Fatal(err) - } - sessions := session.New(5*time.Minute, 100, 0) - defer sessions.Close() - - srv, err := forwardproxy.NewServer(pipeline.NewHolder(pipe), sessions, nil) - if err != nil { - t.Fatal(err) - } - proxy := httptest.NewServer(srv.Handler()) - defer proxy.Close() - - backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { - w.WriteHeader(http.StatusOK) - _, _ = w.Write([]byte("ok")) - })) - defer backend.Close() - - proxyURL, _ := url.Parse(proxy.URL) - client := &http.Client{Transport: &http.Transport{Proxy: http.ProxyURL(proxyURL)}} - - req, _ := http.NewRequest(http.MethodGet, backend.URL+"/v1/chat/completions", nil) - resp, err := client.Do(req) - if err != nil { - t.Fatalf("request through proxy: %v", err) - } - defer resp.Body.Close() - if resp.StatusCode != http.StatusOK { - body, _ := io.ReadAll(resp.Body) - t.Fatalf("expected 200, got %d: %s", resp.StatusCode, body) - } -} - -// TestE2E_AccumulateAndDeny verifies the full lifecycle: accumulate -// tokens via OnResponseFrame, then OnRequest denies with a 403. -func TestE2E_AccumulateAndDeny(t *testing.T) { - p := newE2EPlugin(t, 150, newMemStore()) - - for i := 0; i < 3; i++ { - respond(p, "sess", 60) - } - - action := request(p, "sess") - if action.Type != pipeline.Reject { - t.Fatalf("expected Reject, got %v", action.Type) - } - status, _, body := action.Violation.Render() - if status != http.StatusForbidden { - t.Errorf("status = %d, want 403", status) - } - var parsed map[string]any - if err := json.Unmarshal(body, &parsed); err != nil { - t.Fatal(err) - } - if parsed["error"] != "budget.exceeded" { - t.Errorf("error = %v, want budget.exceeded", parsed["error"]) - } -} - -// TestE2E_MultiSession verifies independent session budgets. -func TestE2E_MultiSession(t *testing.T) { - p := newE2EPlugin(t, 100, newMemStore()) - - for i := 0; i < 3; i++ { - respond(p, "A", 40) // 120 > 100 - } - respond(p, "B", 20) // 20 < 100 - - if a := request(p, "A"); a.Type != pipeline.Reject { - t.Fatalf("session A: expected Reject, got %v", a.Type) - } - if a := request(p, "B"); a.Type != pipeline.Continue { - t.Fatalf("session B: expected Continue, got %v", a.Type) - } -} - -// TestE2E_LocalCacheEnforcesDuringOutage confirms that a populated -// cache enforces even when the backing store is unreachable. -func TestE2E_LocalCacheEnforcesDuringOutage(t *testing.T) { - p := newE2EPlugin(t, 100, newMemStore()) - p.store = &failingStore{} - - p.mu.Lock() - p.cache["s"] = &counters{tokens: 110, calls: 5, startedAt: time.Now()} - p.mu.Unlock() - - if a := request(p, "s"); a.Type != pipeline.Reject { - t.Fatalf("expected Reject from cache with store down, got %v", a.Type) - } -} - -// TestE2E_RefreshRecovery confirms that refreshCache picks up -// authoritative store values after an outage resolves. -func TestE2E_RefreshRecovery(t *testing.T) { - inner := newMemStore() - cs := &controllableStore{inner: inner} - p := newE2EPlugin(t, 200, newMemStore()) - p.store = cs - - ctx := context.Background() - inner.HashIncr(ctx, "token-budget:s", "tokens", 180) - inner.HashIncr(ctx, "token-budget:s", "calls", 7) - inner.HashSetNX(ctx, "token-budget:s", "started_at", "1700000000") - - p.mu.Lock() - p.cache["s"] = &counters{tokens: 50} - p.mu.Unlock() - - cs.setFailing(true) - p.refreshCache() - p.mu.RLock() - if p.cache["s"].tokens != 50 { - t.Fatalf("during outage: tokens = %d, want 50", p.cache["s"].tokens) - } - p.mu.RUnlock() - - cs.setFailing(false) - p.refreshCache() - p.mu.RLock() - if p.cache["s"].tokens != 180 { - t.Errorf("after recovery: tokens = %d, want 180", p.cache["s"].tokens) - } - p.mu.RUnlock() -} - -// TestE2E_PodRestart verifies that a fresh plugin with an empty cache -// resumes enforcement after refresh picks up pre-existing store counters. -func TestE2E_PodRestart(t *testing.T) { - store := newMemStore() - p := newE2EPlugin(t, 200, store) - - ctx := context.Background() - // Pre-seed Redis above the limit (210 > 200). - store.HashIncr(ctx, "token-budget:s", "tokens", 210) - store.HashIncr(ctx, "token-budget:s", "calls", 8) - store.HashSetNX(ctx, "token-budget:s", "started_at", "1700000000") - - // Cold cache — first request passes (overshoot window). - if a := request(p, "s"); a.Type != pipeline.Continue { - t.Fatalf("cold cache: expected Continue, got %v", a.Type) - } - - // Seed cache entry so refresh discovers this session key. - respond(p, "s", 5) - - // Directly invoke refresh (deterministic, no timing dependency). - p.refreshCache() - - if a := request(p, "s"); a.Type != pipeline.Reject { - t.Fatalf("after refresh: expected Reject, got %v", a.Type) - } -} - -// controllableStore delegates to inner memStore but can be toggled to fail. -type controllableStore struct { - inner *memStore - failing bool - mu sync.Mutex -} - -func (c *controllableStore) setFailing(v bool) { c.mu.Lock(); c.failing = v; c.mu.Unlock() } -func (c *controllableStore) isFailing() bool { c.mu.Lock(); defer c.mu.Unlock(); return c.failing } -func (c *controllableStore) err() error { return context.DeadlineExceeded } - -func (c *controllableStore) Get(ctx context.Context, key string) (string, error) { - if c.isFailing() { return "", c.err() } - return c.inner.Get(ctx, key) -} -func (c *controllableStore) Set(ctx context.Context, key, value string, ttl time.Duration) error { - if c.isFailing() { return c.err() } - return c.inner.Set(ctx, key, value, ttl) -} -func (c *controllableStore) Incr(ctx context.Context, key string, delta int64) (int64, error) { - if c.isFailing() { return 0, c.err() } - return c.inner.Incr(ctx, key, delta) -} -func (c *controllableStore) HashIncr(ctx context.Context, key, field string, delta int64) (int64, error) { - if c.isFailing() { return 0, c.err() } - return c.inner.HashIncr(ctx, key, field, delta) -} -func (c *controllableStore) HashGet(ctx context.Context, key string) (map[string]string, error) { - if c.isFailing() { return nil, c.err() } - return c.inner.HashGet(ctx, key) -} -func (c *controllableStore) HashSetNX(ctx context.Context, key, field, value string) (bool, error) { - if c.isFailing() { return false, c.err() } - return c.inner.HashSetNX(ctx, key, field, value) -} -func (c *controllableStore) Expire(ctx context.Context, key string, ttl time.Duration) error { - if c.isFailing() { return c.err() } - return c.inner.Expire(ctx, key, ttl) -} -func (c *controllableStore) Close() error { return nil } - -// TestE2E_ShadowMode verifies that on_exceed=observe allows requests -// through even when budget is exceeded, while still accumulating. -func TestE2E_ShadowMode(t *testing.T) { - p := New() - cfg, _ := json.Marshal(config{ - RedisURL: "mem://test", - MaxTokens: 150, - OnExceed: "observe", - RefreshInterval: "30ms", - RedisUnavailable: "fail_open", - }) - if err := p.Configure(cfg); err != nil { - t.Fatalf("Configure: %v", err) - } - store := newMemStore() - p.store = store - go p.refreshLoop(30 * time.Millisecond) - t.Cleanup(func() { close(p.stopCh); <-p.stopped }) - - for i := 0; i < 3; i++ { - respond(p, "sess", 60) // 180 total > 150 limit - } - - // In observe mode, request should continue (not reject). - action := request(p, "sess") - if action.Type != pipeline.Continue { - t.Fatalf("shadow mode: expected Continue, got %v", action.Type) - } - - // Counters should still accumulate past the limit. - respond(p, "sess", 20) // 200 total - p.mu.RLock() - c := p.cache["sess"] - p.mu.RUnlock() - if c.tokens != 200 { - t.Errorf("tokens = %d, want 200 (accumulation continues in shadow mode)", c.tokens) - } - - // Subsequent requests also continue. - action = request(p, "sess") - if action.Type != pipeline.Continue { - t.Fatalf("shadow mode (2nd request): expected Continue, got %v", action.Type) - } -} diff --git a/authbridge/authlib/plugins/tokenbudget/plugin.go b/authbridge/authlib/plugins/tokenbudget/plugin.go deleted file mode 100644 index fbc3780c0..000000000 --- a/authbridge/authlib/plugins/tokenbudget/plugin.go +++ /dev/null @@ -1,339 +0,0 @@ -// Package tokenbudget enforces per-session lifetime budgets on tokens, -// inference calls, and wall-clock duration. Must run before inference-parser -// in the declared plugin order (response path is reverse: inference-parser -// finalizes counts first, then this plugin reads them). -package tokenbudget - -import ( - "context" - "encoding/json" - "fmt" - "log/slog" - "strconv" - "sync" - "time" - - "github.com/rossoctl/cortex/authbridge/authlib/pipeline" - "github.com/rossoctl/cortex/authbridge/authlib/plugins" - "github.com/rossoctl/cortex/authbridge/authlib/storage" -) - -type config struct { - RedisURL string `json:"redis_url" required:"true" description:"Redis/Valkey connection URL."` - MaxTokens int64 `json:"max_tokens" description:"Cumulative token ceiling per session. 0 = no limit."` - MaxCalls int64 `json:"max_calls" description:"Max inference calls per session. 0 = no limit."` - MaxDurationSeconds int64 `json:"max_duration_seconds" description:"Wall-clock session lifetime in seconds. 0 = no limit."` - OnExceed string `json:"on_exceed" description:"Action on breach: deny (block) or observe (shadow — log but continue)." default:"deny" enum:"deny,observe"` - SessionTTLSeconds int `json:"session_ttl_seconds" description:"Redis key TTL; should be >= max_duration_seconds." default:"7200"` - RefreshInterval string `json:"refresh_interval" description:"How often to sync local cache from Redis." default:"5s"` - RedisUnavailable string `json:"redis_unavailable" description:"Behavior when Redis is unreachable. Only fail_open is supported; fail_closed is reserved." default:"fail_open"` -} - -type counters struct { - tokens int64 - calls int64 - startedAt time.Time -} - -// TokenBudget is the plugin state. Redis provides cross-pod durability; -// the local cache provides zero-I/O enforcement on the request path. -type TokenBudget struct { - cfg config - store storage.Store - log *slog.Logger - - mu sync.RWMutex - cache map[string]*counters - stopCh chan struct{} - stopped chan struct{} -} - -func New() *TokenBudget { - return &TokenBudget{ - cache: make(map[string]*counters), - stopCh: make(chan struct{}), - stopped: make(chan struct{}), - log: slog.Default().With("plugin", "token-budget"), - } -} - -func init() { - plugins.RegisterPlugin("token-budget", func() pipeline.Plugin { return New() }) -} - -func (p *TokenBudget) Name() string { return "token-budget" } - -func (p *TokenBudget) Capabilities() pipeline.PluginCapabilities { - return pipeline.PluginCapabilities{ - Description: "Enforce per-session token, call, and duration budgets via Redis.", - } -} - -func (p *TokenBudget) Configure(raw json.RawMessage) error { - p.cfg = config{ - OnExceed: "deny", - SessionTTLSeconds: 7200, - RefreshInterval: "5s", - RedisUnavailable: "fail_open", - } - if err := json.Unmarshal(raw, &p.cfg); err != nil { - return fmt.Errorf("token-budget config: %w", err) - } - if p.cfg.RedisURL == "" { - return fmt.Errorf("token-budget: redis_url is required") - } - if p.cfg.MaxTokens <= 0 && p.cfg.MaxCalls <= 0 && p.cfg.MaxDurationSeconds <= 0 { - return fmt.Errorf("token-budget: at least one limit (max_tokens, max_calls, max_duration_seconds) must be > 0") - } - if p.cfg.OnExceed != "deny" && p.cfg.OnExceed != "observe" { - return fmt.Errorf("token-budget: on_exceed must be \"deny\" or \"observe\" (got %q)", p.cfg.OnExceed) - } - if d, err := time.ParseDuration(p.cfg.RefreshInterval); err != nil { - return fmt.Errorf("token-budget: invalid refresh_interval %q: %w", p.cfg.RefreshInterval, err) - } else if d <= 0 { - return fmt.Errorf("token-budget: refresh_interval must be > 0 (got %q)", p.cfg.RefreshInterval) - } - if p.cfg.RedisUnavailable == "fail_closed" { - return fmt.Errorf("token-budget: redis_unavailable=fail_closed is not yet implemented; use fail_open") - } - return nil -} - -func (p *TokenBudget) Init(_ context.Context) error { - // "redis" driver handles both Redis and Valkey (wire-compatible); URL must use redis:// scheme. - store, err := storage.Open("redis", p.cfg.RedisURL) - if err != nil { - return fmt.Errorf("token-budget: redis connect: %w", err) - } - p.store = store - - interval, _ := time.ParseDuration(p.cfg.RefreshInterval) - go p.refreshLoop(interval) - return nil -} - -// In-flight accumulate goroutines get ErrClosed after store.Close — bounded by their 2s ctx. -func (p *TokenBudget) Shutdown(_ context.Context) error { - close(p.stopCh) - <-p.stopped - if p.store != nil { - return p.store.Close() - } - return nil -} - -// OnRequest evaluates cached counters against limits and optimistically reserves -// a call slot so concurrent requests on the same session see each other. No I/O. -func (p *TokenBudget) OnRequest(_ context.Context, pctx *pipeline.Context) pipeline.Action { - sessionID := p.sessionID(pctx) - if sessionID == "" { - return pipeline.Action{Type: pipeline.Continue} - } - - p.mu.Lock() - c, ok := p.cache[sessionID] - if !ok { - p.mu.Unlock() - // Cold cache: session not yet seen by this pod. First request passes; refresh loop - // picks up Redis counters within one interval. Intentional one-request overshoot - // tradeoff for zero-I/O enforcement on the hot path. - return pipeline.Action{Type: pipeline.Continue} - } - snap := *c - if reason := p.evaluate(&snap); reason != "" { - if p.cfg.OnExceed == "observe" { - // Still reserve a call — the request will proceed in shadow mode. - c.calls++ - p.mu.Unlock() - pctx.Observe("shadow_budget_exceeded") - p.log.Warn("budget exceeded (shadow mode)", - "session", sessionID, - "reason", reason, - "tokens", snap.tokens, - "calls", snap.calls) - return pipeline.Action{Type: pipeline.Continue} - } - p.mu.Unlock() - details := map[string]any{ - "spent_tokens": snap.tokens, - "spent_calls": snap.calls, - "token_limit": p.cfg.MaxTokens, - "call_limit": p.cfg.MaxCalls, - } - if p.cfg.MaxDurationSeconds > 0 && !snap.startedAt.IsZero() { - details["duration_seconds"] = int64(time.Since(snap.startedAt).Seconds()) - details["duration_limit"] = p.cfg.MaxDurationSeconds - } - return pipeline.DenyWithDetails("budget.exceeded", reason, details) - } - // Optimistically reserve a call slot so concurrent requests see it. - c.calls++ - p.mu.Unlock() - - return pipeline.Action{Type: pipeline.Continue} -} - -// OnResponse is a no-op; see OnResponseFrame. -func (p *TokenBudget) OnResponse(_ context.Context, _ *pipeline.Context) pipeline.Action { - return pipeline.Action{Type: pipeline.Continue} -} - -// OnResponseFrame accumulates token counts on finalization (last=true). -func (p *TokenBudget) OnResponseFrame(_ context.Context, pctx *pipeline.Context, _ []byte, last bool) pipeline.Action { - if !last { - return pipeline.Action{Type: pipeline.Continue} - } - - sessionID := p.sessionID(pctx) - if sessionID == "" { - return pipeline.Action{Type: pipeline.Continue} - } - - inf := pctx.Extensions.Inference - if inf == nil { - return pipeline.Action{Type: pipeline.Continue} - } - - tokens := int64(inf.TotalTokens) - - go p.accumulate(sessionID, tokens) - - p.mu.Lock() - c, ok := p.cache[sessionID] - if !ok { - // First response without a prior OnRequest (e.g. cold cache path). - c = &counters{startedAt: time.Now(), calls: 1} - p.cache[sessionID] = c - } - c.tokens += tokens - // calls already incremented by OnRequest's optimistic reserve. - p.mu.Unlock() - - return pipeline.Action{Type: pipeline.Continue} -} - -func (p *TokenBudget) evaluate(c *counters) string { - if p.cfg.MaxTokens > 0 && c.tokens >= p.cfg.MaxTokens { - return fmt.Sprintf("token limit reached: %d/%d", c.tokens, p.cfg.MaxTokens) - } - if p.cfg.MaxCalls > 0 && c.calls >= p.cfg.MaxCalls { - return fmt.Sprintf("call limit reached: %d/%d", c.calls, p.cfg.MaxCalls) - } - if p.cfg.MaxDurationSeconds > 0 && !c.startedAt.IsZero() { - elapsed := time.Since(c.startedAt).Seconds() - if int64(elapsed) >= p.cfg.MaxDurationSeconds { - return fmt.Sprintf("duration limit reached: %ds/%ds", int64(elapsed), p.cfg.MaxDurationSeconds) - } - } - return "" -} - -// accumulate writes counters to Redis. On failure, writes are dropped (fail-open). -func (p *TokenBudget) accumulate(sessionID string, tokens int64) { - ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) - defer cancel() - - key := p.redisKey(sessionID) - ttl := time.Duration(p.cfg.SessionTTLSeconds) * time.Second - - if tokens > 0 { - if _, err := p.store.HashIncr(ctx, key, "tokens", tokens); err != nil { - p.log.Warn("redis HashIncr tokens failed", "session", sessionID, "err", err) - } - } - - if _, err := p.store.HashIncr(ctx, key, "calls", 1); err != nil { - p.log.Warn("redis HashIncr calls failed", "session", sessionID, "err", err) - } - - set, _ := p.store.HashSetNX(ctx, key, "started_at", strconv.FormatInt(time.Now().Unix(), 10)) - if set { - _ = p.store.Expire(ctx, key, ttl) - } -} - -func (p *TokenBudget) refreshLoop(interval time.Duration) { - defer close(p.stopped) - ticker := time.NewTicker(interval) - defer ticker.Stop() - - for { - select { - case <-p.stopCh: - return - case <-ticker.C: - p.refreshCache() - } - } -} - -// refreshCache replaces local counters with authoritative Redis values. -func (p *TokenBudget) refreshCache() { - p.mu.RLock() - keys := make([]string, 0, len(p.cache)) - for k := range p.cache { - keys = append(keys, k) - } - p.mu.RUnlock() - - for _, sessionID := range keys { - ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) - fields, err := p.store.HashGet(ctx, p.redisKey(sessionID)) - cancel() - - if err != nil { - p.log.Warn("redis refresh failed", "session", sessionID, "err", err) - continue - } - - if len(fields) == 0 { - p.mu.Lock() - delete(p.cache, sessionID) - p.mu.Unlock() - continue - } - - tokens, _ := strconv.ParseInt(fields["tokens"], 10, 64) - calls, _ := strconv.ParseInt(fields["calls"], 10, 64) - var startedAt time.Time - if ts, err := strconv.ParseInt(fields["started_at"], 10, 64); err == nil { - startedAt = time.Unix(ts, 0) - } - - p.mu.Lock() - if existing, ok := p.cache[sessionID]; ok { - // Take the max of local and Redis to avoid regressing counters when - // in-flight accumulate goroutines haven't committed to Redis yet. - if tokens < existing.tokens { - tokens = existing.tokens - } - if calls < existing.calls { - calls = existing.calls - } - if startedAt.IsZero() && !existing.startedAt.IsZero() { - startedAt = existing.startedAt - } - } - p.cache[sessionID] = &counters{tokens: tokens, calls: calls, startedAt: startedAt} - p.mu.Unlock() - } -} - -func (p *TokenBudget) sessionID(pctx *pipeline.Context) string { - if pctx.Session != nil && pctx.Session.ID != "" { - return pctx.Session.ID - } - return "" -} - -func (p *TokenBudget) redisKey(sessionID string) string { - return "token-budget:" + sessionID -} - -var ( - _ pipeline.Plugin = (*TokenBudget)(nil) - _ pipeline.Configurable = (*TokenBudget)(nil) - _ pipeline.Initializer = (*TokenBudget)(nil) - _ pipeline.Shutdowner = (*TokenBudget)(nil) - _ pipeline.StreamingResponder = (*TokenBudget)(nil) -) diff --git a/authbridge/authlib/plugins/tokenbudget/plugin_test.go b/authbridge/authlib/plugins/tokenbudget/plugin_test.go deleted file mode 100644 index d11e63418..000000000 --- a/authbridge/authlib/plugins/tokenbudget/plugin_test.go +++ /dev/null @@ -1,478 +0,0 @@ -package tokenbudget - -import ( - "context" - "encoding/json" - "fmt" - "net/http" - "sync" - "testing" - "time" - - "github.com/rossoctl/cortex/authbridge/authlib/pipeline" - "github.com/rossoctl/cortex/authbridge/authlib/storage" -) - -// memStore is a minimal in-memory storage.Store for testing. -type memStore struct { - mu sync.Mutex - hashes map[string]map[string]string - kvs map[string]string - ttls map[string]time.Duration -} - -func newMemStore() *memStore { - return &memStore{ - hashes: make(map[string]map[string]string), - kvs: make(map[string]string), - ttls: make(map[string]time.Duration), - } -} - -func (m *memStore) Get(_ context.Context, key string) (string, error) { - m.mu.Lock() - defer m.mu.Unlock() - return m.kvs[key], nil -} - -func (m *memStore) Set(_ context.Context, key, value string, ttl time.Duration) error { - m.mu.Lock() - defer m.mu.Unlock() - m.kvs[key] = value - if ttl > 0 { - m.ttls[key] = ttl - } - return nil -} - -func (m *memStore) Incr(_ context.Context, key string, delta int64) (int64, error) { - m.mu.Lock() - defer m.mu.Unlock() - var cur int64 - if v, ok := m.kvs[key]; ok { - fmt.Sscanf(v, "%d", &cur) - } - cur += delta - m.kvs[key] = fmt.Sprintf("%d", cur) - return cur, nil -} - -func (m *memStore) HashIncr(_ context.Context, key, field string, delta int64) (int64, error) { - m.mu.Lock() - defer m.mu.Unlock() - if m.hashes[key] == nil { - m.hashes[key] = make(map[string]string) - } - var cur int64 - if v, ok := m.hashes[key][field]; ok { - fmt.Sscanf(v, "%d", &cur) - } - cur += delta - m.hashes[key][field] = fmt.Sprintf("%d", cur) - return cur, nil -} - -func (m *memStore) HashGet(_ context.Context, key string) (map[string]string, error) { - m.mu.Lock() - defer m.mu.Unlock() - h := m.hashes[key] - if h == nil { - return map[string]string{}, nil - } - out := make(map[string]string, len(h)) - for k, v := range h { - out[k] = v - } - return out, nil -} - -func (m *memStore) HashSetNX(_ context.Context, key, field, value string) (bool, error) { - m.mu.Lock() - defer m.mu.Unlock() - if m.hashes[key] == nil { - m.hashes[key] = make(map[string]string) - } - if _, exists := m.hashes[key][field]; exists { - return false, nil - } - m.hashes[key][field] = value - return true, nil -} - -func (m *memStore) Expire(_ context.Context, key string, ttl time.Duration) error { - m.mu.Lock() - defer m.mu.Unlock() - m.ttls[key] = ttl - return nil -} - -func (m *memStore) Close() error { return nil } - -var _ storage.Store = (*memStore)(nil) - -// failingStore always returns errors (simulates total store unavailability). -type failingStore struct{} - -func (failingStore) Get(context.Context, string) (string, error) { return "", context.DeadlineExceeded } -func (failingStore) Set(context.Context, string, string, time.Duration) error { return context.DeadlineExceeded } -func (failingStore) Incr(context.Context, string, int64) (int64, error) { return 0, context.DeadlineExceeded } -func (failingStore) HashIncr(context.Context, string, string, int64) (int64, error) { return 0, context.DeadlineExceeded } -func (failingStore) HashGet(context.Context, string) (map[string]string, error) { return nil, context.DeadlineExceeded } -func (failingStore) HashSetNX(context.Context, string, string, string) (bool, error) { return false, context.DeadlineExceeded } -func (failingStore) Expire(context.Context, string, time.Duration) error { return context.DeadlineExceeded } -func (failingStore) Close() error { return nil } - -func init() { - storage.Register("mem", func(_ string) (storage.Store, error) { - return newMemStore(), nil - }) -} - -func newTestPlugin(maxTokens, maxCalls, maxDuration int64) *TokenBudget { - p := New() - cfg := fmt.Sprintf(`{ - "redis_url": "mem://test", - "max_tokens": %d, - "max_calls": %d, - "max_duration_seconds": %d, - "refresh_interval": "100ms" - }`, maxTokens, maxCalls, maxDuration) - if err := p.Configure(json.RawMessage(cfg)); err != nil { - panic(err) - } - store := newMemStore() - p.store = store - return p -} - -func makePctx(sessionID string, totalTokens int) *pipeline.Context { - pctx := &pipeline.Context{ - Direction: pipeline.Outbound, - Headers: http.Header{}, - Session: &pipeline.SessionView{ID: sessionID}, - Extensions: pipeline.Extensions{ - Inference: &pipeline.InferenceExtension{ - TotalTokens: totalTokens, - }, - }, - } - return pctx -} - -func TestOnRequest_UnderLimit(t *testing.T) { - p := newTestPlugin(1000, 0, 0) - pctx := makePctx("sess-1", 0) - - action := p.OnRequest(context.Background(), pctx) - if action.Type != pipeline.Continue { - t.Fatalf("expected Continue, got %v", action.Type) - } -} - - -func TestOnResponseFrame_Accumulates(t *testing.T) { - p := newTestPlugin(1000, 0, 0) - pctx := makePctx("sess-1", 42) - - action := p.OnResponseFrame(context.Background(), pctx, nil, true) - if action.Type != pipeline.Continue { - t.Fatalf("expected Continue, got %v", action.Type) - } - - // Check in-memory cache was updated. - p.mu.RLock() - c := p.cache["sess-1"] - p.mu.RUnlock() - - if c == nil { - t.Fatal("expected cache entry") - } - if c.tokens != 42 { - t.Errorf("tokens = %d, want 42", c.tokens) - } - if c.calls != 1 { - t.Errorf("calls = %d, want 1", c.calls) - } -} - -func TestOnResponseFrame_SkipsNonLast(t *testing.T) { - p := newTestPlugin(1000, 0, 0) - pctx := makePctx("sess-1", 42) - - action := p.OnResponseFrame(context.Background(), pctx, []byte("data"), false) - if action.Type != pipeline.Continue { - t.Fatalf("expected Continue, got %v", action.Type) - } - - p.mu.RLock() - _, ok := p.cache["sess-1"] - p.mu.RUnlock() - if ok { - t.Error("expected no cache entry on non-last frame") - } -} - -func TestOnResponseFrame_NoInference(t *testing.T) { - p := newTestPlugin(1000, 0, 0) - pctx := &pipeline.Context{ - Direction: pipeline.Outbound, - Headers: http.Header{}, - Session: &pipeline.SessionView{ID: "sess-1"}, - } - - action := p.OnResponseFrame(context.Background(), pctx, nil, true) - if action.Type != pipeline.Continue { - t.Fatalf("expected Continue, got %v", action.Type) - } - - p.mu.RLock() - _, ok := p.cache["sess-1"] - p.mu.RUnlock() - if ok { - t.Error("expected no cache entry when no inference data") - } -} - -func TestOnRequest_NoSession(t *testing.T) { - p := newTestPlugin(100, 0, 0) - pctx := &pipeline.Context{ - Direction: pipeline.Outbound, - Headers: http.Header{}, - } - - action := p.OnRequest(context.Background(), pctx) - if action.Type != pipeline.Continue { - t.Fatalf("expected Continue for nil session, got %v", action.Type) - } -} - -func TestAccumulate_WritesToStore(t *testing.T) { - p := newTestPlugin(1000, 0, 0) - store := newMemStore() - p.store = store - - p.accumulate("sess-1", 100) - - fields, _ := store.HashGet(context.Background(), "token-budget:sess-1") - if fields["tokens"] != "100" { - t.Errorf("tokens in store = %q, want 100", fields["tokens"]) - } - if fields["calls"] != "1" { - t.Errorf("calls in store = %q, want 1", fields["calls"]) - } - if fields["started_at"] == "" { - t.Error("started_at not set in store") - } -} - -func TestAccumulate_ZeroTokens(t *testing.T) { - p := newTestPlugin(1000, 10, 0) - store := newMemStore() - p.store = store - - p.accumulate("sess-1", 0) - - fields, _ := store.HashGet(context.Background(), "token-budget:sess-1") - if fields["tokens"] != "" { - t.Errorf("tokens in store = %q, want empty (no HINCRBY for 0)", fields["tokens"]) - } - if fields["calls"] != "1" { - t.Errorf("calls in store = %q, want 1", fields["calls"]) - } - if fields["started_at"] == "" { - t.Error("started_at not set in store") - } -} - -func TestOnResponseFrame_ZeroTokensCountsCalls(t *testing.T) { - p := newTestPlugin(1000, 5, 0) - pctx := makePctx("sess-1", 0) - - action := p.OnResponseFrame(context.Background(), pctx, nil, true) - if action.Type != pipeline.Continue { - t.Fatalf("expected Continue, got %v", action.Type) - } - - p.mu.RLock() - c := p.cache["sess-1"] - p.mu.RUnlock() - if c == nil { - t.Fatal("expected cache entry for zero-token response") - } - if c.calls != 1 { - t.Errorf("calls = %d, want 1", c.calls) - } - if c.tokens != 0 { - t.Errorf("tokens = %d, want 0", c.tokens) - } -} - -func TestConfigure_Validation(t *testing.T) { - tests := []struct { - name string - cfg string - wantErr bool - }{ - {"valid", `{"redis_url":"redis://localhost","max_tokens":100}`, false}, - {"missing redis_url", `{"max_tokens":100}`, true}, - {"no limits", `{"redis_url":"redis://localhost"}`, true}, - {"invalid json", `{broken}`, true}, - {"zero refresh_interval", `{"redis_url":"redis://localhost","max_tokens":100,"refresh_interval":"0s"}`, true}, - {"negative refresh_interval", `{"redis_url":"redis://localhost","max_tokens":100,"refresh_interval":"-1s"}`, true}, - {"unparseable refresh_interval", `{"redis_url":"redis://localhost","max_tokens":100,"refresh_interval":"abc"}`, true}, - {"fail_closed rejected", `{"redis_url":"redis://localhost","max_tokens":100,"redis_unavailable":"fail_closed"}`, true}, - {"invalid on_exceed", `{"redis_url":"redis://localhost","max_tokens":100,"on_exceed":"block"}`, true}, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - p := New() - err := p.Configure(json.RawMessage(tt.cfg)) - if (err != nil) != tt.wantErr { - t.Errorf("Configure() error = %v, wantErr = %v", err, tt.wantErr) - } - }) - } -} - -func TestOnRequest_ShadowMode(t *testing.T) { - p := New() - cfg := `{ - "redis_url": "mem://test", - "max_tokens": 100, - "on_exceed": "observe", - "refresh_interval": "100ms" - }` - if err := p.Configure(json.RawMessage(cfg)); err != nil { - t.Fatalf("Configure: %v", err) - } - p.store = newMemStore() - - p.OnResponseFrame(context.Background(), makePctx("sess-1", 60), nil, true) - p.OnResponseFrame(context.Background(), makePctx("sess-1", 60), nil, true) - - p.mu.RLock() - c := p.cache["sess-1"] - p.mu.RUnlock() - if c.tokens != 120 { - t.Errorf("tokens = %d, want 120", c.tokens) - } - - action := p.OnRequest(context.Background(), makePctx("sess-1", 0)) - if action.Type != pipeline.Continue { - t.Fatalf("shadow mode: expected Continue past limit, got %v", action.Type) - } - - // Verify observe mode still reserves a call slot. - p.mu.RLock() - calls := p.cache["sess-1"].calls - p.mu.RUnlock() - if calls != 2 { - t.Errorf("calls after shadow OnRequest = %d, want 2 (1 from cold-cache response + 1 observe reservation)", calls) - } -} - -func TestOnRequest_OptimisticReservation(t *testing.T) { - p := newTestPlugin(1000, 10, 0) - - // Seed cache so OnRequest finds the session. - p.mu.Lock() - p.cache["sess-1"] = &counters{tokens: 50, calls: 7, startedAt: time.Now()} - p.mu.Unlock() - - // Three sequential requests each reserve a slot; calls should reach the limit. - for i := 0; i < 3; i++ { - action := p.OnRequest(context.Background(), makePctx("sess-1", 0)) - if action.Type != pipeline.Continue { - t.Fatalf("request %d: expected Continue, got %v", i+1, action.Type) - } - } - - // Fourth request should be denied (7+3 = 10, limit is 10). - action := p.OnRequest(context.Background(), makePctx("sess-1", 0)) - if action.Type != pipeline.Reject { - t.Fatalf("request 4: expected Reject at limit, got %v", action.Type) - } - - p.mu.RLock() - calls := p.cache["sess-1"].calls - p.mu.RUnlock() - if calls != 10 { - t.Errorf("calls = %d, want 10", calls) - } -} - -// TestOnRequest_ConcurrentCallLimit verifies that concurrent goroutines racing -// through OnRequest cannot exceed the call limit. This is the scenario the -// RLock→Lock upgrade and optimistic reservation were designed to prevent. -// Distinct from TestOnRequest_OptimisticReservation which tests serial ordering. -func TestOnRequest_ConcurrentCallLimit(t *testing.T) { - p := newTestPlugin(1000, 10, 0) - - p.mu.Lock() - p.cache["sess-1"] = &counters{tokens: 0, calls: 0, startedAt: time.Now()} - p.mu.Unlock() - - const n = 20 - var wg sync.WaitGroup - results := make([]pipeline.ActionType, n) - wg.Add(n) - for i := 0; i < n; i++ { - go func(idx int) { - defer wg.Done() - action := p.OnRequest(context.Background(), makePctx("sess-1", 0)) - results[idx] = action.Type - }(i) - } - wg.Wait() - - var continued, rejected int - for _, r := range results { - switch r { - case pipeline.Continue: - continued++ - case pipeline.Reject: - rejected++ - } - } - if continued != 10 { - t.Errorf("continued = %d, want exactly 10 (call limit)", continued) - } - if rejected != 10 { - t.Errorf("rejected = %d, want 10", rejected) - } - - p.mu.RLock() - calls := p.cache["sess-1"].calls - p.mu.RUnlock() - if calls != 10 { - t.Errorf("final calls = %d, want 10", calls) - } -} - -func TestEvaluate_MultipleLimits(t *testing.T) { - p := newTestPlugin(100, 10, 60) - - tests := []struct { - name string - c *counters - wantDeny bool - }{ - {"all under", &counters{tokens: 50, calls: 5, startedAt: time.Now()}, false}, - {"tokens over", &counters{tokens: 100, calls: 5, startedAt: time.Now()}, true}, - {"calls over", &counters{tokens: 50, calls: 10, startedAt: time.Now()}, true}, - {"duration over", &counters{tokens: 50, calls: 5, startedAt: time.Now().Add(-90 * time.Second)}, true}, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - reason := p.evaluate(tt.c) - if tt.wantDeny && reason == "" { - t.Error("expected denial reason, got empty") - } - if !tt.wantDeny && reason != "" { - t.Errorf("expected no denial, got %q", reason) - } - }) - } -} diff --git a/authbridge/cmd/authbridge-envoy/plugins_sessionbudget.go b/authbridge/cmd/authbridge-envoy/plugins_sessionbudget.go new file mode 100644 index 000000000..bd3dd9c4b --- /dev/null +++ b/authbridge/cmd/authbridge-envoy/plugins_sessionbudget.go @@ -0,0 +1,10 @@ +//go:build include_plugin_sessionbudget + +// session-budget is opt-IN: it pulls the storage/redis module (go-redis) +// into the binary. Build with -tags include_plugin_sessionbudget to link it. +package main + +import ( + _ "github.com/rossoctl/cortex/authbridge/authlib/plugins/sessionbudget" + _ "github.com/rossoctl/cortex/authbridge/storage/redis" +) diff --git a/authbridge/cmd/authbridge-envoy/plugins_tokenbudget.go b/authbridge/cmd/authbridge-envoy/plugins_tokenbudget.go deleted file mode 100644 index 93a74754b..000000000 --- a/authbridge/cmd/authbridge-envoy/plugins_tokenbudget.go +++ /dev/null @@ -1,10 +0,0 @@ -//go:build include_plugin_tokenbudget - -// token-budget is opt-IN: it pulls the storage/redis module (go-redis) -// into the binary. Build with -tags include_plugin_tokenbudget to link it. -package main - -import ( - _ "github.com/rossoctl/cortex/authbridge/authlib/plugins/tokenbudget" - _ "github.com/rossoctl/cortex/authbridge/storage/redis" -) diff --git a/authbridge/cmd/authbridge-proxy/plugins_sessionbudget.go b/authbridge/cmd/authbridge-proxy/plugins_sessionbudget.go new file mode 100644 index 000000000..bd3dd9c4b --- /dev/null +++ b/authbridge/cmd/authbridge-proxy/plugins_sessionbudget.go @@ -0,0 +1,10 @@ +//go:build include_plugin_sessionbudget + +// session-budget is opt-IN: it pulls the storage/redis module (go-redis) +// into the binary. Build with -tags include_plugin_sessionbudget to link it. +package main + +import ( + _ "github.com/rossoctl/cortex/authbridge/authlib/plugins/sessionbudget" + _ "github.com/rossoctl/cortex/authbridge/storage/redis" +) diff --git a/authbridge/cmd/authbridge-proxy/plugins_tokenbudget.go b/authbridge/cmd/authbridge-proxy/plugins_tokenbudget.go deleted file mode 100644 index 93a74754b..000000000 --- a/authbridge/cmd/authbridge-proxy/plugins_tokenbudget.go +++ /dev/null @@ -1,10 +0,0 @@ -//go:build include_plugin_tokenbudget - -// token-budget is opt-IN: it pulls the storage/redis module (go-redis) -// into the binary. Build with -tags include_plugin_tokenbudget to link it. -package main - -import ( - _ "github.com/rossoctl/cortex/authbridge/authlib/plugins/tokenbudget" - _ "github.com/rossoctl/cortex/authbridge/storage/redis" -) diff --git a/authbridge/demos/README.md b/authbridge/demos/README.md index 690a89be2..f44bf8fe3 100644 --- a/authbridge/demos/README.md +++ b/authbridge/demos/README.md @@ -20,6 +20,7 @@ more AuthBridge capabilities. | **[GitHub Issue Agent](github-issue/demo.md)** | Intermediate | Inbound validation + outbound token exchange + scope-based access control | [UI](github-issue/demo-ui.md) or [Manual](github-issue/demo-manual.md) | | **[Token-Exchange Routes](token-exchange-routes/README.md)** | Reference | How to write `authproxy-routes` for single- and multi-target token exchange | Configuration only | | **[MCP Parser Plugin](mcp-parser/README.md)** | Reference | Enable the `mcp-parser` plugin to surface tool calls / resource reads in session events | Configuration only | +| **[Session Budget](session-budget/README.md)** | Reference | Test assets for the `session-budget` plugin, including a pause-mode webhook stub | kubectl | | **[abctl Walkthrough](weather-agent/demo-with-abctl.md)** | Reference | Watch the AuthBridge plugin pipeline live with the `abctl` TUI | Tooling only | | **[IBAC](ibac/README.md)** | Intermediate | Intent-Based Access Control: LLM judge denies outbound HTTP that doesn't align with the user's recorded intent. Reproduces the email-poison / prompt-injection attack from `huang195/ibac`; chat with the agent through the rossoctl UI and see the exfiltration blocked, then `make show-result` for a pipeline-level forensic | UI + kubectl | | **[SPARC (finance)](finance-sparc/README.md)** | Intermediate | SPARC pre-tool reflection: the `sparc` plugin blocks a hallucinated/ungrounded tool argument (an invented transaction id) before it executes and transparently asks the user to clarify, then approves the corrected call. Complements IBAC — SPARC verifies argument grounding, IBAC verifies intent alignment | UI + kubectl | diff --git a/authbridge/demos/session-budget/README.md b/authbridge/demos/session-budget/README.md new file mode 100644 index 000000000..1a58355a7 --- /dev/null +++ b/authbridge/demos/session-budget/README.md @@ -0,0 +1,116 @@ +# session-budget demo assets + +Deployable helpers for exercising the `session-budget` plugin. For plugin +configuration and mode semantics, see +[`../../docs/session-budget-plugin.md`](../../docs/session-budget-plugin.md). + +## `k8s/pause-webhook-stub.yaml` + +Minimal HITL webhook that returns `{"action":"approve"}` for every POST — +enough to smoke-test `on_exceed: pause` end-to-end. Also logs each +incoming request body so you can see exactly what session-budget sent. + +```bash +kubectl apply -f k8s/pause-webhook-stub.yaml +``` + +To exercise the deny path, edit the inline Python in the manifest to +return `{"action":"deny"}` and re-apply. + +Follow the webhook stub: + +```bash +kubectl logs -n "$NS" deploy/pause-webhook-stub -f +``` + +## Prerequisites + +- **A Redis-wire-compatible store** reachable from the agent pod. Any + Valkey/Redis deployment works; point `redis_url` at its Service. +- **`a2a-parser` on the inbound pipeline** — see the note below. + +**Ambient-mesh note:** if your namespace has +`istio.io/dataplane-mode: ambient`, the datastore pod needs the +pod-level label `istio.io/dataplane-mode: none`. Ambient's ztunnel drops +non-HBONE connections with `Connection reset by peer`, and Redis RESP is +raw TCP — it can't ride HBONE. The pause webhook stub manifest already +carries the exemption. + +## Configuring the plugin + +Minimum pipeline for session-budget: + +```yaml +pipeline: + inbound: + plugins: + - name: a2a-parser # REQUIRED — parses contextId → Session.ID + outbound: + plugins: + - name: session-budget + config: + redis_url: "redis://valkey.team1.svc:6379" + max_calls: 3 + max_duration_seconds: 1800 + on_exceed: pause + pause_webhook: "http://pause-webhook-stub.team1.svc.cluster.local" + pause_timeout: 10s + pause_timeout_action: deny + pause_grace_period: 5m + - name: inference-parser # supplies token counts to session-budget +``` + +**`a2a-parser` on inbound is not optional.** Without it, every request +lands in the `default` session bucket (no `Rekey` from `contextId`), so +session-budget can never distinguish sessions and cold-cache hydrate +looks for the wrong key. + +## Try it end-to-end (pause mode) + +Assumes an authbridge-sidecar'd agent is already running in `${NS}` with +`session-budget` (pause mode) + `a2a-parser` inbound configured per +above. Substitute your own agent, session id, and A2A payload. + +```bash +NS=team1 +AGENT_POD=$(kubectl -n "$NS" get pod -l app.kubernetes.io/name= \ + -o jsonpath='{.items[0].metadata.name}') +SESSION=demo-$RANDOM + +# Substitute your Redis/Valkey pod + CLI. E.g. REDIS_POD=valkey and +# REDIS_CLI=valkey-cli, or REDIS_POD=redis-0 and REDIS_CLI=redis-cli. +REDIS_POD=valkey +REDIS_CLI=valkey-cli + +# 1. Seed Redis so this session is already over budget. +kubectl -n "$NS" exec "$REDIS_POD" -- "$REDIS_CLI" HSET \ + "session-budget:$SESSION" calls 99 started_at "$(date +%s)" + +# 2. Fire one A2A request with contextId = seeded session. +# (Any callable agent works; adjust auth + payload to fit yours.) +# $TOKEN is a Keycloak-issued bearer for the agent's inbound audience — +# obtain via the same setup script you use for the rest of your demos +# (see e.g. authbridge/demos/weather-agent). If your agent's inbound +# plugin chain has no jwt-validation, omit the Authorization header. +kubectl -n "$NS" port-forward pod/"$AGENT_POD" 8000:8000 & +curl -sS -X POST http://localhost:8000/ \ + -H "Authorization: Bearer $TOKEN" -H "Content-Type: application/json" \ + -d '{"jsonrpc":"2.0","id":"1","method":"message/send","params":{ + "message":{"messageId":"m1","role":"user", + "parts":[{"kind":"text","text":"hi"}], + "contextId":"'"$SESSION"'"}}}' + +# 3. Confirm the webhook was called with the right session_id. +kubectl -n "$NS" logs deploy/pause-webhook-stub | grep "$SESSION" +``` + +Expected: the request returns 200 (stub approves), and the webhook log +shows one POST body with `"session_id":""` and `"reason": +"call limit reached: ..."`. + +**Try the other modes:** swap `on_exceed: pause` for `deny` or +`observe` in the plugin config, redeploy, and repeat step 2. `deny` +returns 403 with a `budget.exceeded` body once the local cache catches +up (one request may pass first — see the cold-cache note in the +reference doc). `observe` never blocks; grep the authbridge-proxy logs +for `"budget exceeded (shadow mode)"` to see breaches. diff --git a/authbridge/demos/session-budget/k8s/pause-webhook-stub.yaml b/authbridge/demos/session-budget/k8s/pause-webhook-stub.yaml new file mode 100644 index 000000000..878d715eb --- /dev/null +++ b/authbridge/demos/session-budget/k8s/pause-webhook-stub.yaml @@ -0,0 +1,91 @@ +# Minimal pause-mode webhook stub for the session-budget plugin. +# Returns {"action":"approve"} for every POST — use to smoke-test the +# on_exceed: pause flow end-to-end. For deny-path testing, edit the inline +# Python to write {"action":"deny"} instead. +# +# Usage: +# kubectl apply -f pause-webhook-stub.yaml +# # In session-budget config: +# # pause_webhook: "http://pause-webhook-stub.team1.svc.cluster.local" +# +# Logs: +# kubectl logs -n team1 deploy/pause-webhook-stub +--- +apiVersion: apps/v1 +kind: Deployment +metadata: + name: pause-webhook-stub + namespace: team1 + labels: + app: pause-webhook-stub +spec: + replicas: 1 + selector: + matchLabels: + app: pause-webhook-stub + template: + metadata: + labels: + app: pause-webhook-stub + # Opt out of ambient mesh: the plain HTTP stub isn't part of the + # authbridge zero-trust boundary and ambient interception drops + # non-HBONE peers with 'Connection reset'. + istio.io/dataplane-mode: none + annotations: + ambient.istio.io/redirection: disabled + spec: + securityContext: + runAsNonRoot: true + runAsUser: 65532 + runAsGroup: 65532 + seccompProfile: + type: RuntimeDefault + containers: + - name: stub + image: python:3.12-alpine@sha256:d09d15e60962ca365d1cd544a48773bac9d33f2fb1b00f2aa0deec78ade7dc31 + securityContext: + allowPrivilegeEscalation: false + readOnlyRootFilesystem: true + capabilities: + drop: ["ALL"] + command: ["python", "-c"] + args: + - | + import http.server, json, sys + class Handler(http.server.BaseHTTPRequestHandler): + def do_POST(self): + n = int(self.headers.get("Content-Length", 0)) + body = self.rfile.read(n).decode("utf-8", "replace") + print(f"[pause-webhook] POST {self.path} body={body}", flush=True) + resp = b'{"action":"approve"}' + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(resp))) + self.end_headers() + self.wfile.write(resp) + def log_message(self, *a, **k): + pass + http.server.HTTPServer(("", 8080), Handler).serve_forever() + ports: + - name: http + containerPort: 8080 + resources: + requests: + cpu: 10m + memory: 32Mi + limits: + cpu: 100m + memory: 64Mi +--- +apiVersion: v1 +kind: Service +metadata: + name: pause-webhook-stub + namespace: team1 +spec: + selector: + app: pause-webhook-stub + ports: + - name: http + port: 80 + targetPort: 8080 diff --git a/authbridge/docs/plugin-catalog.md b/authbridge/docs/plugin-catalog.md index f533e25c5..ac7c23e33 100644 --- a/authbridge/docs/plugin-catalog.md +++ b/authbridge/docs/plugin-catalog.md @@ -33,8 +33,8 @@ AuthBridge pipeline YAML, not whether it is compiled into the binary | [`opa`](#opa) | OPA policy enforcement for inbound and outbound requests. | Alpha | Both | No | | [`sparc`](#sparc) | Pre-tool reflection: blocks ungrounded/hallucinated tool calls. | Alpha | Outbound | No | | [`static-inject`](#static-inject) | Swaps a placeholder credential for a real static credential on outbound requests. | Alpha | Outbound | No | +| [`session-budget`](#session-budget) | Enforces per-session token, call, and duration budgets via Redis. | Alpha | Outbound | No | | [`token-broker`](#token-broker) | Exchanges incoming tokens against a configured IdP via a broker service. | Alpha | Outbound | No | -| [`token-budget`](#token-budget) | Enforces per-session token, call, and duration budgets via Redis. | Coming Soon | Outbound | No | | [`token-exchange`](#token-exchange) | RFC 8693 outbound token exchange per route. | Ready | Outbound | YES | ## `a2a-parser` @@ -169,6 +169,28 @@ outbound requests, so the workload never holds the real secret. - `placeholder` (string) — if set, the inbound bearer must exactly equal this value before injection proceeds. - `inject_header` (string) — header to inject the credential into. Default `Authorization` (writes `Bearer `); any other value writes the raw credential and drops the inbound `Authorization` header. +## `session-budget` + +Enforces per-session token, call-count, and duration budgets via Redis. Opt-in at build time (`-tags include_plugin_sessionbudget`). + +- `redis_url` (string) — Redis/Valkey connection URL; required. +- `max_tokens` (int64) — cumulative token ceiling per session. `0` = no limit. +- `max_calls` (int64) — max inference calls per session. `0` = no limit. +- `max_duration_seconds` (int64) — wall-clock session lifetime. `0` = no limit. +- `on_exceed` (string) — `deny` (default, block), `observe` (log only), or `pause` (HITL webhook approval). +- `pause_webhook` (string) — URL to POST for approval when `on_exceed=pause`. Required in pause mode. +- `pause_timeout` (string) — how long to wait for webhook response. Default `30s`. +- `pause_timeout_action` (string) — fallback on timeout/error: `deny` (default) or `allow`. +- `pause_grace_period` (string) — suppress repeated webhooks after approval. Default `5m`. +- `session_ttl_seconds` (int) — Redis key TTL; should be ≥ `max_duration_seconds`. Default 7200. +- `refresh_interval` (string) — how often the local cache syncs from Redis. Default `5s`. +- `redis_unavailable` (string) — only `fail_open` (default) is implemented; `fail_closed` is rejected at Configure time. + +Cold-cache behavior is mode-dependent; see +[session-budget-plugin.md](session-budget-plugin.md#cold-cache-behavior) +for details. + + ## `token-broker` Exchanges incoming tokens against a configured IdP through an external @@ -182,21 +204,6 @@ token broker service, per host-based routing rules. - `action` — `broker` (default) or `passthrough`. - `authorization_endpoint` / `token_endpoint` — per-route OAuth endpoint overrides sent to the broker. -## `token-budget` - -Enforces per-session token, call-count, and duration budgets via Redis. -Opt-in at build time (`-tags include_plugin_tokenbudget`). - -- `redis_url` (string) — Redis/Valkey connection URL; required. -- `max_tokens` (int64) — cumulative token ceiling per session. `0` = no limit. -- `max_calls` (int64) — max inference calls per session. `0` = no limit. -- `max_duration_seconds` (int64) — wall-clock session lifetime. `0` = no limit. -- `on_exceed` (string) — `deny` (default, block) or `observe` (log only). -- `session_ttl_seconds` (int) — Redis key TTL; should be ≥ `max_duration_seconds`. Default 7200. -- `refresh_interval` (string) — how often the local cache syncs from Redis. Default `5s`. -- `redis_unavailable` (string) — only `fail_open` (default) is implemented; `fail_closed` is rejected at Configure time. - -Note that enforcement uses a zero-I/O local cache. After a pod restart, the first request can pass while the cache is cold, before Redis refresh restores the counters. ## `token-exchange` diff --git a/authbridge/docs/session-budget-plugin.md b/authbridge/docs/session-budget-plugin.md new file mode 100644 index 000000000..da28ae096 --- /dev/null +++ b/authbridge/docs/session-budget-plugin.md @@ -0,0 +1,258 @@ +# session-budget Plugin + +Enforces per-session budgets on tokens, LLM/inference calls, and wall-clock +duration. `max_calls` counts inference calls only — MCP tool calls, A2A +messages, and other outbound traffic do not count toward it. +Supports three `on_exceed` modes: + +- `deny` — return 403 (default) +- `observe` — shadow mode: log without blocking, useful for calibrating limits +- `pause` — HITL: POST to a webhook for approval before continuing + +A "session" is the AuthBridge session ID (typically one A2A conversation +or agent task invocation). Redis holds durable counters across pods; a +local cache serves the hot path with zero I/O. + +## Build + +Opt-in — build with `-tags include_plugin_sessionbudget`: + +```bash +docker build -f cmd/authbridge-proxy/Dockerfile \ + --build-arg GO_BUILD_TAGS="include_plugin_sessionbudget" \ + -t authbridge:latest . +``` + +Same tag works for `cmd/authbridge-envoy/Dockerfile`. + +## Configuration + +```yaml +pipeline: + outbound: + plugins: + - name: token-exchange + config: { ... } + - name: session-budget + config: + redis_url: "redis://valkey.infra.svc:6379" + max_tokens: 50000 + max_calls: 100 + max_duration_seconds: 1800 + - name: inference-parser +``` + +| Field | Default | Description | +|-------|---------|-------------| +| `redis_url` | — (required) | Redis/Valkey URL | +| `max_tokens` | 0 | Token ceiling (0 = no limit) | +| `max_calls` | 0 | LLM/inference call cap. Counts calls surfaced by `inference-parser`; MCP, A2A, and other outbound traffic do NOT count toward the counter. 0 = no limit. **The limit check runs on every outbound request**, though — once the LLM counter crosses `max_calls`, the next outbound of any kind (MCP tool call, A2A message, etc.) is the one that gets the 403. | +| `max_duration_seconds` | 0 | Session lifetime cap (0 = no limit) | +| `on_exceed` | `deny` | `deny` (403), `observe` (log only), or `pause` (webhook) | +| `pause_webhook` | — | URL to POST on breach (required when `on_exceed=pause`) | +| `pause_timeout` | `30s` | Max wait for webhook response | +| `pause_timeout_action` | `deny` | Fallback on timeout/error: `deny` or `allow` | +| `pause_grace_period` | `5m` | Suppress repeat webhooks after approval. `0s` disables the grace window (webhook fires on every breach). | +| `session_ttl_seconds` | 7200 | Redis key TTL; must be ≥ `max_duration_seconds` | +| `refresh_interval` | `5s` | Local-cache sync interval | +| `redis_unavailable` | `fail_open` | Only `fail_open` supported today | + +At least one of `max_tokens`, `max_calls`, `max_duration_seconds` must be > 0. + +**Pipeline position:** must appear **before** `inference-parser` on the outbound +pipeline. Both must be present for token counting (inference-parser supplies the +token counts session-budget accumulates). + +## Modes + +### `deny` (default) + +Returns 403 with a JSON body: + +```json +{ + "error": "budget.exceeded", + "message": "token limit reached: 50200/50000", + "details": { + "spent_tokens": 50200, + "spent_calls": 42, + "token_limit": 50000, + "call_limit": 100, + "duration_seconds": 1205, + "duration_limit": 1800 + } +} +``` + +`duration_seconds` / `duration_limit` are included only when +`max_duration_seconds` is set. + +### `observe` (shadow mode) + +Counters still accumulate and limits are still evaluated, but breaches only emit +a WARN log (`"budget exceeded (shadow mode)"`) and the request continues. Use to +calibrate limits before enforcing: + +1. Deploy with `on_exceed: observe` and conservative limits. +2. Watch logs for shadow-mode entries. +3. Adjust `max_tokens` / `max_calls` / `max_duration_seconds` to fit real + workloads. +4. Flip to `on_exceed: deny` (or `pause`) once confident. + +Call accounting is response-driven: `max_calls` is checked against the +count of completed inference responses, so under bursty concurrency +`max_calls` breach logs may lag actual in-flight calls by up to one per +concurrent request. + +### `pause` (HITL webhook) + +On breach, POST to `pause_webhook` and block the request until the webhook +responds or `pause_timeout` fires. + +The webhook is any HTTP endpoint you build that speaks the contract +below — an in-cluster Service, a workflow entrypoint (Temporal, Argo, +Slack middleware), or a local stub. It must be reachable from the +AuthBridge pod and respond within `pause_timeout` (default `30s`), or +the plugin falls back to `pause_timeout_action`. + +**Contract:** + +| Aspect | Requirement | +|--------|-------------| +| Method | `POST` | +| URL | Exactly the `pause_webhook` value (no path templating) | +| Request `Content-Type` | `application/json` | +| Request body | See below — always the same schema | +| Success response | HTTP `200` with `application/json` body containing an `action` field | +| Response body cap | 4 KiB (larger responses are truncated at decode) | +| Latency budget | Must respond within `pause_timeout`; slow webhooks block the caller | +| Auth | None injected by the plugin — add your own (mTLS, network policy, IP allowlist) at the transport layer | +| Retries | None — the plugin calls once per breach | + +**Request body:** +```json +{ + "session_id": "abc-123", + "reason": "call limit reached: 50/50", + "spent_tokens": 48200, + "spent_calls": 50, + "token_limit": 100000, + "call_limit": 50, + "duration_seconds": 1205, + "duration_limit": 1800 +} +``` + +**Expected response:** +```json +{"action": "approve"} +``` +or +```json +{"action": "deny", "reason": "operator rejected"} +``` + +**On approval:** the request continues; subsequent requests from the +same session skip the webhook for `pause_grace_period` (default `5m`, +pod-local — each pod fires one webhook before its own grace kicks in). +Concurrent breaches during an in-flight webhook wait on the pending +call and honor its outcome — all approved together, or all denied +together. + +**On timeout / non-200 / bad JSON / unreachable:** falls back to +`pause_timeout_action` (`deny` returns 403; `allow` continues). If your +webhook can be unhealthy and `pause_timeout_action: deny`, an outage +turns budget breaches into hard 403s. + +If a human is in the loop, bump `pause_timeout` to minutes so the +request can wait for a real approval decision. + +## Failure Modes + +| Scenario | Behavior | +|----------|----------| +| Redis down at startup | Fail-open until refresh populates cache | +| Redis fails mid-session | Local cache keeps enforcing; writes dropped | +| Pod restart, `pause` | Request #1 hydrates from Redis synchronously | +| Pod restart, `deny` / `observe` | Request #1 skips (`cold_cache`); see below | +| Webhook unreachable | Falls back to `pause_timeout_action` | + +### Cold-cache behavior + +When a request arrives with no local cache entry for the session: + +- **`pause`** — hydrates from Redis synchronously, so an over-budget + session fires the webhook on request #1. +- **`deny` / `observe`** — skip with `reason=cold_cache` and continue. + Counters populate as inference responses stream back and via the + background refresh loop. A pre-existing over-budget session may pass + **up to one request per pod** before enforcement resumes. Keeps Redis + off the hot path. + +Redis unavailability degrades enforcement to local-cache-only rather +than blocking requests. + +**Token counting requires `usage` in provider responses.** Providers that omit +`usage` from streaming chunks (e.g. Anthropic via LiteLLM) will show +`promptTokens=0` in inference-parser logs — `max_tokens` enforcement won't +trigger, but `max_calls` and `max_duration_seconds` still apply. Ollama, +OpenAI, and Azure OpenAI include usage in streaming responses and work fully. + +## Redis Keys + +```text +session-budget: (Hash, TTL = session_ttl_seconds) + tokens cumulative token count + calls inference call count + started_at first-call unix timestamp +``` + +## Local Development + +**Redis / Valkey:** + +```bash +docker run -d --name valkey -p 6379:6379 valkey/valkey:latest +# redis_url: redis://localhost:6379 (or host.docker.internal from a container) +``` + +**Pause-mode webhook stub.** A one-liner that returns `approve` for every +request — enough to smoke-test `on_exceed: pause` end-to-end: + +```bash +docker run -d --name pause-webhook -p 8888:8888 \ + python:3.12-alpine@sha256:d09d15e60962ca365d1cd544a48773bac9d33f2fb1b00f2aa0deec78ade7dc31 \ + python -c "import http.server,json; \ +h=type('H',(http.server.BaseHTTPRequestHandler,),{ \ +'do_POST':lambda s:(s.send_response(200), \ +s.send_header('Content-Type','application/json'),s.end_headers(), \ +s.wfile.write(b'{\"action\":\"approve\"}'))}); \ +http.server.HTTPServer(('',8888),h).serve_forever()" + +# pause_webhook: http://localhost:8888 (or http://host.docker.internal:8888) +``` + +Swap `approve` for `deny` to test the reject path. Logs land in +`docker logs pause-webhook`. + +## In-cluster deployment note + +**If your namespace runs Istio ambient mesh** (label +`istio.io/dataplane-mode: ambient`), the Valkey pod and any plain-HTTP +pause webhook need to opt out with the pod-level label +`istio.io/dataplane-mode: none`. Ambient's ztunnel only accepts HBONE +(HTTP/2 CONNECT), and Redis RESP is raw TCP — the connection is closed +before reaching Valkey. Symptom: `session-budget action=skip +reason=cold_cache` on every request even though the Redis key exists. + +For an in-cluster stub, apply +`authbridge/demos/session-budget/k8s/pause-webhook-stub.yaml` and set +`pause_webhook: http://pause-webhook-stub.team1.svc.cluster.local`. See +[`../demos/session-budget/README.md`](../demos/session-budget/README.md). + +**Run the plugin tests:** + +```bash +cd authbridge/authlib +go test ./plugins/sessionbudget/... -v -count=1 +``` diff --git a/authbridge/docs/token-budget-plugin.md b/authbridge/docs/token-budget-plugin.md deleted file mode 100644 index d2bbaab2f..000000000 --- a/authbridge/docs/token-budget-plugin.md +++ /dev/null @@ -1,141 +0,0 @@ -# token-budget Plugin - -Enforces per-session lifetime budgets on tokens, inference calls, and -wall-clock duration. Supports `observe` (shadow — log without blocking) and -`deny` (return 403) modes. Uses Redis for cross-pod durable counters; -evaluates limits from a local cache with zero I/O on the hot path. - -A "session" maps to the AuthBridge session ID (typically one A2A conversation -or agent task invocation). - -## Build Tag - -This plugin is **opt-IN**. Build with `-tags include_plugin_tokenbudget` -to include it (and its `storage/redis` dependency) in the binary: - -```bash -cd authbridge -docker build -f cmd/authbridge-proxy/Dockerfile \ - --build-arg GO_BUILD_TAGS="include_plugin_tokenbudget" \ - -t authbridge:latest . -``` - -Without the tag, neither token-budget nor go-redis are linked. - -The same build tag works for the envoy-sidecar image: - -```bash -docker build -f cmd/authbridge-envoy/Dockerfile \ - --build-arg GO_BUILD_TAGS="include_plugin_tokenbudget" \ - -t authbridge-envoy:latest . -``` - -## Configuration - -```yaml -pipeline: - outbound: - plugins: - - name: token-exchange - config: { ... } - - name: token-budget - config: - redis_url: "redis://valkey.infra.svc:6379" - max_tokens: 50000 - max_calls: 100 - max_duration_seconds: 1800 - - name: inference-parser -``` - -| Field | Required | Default | Description | -|-------|----------|---------|-------------| -| `redis_url` | yes | — | Redis/Valkey connection URL | -| `max_tokens` | no | 0 | Cumulative token ceiling per session (0 = no limit) | -| `max_calls` | no | 0 | Max inference calls per session (0 = no limit) | -| `max_duration_seconds` | no | 0 | Wall-clock session lifetime in seconds (0 = no limit) | -| `on_exceed` | no | "deny" | `deny` (block with 403) or `observe` (shadow — log but continue) | -| `session_ttl_seconds` | no | 7200 | Redis key TTL; should be >= `max_duration_seconds` | -| `refresh_interval` | no | "5s" | How often to sync local cache from Redis | -| `redis_unavailable` | no | "fail_open" | Only `fail_open` supported (stale cache retained on failure). `fail_closed` reserved. | - -At least one of `max_tokens`, `max_calls`, or `max_duration_seconds` must be > 0. - -**Local Redis/Valkey** for development: `docker run -d --name valkey -p 6379:6379 valkey/valkey:latest`. -Use `redis://localhost:6379` or `redis://host.docker.internal:6379` (container mode). - -## Shadow Mode - -Set `on_exceed: "observe"` to run the plugin in shadow mode. The plugin -still accumulates counters and evaluates limits, but instead of blocking -requests it logs a WARN and continues the pipeline. Use this to calibrate -limits under real workloads before enabling enforcement. - -Rollout workflow: -1. Deploy with `on_exceed: "observe"` and conservative limits -2. Monitor logs for `"budget exceeded (shadow mode)"` entries -3. Adjust `max_tokens` / `max_calls` / `max_duration_seconds` based on observed patterns -4. Flip to `on_exceed: "deny"` when confident in the thresholds - -## Pipeline Position - -Must be declared **before** `inference-parser` in the outbound plugin list. -The response path runs in reverse order, so inference-parser finalizes token -counts first, then token-budget reads them. Both plugins implement -`StreamingResponder` so they work with streaming SSE responses (Ollama, -LiteLLM, OpenAI) and buffered JSON responses alike. - -## Response Format (403) - -```json -{ - "error": "budget.exceeded", - "message": "token limit reached: 50200/50000", - "details": { - "spent_tokens": 50200, - "spent_calls": 42, - "token_limit": 50000, - "call_limit": 100, - "duration_seconds": 1205, - "duration_limit": 1800 - } -} -``` - -`duration_seconds` and `duration_limit` are included only when `max_duration_seconds` is configured. - -## Redis Key Schema - -```text -token-budget: (Hash, TTL = session_ttl_seconds) - tokens int cumulative TotalTokens - calls int inference call count - started_at unix first-call timestamp (set-if-not-exists) -``` - -## Failure Modes - -| Scenario | Behavior | -|----------|----------| -| Redis down at startup | `Init` succeeds (no connectivity check); enforcement fail-open until first refresh populates cache | -| Redis fails mid-session | Local cache continues enforcing; writes dropped silently | -| Pod restarts | First request passes (cold cache); refresh picks up Redis counters within one interval | -| Provider returns no usage data | `max_tokens` not enforced; `max_calls` and `max_duration_seconds` still work | - -**Fail-open guarantee:** The plugin never blocks requests due to its own infrastructure failures. Redis unavailability degrades enforcement (local cache only, no cross-pod consistency) but never causes false denials. - -**Note on token counting:** Token accumulation requires the LLM provider to -return `usage` (prompt/completion token counts) in responses. Providers that -omit usage from streaming chunks (e.g. Anthropic via LiteLLM) will show -`promptTokens=0` in inference-parser logs — `max_tokens` enforcement won't -trigger for these providers, but `max_calls` and `max_duration_seconds` still -apply. Ollama, OpenAI, and Azure OpenAI include usage in streaming responses -and work fully. - -## Testing - -```bash -cd authbridge/authlib -go test ./plugins/tokenbudget/... -v -count=1 -``` - -No external dependencies — tests use in-memory stores.