1use std::collections::{HashSet, VecDeque};
2
3use serde_json::Value;
4use tokio::sync::broadcast;
5
6use crate::events::{RunEvent, ToolStatus};
7use crate::result::RunResult;
8use crate::run_handle::SharedRunResult;
9
10pub struct RunEventStream {
11 pending: VecDeque<RunEvent>,
12 seen_event_ids: HashSet<String>,
13 receiver: Option<broadcast::Receiver<RunEvent>>,
14 shared_result: Option<SharedRunResult>,
15}
16
17impl RunEventStream {
18 pub(crate) fn from_live(
19 receiver: Option<broadcast::Receiver<RunEvent>>,
20 result: Option<SharedRunResult>,
21 backlog: Vec<RunEvent>,
22 ) -> Self {
23 let seen_event_ids = backlog
24 .iter()
25 .map(|event| event.event_id().as_str().to_string())
26 .collect();
27 Self {
28 pending: VecDeque::from(backlog),
29 seen_event_ids,
30 receiver,
31 shared_result: result,
32 }
33 }
34
35 pub async fn next(&mut self) -> Option<Result<RunEvent, String>> {
36 if let Some(event) = self.pending.pop_front() {
37 return Some(Ok(event));
38 }
39 let receiver = self.receiver.as_mut()?;
40 loop {
41 match receiver.recv().await {
42 Ok(event) => {
43 if self
44 .seen_event_ids
45 .insert(event.event_id().as_str().to_string())
46 {
47 return Some(Ok(event));
48 }
49 }
50 Err(broadcast::error::RecvError::Closed) => return None,
51 Err(broadcast::error::RecvError::Lagged(count)) => {
52 return Some(Err(format!(
53 "run event stream lagged and dropped {count} events"
54 )));
55 }
56 }
57 }
58 }
59
60 pub async fn into_result(mut self) -> Result<RunResult, String> {
61 if let Some(result) = self.shared_result.take() {
62 return result.wait().await;
63 }
64 Err("stream result already taken".to_string())
65 }
66}
67
68pub(super) fn map_runtime_event(
69 event: &str,
70 payload: &std::collections::BTreeMap<String, Value>,
71) -> Option<RunEvent> {
72 let run_id = payload
73 .get("task_id")
74 .and_then(Value::as_str)
75 .or_else(|| payload.get("run_id").and_then(Value::as_str))
76 .unwrap_or("run")
77 .to_string();
78 let trace_id = payload
79 .get("trace_id")
80 .and_then(Value::as_str)
81 .unwrap_or(&run_id)
82 .to_string();
83 let agent_name = payload
84 .get("agent_name")
85 .and_then(Value::as_str)
86 .unwrap_or("agent")
87 .to_string();
88 match event {
89 "run_started" => Some(RunEvent::run_started(
90 run_id,
91 trace_id,
92 agent_name,
93 payload
94 .get("input")
95 .and_then(Value::as_str)
96 .unwrap_or_default(),
97 )),
98 "cycle_started" => Some(RunEvent::cycle_started(
99 run_id,
100 trace_id,
101 agent_name,
102 payload
103 .get("cycle")
104 .and_then(Value::as_u64)
105 .unwrap_or_default() as u32,
106 )),
107 "cycle_llm_response" => {
108 let cycle_index = payload
109 .get("cycle")
110 .and_then(Value::as_u64)
111 .unwrap_or_default() as u32;
112 payload
113 .get("assistant_message")
114 .and_then(Value::as_str)
115 .filter(|delta| !delta.is_empty())
116 .map(|delta| {
117 RunEvent::assistant_delta(run_id, trace_id, agent_name, cycle_index, delta)
118 })
119 }
120 "tool_call_started" => Some(RunEvent::tool_call_started(
121 run_id,
122 trace_id,
123 agent_name,
124 payload
125 .get("cycle")
126 .and_then(Value::as_u64)
127 .unwrap_or_default() as u32,
128 payload
129 .get("tool_call_id")
130 .and_then(Value::as_str)
131 .unwrap_or_default(),
132 payload
133 .get("tool_name")
134 .and_then(Value::as_str)
135 .unwrap_or_default(),
136 payload
137 .get("tool_arguments")
138 .cloned()
139 .unwrap_or(Value::Null),
140 )),
141 "approval_requested" => Some(RunEvent::approval_requested(
142 run_id,
143 trace_id,
144 agent_name,
145 payload
146 .get("request_id")
147 .and_then(Value::as_str)
148 .unwrap_or_default(),
149 payload
150 .get("tool_call_id")
151 .and_then(Value::as_str)
152 .unwrap_or_default(),
153 payload
154 .get("tool_name")
155 .and_then(Value::as_str)
156 .unwrap_or_default(),
157 payload
158 .get("preview")
159 .and_then(Value::as_str)
160 .unwrap_or_default(),
161 )),
162 "approval_resolved" => Some(RunEvent::new(
163 run_id,
164 trace_id,
165 agent_name,
166 payload
167 .get("cycle")
168 .and_then(Value::as_u64)
169 .map(|cycle| cycle as u32),
170 crate::events::RunEventPayload::ApprovalResolved {
171 request_id: payload
172 .get("request_id")
173 .and_then(Value::as_str)
174 .unwrap_or_default()
175 .to_string(),
176 tool_call_id: payload
177 .get("tool_call_id")
178 .and_then(Value::as_str)
179 .unwrap_or_default()
180 .to_string(),
181 tool_name: payload
182 .get("tool_name")
183 .and_then(Value::as_str)
184 .unwrap_or_default()
185 .to_string(),
186 approved: payload
187 .get("approved")
188 .and_then(Value::as_bool)
189 .unwrap_or(false),
190 },
191 )),
192 "sub_run_started" => Some(
193 RunEvent::new(
194 payload
195 .get("task_id_hint")
196 .and_then(Value::as_str)
197 .unwrap_or("sub_run"),
198 trace_id,
199 agent_name,
200 payload
201 .get("cycle")
202 .and_then(Value::as_u64)
203 .map(|cycle| cycle as u32),
204 crate::events::RunEventPayload::SubRunStarted {
205 parent_tool_call_id: payload
206 .get("parent_tool_call_id")
207 .and_then(Value::as_str)
208 .unwrap_or_default()
209 .to_string(),
210 child_session_id: payload
211 .get("child_session_id")
212 .and_then(Value::as_str)
213 .map(str::to_string),
214 task_id: payload
215 .get("task_id_hint")
216 .and_then(Value::as_str)
217 .map(str::to_string),
218 },
219 )
220 .with_parent_run_id(
221 payload
222 .get("parent_run_id")
223 .and_then(Value::as_str)
224 .unwrap_or(&run_id),
225 ),
226 ),
227 "sub_run_completed" => Some(
228 RunEvent::new(
229 format!(
230 "sub_run:{}",
231 payload
232 .get("parent_tool_call_id")
233 .and_then(Value::as_str)
234 .unwrap_or_default()
235 ),
236 trace_id,
237 agent_name,
238 payload
239 .get("cycle")
240 .and_then(Value::as_u64)
241 .map(|cycle| cycle as u32),
242 crate::events::RunEventPayload::SubRunCompleted {
243 parent_tool_call_id: payload
244 .get("parent_tool_call_id")
245 .and_then(Value::as_str)
246 .unwrap_or_default()
247 .to_string(),
248 status: crate::types::AgentStatus::Completed,
249 final_output: payload
250 .get("final_output")
251 .and_then(Value::as_str)
252 .map(str::to_string),
253 },
254 )
255 .with_parent_run_id(
256 payload
257 .get("parent_run_id")
258 .and_then(Value::as_str)
259 .unwrap_or(&run_id),
260 ),
261 ),
262 "tool_result" => {
263 let metadata = payload.get("metadata").and_then(Value::as_object);
264 if let Some(interruption_id) = metadata
265 .and_then(|metadata| metadata.get("approval_interruption_id"))
266 .and_then(Value::as_str)
267 {
268 return Some(RunEvent::approval_requested(
269 run_id,
270 trace_id,
271 agent_name,
272 interruption_id,
273 payload
274 .get("tool_call_id")
275 .and_then(Value::as_str)
276 .unwrap_or_default(),
277 payload
278 .get("tool_name")
279 .and_then(Value::as_str)
280 .unwrap_or_default(),
281 "Approval required for tool call.",
282 ));
283 }
284 let status = match payload
285 .get("status")
286 .and_then(Value::as_str)
287 .unwrap_or_default()
288 {
289 "error" | "ERROR" => ToolStatus::Error,
290 "wait_response" | "WAIT_RESPONSE" => ToolStatus::WaitResponse,
291 _ => ToolStatus::Success,
292 };
293 Some(RunEvent::tool_call_completed(
294 run_id,
295 trace_id,
296 agent_name,
297 None,
298 payload
299 .get("tool_call_id")
300 .and_then(Value::as_str)
301 .unwrap_or_default(),
302 payload
303 .get("tool_name")
304 .and_then(Value::as_str)
305 .unwrap_or_default(),
306 status,
307 ))
308 }
309 "run_completed" => Some(RunEvent::run_completed(
310 run_id,
311 trace_id,
312 agent_name,
313 crate::types::AgentStatus::Completed,
314 )),
315 "run_wait_user" => Some(RunEvent::run_completed(
316 run_id,
317 trace_id,
318 agent_name,
319 crate::types::AgentStatus::WaitUser,
320 )),
321 "run_cancelled" => Some(RunEvent::new(
322 run_id,
323 trace_id,
324 agent_name,
325 None,
326 crate::events::RunEventPayload::RunCancelled {
327 reason: payload
328 .get("reason")
329 .and_then(Value::as_str)
330 .or_else(|| payload.get("error").and_then(Value::as_str))
331 .unwrap_or("run cancelled")
332 .to_string(),
333 },
334 )),
335 "run_max_cycles" => Some(RunEvent::run_completed(
336 run_id,
337 trace_id,
338 agent_name,
339 crate::types::AgentStatus::MaxCycles,
340 )),
341 "cycle_failed" => Some(RunEvent::run_failed(
342 run_id,
343 trace_id,
344 agent_name,
345 crate::events::AgentErrorPayload::new(
346 payload
347 .get("error")
348 .and_then(Value::as_str)
349 .unwrap_or("cycle failed"),
350 ),
351 )),
352 _ => None,
353 }
354}