Skip to content

Commit a1857ea

Browse files
jboldjbold
authored
Fix streaming deadlock and harden websocket/SSE chat flow (#7)
* test: add full rust/wasm/e2e suite with ci gates * test: build mock-channel wasm fixture before rust suite * test: build wasm fixture plugins for rust ci * Fix streaming deadlock and harden websocket/sse chat flow
1 parent dd97b30 commit a1857ea

11 files changed

Lines changed: 574 additions & 107 deletions

File tree

‎Cargo.lock‎

Lines changed: 22 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎src/agent/mod.rs‎

Lines changed: 144 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -91,52 +91,84 @@ impl AgentRunner {
9191
return Ok(());
9292
}
9393

94-
// Create an internal channel to collect events from this LLM call
95-
let (inner_tx, mut inner_rx) = mpsc::channel::<AgentEvent>(32);
96-
97-
provider
98-
.call_streaming(&current_messages, tools, system_prompt, inner_tx)
99-
.await?;
100-
101-
// Collect events, forwarding text/usage/error to client,
102-
// collecting tool_use calls for dispatch
103-
let mut tool_calls: Vec<(String, String, serde_json::Value)> = Vec::new();
94+
let tool_calls: Vec<(String, String, serde_json::Value)> = {
95+
// Create an internal channel to collect events from this LLM call.
96+
// We must drain this channel while provider streaming is in-flight;
97+
// otherwise long responses can fill the buffer and deadlock.
98+
let (inner_tx, mut inner_rx) = mpsc::channel::<AgentEvent>(32);
99+
let mut tool_calls: Vec<(String, String, serde_json::Value)> = Vec::new();
100+
let mut provider_done = false;
101+
let mut provider_error: Option<anyhow::Error> = None;
102+
103+
let provider_call =
104+
provider.call_streaming(&current_messages, tools, system_prompt, inner_tx);
105+
tokio::pin!(provider_call);
106+
107+
loop {
108+
tokio::select! {
109+
result = &mut provider_call, if !provider_done => {
110+
provider_done = true;
111+
if let Err(e) = result {
112+
provider_error = Some(e);
113+
}
114+
}
115+
maybe_event = inner_rx.recv() => {
116+
let Some(event) = maybe_event else {
117+
if !provider_done {
118+
if let Err(e) = (&mut provider_call).await {
119+
provider_error = Some(e);
120+
}
121+
}
122+
break;
123+
};
104124

105-
while let Some(event) = inner_rx.recv().await {
106-
match event {
107-
AgentEvent::Text(ref _t) => {
108-
let _ = tx.send(event).await;
109-
}
110-
AgentEvent::ToolUse {
111-
ref id,
112-
ref name,
113-
ref input,
114-
} => {
115-
// Forward to client so they can observe
116-
let _ = tx
117-
.send(AgentEvent::ToolUse {
118-
id: id.clone(),
119-
name: name.clone(),
120-
input: input.clone(),
121-
})
122-
.await;
123-
tool_calls.push((id.clone(), name.clone(), input.clone()));
124-
}
125-
AgentEvent::Usage { .. } => {
126-
let _ = tx.send(event).await;
127-
}
128-
AgentEvent::Error(ref _e) => {
129-
let _ = tx.send(event).await;
130-
}
131-
AgentEvent::Done => {
132-
// Don't forward Done yet — we may need to continue the loop
125+
match event {
126+
AgentEvent::Text(ref _t) => {
127+
let _ = tx.send(event).await;
128+
}
129+
AgentEvent::ToolUse {
130+
ref id,
131+
ref name,
132+
ref input,
133+
} => {
134+
// Forward to client so they can observe
135+
let _ = tx
136+
.send(AgentEvent::ToolUse {
137+
id: id.clone(),
138+
name: name.clone(),
139+
input: input.clone(),
140+
})
141+
.await;
142+
tool_calls.push((id.clone(), name.clone(), input.clone()));
143+
}
144+
AgentEvent::Usage { .. } => {
145+
let _ = tx.send(event).await;
146+
}
147+
AgentEvent::Error(ref _e) => {
148+
let _ = tx.send(event).await;
149+
}
150+
AgentEvent::Done => {
151+
// Don't forward Done yet — we may need to continue the loop
152+
}
153+
AgentEvent::ToolResult { .. } => {
154+
// Shouldn't come from provider, but forward if it does
155+
let _ = tx.send(event).await;
156+
}
157+
}
158+
}
133159
}
134-
AgentEvent::ToolResult { .. } => {
135-
// Shouldn't come from provider, but forward if it does
136-
let _ = tx.send(event).await;
160+
161+
if provider_done && inner_rx.is_closed() && inner_rx.is_empty() {
162+
break;
137163
}
138164
}
139-
}
165+
166+
if let Some(e) = provider_error {
167+
return Err(e);
168+
}
169+
170+
tool_calls
171+
};
140172

141173
// If no tool calls, we're done
142174
if tool_calls.is_empty() {
@@ -366,3 +398,73 @@ impl Default for AgentRunner {
366398
Self::new()
367399
}
368400
}
401+
402+
#[cfg(test)]
403+
mod tests {
404+
use super::{AgentEvent, AgentRunner, providers};
405+
use crate::sandbox::PluginHost;
406+
use async_trait::async_trait;
407+
use std::sync::Arc;
408+
use tokio::sync::{RwLock, mpsc};
409+
use tokio::time::{Duration, timeout};
410+
411+
struct BurstProvider {
412+
chunks: usize,
413+
}
414+
415+
#[async_trait]
416+
impl providers::LlmProvider for BurstProvider {
417+
async fn call_streaming(
418+
&self,
419+
_messages: &[serde_json::Value],
420+
_tools: &[serde_json::Value],
421+
_system_prompt: Option<&str>,
422+
tx: mpsc::Sender<AgentEvent>,
423+
) -> anyhow::Result<()> {
424+
for i in 0..self.chunks {
425+
tx.send(AgentEvent::Text(format!("chunk-{i}"))).await?;
426+
}
427+
tx.send(AgentEvent::Done).await?;
428+
Ok(())
429+
}
430+
}
431+
432+
#[tokio::test]
433+
async fn run_with_tools_drains_stream_while_provider_is_running() {
434+
let runner = AgentRunner::new();
435+
let provider = BurstProvider { chunks: 64 };
436+
let plugins = Arc::new(RwLock::new(PluginHost::new()));
437+
let (tx, mut rx) = mpsc::channel::<AgentEvent>(256);
438+
439+
timeout(
440+
Duration::from_secs(5),
441+
runner.run_with_tools(
442+
&provider,
443+
vec![serde_json::json!({
444+
"role": "user",
445+
"content": "hello",
446+
})],
447+
&[],
448+
None,
449+
&plugins,
450+
tx,
451+
),
452+
)
453+
.await
454+
.expect("runner should not deadlock")
455+
.expect("runner should succeed");
456+
457+
let mut text_count = 0usize;
458+
let mut saw_done = false;
459+
while let Ok(event) = rx.try_recv() {
460+
match event {
461+
AgentEvent::Text(_) => text_count += 1,
462+
AgentEvent::Done => saw_done = true,
463+
_ => {}
464+
}
465+
}
466+
467+
assert_eq!(text_count, 64);
468+
assert!(saw_done);
469+
}
470+
}

0 commit comments

Comments
 (0)