Skip to content

feat: add rate limit - #1

Open
odelya-krief wants to merge 1 commit into
mainfrom
add-rate-limit
Open

odelya-krief wants to merge 1 commit into
mainfrom
add-rate-limit

Conversation

@odelya-krief

@odelya-krief odelya-krief commented May 19, 2026 •

Copy link
Copy Markdown
Owner

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:

[rate_limit]
enabled = true
bytes_per_window = 104857600  # 100 MB
window_secs = 3600            # 1 hour (default if omitted)

Design:

The gateway already measures bytes per tunnel, copy_bidirectional returns (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:

  • Check right after the policy engine allows a connection (before the tunnel opens)
  • Record right after the tunnel closes

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 a Mutex is sufficient and adds no external dependencies.

The window resets per identity on the next record_bytes call 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

File What changed
src/rate_limit.rs New module — RateLimiter struct with is_allowed and record_bytes, apply() function that builds the limiter from config
src/config.rs New RateLimitConfig struct (enabled, bytes_per_window, window_secs), defaults to disabled
src/proxy.rs MakeProxyService holds Arc<RateLimiter>, checks limit after policy allows, records bytes after tunnel closes, denial routed through existing log_denial for consistent structured logging
src/main.rs One added line: let rate_limiter = rate_limit::apply(&config.rate_limit)
tests/e2e.rs New test rate_limit_blocks_after_limit_exceeded
tests/common/mod.rs AllowAllPolicyEngine and start_proxy_with_rate_limit test helpers

Testing

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_exceeded

All existing tests continue to pass.


I used Claude Code (Anthropic's AI coding assistant) throughout this task

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant