mirror of
https://github.com/block/buzz.git
synced 2026-08-18 06:50:31 +02:00
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:
co-authored by
Will Pfleger
parent
4b676908bd
commit
d22ba793a6
@@ -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 {
|
||||
|
||||
@@ -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());
|
||||
}
|
||||
|
||||
+1224
-11
File diff suppressed because it is too large
Load Diff
@@ -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;
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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",
|
||||
|
||||
Reference in New Issue
Block a user