feat: add proactive rate-limit - #47
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughWalkthroughAdds endpoint-keyed token-bucket rate limiting, exposes provider base URLs through ChangesProactive stream rate limiting
Sequence Diagram(s)sequenceDiagram
participant Client
participant StreamHandler
participant ApiClient
participant RateLimiter
participant TokenBucket
participant SSE
Client->>StreamHandler: start stream turn
StreamHandler->>ApiClient: get base_url()
StreamHandler->>RateLimiter: acquire(base_url)
RateLimiter->>TokenBucket: take token
TokenBucket-->>RateLimiter: permit or wait duration
RateLimiter-->>StreamHandler: permit or wait
StreamHandler->>SSE: open stream
Possibly related PRs
🚥 Pre-merge checks | ✅ 4✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 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: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
src/stream/handler.rs (2)
1330-1390: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
gate_on_rate_limitignorestotal_deadline, so a turn can overrun its configuredtotal_stream_timeout.
try_streamalready receivestotal_deadlineand threads it intoprocess_events, butgate_on_rate_limit(called first, at Line 1339) doesn't receive it at all — its wait loop is bounded only by the limiter's ownmax_wait(default 30s). Iftotal_stream_timeoutis shorter thanmax_wait(or the gate wait is large relative to remaining budget), the turn can block in the gate well past the deadline beforeprocess_eventsgets a chance to fail withTotalTimeout. This is inconsistent with the precedent already established in this same file for rate-limit backoff sleeps, which are explicitly clamped to the deadline so a largeRetry-After"cannot overrun the turn budget."♻️ Proposed fix: thread `total_deadline` through the gate
async fn try_stream<C: ApiClient>( &self, client: &C, conversation: Vec<Message>, system: Option<String>, tool_schemas: Option<Vec<ToolSchema>>, cancel: &Arc<CancelSignal>, total_deadline: Option<Instant>, ) -> Result<StreamTurnResult, StreamHandlerError> { - self.gate_on_rate_limit(client, cancel).await?; + self.gate_on_rate_limit(client, cancel, total_deadline).await?; let stream = client.stream_messages(conversation, system, tool_schemas); self.process_events(stream, cancel, total_deadline).await } async fn gate_on_rate_limit<C: ApiClient>( &self, client: &C, cancel: &Arc<CancelSignal>, + total_deadline: Option<Instant>, ) -> Result<(), StreamHandlerError> { let Some(limiter) = &self.rate_limiter else { return Ok(()); }; let key = client.base_url(); let max_wait = limiter.max_wait(); let mut waited = Duration::ZERO; loop { match limiter.acquire(&key) { Ok(()) => return Ok(()), Err(wait) => { - if waited >= max_wait { + if waited >= max_wait || Self::deadline_exceeded(total_deadline) { return Ok(()); } let remaining = max_wait.checked_sub(waited).unwrap_or(Duration::ZERO); - let capped = wait.min(remaining); + let mut capped = wait.min(remaining); + if let Some(deadline) = total_deadline { + if let Some(left) = deadline.checked_duration_since(Instant::now()) { + capped = capped.min(left); + } + } tokio::select! { () = tokio::time::sleep(capped) => {} () = cancel.notified() => return Err(StreamHandlerError::Cancelled), } waited = waited.saturating_add(capped); } } } }🤖 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 `@src/stream/handler.rs` around lines 1330 - 1390, Update try_stream and gate_on_rate_limit to accept and enforce total_deadline while waiting for rate-limit tokens. Clamp each rate-limit sleep to the remaining turn budget, preserve cancellation handling, and return the existing TotalTimeout error when the deadline expires so the gate cannot overrun total_stream_timeout.
1006-1014: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
RateLimitConfig::requests_per_minuteno longer does proactive throttling
This field’s doc still says it “proactively throttle[s] outgoing requests,” but the runtime only gates requests throughwith_rate_limiter(). Either wirenew()/with_rate_limit_config()to build a limiter from this value, or update the doc to point users atwith_rate_limiter()instead.🤖 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 `@src/stream/handler.rs` around lines 1006 - 1014, Update the documentation for RateLimitConfig::requests_per_minute to describe its actual behavior and direct users to StreamHandler::with_rate_limiter() for proactive throttling. Do not claim that the field itself proactively throttles requests, and leave the existing runtime configuration unchanged.
🤖 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 `@src/stream/rate_limit.rs`:
- Around line 57-83: Update TokenBucket::take_at (and the take path if it
contains equivalent logic) to detect refill_per_sec <= 0.0 before calculating
the wait duration, returning the existing Err result for an empty bucket instead
of passing an infinite value to Duration::from_secs_f64. Preserve normal refill
behavior for positive-capacity buckets.
---
Outside diff comments:
In `@src/stream/handler.rs`:
- Around line 1330-1390: Update try_stream and gate_on_rate_limit to accept and
enforce total_deadline while waiting for rate-limit tokens. Clamp each
rate-limit sleep to the remaining turn budget, preserve cancellation handling,
and return the existing TotalTimeout error when the deadline expires so the gate
cannot overrun total_stream_timeout.
- Around line 1006-1014: Update the documentation for
RateLimitConfig::requests_per_minute to describe its actual behavior and direct
users to StreamHandler::with_rate_limiter() for proactive throttling. Do not
claim that the field itself proactively throttles requests, and leave the
existing runtime configuration unchanged.
🪄 Autofix (Beta)
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: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 73db8823-39fb-4c4c-a1ea-2150cfc30d67
📒 Files selected for processing (8)
CHANGELOG.mdsrc/api.rssrc/provider/anthropic.rssrc/provider/gemini.rssrc/provider/openai.rssrc/stream.rssrc/stream/handler.rssrc/stream/rate_limit.rs
feat: add proactive rate-limit
1 participant
No description provided.