Skip to main content

agent_engine/engine/
stream.rs

1//! Engine-level stream processing — TUI-agnostic event handling.
2//!
3//! Processes StreamEvent variants, tracks subagent state and usage,
4//! and returns renderer-agnostic actions.
5
6use crate::{StreamEvent, LlmEvent, SessionEvent, AgentEvent};
7
8/// What happened during a stream event — renderer decides how to display.
9#[derive(Debug)]
10pub enum EngineStreamEvent {
11    /// Thinking text arrived.
12    Thinking(String),
13    /// Response text arrived.
14    Text(String),
15    /// Tool use started.
16    ToolStart { tool_id: String, tool_name: String },
17    /// Tool use input delta.
18    ToolDelta { tool_id: String, delta: String },
19    /// Tool use finalized.
20    ///
21    /// `input` is the parsed JSON value, not a stringified version. Renderers
22    /// that need a string preview (chat.rs, server's HistoryEntry::ToolUse)
23    /// can call `serde_json::to_string` themselves; the wire-format ToolUse
24    /// in server mode passes the Value through directly.
25    ToolFinalized { tool_id: String, tool_name: String, input: serde_json::Value },
26    /// Tool result delta.
27    ToolResultDelta { tool_id: String, delta: String },
28    /// Tool result complete.
29    ToolResult { tool_id: String, result: String },
30    /// Subagent dispatched.
31    SubagentStart { id: u64, name: String, task: String },
32    /// Subagent status update.
33    SubagentUpdate { id: u64, status: String },
34    /// Subagent finished.
35    SubagentDone { id: u64, status: String, duration_secs: f64 },
36    /// Steering message was delivered.
37    SteeringDelivered { message: String },
38    /// Usage stats for this turn.
39    Usage {
40        input_tokens: u64,
41        output_tokens: u64,
42        cache_read: u64,
43        cache_creation: u64,
44        /// Cache-write TTL split — `None` when the API omitted the breakdown.
45        cache_creation_5m: Option<u64>,
46        cache_creation_1h: Option<u64>,
47        model: Option<String>,
48    },
49    /// Internal bookkeeping — no visual output.
50    Noop,
51    /// Display-only status notice (e.g. API retry). Not part of the transcript.
52    Notice(String),
53    /// Stream completed.
54    Done,
55    /// Stream errored.
56    Error(String),
57}
58
59/// Subagent tracking state.
60#[derive(Debug, Clone)]
61pub struct SubagentTracker {
62    pub id: u64,
63    pub name: String,
64    pub status: String,
65    pub start_time: std::time::Instant,
66    pub done: bool,
67    pub duration_secs: Option<f64>,
68}
69
70/// What the caller should do after stream completion.
71#[derive(Debug)]
72pub enum StreamCompletion {
73    /// Stream is still going.
74    Continue,
75    /// Stream done — auto-send this queued message.
76    AutoSendQueued(String),
77    /// Stream done — pending events need a new model turn.
78    AutoTriggerEvents,
79    /// Stream done — nothing special.
80    Done,
81    /// Stream errored.
82    Error(String),
83}
84
85/// Convert a raw StreamEvent into an EngineStreamEvent.
86/// Also handles message history capture and returns completion signals.
87///
88/// `messages` — the conversation history (updated in place on MessageHistory)
89/// `subagents` — tracked subagent states (updated in place)
90/// `queued_message` — message queued during streaming (taken if stream completes)
91/// `pending_events` — events buffered during streaming (drained on completion)
92pub fn process_stream_event(
93    event: StreamEvent,
94    messages: &mut Vec<crate::SharedMessage>,
95    subagents: &mut Vec<SubagentTracker>,
96    queued_message: &mut Option<String>,
97    pending_events: &mut Vec<String>,
98) -> (EngineStreamEvent, StreamCompletion) {
99    match event {
100        StreamEvent::Llm(LlmEvent::Thinking(text)) => {
101            (EngineStreamEvent::Thinking(text), StreamCompletion::Continue)
102        }
103        StreamEvent::Llm(LlmEvent::Text(text)) => {
104            (EngineStreamEvent::Text(text), StreamCompletion::Continue)
105        }
106        StreamEvent::Llm(LlmEvent::ToolUseStart { tool_name, tool_id }) => {
107            (EngineStreamEvent::ToolStart { tool_id, tool_name }, StreamCompletion::Continue)
108        }
109        StreamEvent::Llm(LlmEvent::ToolUseDelta { tool_id, delta }) => {
110            (EngineStreamEvent::ToolDelta { tool_id, delta }, StreamCompletion::Continue)
111        }
112        StreamEvent::Llm(LlmEvent::ToolUse { tool_name, tool_id, input }) => {
113            (EngineStreamEvent::ToolFinalized { tool_id, tool_name, input }, StreamCompletion::Continue)
114        }
115        StreamEvent::Llm(LlmEvent::ToolResultDelta { delta, tool_id }) => {
116            (EngineStreamEvent::ToolResultDelta { tool_id, delta }, StreamCompletion::Continue)
117        }
118        StreamEvent::Llm(LlmEvent::ToolResult { result, tool_id }) => {
119            (EngineStreamEvent::ToolResult { tool_id, result }, StreamCompletion::Continue)
120        }
121        StreamEvent::Session(SessionEvent::MessageHistory(history)) => {
122            *messages = history;
123            (EngineStreamEvent::Noop, StreamCompletion::Continue)
124        }
125        StreamEvent::Agent(AgentEvent::SubagentStart { subagent_id, agent_name, task_preview }) => {
126            subagents.push(SubagentTracker {
127                id: subagent_id,
128                name: agent_name.clone(),
129                status: format!("starting: {}", task_preview),
130                start_time: std::time::Instant::now(),
131                done: false,
132                duration_secs: None,
133            });
134            (EngineStreamEvent::SubagentStart { id: subagent_id, name: agent_name, task: task_preview }, StreamCompletion::Continue)
135        }
136        StreamEvent::Agent(AgentEvent::SubagentUpdate { subagent_id, status, .. }) => {
137            if let Some(sa) = subagents.iter_mut().find(|s| s.id == subagent_id) {
138                sa.status = status.clone();
139            }
140            (EngineStreamEvent::SubagentUpdate { id: subagent_id, status }, StreamCompletion::Continue)
141        }
142        StreamEvent::Agent(AgentEvent::SubagentDone { subagent_id, result_preview, duration_secs, .. }) => {
143            let status = if result_preview.starts_with("[TIMED OUT") {
144                "\u{26a0} timed out".to_string()
145            } else if result_preview.starts_with("ERROR") {
146                let preview: String = result_preview.chars().take(40).collect();
147                format!("\u{2718} {}", preview)
148            } else {
149                let preview: String = result_preview.chars().take(40).collect();
150                format!("\u{2714} {}", preview)
151            };
152            if let Some(sa) = subagents.iter_mut().find(|s| s.id == subagent_id) {
153                sa.done = true;
154                sa.duration_secs = Some(duration_secs);
155                sa.status = status.clone();
156            }
157            (EngineStreamEvent::SubagentDone { id: subagent_id, status, duration_secs }, StreamCompletion::Continue)
158        }
159        StreamEvent::Agent(AgentEvent::SteeringDelivered { message }) => {
160            if queued_message.as_ref() == Some(&message) {
161                *queued_message = None;
162            }
163            (EngineStreamEvent::SteeringDelivered { message }, StreamCompletion::Continue)
164        }
165        StreamEvent::Session(SessionEvent::Usage {
166            input_tokens,
167            output_tokens,
168            cache_read_input_tokens,
169            cache_creation_input_tokens,
170            cache_creation_5m,
171            cache_creation_1h,
172            model,
173        }) => {
174            (EngineStreamEvent::Usage {
175                input_tokens,
176                output_tokens,
177                cache_read: cache_read_input_tokens,
178                cache_creation: cache_creation_input_tokens,
179                cache_creation_5m,
180                cache_creation_1h,
181                model,
182            }, StreamCompletion::Continue)
183        }
184        StreamEvent::Session(SessionEvent::Done) => {
185            subagents.clear();
186
187            // Drain pending events into messages
188            let had_pending = !pending_events.is_empty();
189            for formatted in pending_events.drain(..) {
190                messages.push(std::sync::Arc::new(serde_json::json!({
191                    "role": "user",
192                    "content": formatted
193                })));
194            }
195
196            // Check for queued message
197            if let Some(queued) = queued_message.take() {
198                (EngineStreamEvent::Done, StreamCompletion::AutoSendQueued(queued))
199            } else if had_pending {
200                (EngineStreamEvent::Done, StreamCompletion::AutoTriggerEvents)
201            } else {
202                (EngineStreamEvent::Done, StreamCompletion::Done)
203            }
204        }
205        StreamEvent::Session(SessionEvent::Notice(text)) => {
206            // Display-only status (e.g. retry notice) — surface as a system
207            // line, never recorded into message history.
208            (EngineStreamEvent::Notice(text), StreamCompletion::Continue)
209        }
210        StreamEvent::Session(SessionEvent::Error(err)) => {
211            subagents.clear();
212            // Drop trailing unmatched messages
213            if let Some(last) = messages.last() {
214                let role = last["role"].as_str().unwrap_or("");
215                let is_text_user = role == "user" && last["content"].is_string();
216                let is_assistant = role == "assistant";
217                if is_text_user || is_assistant {
218                    messages.pop();
219                }
220            }
221            (EngineStreamEvent::Error(err.clone()), StreamCompletion::Error(err))
222        }
223    }
224}