Skip to content

fix: 🐛 Fallback for default session ID #807

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
evaline-ju wants to merge 4 commits into
base: main
Choose a base branch
from
+92 −14
Open

fix: 🐛 Fallback for default session ID #807

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
8e7130c
:bug: Fallback for default session ID
evaline-ju Aug 25, 2026
491bcb6
:white_check_mark: Add test for default session
evaline-ju Aug 25, 2026
84e6616
:wrench::white_check_mark: Gate fallback
evaline-ju Aug 25, 2026
1abbe1a
:art: Format struct
evaline-ju Aug 25, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 20 additions & 12 deletions authbridge/authlib/plugins/sessionbudget/plugin.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 ""
}

Expand Down
74 changes: 72 additions & 2 deletions authbridge/authlib/plugins/sessionbudget/plugin_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)

Expand Down Expand Up @@ -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)
}
Expand Down Expand Up @@ -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) {

@clawgenti clawgenti Aug 25, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: The test name TestOnResponseFrame_NoSession_UsesDefaultBucket is clear, but consider also checking that the redisKey round-trip works correctly for DefaultSessionID if that path is exercised in integration tests. Not blocking — the unit test assertions here are solid.

Comment thread
evaline-ju marked this conversation as resolved.
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) {
Expand Down
Loading