feat: ✨ Session budget enforcement plugin - #723
Conversation
Implements per-session lifetime budgets on tokens, inference calls, and wall-clock duration. Uses Redis for cross-pod durable counters with a local in-memory cache for zero-latency evaluation on the request path. Fail-open by default when Redis is unavailable. Includes unit tests, e2e tests (forward proxy round-trip, accumulate and deny, multi-session isolation, Redis failure/recovery, pod restart), storage interface with driver registry, Redis driver, and user-facing plugin documentation. Assisted-By: Claude (Anthropic AI) <noreply@anthropic.com> Signed-off-by: Evaline Ju <69598118+evaline-ju@users.noreply.github.com>
Signed-off-by: Evaline Ju <69598118+evaline-ju@users.noreply.github.com>
Signed-off-by: Evaline Ju <69598118+evaline-ju@users.noreply.github.com>
Signed-off-by: Evaline Ju <69598118+evaline-ju@users.noreply.github.com>
Signed-off-by: Evaline Ju <69598118+evaline-ju@users.noreply.github.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (6)
🚧 Files skipped from review as they are similar to previous changes (5)
📝 WalkthroughWalkthroughAdds an opt-in token-budget plugin with Redis-backed per-session token, call, and duration limits. It supports enforcement and observe modes, cache refresh, failure handling, build integration, documentation, and comprehensive tests. ChangesToken budget plugin
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Pipeline
participant TokenBudget
participant Redis
Pipeline->>TokenBudget: Check cached session budget
TokenBudget-->>Pipeline: Allow or budget.exceeded denial
Pipeline->>TokenBudget: Process final response frame
TokenBudget->>Redis: Persist token and call counters
TokenBudget->>TokenBudget: Refresh cached session state
Suggested reviewers: 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 13
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@authbridge/authlib/listener/forwardproxy/server.go`:
- Around line 610-615: Update the finalization flow around RunResponseFrame to
wrap context.WithoutCancel(r.Context()) with a short context.WithTimeout before
dispatch. Use the resulting context for the final last=true call and ensure its
cancel function is released after dispatch, while preserving detached
cancellation behavior.
In `@authbridge/authlib/plugins/tokenbudget/plugin.go`:
- Around line 255-261: The RedisUnavailable configuration accepts fail_closed
even though refresh failures still allow cache-miss sessions to continue. In the
configuration validation path, reject the fail_closed value until enforcement is
implemented, while preserving supported modes such as fail_open and the existing
refresh behavior.
- Around line 72-103: Validate p.cfg.RefreshInterval in TokenBudget.Configure by
parsing it with time.ParseDuration and returning a configuration error when the
duration is zero or negative; preserve the existing default and valid-duration
behavior. Add validation tests covering "0s" and "-1s", ensuring both are
rejected before Init can start refreshLoop.
- Around line 166-183: Update the inference handling around
pctx.Extensions.Inference to return early only when inf is nil. For completed
responses with TotalTokens equal to zero, retain the zero token delta while
still initializing the session counters, incrementing calls, and preserving the
existing cache update flow.
- Around line 126-128: Update the cache-miss path around the visible !ok check
in the token-budget request flow to load the session’s persisted counters from
Redis before returning pipeline.Continue; do not let an uncached session bypass
budget enforcement, while preserving normal refresh behavior for existing keys.
Add or update the preload logic needed for first requests and revise
TestE2E_PodRestart to assert that over-budget sessions remain blocked after a
replica restart.
- Around line 122-145: The cache lookup in the token-budget evaluation path must
use a stable counter snapshot. In the code around p.evaluate and the shadow-mode
log, copy tokens, calls, and startedAt from the cached counters while holding
p.mu.RLock(), then unlock and evaluate/log the copied value instead of the
shared c instance.
- Around line 106-110: Update TokenBudget.Shutdown and the accumulate goroutine
to track pending persistence work, wait for all queued writes using a bounded
shutdown context, and only then close p.store. Ensure shutdown proceeds after
the timeout while preventing accumulate from writing against a closed store.
In `@authbridge/cmd/authbridge-envoy/go.mod`:
- Around line 5-13: Align the authbridge-envoy Docker build context with the
storage/redis module replacement in go.mod by ensuring the Dockerfile makes
storage/ available at the expected relative path, either through a matching COPY
storage/ storage/ directive or an equivalent build context that includes it.
Preserve the existing authlib and authbridge-envoy build inputs.
In `@authbridge/docs/token-budget-plugin.md`:
- Around line 11-19: Update token-budget plugin documentation to clarify its
deployment scope: either add the corresponding Envoy Docker build procedure
using include_plugin_tokenbudget, or explicitly state that the plugin is
supported only by the proxy and not Envoy. Ensure the documented guidance
matches the behavior enabled by plugins_tokenbudget.go.
- Around line 104-105: Update the token-budget enforcement flow documented in
the pod-restart and fail_closed rows so requests cannot bypass exhausted
Redis-backed session budgets: load current session state before admitting the
request, and when refresh fails under fail_closed, deny the request rather than
retaining stale state. If stale-cache enforcement is intentional, rename
fail_closed and document that behavior instead.
- Around line 91-96: Add the `text` language tag to the schema code fence in the
token budget documentation, changing the opening fence to specify `text` while
leaving the schema content unchanged.
- Line 102: Update the Redis down at startup row in the token-budget behavior
table to reflect that token-budget.Init succeeds without connectivity
validation, while failure occurs later during lazy access or refresh operations.
Remove the incorrect Pod Init error outcome and describe the actual refresh-path
behavior.
In `@authbridge/storage/redis/redis_test.go`:
- Line 18: Update the Redis tests to check and fail on the error returned by
c.Close in t.Cleanup, and validate every HashIncr result before proceeding,
especially before calling Expire in TestExpire. Ensure a failed HashIncr cannot
allow Expire to continue as a no-op and produce a false passing test.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 4f19bf22-3c03-4cee-a995-66532013ca4b
⛔ Files ignored due to path filters (4)
authbridge/cmd/authbridge-envoy/go.sumis excluded by!**/*.sumauthbridge/cmd/authbridge-proxy/go.sumis excluded by!**/*.sumauthbridge/go.workis excluded by!**/*.workauthbridge/storage/redis/go.sumis excluded by!**/*.sum
📒 Files selected for processing (18)
authbridge/authlib/listener/forwardproxy/server.goauthbridge/authlib/pipeline/action.goauthbridge/authlib/plugins/tokenbudget/e2e_test.goauthbridge/authlib/plugins/tokenbudget/lifecycle_test.goauthbridge/authlib/plugins/tokenbudget/plugin.goauthbridge/authlib/plugins/tokenbudget/plugin_test.goauthbridge/authlib/storage/provider.goauthbridge/authlib/storage/provider_test.goauthbridge/authlib/storage/store.goauthbridge/cmd/authbridge-envoy/go.modauthbridge/cmd/authbridge-envoy/plugins_tokenbudget.goauthbridge/cmd/authbridge-proxy/Dockerfileauthbridge/cmd/authbridge-proxy/go.modauthbridge/cmd/authbridge-proxy/plugins_tokenbudget.goauthbridge/docs/token-budget-plugin.mdauthbridge/storage/redis/go.modauthbridge/storage/redis/redis.goauthbridge/storage/redis/redis_test.go
| // 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 | ||
| // need their last=true dispatch to finalize state. | ||
| finalCtx := context.WithoutCancel(r.Context()) | ||
| finalAction := s.OutboundPipeline.RunResponseFrame(finalCtx, pctx, nil, true) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Bound detached response finalization.
context.WithoutCancel removes both cancellation and the request deadline. If a streaming responder blocks during RunResponseFrame, this handler can block indefinitely after the client disconnects.
Wrap the detached context with a short context.WithTimeout before final dispatch.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@authbridge/authlib/listener/forwardproxy/server.go` around lines 610 - 615,
Update the finalization flow around RunResponseFrame to wrap
context.WithoutCancel(r.Context()) with a short context.WithTimeout before
dispatch. Use the resulting context for the final last=true call and ensure its
cancel function is released after dispatch, while preserving detached
cancellation behavior.
| func (p *TokenBudget) Shutdown(_ context.Context) error { | ||
| close(p.stopCh) | ||
| <-p.stopped | ||
| if p.store != nil { | ||
| return p.store.Close() |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
Drain persistence work before closing the store.
Line 173 starts an untracked accumulate goroutine. Shutdown waits only for refreshLoop, then closes p.store at line 110. A pending accumulation can run after close and lose a completed request's counters.
Track persistence tasks and wait for them with a bounded shutdown context. Do not close the store until queued writes finish or time out.
Also applies to: 173-173
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@authbridge/authlib/plugins/tokenbudget/plugin.go` around lines 106 - 110,
Update TokenBudget.Shutdown and the accumulate goroutine to track pending
persistence work, wait for all queued writes using a bounded shutdown context,
and only then close p.store. Ensure shutdown proceeds after the timeout while
preventing accumulate from writing against a closed store.
| if !ok { | ||
| return pipeline.Action{Type: pipeline.Continue} | ||
| } |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
Load persisted counters before allowing a cache-miss request.
Line 127 allows every session that is absent from the local cache. refreshCache only refreshes existing cache keys. A fresh replica therefore allows requests for a session that already exceeded its Redis budget.
Load the session counters on the first request, or maintain a safe preload mechanism. Update TestE2E_PodRestart so it does not accept this enforcement bypass.
Also applies to: 241-281
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@authbridge/authlib/plugins/tokenbudget/plugin.go` around lines 126 - 128,
Update the cache-miss path around the visible !ok check in the token-budget
request flow to load the session’s persisted counters from Redis before
returning pipeline.Continue; do not let an uncached session bypass budget
enforcement, while preserving normal refresh behavior for existing keys. Add or
update the preload logic needed for first requests and revise TestE2E_PodRestart
to assert that over-budget sessions remain blocked after a replica restart.
| | Pod restarts | First request passes (cold cache); refresh picks up Redis counters within one interval | | ||
| | `fail_closed` + refresh failure | Stale cache retained; enforcement lags until Redis recovers | |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
Prevent stale cache state from bypassing exhausted budgets.
After a restart, Line 104 says that the first request passes before Redis refreshes the local cache. Line 105 also says that fail_closed retains stale cache state during refresh failures. If Redis already records an exhausted session, either path can admit traffic beyond the configured limit. Load session state before admitting the request. For true fail_closed behavior, deny while refresh fails, or rename and document the mode as stale-cache enforcement.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@authbridge/docs/token-budget-plugin.md` around lines 104 - 105, Update the
token-budget enforcement flow documented in the pod-restart and fail_closed rows
so requests cannot bypass exhausted Redis-backed session budgets: load current
session state before admitting the request, and when refresh fails under
fail_closed, deny the request rather than retaining stale state. If stale-cache
enforcement is intentional, rename fail_closed and document that behavior
instead.
rubambiza
added
the
ready-for-ai-review
clawgenti
left a comment
clawgenti
left a comment
There was a problem hiding this comment.
New opt-in token-budget plugin enforcing per-session token, call, and duration limits via Redis, with shadow mode and a pluggable storage abstraction — solid design overall. Two findings worth addressing before merge.
- Goroutine leak on shutdown:
accumulateis fired as a detached goroutine with no WaitGroup;Shutdowncloses the store while in-flight accumulate goroutines may still be writing to it. fail_closedis not implemented: the config acceptsfail_closedand docs describe it as "log warnings; stale cache retained", but the TODO comment confirms it currently behaves identically tofail_open. Either remove the option or implement it before merge to avoid user confusion.
Reviewed by clawgenti using the github-pr-review skill
| tokens := int64(inf.TotalTokens) | ||
|
|
||
| go p.accumulate(sessionID, tokens) | ||
|
|
There was a problem hiding this comment.
Goroutine leak on shutdown: accumulate is fire-and-forget with no WaitGroup. When Shutdown is called, p.store.Close() may race with in-flight goroutines still writing to the store. Consider adding a sync.WaitGroup to track outstanding accumulate calls and waiting on it before closing the store.
|
|
||
| if err != nil { | ||
| // TODO: fail_closed should deny requests when Redis is unreachable | ||
| // and the local cache has no data. Currently both modes retain stale cache. |
There was a problem hiding this comment.
Unimplemented fail_closed: The config option is accepted and documented, but the TODO comment confirms it currently behaves the same as fail_open (stale cache retained, no enforcement difference). Remove the option from the public API or implement the intended behavior (e.g., deny requests when the store is unreachable) before shipping.
Signed-off-by: Evaline Ju <69598118+evaline-ju@users.noreply.github.com>
clawgenti
left a comment
clawgenti
left a comment
There was a problem hiding this comment.
This PR introduces a solid, well-structured token-budget plugin with clean storage abstraction, good test coverage across unit/e2e/lifecycle scenarios, and a sensible fail-open default for Redis outages. Two findings worth addressing before merge.
Findings:
Init()hardcodesstorage.Open("redis", ...)regardless of the URL scheme, breaking the pluggable registry for any non-redis://scheme (e.g.valkey://).Configure()silently ignoresrefresh_intervalparse errors (theerr == nilguard means invalid strings pass validation;Init()then silently falls back to 5s with no user-visible error).
Reviewed by clawgenti using the github-pr-review skill
| } | ||
|
|
||
| func (p *TokenBudget) Init(_ context.Context) error { | ||
| store, err := storage.Open("redis", p.cfg.RedisURL) |
There was a problem hiding this comment.
storage.Open("redis", ...) hardcodes the scheme instead of deriving it from the URL. This bypasses the pluggable registry — a valkey:// URL (which the docs advertise as supported) would fail with unknown scheme. Consider parsing the scheme from p.cfg.RedisURL via url.Parse and passing it to Open, or documenting that only redis:// is valid here.
| 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 && d <= 0 { |
There was a problem hiding this comment.
The guard if d, err := time.ParseDuration(...); err == nil && d <= 0 only rejects non-positive durations when the string is valid. An invalid string (e.g. "abc") silently passes Configure() and Init() then falls back to 5s without any error or warning. Consider returning an error on parse failure here too: if _, err := time.ParseDuration(p.cfg.RefreshInterval); err != nil { return fmt.Errorf(...) }.
3 participants
Summary
include_plugin_tokenbudget) token-budget plugin for per-session lifetime enforcement for tokens, calls, and wall-clock duration, working with streaming SSE responses and buffered responsestoken-exchangeplugin, Cortex as a sidecar will not be able to surface these errors clearly through agents themselvesRelated issue(s)
Closes #708 , future: extend with HITL
Testing Instructions
helm install valkey bitnami/valkey -n <ns>) or any Redis-compatible storecd authbridge && docker build -f cmd/authbridge-proxy/Dockerfile \ --build-arg GO_BUILD_TAGS="include_plugin_tokenbudget" \ -t authbridge:latest .and load into clusterauthbridge-runtime-config(or seedocs/token-budget-plugin.md), roll pod or wait for hot-reloadon_exceed: "observe", send agent requests past limit, these requests should succeed and info can be seen in logs/through abctlon_exceed: "deny", edit the authbridge configmapAssisted-By: Claude (Anthropic AI)
Summary by CodeRabbit
New Features
Bug Fixes
Documentation
Tests