diff --git a/authbridge/authlib/plugins/sessionbudget/plugin.go b/authbridge/authlib/plugins/sessionbudget/plugin.go index da069c490..0cd5cfcb0 100644 --- a/authbridge/authlib/plugins/sessionbudget/plugin.go +++ b/authbridge/authlib/plugins/sessionbudget/plugin.go @@ -18,23 +18,25 @@ import ( "github.com/rossoctl/cortex/authbridge/authlib/pipeline" "github.com/rossoctl/cortex/authbridge/authlib/plugins" + "github.com/rossoctl/cortex/authbridge/authlib/session" "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. Only inference-parser output increments this counter; MCP tool calls and other outbound traffic do not. Once the limit is reached, all subsequent outbound requests (including MCP tool calls) are blocked until the session resets. 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"` + 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. Only inference-parser output increments this counter; MCP tool calls and other outbound traffic do not. Once the limit is reached, all subsequent outbound requests (including MCP tool calls) are blocked until the session resets. 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"` + DefaultSessionFallback bool `json:"default_session_fallback" description:"Pool sessionless traffic into a shared 'default' bucket. Off by default. Single-workload only." default:"false"` } // approvalFlight carries the outcome of one webhook call. The leader writes @@ -665,6 +667,12 @@ func (p *SessionBudget) sessionID(pctx *pipeline.Context) string { if pctx.Session != nil && pctx.Session.ID != "" { return pctx.Session.ID } + // Opt-in fallback for single-workload deployments where all sessionless + // egress should share one bucket. Off by default; callers with no + // session then skip enforcement (existing no_session_id path). + if p.cfg.DefaultSessionFallback { + return session.DefaultSessionID + } return "" } diff --git a/authbridge/authlib/plugins/sessionbudget/plugin_test.go b/authbridge/authlib/plugins/sessionbudget/plugin_test.go index 2dac6052b..cfb3a6d34 100644 --- a/authbridge/authlib/plugins/sessionbudget/plugin_test.go +++ b/authbridge/authlib/plugins/sessionbudget/plugin_test.go @@ -14,6 +14,7 @@ import ( "time" "github.com/rossoctl/cortex/authbridge/authlib/pipeline" + "github.com/rossoctl/cortex/authbridge/authlib/session" "github.com/rossoctl/cortex/authbridge/authlib/storage" ) @@ -145,14 +146,19 @@ func init() { } func newTestPlugin(maxTokens, maxCalls, maxDuration int64) *SessionBudget { + return newTestPluginWithFallback(maxTokens, maxCalls, maxDuration, false) +} + +func newTestPluginWithFallback(maxTokens, maxCalls, maxDuration int64, defaultSessionFallback bool) *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) + "refresh_interval": "100ms", + "default_session_fallback": %t + }`, maxTokens, maxCalls, maxDuration, defaultSessionFallback) if err := p.Configure(json.RawMessage(cfg)); err != nil { panic(err) } @@ -259,6 +265,70 @@ func TestOnRequest_NoSession(t *testing.T) { if action.Type != pipeline.Continue { t.Fatalf("expected Continue for nil session, got %v", action.Type) } + + // Pin the mechanism, not just the outcome: with fallback off (the + // default), sessionless traffic must record a no_session_id skip so + // operators can distinguish it from real-session traffic. + if pctx.Extensions.Invocations == nil { + t.Fatal("expected invocation recorded, got none") + } + out := pctx.Extensions.Invocations.Outbound + if len(out) != 1 { + t.Fatalf("expected 1 outbound invocation, got %d", len(out)) + } + if out[0].Action != pipeline.ActionSkip || out[0].Reason != "no_session_id" { + t.Errorf("invocation = {%s, %s}, want {skip, no_session_id}", out[0].Action, out[0].Reason) + } +} + +// With default_session_fallback enabled, sessionless response frames +// accumulate under session.DefaultSessionID instead of being skipped. +func TestOnResponseFrame_NoSession_UsesDefaultBucket(t *testing.T) { + p := newTestPluginWithFallback(1000, 0, 0, true) + pctx := &pipeline.Context{ + Direction: pipeline.Outbound, + Headers: http.Header{}, + Extensions: pipeline.Extensions{ + Inference: &pipeline.InferenceExtension{TotalTokens: 42}, + }, + } + + p.OnResponseFrame(context.Background(), pctx, nil, true) + + p.mu.RLock() + c := p.cache[session.DefaultSessionID] + p.mu.RUnlock() + if c == nil { + t.Fatalf("expected cache entry under %q, got none", session.DefaultSessionID) + } + if c.tokens != 42 { + t.Errorf("tokens = %d, want 42", c.tokens) + } + if c.calls != 1 { + t.Errorf("calls = %d, want 1", c.calls) + } +} + +// With default_session_fallback disabled (the default), sessionless response +// frames are skipped and nothing accumulates. +func TestOnResponseFrame_NoSession_FallbackOff_Skips(t *testing.T) { + p := newTestPlugin(1000, 0, 0) + pctx := &pipeline.Context{ + Direction: pipeline.Outbound, + Headers: http.Header{}, + Extensions: pipeline.Extensions{ + Inference: &pipeline.InferenceExtension{TotalTokens: 42}, + }, + } + + p.OnResponseFrame(context.Background(), pctx, nil, true) + + p.mu.RLock() + _, ok := p.cache[session.DefaultSessionID] + p.mu.RUnlock() + if ok { + t.Errorf("expected no cache entry under %q when fallback is off", session.DefaultSessionID) + } } func TestAccumulate_WritesToStore(t *testing.T) { diff --git a/authbridge/docs/session-budget-plugin.md b/authbridge/docs/session-budget-plugin.md index 42634e46f..005969c91 100644 --- a/authbridge/docs/session-budget-plugin.md +++ b/authbridge/docs/session-budget-plugin.md @@ -56,6 +56,7 @@ pipeline: | `session_ttl_seconds` | 7200 | Redis key TTL. Must be ≥ `max_duration_seconds` when the latter is set (rejected at Configure time otherwise). | | `refresh_interval` | `5s` | Local-cache sync interval | | `redis_unavailable` | `fail_open` | Only `fail_open` supported today | +| `default_session_fallback` | `false` | Pool sessionless traffic into a shared `"default"` bucket. Single-workload only — one caller exhausting the budget denies the rest. Under `max_duration_seconds`, continuous traffic refreshes the TTL, so once elapsed exceeds the limit requests stay denied until the key expires or is deleted. | At least one of `max_tokens`, `max_calls`, `max_duration_seconds` must be > 0.