1#![allow(missing_docs)]
2
3use std::io::{self, Write};
4
5use agent_base::{AgentResult, RuntimeEvent, UserEvent};
6use serde_json::{Value, json};
7
8use crate::render::EventRenderer;
9
10pub struct JsonStreamRenderer {
13 writer: Box<dyn Write + Send>,
14 turn_start: Option<std::time::Instant>,
15 tool_call_count: u32,
16 last_assistant_text: String,
17}
18
19impl JsonStreamRenderer {
20 pub fn new(writer: Box<dyn Write + Send>) -> Self {
22 Self { writer, turn_start: None, tool_call_count: 0, last_assistant_text: String::new() }
23 }
24
25 pub fn stdout() -> Self {
27 Self::new(Box::new(io::stdout()))
28 }
29
30 fn emit(&mut self, value: &Value) -> AgentResult<()> {
31 let line = serde_json::to_string(value)
32 .map_err(|e| agent_base::AgentError::internal(format!("JSON serialize error: {e}")))?;
33 writeln!(self.writer, "{}", line).map_err(|e| agent_base::AgentError::internal(format!("write error: {e}")))?;
34 Ok(())
35 }
36
37 fn emit_event(&mut self, event: &RuntimeEvent, mut value: Value) -> AgentResult<()> {
39 if let Some(agent_id) = event.agent_id() {
40 value["agent_id"] = json!(agent_id);
41 }
42 self.emit(&value)
43 }
44}
45
46impl EventRenderer for JsonStreamRenderer {
47 fn render(&mut self, event: RuntimeEvent) -> AgentResult<()> {
48 if self.turn_start.is_none() {
49 self.turn_start = Some(std::time::Instant::now());
50 }
51
52 match &event {
53 RuntimeEvent::ThoughtDelta { text, .. } => {
54 self.emit_event(&event, json!({ "type": "thought_delta", "text": text }))?;
55 },
56 RuntimeEvent::TextDelta { text, .. } => {
57 if event.agent_id().is_none() {
58 self.last_assistant_text.push_str(text);
59 }
60 self.emit_event(&event, json!({ "type": "text_delta", "text": text }))?;
61 },
62 RuntimeEvent::ToolCallStarted { tool_name, args_json, .. } => {
63 self.tool_call_count += 1;
64 let args: Value = serde_json::from_str(args_json).unwrap_or(Value::Null);
65 self.emit_event(
66 &event,
67 json!({
68 "type": "tool_call_started",
69 "tool": tool_name,
70 "args": args,
71 }),
72 )?;
73 },
74 RuntimeEvent::ToolCallFinished { tool_name, summary, denied, .. } => {
75 self.emit_event(
76 &event,
77 json!({
78 "type": "tool_call_finished",
79 "tool": tool_name,
80 "summary": summary,
81 "denied": denied,
82 }),
83 )?;
84 },
85 RuntimeEvent::AwaitingApproval { request, .. } => {
86 self.emit_event(
87 &event,
88 json!({
89 "type": "approval_request",
90 "title": request.title,
91 "risk": format!("{:?}", request.risk_level),
92 "message": request.message,
93 }),
94 )?;
95 },
96 RuntimeEvent::PlanUpdated { explanation, plan, .. } => {
97 self.emit_event(
98 &event,
99 json!({
100 "type": "plan_updated",
101 "explanation": explanation,
102 "plan": plan,
103 }),
104 )?;
105 },
106 RuntimeEvent::UserEvent { event: UserEvent::Structured { event_type, data }, .. } => {
107 self.emit_event(
108 &event,
109 json!({
110 "type": "user_event",
111 "event_type": event_type,
112 "data": data,
113 }),
114 )?;
115 },
116 RuntimeEvent::UserEvent { .. } => {},
117 RuntimeEvent::Checkpoint { .. } => {},
118 RuntimeEvent::RunFinished { .. } => {},
119 RuntimeEvent::RunCancelled { .. } => {
120 self.emit_event(&event, json!({ "type": "run_cancelled" }))?;
121 },
122 }
123
124 Ok(())
125 }
126
127 fn finish_turn(&mut self) -> AgentResult<()> {
128 let duration_ms = self.turn_start.map(|s| s.elapsed().as_millis() as u64).unwrap_or(0);
129
130 self.emit(&json!({
131 "type": "turn_finished",
132 "duration_ms": duration_ms,
133 "tool_call_count": self.tool_call_count,
134 "assistant_text": self.last_assistant_text.trim(),
135 }))?;
136
137 self.turn_start = None;
138 self.tool_call_count = 0;
139 self.last_assistant_text.clear();
140
141 Ok(())
142 }
143}
144
145#[cfg(test)]
146mod tests {
147 use super::*;
148 use agent_base::{ApprovalRequest, PlanItem, PlanStepStatus, RiskLevel, SessionId, UserEvent};
149 use std::io::Write;
150 use std::sync::{Arc, Mutex};
151
152 struct SharedWriter {
153 inner: Arc<Mutex<Vec<u8>>>,
154 }
155
156 impl Write for SharedWriter {
157 fn write(&mut self, data: &[u8]) -> std::io::Result<usize> {
158 self.inner.lock().unwrap().extend_from_slice(data);
159 Ok(data.len())
160 }
161 fn flush(&mut self) -> std::io::Result<()> {
162 Ok(())
163 }
164 }
165
166 impl SharedWriter {
167 fn new() -> (Self, Arc<Mutex<Vec<u8>>>) {
168 let inner = Arc::new(Mutex::new(Vec::new()));
169 (Self { inner: inner.clone() }, inner)
170 }
171 }
172
173 fn session_id() -> SessionId {
174 SessionId { id: 1, external_id: None }
175 }
176
177 fn render_one(event: RuntimeEvent) -> Vec<String> {
178 let (writer, buf) = SharedWriter::new();
179 let mut r = JsonStreamRenderer::new(Box::new(writer));
180 r.render(event).unwrap();
181 drop(r);
182 let text = String::from_utf8(buf.lock().unwrap().clone()).unwrap();
183 if text.is_empty() { vec![] } else { text.lines().map(|l| l.to_string()).collect() }
184 }
185
186 fn render_and_finish(events: &[RuntimeEvent]) -> Vec<String> {
187 let (writer, buf) = SharedWriter::new();
188 let mut r = JsonStreamRenderer::new(Box::new(writer));
189 for e in events {
190 r.render(e.clone()).unwrap();
191 }
192 r.finish_turn().unwrap();
193 drop(r);
194 let text = String::from_utf8(buf.lock().unwrap().clone()).unwrap();
195 text.lines().map(|l| l.to_string()).collect()
196 }
197
198 #[test]
199 fn test_text_delta_produces_valid_json() {
200 let lines = render_one(RuntimeEvent::TextDelta {
201 session_id: session_id(),
202 text: "hello".into(),
203 agent_id: None,
204 trace_id: None,
205 });
206 assert_eq!(lines.len(), 1);
207 let v: serde_json::Value = serde_json::from_str(&lines[0]).unwrap();
208 assert_eq!(v["type"], "text_delta");
209 assert_eq!(v["text"], "hello");
210 }
211
212 #[test]
213 fn test_thought_delta_produces_valid_json() {
214 let lines = render_one(RuntimeEvent::ThoughtDelta {
215 session_id: session_id(),
216 text: "thinking...".into(),
217 agent_id: None,
218 trace_id: None,
219 });
220 let v: serde_json::Value = serde_json::from_str(&lines[0]).unwrap();
221 assert_eq!(v["type"], "thought_delta");
222 }
223
224 #[test]
225 fn test_tool_call_started_parses_args() {
226 let lines = render_one(RuntimeEvent::ToolCallStarted {
227 session_id: session_id(),
228 tool_name: "shell".into(),
229 args_json: r#"{"cmd":"ls"}"#.into(),
230 agent_id: None,
231 trace_id: None,
232 });
233 let v: serde_json::Value = serde_json::from_str(&lines[0]).unwrap();
234 assert_eq!(v["type"], "tool_call_started");
235 assert_eq!(v["tool"], "shell");
236 assert_eq!(v["args"]["cmd"], "ls");
237 }
238
239 #[test]
240 fn test_tool_call_finished_produces_valid_json() {
241 let lines = render_one(RuntimeEvent::ToolCallFinished {
242 session_id: session_id(),
243 tool_name: "shell".into(),
244 summary: "done".into(),
245 agent_id: None,
246 trace_id: None,
247 denied: false,
248 });
249 let v: serde_json::Value = serde_json::from_str(&lines[0]).unwrap();
250 assert_eq!(v["type"], "tool_call_finished");
251 assert_eq!(v["tool"], "shell");
252 assert_eq!(v["summary"], "done");
253 assert_eq!(v["denied"], false);
254 }
255
256 #[test]
257 fn test_awaiting_approval_produces_valid_json() {
258 let lines = render_one(RuntimeEvent::AwaitingApproval {
259 session_id: session_id(),
260 request: ApprovalRequest {
261 title: "Delete".into(),
262 message: "Dangerous".into(),
263 action_key: None,
264 risk_level: RiskLevel::Destructive,
265 raw: None,
266 },
267 agent_id: None,
268 trace_id: None,
269 });
270 let v: serde_json::Value = serde_json::from_str(&lines[0]).unwrap();
271 assert_eq!(v["type"], "approval_request");
272 assert_eq!(v["title"], "Delete");
273 }
274
275 #[test]
276 fn test_plan_updated_produces_valid_json() {
277 let lines = render_one(RuntimeEvent::PlanUpdated {
278 session_id: session_id(),
279 objective: "test".into(),
280 explanation: Some("step 1 done".into()),
281 plan: vec![PlanItem { step: "Step 1".into(), status: PlanStepStatus::Completed }],
282 agent_id: None,
283 trace_id: None,
284 });
285 let v: serde_json::Value = serde_json::from_str(&lines[0]).unwrap();
286 assert_eq!(v["type"], "plan_updated");
287 }
288
289 #[test]
290 fn test_user_event_structured() {
291 let lines = render_one(RuntimeEvent::UserEvent {
292 session_id: session_id(),
293 event: UserEvent::Structured { event_type: "custom".into(), data: serde_json::json!({"key": "value"}) },
294 agent_id: None,
295 trace_id: None,
296 });
297 let v: serde_json::Value = serde_json::from_str(&lines[0]).unwrap();
298 assert_eq!(v["type"], "user_event");
299 assert_eq!(v["event_type"], "custom");
300 assert_eq!(v["data"]["key"], "value");
301 }
302
303 #[test]
304 fn test_user_event_progress_ignored() {
305 let lines = render_one(RuntimeEvent::UserEvent {
306 session_id: session_id(),
307 event: UserEvent::Progress { text: "loading...".into() },
308 agent_id: None,
309 trace_id: None,
310 });
311 assert!(lines.is_empty());
312 }
313
314 #[test]
315 fn test_run_finished_no_output() {
316 let lines = render_one(RuntimeEvent::RunFinished { session_id: session_id(), agent_id: None, trace_id: None });
317 assert!(lines.is_empty());
318 }
319
320 #[test]
321 fn test_run_cancelled_produces_valid_json() {
322 let lines = render_one(RuntimeEvent::RunCancelled { session_id: session_id(), agent_id: None, trace_id: None });
323 let v: serde_json::Value = serde_json::from_str(&lines[0]).unwrap();
324 assert_eq!(v["type"], "run_cancelled");
325 }
326
327 #[test]
328 fn test_finish_turn_emits_summary() {
329 let lines = render_and_finish(&[RuntimeEvent::TextDelta {
330 session_id: session_id(),
331 text: "hello".into(),
332 agent_id: None,
333 trace_id: None,
334 }]);
335 let last: serde_json::Value = serde_json::from_str(lines.last().unwrap()).unwrap();
336 assert_eq!(last["type"], "turn_finished");
337 assert!(last["duration_ms"].as_u64().is_some());
338 assert_eq!(last["tool_call_count"], 0);
339 }
340
341 #[test]
342 fn test_tool_call_count_incremented() {
343 let lines = render_and_finish(&[
344 RuntimeEvent::ToolCallStarted {
345 session_id: session_id(),
346 tool_name: "a".into(),
347 args_json: "{}".into(),
348 agent_id: None,
349 trace_id: None,
350 },
351 RuntimeEvent::ToolCallStarted {
352 session_id: session_id(),
353 tool_name: "b".into(),
354 args_json: "{}".into(),
355 agent_id: None,
356 trace_id: None,
357 },
358 ]);
359 let last: serde_json::Value = serde_json::from_str(lines.last().unwrap()).unwrap();
360 assert_eq!(last["tool_call_count"], 2);
361 }
362
363 #[test]
364 fn test_assistant_text_accumulated() {
365 let lines = render_and_finish(&[
366 RuntimeEvent::TextDelta { session_id: session_id(), text: "Hello ".into(), agent_id: None, trace_id: None },
367 RuntimeEvent::TextDelta { session_id: session_id(), text: "World".into(), agent_id: None, trace_id: None },
368 ]);
369 let last: serde_json::Value = serde_json::from_str(lines.last().unwrap()).unwrap();
370 assert_eq!(last["assistant_text"], "Hello World");
371 }
372
373 #[test]
376 fn test_subagent_text_delta_includes_agent_id() {
377 let lines = render_one(RuntimeEvent::TextDelta {
378 session_id: session_id(),
379 text: "found 3 items".into(),
380 agent_id: Some("root/searcher".into()),
381 trace_id: None,
382 });
383 assert_eq!(lines.len(), 1);
384 let v: serde_json::Value = serde_json::from_str(&lines[0]).unwrap();
385 assert_eq!(v["type"], "text_delta");
386 assert_eq!(v["text"], "found 3 items");
387 assert_eq!(v["agent_id"], "root/searcher");
388 }
389
390 #[test]
391 fn test_subagent_tool_call_includes_agent_id() {
392 let lines = render_one(RuntimeEvent::ToolCallStarted {
393 session_id: session_id(),
394 tool_name: "shell".into(),
395 args_json: r#"{"cmd":"ls"}"#.into(),
396 agent_id: Some("root/worker".into()),
397 trace_id: None,
398 });
399 let v: serde_json::Value = serde_json::from_str(&lines[0]).unwrap();
400 assert_eq!(v["type"], "tool_call_started");
401 assert_eq!(v["tool"], "shell");
402 assert_eq!(v["agent_id"], "root/worker");
403 }
404
405 #[test]
406 fn test_root_event_omits_agent_id() {
407 let lines = render_one(RuntimeEvent::TextDelta {
408 session_id: session_id(),
409 text: "hello".into(),
410 agent_id: None,
411 trace_id: None,
412 });
413 let v: serde_json::Value = serde_json::from_str(&lines[0]).unwrap();
414 assert!(v.get("agent_id").is_none());
415 }
416
417 #[test]
418 fn test_subagent_text_does_not_pollute_assistant_text() {
419 let lines = render_and_finish(&[
420 RuntimeEvent::TextDelta { session_id: session_id(), text: "root ".into(), agent_id: None, trace_id: None },
421 RuntimeEvent::TextDelta {
422 session_id: session_id(),
423 text: "child".into(),
424 agent_id: Some("root/child".into()),
425 trace_id: None,
426 },
427 ]);
428 let last: serde_json::Value = serde_json::from_str(lines.last().unwrap()).unwrap();
429 assert_eq!(last["assistant_text"], "root");
430 }
431}