feat(streaming): stream assistant tokens on-device and over the cloud
Thread an optional TokenSink through the InferenceBackend trait so a turn streams text as it is generated while still returning the full result. - LocalBackend streams via onde stream_message / stream_tool_results when no tools are offered. onde's tool-aware path must buffer to detect tool calls, so streaming covers the tools-disabled and final forced-text rounds. - OpenAiBackend streams via SSE (stream:true), reassembling content and index-keyed tool_calls; adds the reqwest "stream" feature. - TUI renders live tokens with <think> reasoning hidden. - ACP emits AgentMessageChunk deltas live, stripping <think> across chunk boundaries. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
paydii committed
Jun 27, 2026 at 08:38 UTC
33432e5e59517a4bf5309af30839d67fae8fd3e6
4 files changed
+527
-71
Cargo.toml
+1
-1
index 42b2d69..e79196d 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -43,6 +43,6 @@ serde_json = "1"
toml = "0.8"
async-trait = "0.1"
regex = "1"
-reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls"] }
+reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls", "stream"] }
uuid = { version = "1", features = ["v4"] }
rpassword = "7"
src/backend.rs
+311
-6
index 0692e6d..c99e90d 100644
--- a/src/backend.rs
+++ b/src/backend.rs
@@ -62,24 +62,40 @@ pub struct TurnResult {
/// Backend errors are plain strings. Callers map them to ACP errors.
pub type BackendError = String;
+/// A sink for streaming assistant text deltas to the UI as they are produced.
+///
+/// When a caller passes `Some(sink)`, a streaming-capable backend forwards each
+/// text fragment through it as the model emits it; the returned [`TurnResult`]
+/// still carries the fully assembled text (and any tool calls). When the sink is
+/// `None`, the backend runs in non-streaming mode. Unbounded so the inference
+/// task never blocks on a slow consumer.
+pub type TokenSink = tokio::sync::mpsc::UnboundedSender<String>;
+
// ── The trait ───────────────────────────────────────────────────────────────────
/// A swappable inference backend driving siGit Code's agent loop.
#[async_trait]
pub trait InferenceBackend: Send + Sync {
/// Start an assistant turn from a new user message, offering `tools`.
+ ///
+ /// If `sink` is `Some`, text is streamed through it as it is generated. A
+ /// backend may decline to stream a given round (for example, on-device
+ /// inference cannot stream while it is still deciding whether to call a
+ /// tool); in that case the text is delivered only via the returned result.
async fn send_message_with_tools(
&self,
text: &str,
tools: &[ToolSpec],
+ sink: Option<&TokenSink>,
) -> Result<TurnResult, BackendError>;
/// Continue the turn by returning tool results. `tools` may be `None` on the
- /// final round to force a text answer.
+ /// final round to force a text answer. `sink` streams that text when set.
async fn send_tool_results(
&self,
results: Vec<ToolResult>,
tools: Option<&[ToolSpec]>,
+ sink: Option<&TokenSink>,
) -> Result<TurnResult, BackendError>;
/// Whether inference runs over the network (a configured provider) rather
@@ -118,7 +134,22 @@ impl InferenceBackend for LocalBackend {
&self,
text: &str,
tools: &[ToolSpec],
+ sink: Option<&TokenSink>,
) -> Result<TurnResult, BackendError> {
+ // onde's tool-aware path is non-streaming: it has to buffer the whole
+ // reply to detect tool calls. We can only stream when no tools are on
+ // offer (a plain answer), which is exactly the tools-disabled case.
+ if let Some(sink) = sink
+ && tools.is_empty()
+ {
+ let rx = self
+ .engine
+ .stream_message(text)
+ .await
+ .map_err(|error| error.to_string())?;
+ return drain_onde_stream(rx, sink).await;
+ }
+
let onde_tools = to_onde_tools(tools);
let result = self
.engine
@@ -132,6 +163,7 @@ impl InferenceBackend for LocalBackend {
&self,
results: Vec<ToolResult>,
tools: Option<&[ToolSpec]>,
+ sink: Option<&TokenSink>,
) -> Result<TurnResult, BackendError> {
let onde_results: Vec<onde::inference::ToolResult> = results
.into_iter()
@@ -140,6 +172,20 @@ impl InferenceBackend for LocalBackend {
content: result.content,
})
.collect();
+
+ // The final round passes `tools = None` to force a text answer; that's
+ // the only round onde can stream, since no further tool calls are parsed.
+ if let Some(sink) = sink
+ && tools.is_none()
+ {
+ let rx = self
+ .engine
+ .stream_tool_results(onde_results, None)
+ .await
+ .map_err(|error| error.to_string())?;
+ return drain_onde_stream(rx, sink).await;
+ }
+
let onde_tools = tools.map(to_onde_tools);
let result = self
.engine
@@ -154,6 +200,38 @@ impl InferenceBackend for LocalBackend {
}
}
+/// Drain an onde streaming receiver, forwarding each token to `sink` and
+/// assembling the full text. onde reports stream failures as a final chunk whose
+/// `finish_reason` is `"error: …"`; surface those as a backend error.
+async fn drain_onde_stream(
+ mut rx: tokio::sync::mpsc::Receiver<onde::inference::StreamChunk>,
+ sink: &TokenSink,
+) -> Result<TurnResult, BackendError> {
+ let mut text = String::new();
+ while let Some(chunk) = rx.recv().await {
+ if !chunk.delta.is_empty() {
+ text.push_str(&chunk.delta);
+ // The receiver is the UI; if it's gone the turn is being cancelled,
+ // so stop assembling rather than spinning the model to completion.
+ if sink.send(chunk.delta).is_err() {
+ break;
+ }
+ }
+ if chunk.done {
+ if let Some(reason) = chunk.finish_reason
+ && let Some(message) = reason.strip_prefix("error: ")
+ {
+ return Err(message.to_string());
+ }
+ break;
+ }
+ }
+ Ok(TurnResult {
+ text,
+ tool_calls: Vec::new(),
+ })
+}
+
/// Convert an `onde` tool-aware result into the neutral [`TurnResult`].
fn onde_result_to_turn(result: onde::inference::ToolAwareResult) -> TurnResult {
TurnResult {
@@ -230,14 +308,20 @@ impl OpenAiBackend {
}
/// POST the current history (plus `tools`) and apply the assistant reply to
- /// history, returning the neutral turn result.
- async fn complete(&self, tools: Option<&[ToolSpec]>) -> Result<TurnResult, BackendError> {
+ /// history, returning the neutral turn result. Streams via SSE when `sink`
+ /// is set; otherwise reads a single JSON response.
+ async fn complete(
+ &self,
+ tools: Option<&[ToolSpec]>,
+ sink: Option<&TokenSink>,
+ ) -> Result<TurnResult, BackendError> {
let url = format!("{}/chat/completions", self.base_url.trim_end_matches('/'));
+ let streaming = sink.is_some();
let mut body = serde_json::json!({
"model": self.model,
"messages": *self.history.lock().await,
- "stream": false,
+ "stream": streaming,
});
if let Some(tools) = tools
&& !tools.is_empty()
@@ -260,6 +344,15 @@ impl OpenAiBackend {
return Err(format!("endpoint returned {status}: {detail}"));
}
+ if let Some(sink) = sink {
+ self.consume_stream(response, sink).await
+ } else {
+ self.consume_json(response).await
+ }
+ }
+
+ /// Parse a single non-streaming chat-completion response.
+ async fn consume_json(&self, response: reqwest::Response) -> Result<TurnResult, BackendError> {
let parsed: ChatCompletion = response
.json()
.await
@@ -289,6 +382,150 @@ impl OpenAiBackend {
Ok(TurnResult { text, tool_calls })
}
+
+ /// Consume an OpenAI Server-Sent Events stream, forwarding content deltas to
+ /// `sink` and reassembling any tool calls (which arrive fragmented across
+ /// chunks, keyed by `index`).
+ async fn consume_stream(
+ &self,
+ response: reqwest::Response,
+ sink: &TokenSink,
+ ) -> Result<TurnResult, BackendError> {
+ use futures::StreamExt;
+
+ let mut stream = response.bytes_stream();
+ // Newlines are ASCII, so splitting raw bytes on `\n` never bisects a
+ // multibyte UTF-8 sequence; we only lossily decode whole lines.
+ let mut buffer: Vec<u8> = Vec::new();
+ let mut text = String::new();
+ let mut tool_accum: Vec<StreamingToolCall> = Vec::new();
+ let mut done = false;
+
+ while let Some(item) = stream.next().await {
+ let bytes = item.map_err(|error| format!("stream read error: {error}"))?;
+ buffer.extend_from_slice(&bytes);
+
+ while let Some(pos) = buffer.iter().position(|&b| b == b'\n') {
+ let line: Vec<u8> = buffer.drain(..=pos).collect();
+ let line = String::from_utf8_lossy(&line);
+ let line = line.trim();
+
+ let Some(data) = line.strip_prefix("data:") else {
+ continue;
+ };
+ let data = data.trim();
+ if data == "[DONE]" {
+ done = true;
+ break;
+ }
+ if data.is_empty() {
+ continue;
+ }
+
+ let chunk: StreamCompletion = match serde_json::from_str(data) {
+ Ok(chunk) => chunk,
+ // Skip keep-alive comments and anything we can't parse rather
+ // than aborting a turn over one malformed frame.
+ Err(_) => continue,
+ };
+
+ let Some(choice) = chunk.choices.into_iter().next() else {
+ continue;
+ };
+ if let Some(content) = choice.delta.content
+ && !content.is_empty()
+ {
+ text.push_str(&content);
+ if sink.send(content).is_err() {
+ // Consumer dropped (turn cancelled) — stop reading.
+ done = true;
+ break;
+ }
+ }
+ for delta in choice.delta.tool_calls.into_iter().flatten() {
+ let index = delta.index.unwrap_or(0) as usize;
+ if tool_accum.len() <= index {
+ tool_accum.resize_with(index + 1, StreamingToolCall::default);
+ }
+ let slot = &mut tool_accum[index];
+ if let Some(id) = delta.id {
+ slot.id = id;
+ }
+ if let Some(function) = delta.function {
+ if let Some(name) = function.name {
+ slot.name = name;
+ }
+ if let Some(arguments) = function.arguments {
+ slot.arguments.push_str(&arguments);
+ }
+ }
+ }
+ }
+
+ if done {
+ break;
+ }
+ }
+
+ let tool_calls: Vec<ToolCall> = tool_accum
+ .iter()
+ .filter(|call| !call.name.is_empty())
+ .enumerate()
+ .map(|(index, call)| ToolCall {
+ id: if call.id.is_empty() {
+ format!("call_{index}")
+ } else {
+ call.id.clone()
+ },
+ name: call.name.clone(),
+ arguments: call.arguments.clone(),
+ })
+ .collect();
+
+ // Record the assistant turn so later tool results have context.
+ self.history
+ .lock()
+ .await
+ .push(streamed_assistant_history(&text, &tool_calls));
+
+ Ok(TurnResult { text, tool_calls })
+ }
+}
+
+/// One tool call being reassembled from streamed deltas.
+#[derive(Default)]
+struct StreamingToolCall {
+ id: String,
+ name: String,
+ arguments: String,
+}
+
+/// Rebuild the assistant message for replay in history after a streamed turn,
+/// preserving any tool calls so the follow-up request is well-formed. Mirrors
+/// [`ResponseMessage::into_history_value`] for the non-streaming path.
+fn streamed_assistant_history(text: &str, tool_calls: &[ToolCall]) -> serde_json::Value {
+ let mut message = serde_json::json!({ "role": "assistant" });
+ message["content"] = if text.is_empty() {
+ serde_json::Value::Null
+ } else {
+ serde_json::Value::String(text.to_string())
+ };
+ if !tool_calls.is_empty() {
+ message["tool_calls"] = serde_json::json!(
+ tool_calls
+ .iter()
+ .map(|call| serde_json::json!({
+ "id": call.id,
+ "type": "function",
+ "function": {
+ "name": call.name,
+ "arguments": call.arguments,
+ }
+ }))
+ .collect::<Vec<_>>()
+ );
+ }
+ message
}
#[async_trait]
@@ -297,18 +534,20 @@ impl InferenceBackend for OpenAiBackend {
&self,
text: &str,
tools: &[ToolSpec],
+ sink: Option<&TokenSink>,
) -> Result<TurnResult, BackendError> {
self.history
.lock()
.await
.push(serde_json::json!({ "role": "user", "content": text }));
- self.complete(Some(tools)).await
+ self.complete(Some(tools), sink).await
}
async fn send_tool_results(
&self,
results: Vec<ToolResult>,
tools: Option<&[ToolSpec]>,
+ sink: Option<&TokenSink>,
) -> Result<TurnResult, BackendError> {
{
let mut history = self.history.lock().await;
@@ -320,7 +559,7 @@ impl InferenceBackend for OpenAiBackend {
}));
}
}
- self.complete(tools).await
+ self.complete(tools, sink).await
}
fn is_remote(&self) -> bool {
@@ -390,6 +629,46 @@ struct ResponseFunction {
arguments: String,
}
+// ── OpenAI streaming (SSE) chunk shapes ─────────────────────────────────────────
+
+#[derive(Debug, Deserialize)]
+struct StreamCompletion {
+ #[serde(default)]
+ choices: Vec<StreamChoice>,
+}
+
+#[derive(Debug, Deserialize)]
+struct StreamChoice {
+ #[serde(default)]
+ delta: StreamDelta,
+}
+
+#[derive(Debug, Default, Deserialize)]
+struct StreamDelta {
+ #[serde(default)]
+ content: Option<String>,
+ #[serde(default)]
+ tool_calls: Option<Vec<StreamToolCallDelta>>,
+}
+
+#[derive(Debug, Deserialize)]
+struct StreamToolCallDelta {
+ #[serde(default)]
+ index: Option<u32>,
+ #[serde(default)]
+ id: Option<String>,
+ #[serde(default)]
+ function: Option<StreamFunctionDelta>,
+}
+
+#[derive(Debug, Deserialize)]
+struct StreamFunctionDelta {
+ #[serde(default)]
+ name: Option<String>,
+ #[serde(default)]
+ arguments: Option<String>,
+}
+
#[cfg(test)]
mod tests {
use super::*;
@@ -422,6 +701,32 @@ mod tests {
assert_eq!(json[0]["function"]["parameters"]["type"], "object");
}
+ #[test]
+ fn streamed_assistant_history_omits_empty_tool_calls() {
+ let value = streamed_assistant_history("hello", &[]);
+ assert_eq!(value["role"], "assistant");
+ assert_eq!(value["content"], "hello");
+ assert!(value.get("tool_calls").is_none());
+ }
+
+ #[test]
+ fn streamed_assistant_history_preserves_tool_calls() {
+ let calls = vec![ToolCall {
+ id: "call_0".to_string(),
+ name: "read_file".to_string(),
+ arguments: r#"{"path":"a.rs"}"#.to_string(),
+ }];
+ let value = streamed_assistant_history("", &calls);
+ assert!(value["content"].is_null());
+ assert_eq!(value["tool_calls"][0]["id"], "call_0");
+ assert_eq!(value["tool_calls"][0]["type"], "function");
+ assert_eq!(value["tool_calls"][0]["function"]["name"], "read_file");
+ assert_eq!(
+ value["tool_calls"][0]["function"]["arguments"],
+ r#"{"path":"a.rs"}"#
+ );
+ }
+
#[test]
fn assistant_message_with_tool_calls_round_trips() {
let message = ResponseMessage {
src/chat.rs
+89
-35
index df35406..f896c70 100644
--- a/src/chat.rs
+++ b/src/chat.rs
@@ -75,7 +75,7 @@ mod tui {
use anyhow::Result;
use crossterm::event::{Event, EventStream, KeyCode, KeyEvent, KeyEventKind, KeyModifiers};
use futures::StreamExt;
- use onde::inference::{ChatEngine, SamplingConfig, StreamChunk};
+ use onde::inference::{ChatEngine, SamplingConfig};
use crate::backend::{InferenceBackend, LocalBackend, OpenAiBackend, ToolResult, ToolSpec};
use crate::models::{ModelCacheHealth, ModelPickerItem, ModelSource, build_model_picker_items};
@@ -148,6 +148,11 @@ mod tui {
enum InferenceUpdate {
/// show tool name in chat while it runs
ToolUse(String),
+ /// a streamed token fragment of the assistant's reply
+ Delta(String),
+ /// the streamed reply is complete; commit the accumulated buffer
+ StreamEnd,
+ /// a complete (non-streamed) assistant reply
Response(String),
Error(String),
}
@@ -164,7 +169,8 @@ mod tui {
input: String,
cursor: usize,
scroll_offset: u16,
- stream_rx: Option<mpsc::Receiver<StreamChunk>>,
+ /// true while assistant tokens are streaming into `stream_buf`
+ streaming: bool,
stream_buf: String,
inference_rx: Option<mpsc::Receiver<InferenceUpdate>>,
model_load_rx: Option<mpsc::Receiver<ModelLoadUpdate>>,
@@ -260,7 +266,7 @@ mod tui {
input: String::new(),
cursor: 0,
scroll_offset: 0,
- stream_rx: None,
+ streaming: false,
stream_buf: String::new(),
inference_rx: None,
model_load_rx: None,
@@ -298,11 +304,11 @@ mod tui {
}
fn is_streaming(&self) -> bool {
- self.stream_rx.is_some()
+ self.streaming
}
fn finalize_stream(&mut self) {
- self.stream_rx = None;
+ self.streaming = false;
if !self.stream_buf.is_empty() {
let text = std::mem::take(&mut self.stream_buf);
self.messages.push(ChatMessage::assistant(text));
@@ -311,11 +317,23 @@ mod tui {
}
fn push_stream_delta(&mut self, delta: &str) {
+ self.streaming = true;
self.stream_buf.push_str(delta);
+ // Hide reasoning the way the rest of the app does: keep the "thinking"
+ // spinner until visible (non-<think>) text appears, then show the
+ // live reply. Don't call stop_thinking() — that drops the channel.
+ let (_think, visible) = super::strip_think_blocks(&self.stream_buf);
+ self.thinking = visible.trim().is_empty();
self.blink_counter = self.blink_counter.wrapping_add(1);
self.blink_on = self.blink_counter % 4 < 2;
}
+ /// The portion of the streaming buffer to show live, with reasoning hidden.
+ fn visible_stream(&self) -> String {
+ let (_think, visible) = super::strip_think_blocks(&self.stream_buf);
+ visible
+ }
+
fn start_thinking(&mut self) {
self.thinking = true;
self.thinking_tick = 0;
@@ -444,8 +462,9 @@ mod tui {
for msg in &self.messages {
lines += wrapped_line_count(&msg.text, msg.role, w);
}
- if !self.stream_buf.is_empty() {
- lines += wrapped_line_count(&self.stream_buf, Role::Assistant, w);
+ let visible = self.visible_stream();
+ if !visible.is_empty() {
+ lines += wrapped_line_count(&visible, Role::Assistant, w);
}
if self.thinking || self.switching_model {
lines += 1;
@@ -851,10 +870,11 @@ mod tui {
render_chat_message(&mut lines, msg, inner_width as usize);
}
- if !app.stream_buf.is_empty() {
+ let streamed_visible = app.visible_stream();
+ if !streamed_visible.is_empty() {
let fake = ChatMessage {
role: Role::Assistant,
- text: app.stream_buf.clone(),
+ text: streamed_visible,
think_block: None,
};
render_chat_message(&mut lines, &fake, inner_width as usize);
@@ -1387,7 +1407,36 @@ mod tui {
vec![]
};
- let mut result = match backend.send_message_with_tools(&text, &tools).await {
+ // Bridge the backend's token sink (plain strings) onto the UI update
+ // channel as `Delta` messages. The forwarder lives for the whole turn.
+ let (delta_tx, mut delta_rx) = mpsc::unbounded_channel::<String>();
+ let forward_tx = tx.clone();
+ let forwarder = tokio::spawn(async move {
+ while let Some(piece) = delta_rx.recv().await {
+ if forward_tx
+ .send(InferenceUpdate::Delta(piece))
+ .await
+ .is_err()
+ {
+ break;
+ }
+ }
+ });
+
+ // The first round offers tools, so on-device inference can't stream it
+ // (it must buffer to detect tool calls). With tools disabled there are
+ // none to offer, so it streams directly.
+ let first_sink = if tools.is_empty() {
+ Some(&delta_tx)
+ } else {
+ None
+ };
+ let mut streamed = first_sink.is_some();
+
+ let mut result = match backend
+ .send_message_with_tools(&text, &tools, first_sink)
+ .await
+ {
Ok(r) => r,
Err(err) => {
let _ = tx.send(InferenceUpdate::Error(err)).await;
@@ -1398,6 +1447,8 @@ mod tui {
let mut round = 0;
while !result.tool_calls.is_empty() && round < MAX_TOOL_ROUNDS {
+ // any tool call means the first round didn't produce a final answer
+ streamed = false;
round += 1;
log::info!("tool round {} — {} call(s)", round, result.tool_calls.len());
@@ -1421,14 +1472,24 @@ mod tui {
});
}
- // on the last round, pass no tools so the model must produce text
+ // on the last round, pass no tools so the model must produce text —
+ // that's also the round we can stream on-device.
let next_tools = if round < MAX_TOOL_ROUNDS {
Some(tools.as_slice())
} else {
None
};
+ let sink = if next_tools.is_none() {
+ streamed = true;
+ Some(&delta_tx)
+ } else {
+ None
+ };
- match backend.send_tool_results(tool_results, next_tools).await {
+ match backend
+ .send_tool_results(tool_results, next_tools, sink)
+ .await
+ {
Ok(r) => result = r,
Err(err) => {
let _ = tx.send(InferenceUpdate::Error(err)).await;
@@ -1437,6 +1498,11 @@ mod tui {
}
}
+ // Drop the sink so the forwarder finishes draining any buffered tokens
+ // before we commit the reply.
+ drop(delta_tx);
+ let _ = forwarder.await;
+
if result.tool_calls.is_empty() {
if result.text.is_empty() {
log::warn!(
@@ -1449,6 +1515,9 @@ mod tui {
.to_string(),
))
.await;
+ } else if streamed {
+ // tokens already went out as deltas; just commit the buffer
+ let _ = tx.send(InferenceUpdate::StreamEnd).await;
} else {
let _ = tx.send(InferenceUpdate::Response(result.text)).await;
}
@@ -1588,29 +1657,6 @@ mod tui {
app.tick();
}
- // ── Streaming LLM tokens ──────────────────────────────────────
- chunk = async {
- match app.stream_rx.as_mut() {
- Some(rx) => rx.recv().await,
- None => pending().await,
- }
- } => {
- match chunk {
- Some(chunk) => {
- if !chunk.delta.is_empty() {
- app.push_stream_delta(&chunk.delta);
- }
- if chunk.done {
- app.finalize_stream();
- }
- }
- // sender dropped without done=true
- None => {
- app.finalize_stream();
- }
- }
- }
-
// ── inference updates from background task ───────────────────
update = async {
match app.inference_rx.as_mut() {
@@ -1622,16 +1668,24 @@ mod tui {
Some(InferenceUpdate::ToolUse(name)) => {
app.messages.push(ChatMessage::system(format!("🔧 {name}")));
}
+ Some(InferenceUpdate::Delta(delta)) => {
+ app.push_stream_delta(&delta);
+ }
+ Some(InferenceUpdate::StreamEnd) => {
+ app.finalize_stream();
+ }
Some(InferenceUpdate::Response(text)) => {
app.stop_thinking();
app.messages.push(ChatMessage::assistant(text));
}
Some(InferenceUpdate::Error(msg)) => {
+ app.finalize_stream();
app.stop_thinking();
app.messages.push(ChatMessage::system(format!("error: {msg}")));
}
None => {
// task finished, possibly with no text to show
+ app.finalize_stream();
app.stop_thinking();
}
}
src/main.rs
+126
-29
index 9f89ba1..8d6de8e 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -61,6 +61,7 @@ use onde::inference::{ChatEngine, GgufModelConfig};
use crate::backend::{
InferenceBackend, LocalBackend, OpenAiBackend, ToolResult as BackendToolResult, ToolSpec,
+ TurnResult,
};
use std::path::PathBuf;
use std::sync::atomic::{AtomicBool, Ordering};
@@ -568,6 +569,72 @@ impl SiGitAgent {
))
}
+ /// Run one inference turn (`fut`) while concurrently forwarding any streamed
+ /// tokens to the editor. The sink receiver is drained as the future runs, so
+ /// chunks reach the client live rather than all at once when it resolves.
+ ///
+ /// `assembled`/`sent`/`streamed_any` persist across the turns of a single
+ /// prompt so reasoning is stripped consistently and we never re-send text.
+ #[allow(clippy::too_many_arguments)]
+ async fn drain_turn<F>(
+ &self,
+ cx: &ConnectionTo<Client>,
+ session_id: &SessionId,
+ fut: F,
+ sink_rx: &mut tokio::sync::mpsc::UnboundedReceiver<String>,
+ assembled: &mut String,
+ sent: &mut String,
+ streamed_any: &mut bool,
+ ) -> Result<TurnResult, backend::BackendError>
+ where
+ F: std::future::Future<Output = Result<TurnResult, backend::BackendError>>,
+ {
+ tokio::pin!(fut);
+ let result = loop {
+ tokio::select! {
+ done = &mut fut => break done,
+ Some(piece) = sink_rx.recv() => {
+ self.emit_visible_chunk(cx, session_id, &piece, assembled, sent, streamed_any);
+ }
+ }
+ };
+ // Flush tokens that landed between the last poll and the future resolving.
+ while let Ok(piece) = sink_rx.try_recv() {
+ self.emit_visible_chunk(cx, session_id, &piece, assembled, sent, streamed_any);
+ }
+ result
+ }
+
+ /// Append a streamed fragment, strip `<think>` reasoning from the running
+ /// text, and send only the newly revealed visible suffix as a chunk. Tracking
+ /// the assembled text (not just deltas) keeps think-block stripping correct
+ /// even when a tag spans chunk boundaries.
+ fn emit_visible_chunk(
+ &self,
+ cx: &ConnectionTo<Client>,
+ session_id: &SessionId,
+ piece: &str,
+ assembled: &mut String,
+ sent: &mut String,
+ streamed_any: &mut bool,
+ ) {
+ assembled.push_str(piece);
+ let (_think, visible) = chat::strip_think_blocks(assembled);
+ match visible.strip_prefix(sent.as_str()) {
+ Some(extra) if !extra.is_empty() => {
+ let extra = extra.to_string();
+ *sent = visible;
+ *streamed_any = true;
+ self.send_assistant_message(cx, session_id.clone(), extra)
+ .ok();
+ }
+ // No new visible text, or the visible prefix changed retroactively
+ // (rare, e.g. a late-closing think tag): just resync without
+ // resending what's already on the wire.
+ _ => *sent = visible,
+ }
+ }
+
fn send_tool_call_update(
&self,
cx: &ConnectionTo<Client>,
@@ -1076,8 +1143,26 @@ impl SiGitAgent {
let tools = agent_tools_as_specs();
- let mut result = backend
- .send_message_with_tools(&user_text, &tools)
+ // Token sink: backends stream assistant text through this while a turn
+ // runs. We forward the visible portion to the editor as agent-message
+ // chunks live (see `drain_turn` / `emit_visible_chunk`). The sink stays
+ // alive for the whole prompt so `recv()` only ends when a turn future
+ // resolves, never because every sender was dropped.
+ let (sink, mut sink_rx) = tokio::sync::mpsc::unbounded_channel::<String>();
+ let mut assembled = String::new();
+ let mut sent = String::new();
+ let mut streamed_any = false;
+
+ let mut result = self
+ .drain_turn(
+ cx,
+ &session_id,
+ backend.send_message_with_tools(&user_text, &tools, Some(&sink)),
+ &mut sink_rx,
+ &mut assembled,
+ &mut sent,
+ &mut streamed_any,
+ )
.await
.map_err(|error| {
log::error!("send_message_with_tools failed: {error}");
@@ -1120,39 +1205,51 @@ impl SiGitAgent {
None // last round: force text
};
- result = backend
- .send_tool_results(tool_results, next_tools)
+ result = self
+ .drain_turn(
+ cx,
+ &session_id,
+ backend.send_tool_results(tool_results, next_tools, Some(&sink)),
+ &mut sink_rx,
+ &mut assembled,
+ &mut sent,
+ &mut streamed_any,
+ )
.await
.map_err(|e| agent_client_protocol::Error::new(-32603, e.to_string()))?;
}
- // ── Send the final text response ─────────────────────────────────
- let reply_text = result.text.trim().to_string();
-
- let final_text = if reply_text.is_empty() {
- if round > 0 {
- log::warn!(
- "prompt({}) — model returned empty reply after {} tool round(s)",
- session_id,
- round
- );
- "Something went wrong — the edits didn't go through. Try rephrasing what you need, or point me at the specific lines.".to_string()
+ // ── Final text response ───────────────────────────────────────────
+ // If anything streamed, the visible reply is already on the wire; only
+ // send a trailing block for the non-streamed path (e.g. on-device direct
+ // answers, which onde can't stream while tools are on offer).
+ if !streamed_any {
+ let reply_text = result.text.trim().to_string();
+ let final_text = if reply_text.is_empty() {
+ if round > 0 {
+ log::warn!(
+ "prompt({}) — model returned empty reply after {} tool round(s)",
+ session_id,
+ round
+ );
+ "Something went wrong — the edits didn't go through. Try rephrasing what you need, or point me at the specific lines.".to_string()
+ } else {
+ log::warn!(
+ "prompt({}) — model returned empty reply (no tool rounds)",
+ session_id
+ );
+ String::new()
+ }
} else {
- log::warn!(
- "prompt({}) — model returned empty reply (no tool rounds)",
- session_id
- );
- String::new()
- }
- } else {
- // strip <think> blocks so reasoning tokens stay hidden
- let (_think, visible) = chat::strip_think_blocks(&reply_text);
- visible
- };
+ // strip <think> blocks so reasoning tokens stay hidden
+ let (_think, visible) = chat::strip_think_blocks(&reply_text);
+ visible
+ };
- if !final_text.is_empty() {
- self.send_assistant_message(cx, session_id.clone(), final_text)
- .ok();
+ if !final_text.is_empty() {
+ self.send_assistant_message(cx, session_id.clone(), final_text)
+ .ok();
+ }
}
log::info!("prompt({}) complete — {} tool round(s)", session_id, round);