From 8e7130c4f7bebb6de7070d2d9bc57b4ba313b568 Mon Sep 17 00:00:00 2001 From: Evaline Ju <69598118+evaline-ju@users.noreply.github.com> Date: Tue, 25 Aug 2026 13:25:06 -0600 Subject: [PATCH 1/4] :bug: Fallback for default session ID Signed-off-by: Evaline Ju <69598118+evaline-ju@users.noreply.github.com> --- authbridge/authlib/plugins/sessionbudget/plugin.go | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/authbridge/authlib/plugins/sessionbudget/plugin.go b/authbridge/authlib/plugins/sessionbudget/plugin.go index da069c49..a9faee8b 100644 --- a/authbridge/authlib/plugins/sessionbudget/plugin.go +++ b/authbridge/authlib/plugins/sessionbudget/plugin.go @@ -18,6 +18,7 @@ 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" ) @@ -665,7 +666,11 @@ func (p *SessionBudget) sessionID(pctx *pipeline.Context) string { if pctx.Session != nil && pctx.Session.ID != "" { return pctx.Session.ID } - return "" + // Forward-proxy egress with no prior inbound A2A leaves pctx.Session + // nil. Fall back to the well-known default bucket so single-workload + // demos and any traffic ahead of the first inbound request still get + // budgeted, matching the listener's own DefaultSessionID fallback. + return session.DefaultSessionID } func (p *SessionBudget) redisKey(sessionID string) string { From 491bcb6e20e070fe6c9e6f5d474bd98fa576ddd1 Mon Sep 17 00:00:00 2001 From: Evaline Ju <69598118+evaline-ju@users.noreply.github.com> Date: Tue, 25 Aug 2026 14:41:23 -0600 Subject: [PATCH 2/4] :white_check_mark: Add test for default session Signed-off-by: Evaline Ju <69598118+evaline-ju@users.noreply.github.com> --- .../plugins/sessionbudget/plugin_test.go | 31 +++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/authbridge/authlib/plugins/sessionbudget/plugin_test.go b/authbridge/authlib/plugins/sessionbudget/plugin_test.go index 2dac6052..b3e0e725 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" ) @@ -261,6 +262,36 @@ func TestOnRequest_NoSession(t *testing.T) { } } +// Pins the default-session fallback on the response path: forward-proxy egress +// with no inbound A2A leaves pctx.Session nil, and the plugin must still +// accumulate under session.DefaultSessionID so budgets are enforced instead of +// silently skipped. +func TestOnResponseFrame_NoSession_UsesDefaultBucket(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() + 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) + } +} + func TestAccumulate_WritesToStore(t *testing.T) { p := newTestPlugin(1000, 0, 0) store := newMemStore() From 84e6616f3d5343a35da46ca4d0898416029a5b3d Mon Sep 17 00:00:00 2001 From: Evaline Ju <69598118+evaline-ju@users.noreply.github.com> Date: Tue, 25 Aug 2026 15:53:26 -0600 Subject: [PATCH 3/4] :wrench::white_check_mark: Gate fallback Signed-off-by: Evaline Ju <69598118+evaline-ju@users.noreply.github.com> --- .../authlib/plugins/sessionbudget/plugin.go | 13 +++-- .../plugins/sessionbudget/plugin_test.go | 53 ++++++++++++++++--- 2 files changed, 54 insertions(+), 12 deletions(-) diff --git a/authbridge/authlib/plugins/sessionbudget/plugin.go b/authbridge/authlib/plugins/sessionbudget/plugin.go index a9faee8b..57a78baf 100644 --- a/authbridge/authlib/plugins/sessionbudget/plugin.go +++ b/authbridge/authlib/plugins/sessionbudget/plugin.go @@ -36,6 +36,7 @@ type config struct { 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 @@ -666,11 +667,13 @@ func (p *SessionBudget) sessionID(pctx *pipeline.Context) string { if pctx.Session != nil && pctx.Session.ID != "" { return pctx.Session.ID } - // Forward-proxy egress with no prior inbound A2A leaves pctx.Session - // nil. Fall back to the well-known default bucket so single-workload - // demos and any traffic ahead of the first inbound request still get - // budgeted, matching the listener's own DefaultSessionID fallback. - return session.DefaultSessionID + // 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 "" } func (p *SessionBudget) redisKey(sessionID string) string { diff --git a/authbridge/authlib/plugins/sessionbudget/plugin_test.go b/authbridge/authlib/plugins/sessionbudget/plugin_test.go index b3e0e725..cfb3a6d3 100644 --- a/authbridge/authlib/plugins/sessionbudget/plugin_test.go +++ b/authbridge/authlib/plugins/sessionbudget/plugin_test.go @@ -146,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) } @@ -260,14 +265,26 @@ 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) + } } -// Pins the default-session fallback on the response path: forward-proxy egress -// with no inbound A2A leaves pctx.Session nil, and the plugin must still -// accumulate under session.DefaultSessionID so budgets are enforced instead of -// silently skipped. +// With default_session_fallback enabled, sessionless response frames +// accumulate under session.DefaultSessionID instead of being skipped. func TestOnResponseFrame_NoSession_UsesDefaultBucket(t *testing.T) { - p := newTestPlugin(1000, 0, 0) + p := newTestPluginWithFallback(1000, 0, 0, true) pctx := &pipeline.Context{ Direction: pipeline.Outbound, Headers: http.Header{}, @@ -292,6 +309,28 @@ func TestOnResponseFrame_NoSession_UsesDefaultBucket(t *testing.T) { } } +// 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) { p := newTestPlugin(1000, 0, 0) store := newMemStore() From 1abbe1ae37d51b8ad2cfcf3480e0e42ccff7f643 Mon Sep 17 00:00:00 2001 From: Evaline Ju <69598118+evaline-ju@users.noreply.github.com> Date: Tue, 25 Aug 2026 16:00:00 -0600 Subject: [PATCH 4/4] :art: Format struct Signed-off-by: Evaline Ju <69598118+evaline-ju@users.noreply.github.com> --- .../authlib/plugins/sessionbudget/plugin.go | 26 +++++++++---------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/authbridge/authlib/plugins/sessionbudget/plugin.go b/authbridge/authlib/plugins/sessionbudget/plugin.go index 57a78baf..0cd5cfcb 100644 --- a/authbridge/authlib/plugins/sessionbudget/plugin.go +++ b/authbridge/authlib/plugins/sessionbudget/plugin.go @@ -24,19 +24,19 @@ import ( ) 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"` - DefaultSessionFallback bool `json:"default_session_fallback" description:"Pool sessionless traffic into a shared 'default' bucket. Off by default. Single-workload only." default:"false"` + 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