feat: add rate limit - #1
Open
odelya-krief wants to merge 1 commit into
Open
odelya-krief wants to merge 1 commit into
odelya-krief wants to merge 1 commit into
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Agent Gateway – Per-Identity Rate Limit
A per-identity byte rate limit on data proxied through the gateway. Each agent identity gets a configurable byte budget per time window. Once the budget is exhausted, new connection attempts are denied with HTTP 429 until the window resets. The feature is opt-in — the gateway behaves identically to before if not configured.
How to enable:
Add to
config.toml:Design:
The gateway already measures bytes per tunnel,
copy_bidirectionalreturns(bytes_up, bytes_down)after every tunnel closes, and those numbers are already logged. So the infrastructure was already there. The natural enforcement points are:Agents are ephemeral, they are created for a task and deleted when done. There is no value in persisting rate limit counters across gateway restarts. A simple in-memory
HashMap<identity → bytes_used>with aMutexis sufficient and adds no external dependencies.The window resets per identity on the next
record_bytescall after the window has expired. No background task is needed the reset is lazy and happens on access.Tradeoffs:
Best-effort enforcement: Bytes are counted after a tunnel closes, not reserved upfront (we don't know how much data will flow before a tunnel opens). This means two concurrent tunnels from the same identity could together exceed the limit before enforcement kicks in on the next connection attempt.
The alternative would be to close an active tunnel mid-stream once it hits the limit — significantly more complex and more disruptive to the agent. For the use case here, agents are not adversarial toward their own rate limits, and a small overshoot is not a security concern. Simplicity won.
Files changed
src/rate_limit.rsRateLimiterstruct withis_allowedandrecord_bytes,apply()function that builds the limiter from configsrc/config.rsRateLimitConfigstruct (enabled,bytes_per_window,window_secs), defaults to disabledsrc/proxy.rsMakeProxyServiceholdsArc<RateLimiter>, checks limit after policy allows, records bytes after tunnel closes, denial routed through existinglog_denialfor consistent structured loggingsrc/main.rslet rate_limiter = rate_limit::apply(&config.rate_limit)tests/e2e.rsrate_limit_blocks_after_limit_exceededtests/common/mod.rsAllowAllPolicyEngineandstart_proxy_with_rate_limittest helpersTesting
The new e2e test does not require a database — it uses an in-memory allow-all policy engine with a 10-byte limit. It sends 11 bytes through a tunnel, waits for the tunnel to close (so bytes are recorded), then verifies the next connection gets HTTP 429.
cargo test rate_limit_blocks_after_limit_exceededAll existing tests continue to pass.
I used Claude Code (Anthropic's AI coding assistant) throughout this task