feat(agent): SSE streaming for LLM completions with per-provider accumulators

Replace buffered LLM requests with Server-Sent Events streaming. Each
text delta emits an agent_message_chunk session update as it arrives,
providing a natural keepalive that resets the ACP idle clock without
relying solely on the 30s ticker.

Design decisions:
- Two reqwest::Client instances: `http` (with global timeout for
  summarize) and `http_stream` (no global timeout, enforces first-byte
  and inter-chunk timeouts via tokio::time::timeout)
- Two-phase timeout: llm_timeout (120s) until first content delta,
  then stream_chunk_timeout (30s) between subsequent events
- Anthropic: index-keyed HashMap<usize, (String, String)> for parallel
  tool-call accumulation via content_block index
- OpenAI Chat: Vec-indexed accumulation by tool_calls[].index
- Responses API: routes on JSON `type` field inside data payload, uses
  response.output_text.delta and response.function_call_arguments.delta
- MAX_LLM_RESPONSE_BYTES caps accumulated semantic content (text + tool
  args), not raw SSE wire bytes
- SSE parser handles : comments, multi-line data:, id:/retry: fields
- Keepalive ticker (G) preserved as fallback for reasoning models that
  pause before producing content

Co-authored-by: Will Pfleger <pfleger.will@gmail.com>
Signed-off-by: Will Pfleger <pfleger.will@gmail.com>
This commit is contained in:
npub1mn7jgtj4w2pd0g0zeuhxsa6jy6p0rewxz4kujt98my82ahfmp72sxjexk7
2026-06-10 13:45:13 -04:00
co-authored by Will Pfleger
parent 4b676908bd
commit d22ba793a6
6 changed files with 1581 additions and 42 deletions
+7 -5
View File
@@ -6,7 +6,7 @@ use tokio::task::JoinSet;
use crate::config::{Config, MAX_PROMPT_BYTES, MAX_TOOL_CALLS_PER_TURN, MAX_TOOL_RESULT_BYTES};
use crate::handoff::HandoffOutcome;
use crate::llm::Llm;
use crate::llm::{Llm, StreamEmitter};
use crate::mcp::McpRegistry;
use crate::types::{
@@ -91,14 +91,16 @@ impl RunCtx<'_> {
let tools = self.mcp.tools();
round = round.saturating_add(1);
let emitter = StreamEmitter::new(self.wire.clone(), self.session_id.to_owned());
let response = tokio::select! {
biased;
_ = self.cancel.changed() => return Ok(StopReason::Cancelled),
r = self.llm.complete(self.cfg, self.system_prompt, self.history, &tools) => r?,
r = self.llm.complete_stream(self.cfg, self.system_prompt, self.history, &tools, &emitter) => r?,
_ = async {
// Keepalive ticker: emit a lightweight session update every 30s
// while waiting on the LLM provider. This resets the ACP harness
// idle clock so long provider responses don't trigger timeout.
// Keepalive ticker (fallback): fires every 30s while waiting on
// the LLM. During active streaming, text-delta emissions already
// reset the ACP idle clock — this only matters for reasoning
// models that pause before producing content.
let mut interval = tokio::time::interval(std::time::Duration::from_secs(30));
interval.tick().await; // first tick fires immediately — skip it
loop {
+23
View File
@@ -48,6 +48,15 @@ pub struct Config {
pub max_rounds: u32,
pub max_output_tokens: u32,
pub llm_timeout: Duration,
/// Maximum time to wait between consecutive response body chunks from the
/// LLM provider. If no data arrives within this window the request is
/// considered stalled and terminated. This does NOT cap total response
/// time — a stream that keeps producing chunks can run indefinitely.
pub llm_stream_chunk_timeout: Duration,
/// Inter-chunk timeout for SSE streaming after the first content delta
/// arrives. Tighter than `llm_stream_chunk_timeout` because SSE events
/// arrive frequently once content generation starts. Default 30s.
pub stream_chunk_timeout: Duration,
pub tool_timeout: Duration,
pub mcp_init_timeout: Duration,
pub mcp_max_restart_attempts: u32,
@@ -152,6 +161,14 @@ impl Config {
max_rounds: parse_env("SPROUT_AGENT_MAX_ROUNDS", 0)?,
max_output_tokens: parse_env("SPROUT_AGENT_MAX_OUTPUT_TOKENS", 32_768)?,
llm_timeout: Duration::from_secs(parse_env("SPROUT_AGENT_LLM_TIMEOUT_SECS", 120)?),
llm_stream_chunk_timeout: Duration::from_secs(parse_env(
"SPROUT_AGENT_LLM_STREAM_CHUNK_TIMEOUT_SECS",
120,
)?),
stream_chunk_timeout: Duration::from_secs(parse_env(
"SPROUT_AGENT_STREAM_CHUNK_TIMEOUT_SECS",
30,
)?),
tool_timeout: Duration::from_secs(parse_env("SPROUT_AGENT_TOOL_TIMEOUT_SECS", 660)?),
mcp_init_timeout: Duration::from_secs(parse_env(
"SPROUT_AGENT_MCP_INIT_TIMEOUT_SECS",
@@ -211,6 +228,12 @@ impl Config {
if self.llm_timeout < MIN_TIMEOUT {
return Err("config: SPROUT_AGENT_LLM_TIMEOUT_SECS must be >= 1".into());
}
if self.llm_stream_chunk_timeout < MIN_TIMEOUT {
return Err("config: SPROUT_AGENT_LLM_STREAM_CHUNK_TIMEOUT_SECS must be >= 1".into());
}
if self.stream_chunk_timeout < MIN_TIMEOUT {
return Err("config: SPROUT_AGENT_STREAM_CHUNK_TIMEOUT_SECS must be >= 1".into());
}
if self.tool_timeout < MIN_TIMEOUT {
return Err("config: SPROUT_AGENT_TOOL_TIMEOUT_SECS must be >= 1".into());
}
File diff suppressed because it is too large Load Diff
+127 -11
View File
@@ -32,6 +32,7 @@ async fn spawn_fake_llm(responses: Vec<Value>) -> String {
tokio::spawn(async move {
let mut buf = Vec::new();
let mut tmp = [0u8; 4096];
// Read headers
while !buf.windows(4).any(|w| w == b"\r\n\r\n") {
match sock.read(&mut tmp).await {
Ok(0) | Err(_) => return,
@@ -41,17 +42,63 @@ async fn spawn_fake_llm(responses: Vec<Value>) -> String {
return;
}
}
let body = queue
// Extract content-length and read body
let header_str = String::from_utf8_lossy(&buf).to_string();
let header_end = header_str.find("\r\n\r\n").unwrap_or(0) + 4;
let content_length = header_str
.lines()
.find_map(|l| {
let lower = l.to_lowercase();
if lower.starts_with("content-length:") {
lower.split(':').nth(1)?.trim().parse::<usize>().ok()
} else {
None
}
})
.unwrap_or(0);
let already_read = buf.len() - header_end;
let remaining = content_length.saturating_sub(already_read);
if remaining > 0 {
let mut body_buf = vec![0u8; remaining];
let mut read_so_far = 0;
while read_so_far < remaining {
match sock.read(&mut body_buf[read_so_far..]).await {
Ok(0) | Err(_) => break,
Ok(n) => read_so_far += n,
}
}
buf.extend_from_slice(&body_buf[..read_so_far]);
}
let body_bytes = &buf[header_end..];
let is_stream = body_bytes.windows(13).any(|w| w == b"\"stream\":true");
let canned = queue
.lock()
.await
.pop_front()
.unwrap_or_else(|| json!({ "error": "no canned response" }));
let body_s = serde_json::to_string(&body).unwrap();
let resp = format!(
"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
body_s.len(), body_s,
);
let _ = sock.write_all(resp.as_bytes()).await;
if is_stream {
// Convert canned response to SSE streaming format (OpenAI Chat)
let events = openai_to_sse_events(&canned);
let mut sse_body = String::new();
for ev in &events {
sse_body.push_str(&format!("data: {}\n\n", ev));
}
sse_body.push_str("data: [DONE]\n\n");
let resp = format!(
"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nConnection: close\r\n\r\n{}",
sse_body,
);
let _ = sock.write_all(resp.as_bytes()).await;
} else {
let body_s = serde_json::to_string(&canned).unwrap();
let resp = format!(
"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
body_s.len(), body_s,
);
let _ = sock.write_all(resp.as_bytes()).await;
}
let _ = sock.shutdown().await;
});
}
@@ -59,6 +106,70 @@ async fn spawn_fake_llm(responses: Vec<Value>) -> String {
url
}
/// Convert a canned OpenAI Chat Completions response into SSE delta events.
fn openai_to_sse_events(response: &Value) -> Vec<String> {
let mut events = Vec::new();
let choice = &response["choices"][0];
let msg = &choice["message"];
// Text content
if let Some(content) = msg.get("content").and_then(Value::as_str) {
if !content.is_empty() {
events.push(
json!({
"choices": [{"index": 0, "delta": {"content": content}, "finish_reason": null}]
})
.to_string(),
);
}
}
// Tool calls
if let Some(tcs) = msg.get("tool_calls").and_then(Value::as_array) {
for (i, tc) in tcs.iter().enumerate() {
let id = tc.get("id").and_then(Value::as_str).unwrap_or("");
let name = tc["function"]
.get("name")
.and_then(Value::as_str)
.unwrap_or("");
let args = tc["function"]
.get("arguments")
.and_then(Value::as_str)
.unwrap_or("{}");
// First chunk: id + name
events.push(json!({
"choices": [{"index": 0, "delta": {
"tool_calls": [{"index": i, "id": id, "function": {"name": name, "arguments": ""}}]
}, "finish_reason": null}]
}).to_string());
// Second chunk: arguments
events.push(
json!({
"choices": [{"index": 0, "delta": {
"tool_calls": [{"index": i, "function": {"arguments": args}}]
}, "finish_reason": null}]
})
.to_string(),
);
}
}
// Final chunk with finish_reason
let finish = choice
.get("finish_reason")
.and_then(Value::as_str)
.unwrap_or("stop");
events.push(
json!({
"choices": [{"index": 0, "delta": {}, "finish_reason": finish}],
"usage": {"prompt_tokens": 10, "completion_tokens": 5}
})
.to_string(),
);
events
}
// ─── ACP harness ────────────────────────────────────────────────────────────
struct Harness {
@@ -258,11 +369,16 @@ async fn rejects_concurrent_prompts() {
buf.extend_from_slice(&tmp[..n]);
}
tokio::time::sleep(Duration::from_millis(500)).await;
let body = openai_text("done").to_string();
// Return SSE streaming response
let events = openai_to_sse_events(&openai_text("done"));
let mut sse_body = String::new();
for ev in &events {
sse_body.push_str(&format!("data: {}\n\n", ev));
}
sse_body.push_str("data: [DONE]\n\n");
let resp = format!(
"HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
body.len(),
body
"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nConnection: close\r\n\r\n{}",
sse_body,
);
let _ = sock.write_all(resp.as_bytes()).await;
let _ = sock.shutdown().await;
+112 -8
View File
@@ -121,6 +121,7 @@ async fn spawn_fake_llm(responses: Vec<Value>) -> String {
tokio::spawn(async move {
let mut buf = Vec::new();
let mut tmp = [0u8; 4096];
// Read headers
while !buf.windows(4).any(|w| w == b"\r\n\r\n") {
match sock.read(&mut tmp).await {
Ok(0) | Err(_) => return,
@@ -130,18 +131,62 @@ async fn spawn_fake_llm(responses: Vec<Value>) -> String {
return;
}
}
let body = queue
// Extract content-length and read body to detect stream flag
let header_str = String::from_utf8_lossy(&buf).to_string();
let header_end = header_str.find("\r\n\r\n").unwrap_or(0) + 4;
let content_length = header_str
.lines()
.find_map(|l| {
let lower = l.to_lowercase();
if lower.starts_with("content-length:") {
lower.split(':').nth(1)?.trim().parse::<usize>().ok()
} else {
None
}
})
.unwrap_or(0);
let already_read = buf.len() - header_end;
let remaining = content_length.saturating_sub(already_read);
if remaining > 0 {
let mut body_buf = vec![0u8; remaining];
let mut read_so_far = 0;
while read_so_far < remaining {
match sock.read(&mut body_buf[read_so_far..]).await {
Ok(0) | Err(_) => break,
Ok(n) => read_so_far += n,
}
}
buf.extend_from_slice(&body_buf[..read_so_far]);
}
let body_bytes = &buf[header_end..];
let is_stream = body_bytes.windows(13).any(|w| w == b"\"stream\":true");
let canned = queue
.lock()
.await
.pop_front()
.unwrap_or_else(|| json!({ "error": "no canned response" }));
let body_s = serde_json::to_string(&body).unwrap();
let resp = format!(
"HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
body_s.len(),
body_s,
);
let _ = sock.write_all(resp.as_bytes()).await;
if is_stream {
let events = openai_to_sse_events(&canned);
let mut sse_body = String::new();
for ev in &events {
sse_body.push_str(&format!("data: {}\n\n", ev));
}
sse_body.push_str("data: [DONE]\n\n");
let resp = format!(
"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nConnection: close\r\n\r\n{}",
sse_body,
);
let _ = sock.write_all(resp.as_bytes()).await;
} else {
let body_s = serde_json::to_string(&canned).unwrap();
let resp = format!(
"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
body_s.len(), body_s,
);
let _ = sock.write_all(resp.as_bytes()).await;
}
let _ = sock.shutdown().await;
});
}
@@ -149,6 +194,65 @@ async fn spawn_fake_llm(responses: Vec<Value>) -> String {
url
}
/// Convert a canned OpenAI Chat Completions response into SSE delta events.
fn openai_to_sse_events(response: &Value) -> Vec<String> {
let mut events = Vec::new();
let choice = &response["choices"][0];
let msg = &choice["message"];
if let Some(content) = msg.get("content").and_then(Value::as_str) {
if !content.is_empty() {
events.push(
json!({
"choices": [{"index": 0, "delta": {"content": content}, "finish_reason": null}]
})
.to_string(),
);
}
}
if let Some(tcs) = msg.get("tool_calls").and_then(Value::as_array) {
for (i, tc) in tcs.iter().enumerate() {
let id = tc.get("id").and_then(Value::as_str).unwrap_or("");
let name = tc["function"]
.get("name")
.and_then(Value::as_str)
.unwrap_or("");
let args = tc["function"]
.get("arguments")
.and_then(Value::as_str)
.unwrap_or("{}");
events.push(json!({
"choices": [{"index": 0, "delta": {
"tool_calls": [{"index": i, "id": id, "function": {"name": name, "arguments": ""}}]
}, "finish_reason": null}]
}).to_string());
events.push(
json!({
"choices": [{"index": 0, "delta": {
"tool_calls": [{"index": i, "function": {"arguments": args}}]
}, "finish_reason": null}]
})
.to_string(),
);
}
}
let finish = choice
.get("finish_reason")
.and_then(Value::as_str)
.unwrap_or("stop");
events.push(
json!({
"choices": [{"index": 0, "delta": {}, "finish_reason": finish}],
"usage": {"prompt_tokens": 10, "completion_tokens": 5}
})
.to_string(),
);
events
}
fn openai_text(content: &str) -> Value {
json!({
"id": "cc-1", "object": "chat.completion", "model": "fake-model",
+88 -7
View File
@@ -68,17 +68,36 @@ async fn spawn_capturing_llm(responses: Vec<Value>) -> CapturingLlm {
if let Ok(req) = serde_json::from_slice::<Value>(&buf[header_end..]) {
captured.lock().await.push(req);
}
let body = queue
let is_stream = buf[header_end..]
.windows(13)
.any(|w| w == b"\"stream\":true");
let canned = queue
.lock()
.await
.pop_front()
.unwrap_or_else(|| json!({ "error": "no canned response" }));
let body_s = serde_json::to_string(&body).unwrap();
let resp = format!(
"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
body_s.len(), body_s,
);
let _ = sock.write_all(resp.as_bytes()).await;
if is_stream {
let events = openai_to_sse_events(&canned);
let mut sse_body = String::new();
for ev in &events {
sse_body.push_str(&format!("data: {}\n\n", ev));
}
sse_body.push_str("data: [DONE]\n\n");
let resp = format!(
"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nConnection: close\r\n\r\n{}",
sse_body,
);
let _ = sock.write_all(resp.as_bytes()).await;
} else {
let body_s = serde_json::to_string(&canned).unwrap();
let resp = format!(
"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
body_s.len(), body_s,
);
let _ = sock.write_all(resp.as_bytes()).await;
}
let _ = sock.shutdown().await;
});
}
@@ -216,6 +235,68 @@ fn openai_tool_call(id: &str, name: &str, args: Value) -> Value {
})
}
/// Convert a canned OpenAI Chat Completions response into SSE delta events.
fn openai_to_sse_events(response: &Value) -> Vec<String> {
let mut events = Vec::new();
let choice = &response["choices"][0];
let msg = &choice["message"];
if let Some(content) = msg.get("content").and_then(Value::as_str) {
if !content.is_empty() {
events.push(
json!({
"choices": [{"index": 0, "delta": {"content": content}, "finish_reason": null}]
})
.to_string(),
);
}
}
if let Some(tcs) = msg.get("tool_calls").and_then(Value::as_array) {
for (i, tc) in tcs.iter().enumerate() {
let id = tc.get("id").and_then(Value::as_str).unwrap_or("");
let name = tc["function"]
.get("name")
.and_then(Value::as_str)
.unwrap_or("");
let args = tc["function"]
.get("arguments")
.and_then(Value::as_str)
.unwrap_or("{}");
events.push(json!({
"choices": [{"index": 0, "delta": {
"tool_calls": [{"index": i, "id": id, "function": {"name": name, "arguments": ""}}]
}, "finish_reason": null}]
}).to_string());
events.push(
json!({
"choices": [{"index": 0, "delta": {
"tool_calls": [{"index": i, "function": {"arguments": args}}]
}, "finish_reason": null}]
})
.to_string(),
);
}
}
let finish = choice
.get("finish_reason")
.and_then(Value::as_str)
.unwrap_or("stop");
// Carry usage from the original response if present
let mut final_event = json!({
"choices": [{"index": 0, "delta": {}, "finish_reason": finish}],
});
if let Some(usage) = response.get("usage") {
final_event["usage"] = usage.clone();
} else {
final_event["usage"] = json!({"prompt_tokens": 10, "completion_tokens": 5});
}
events.push(final_event.to_string());
events
}
async fn init_session(h: &mut Harness, mcp_servers: Value) -> String {
h.send(
"initialize",