Richard Feldman
29cf14ed2f
The rate limiter's semaphore guard was being held for the entire duration of a turn, including during tool execution. This caused deadlocks when subagents tried to acquire permits while parent requests were waiting for them to complete. ## The Problem In `run_turn_internal`, the stream (which contains the `RateLimitGuard` holding the semaphore permit) was kept alive throughout the entire loop iteration - including during **tool execution**: 1. Parent request acquires permit 2. Parent starts streaming, consumes response 3. Parent starts executing tools (subagents) 4. **Stream/guard still held** while tools execute 5. Subagents try to acquire permits → blocked because parent still holds permit 6. Deadlock if all permits are held by parents waiting for subagent children ## The Fix Two changes were made: 1. **Drop the stream early**: Added an explicit `drop(events)` after the stream is fully consumed but before tool execution begins. This releases the rate limit permit so subagents can acquire it. 2. **Removed the `bypass_rate_limit` workaround**: Since the root cause is now fixed, the bypass mechanism is no longer needed. Note: no release notes because subagents are still feature-flagged, and this rate limiting change isn't actually observable without them. Release Notes: - N/A
78 lines
1.8 KiB
Rust
78 lines
1.8 KiB
Rust
use futures::Stream;
|
|
use smol::lock::{Semaphore, SemaphoreGuardArc};
|
|
use std::{
|
|
future::Future,
|
|
pin::Pin,
|
|
sync::Arc,
|
|
task::{Context, Poll},
|
|
};
|
|
|
|
use crate::LanguageModelCompletionError;
|
|
|
|
#[derive(Clone)]
|
|
pub struct RateLimiter {
|
|
semaphore: Arc<Semaphore>,
|
|
}
|
|
|
|
pub struct RateLimitGuard<T> {
|
|
inner: T,
|
|
_guard: SemaphoreGuardArc,
|
|
}
|
|
|
|
impl<T> Stream for RateLimitGuard<T>
|
|
where
|
|
T: Stream,
|
|
{
|
|
type Item = T::Item;
|
|
|
|
fn poll_next(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
|
|
unsafe { Pin::map_unchecked_mut(self, |this| &mut this.inner).poll_next(cx) }
|
|
}
|
|
}
|
|
|
|
impl RateLimiter {
|
|
pub fn new(limit: usize) -> Self {
|
|
Self {
|
|
semaphore: Arc::new(Semaphore::new(limit)),
|
|
}
|
|
}
|
|
|
|
pub fn run<'a, Fut, T>(
|
|
&self,
|
|
future: Fut,
|
|
) -> impl 'a + Future<Output = Result<T, LanguageModelCompletionError>>
|
|
where
|
|
Fut: 'a + Future<Output = Result<T, LanguageModelCompletionError>>,
|
|
{
|
|
let guard = self.semaphore.acquire_arc();
|
|
async move {
|
|
let guard = guard.await;
|
|
let result = future.await?;
|
|
drop(guard);
|
|
Ok(result)
|
|
}
|
|
}
|
|
|
|
pub fn stream<'a, Fut, T>(
|
|
&self,
|
|
future: Fut,
|
|
) -> impl 'a
|
|
+ Future<
|
|
Output = Result<impl Stream<Item = T::Item> + use<Fut, T>, LanguageModelCompletionError>,
|
|
>
|
|
where
|
|
Fut: 'a + Future<Output = Result<T, LanguageModelCompletionError>>,
|
|
T: Stream,
|
|
{
|
|
let guard = self.semaphore.acquire_arc();
|
|
async move {
|
|
let guard = guard.await;
|
|
let inner = future.await?;
|
|
Ok(RateLimitGuard {
|
|
inner,
|
|
_guard: guard,
|
|
})
|
|
}
|
|
}
|
|
}
|