1use aether_core::core::Prompt;
2use aether_core::events::{AgentMessage, Command};
3use std::io;
4use std::process::ExitCode;
5use tokio::sync::mpsc;
6
7use super::error::CliError;
8use super::{CliEventKind, OutputFormat, RunConfig};
9use crate::runtime::RuntimeBuilder;
10
11pub async fn run(config: RunConfig) -> Result<ExitCode, CliError> {
12 setup_tracing(config.verbose, &config.output);
13
14 let agent = RuntimeBuilder::from_spec(config.cwd.clone(), config.spec)
15 .mcp_sources(config.mcp_config_sources)
16 .build(config.system_prompt.as_deref().map(Prompt::text), None)
17 .await?;
18
19 agent
20 .agent_tx
21 .send(Command::text(&config.prompt))
22 .await
23 .map_err(|e| CliError::AgentError(format!("Failed to send prompt: {e}")))?;
24
25 Ok(stream_output(agent.agent_rx, &config.output, &config.events).await)
26}
27
28async fn stream_output(
29 mut rx: mpsc::Receiver<AgentMessage>,
30 format: &OutputFormat,
31 events: &[CliEventKind],
32) -> ExitCode {
33 while let Some(msg) = rx.recv().await {
34 if should_emit(&msg, events) {
35 match format {
36 OutputFormat::Text => emit_text(&msg),
37 OutputFormat::Pretty | OutputFormat::Json => emit_event(&msg),
38 }
39 }
40 if matches!(msg, AgentMessage::Done) {
41 break;
42 }
43 }
44 ExitCode::SUCCESS
45}
46
47fn should_emit(msg: &AgentMessage, include: &[CliEventKind]) -> bool {
48 if include.is_empty() {
49 return true;
50 }
51 event_kind(msg).is_none_or(|ty| include.contains(&ty))
52}
53
54fn event_kind(msg: &AgentMessage) -> Option<CliEventKind> {
55 match msg {
56 AgentMessage::Text { is_complete: true, .. } => Some(CliEventKind::Text),
57 AgentMessage::Thought { is_complete: true, .. } => Some(CliEventKind::Thought),
58 AgentMessage::ToolCall { .. } => Some(CliEventKind::ToolCall),
59 AgentMessage::ToolResult { .. } => Some(CliEventKind::ToolResult),
60 AgentMessage::ToolError { .. } => Some(CliEventKind::ToolError),
61 AgentMessage::Error { .. } => Some(CliEventKind::Error),
62 AgentMessage::Cancelled { .. } => Some(CliEventKind::Cancelled),
63 AgentMessage::AutoContinue { .. } => Some(CliEventKind::AutoContinue),
64 AgentMessage::Retrying { .. } => Some(CliEventKind::Retrying),
65 AgentMessage::ModelSwitched { .. } => Some(CliEventKind::ModelSwitched),
66 AgentMessage::ToolProgress { .. } => Some(CliEventKind::ToolProgress),
67 AgentMessage::ContextCompactionStarted { .. } => Some(CliEventKind::ContextCompactionStarted),
68 AgentMessage::ContextCompactionResult { .. } => Some(CliEventKind::ContextCompactionResult),
69 AgentMessage::ContextUsageUpdate { .. } => Some(CliEventKind::ContextUsage),
70 AgentMessage::ContextCleared => Some(CliEventKind::ContextCleared),
71 AgentMessage::Text { is_complete: false, .. }
72 | AgentMessage::Thought { is_complete: false, .. }
73 | AgentMessage::ToolCallUpdate { .. }
74 | AgentMessage::Done => None,
75 }
76}
77
78fn format_text(msg: &AgentMessage) -> Option<String> {
79 match msg {
80 AgentMessage::Text { chunk, is_complete: true, .. } => Some(chunk.clone()),
81
82 AgentMessage::Thought { chunk, is_complete: true, .. } => Some(format!("Thought: {chunk}")),
83
84 AgentMessage::ToolCall { request, .. } => Some(format!("Tool call: {}({})", request.name, request.arguments)),
85
86 AgentMessage::ToolResult { result, .. } => Some(format!("Tool result [{}]: {}", result.name, result.result)),
87
88 AgentMessage::ToolError { error, .. } => Some(format!("Tool error [{}]: {}", error.name, error.error)),
89
90 AgentMessage::Error { message } => Some(format!("Error: {message}")),
91
92 AgentMessage::Cancelled { message } => Some(format!("Cancelled: {message}")),
93
94 AgentMessage::AutoContinue { attempt, max_attempts } => {
95 Some(format!("Continuing ({attempt}/{max_attempts})..."))
96 }
97
98 AgentMessage::Retrying { attempt, max_attempts, delay_ms, error } => {
99 Some(format!("Retrying ({attempt}/{max_attempts}) in {delay_ms}ms: {error}"))
100 }
101
102 AgentMessage::ModelSwitched { previous, new } => Some(format!("Model switched: {previous} -> {new}")),
103
104 AgentMessage::ToolProgress { request, progress, total, message } => {
105 let bar = match total {
106 Some(t) => format!("{progress}/{t}"),
107 None => format!("{progress}"),
108 };
109 let suffix = message.as_deref().map(|m| format!(" - {m}")).unwrap_or_default();
110 Some(format!("Tool progress [{}]: {bar}{suffix}", request.name))
111 }
112
113 AgentMessage::ContextCompactionStarted { message_count } => {
114 Some(format!("Context compaction started ({message_count} messages)"))
115 }
116
117 AgentMessage::ContextCompactionResult { summary, messages_removed } => {
118 Some(format!("Context compacted: {messages_removed} messages removed. {summary}"))
119 }
120
121 AgentMessage::ContextUsageUpdate {
122 input_tokens,
123 output_tokens,
124 total_input_tokens,
125 total_output_tokens,
126 ..
127 } => Some(format!(
128 "Tokens: {input_tokens} in, {output_tokens} out (total: {total_input_tokens} in, {total_output_tokens} out)"
129 )),
130
131 AgentMessage::ContextCleared => Some("Context cleared".to_string()),
132
133 AgentMessage::ToolCallUpdate { .. }
134 | AgentMessage::Text { .. }
135 | AgentMessage::Thought { .. }
136 | AgentMessage::Done => None,
137 }
138}
139
140fn emit_text(msg: &AgentMessage) {
141 if let Some(text) = format_text(msg) {
142 if matches!(msg, AgentMessage::Error { .. }) {
143 eprintln!("{text}");
144 } else {
145 println!("{text}");
146 }
147 }
148}
149
150#[allow(clippy::too_many_lines)]
151fn emit_event(msg: &AgentMessage) {
152 let kind = event_kind(msg).map_or("", CliEventKind::as_str);
153 match msg {
154 AgentMessage::Text { chunk, is_complete: true, .. } => {
155 tracing::info!(target: "agent", kind, "{chunk}");
156 }
157
158 AgentMessage::Thought { chunk, is_complete: true, .. } => {
159 tracing::info!(target: "agent", kind, thought = %chunk);
160 }
161
162 AgentMessage::ToolCall { request, .. } => {
163 tracing::info!(
164 target: "agent",
165 kind,
166 tool = %request.name,
167 arguments = %request.arguments,
168 );
169 }
170
171 AgentMessage::ToolResult { result, .. } => {
172 tracing::info!(
173 target: "agent",
174 kind,
175 tool = %result.name,
176 result = %result.result,
177 );
178 }
179
180 AgentMessage::ToolError { error, .. } => {
181 tracing::warn!(
182 target: "agent",
183 kind,
184 tool = %error.name,
185 error = %error.error,
186 );
187 }
188
189 AgentMessage::Error { message } => {
190 tracing::error!(target: "agent", kind, "{message}");
191 }
192
193 AgentMessage::Cancelled { message } => {
194 tracing::info!(target: "agent", kind, cancelled = %message);
195 }
196
197 AgentMessage::AutoContinue { attempt, max_attempts } => {
198 tracing::info!(
199 target: "agent",
200 kind,
201 attempt,
202 max_attempts,
203 "Continuing ({attempt}/{max_attempts})..."
204 );
205 }
206
207 AgentMessage::Retrying { attempt, max_attempts, delay_ms, error } => {
208 tracing::info!(
209 target: "agent",
210 kind,
211 attempt,
212 max_attempts,
213 delay_ms,
214 error = %error,
215 "Retrying ({attempt}/{max_attempts}) in {delay_ms}ms: {error}"
216 );
217 }
218
219 AgentMessage::ModelSwitched { previous, new } => {
220 tracing::info!(
221 target: "agent",
222 kind,
223 previous = %previous,
224 new = %new,
225 "Model switched: {previous} -> {new}"
226 );
227 }
228
229 AgentMessage::ToolProgress { request, progress, total, message } => {
230 tracing::info!(
231 target: "agent",
232 kind,
233 tool = %request.name,
234 progress,
235 total = ?total,
236 message = ?message,
237 );
238 }
239
240 AgentMessage::ContextCompactionStarted { message_count } => {
241 tracing::info!(
242 target: "agent",
243 kind,
244 message_count,
245 "context compaction started"
246 );
247 }
248
249 AgentMessage::ContextCompactionResult { summary, messages_removed } => {
250 tracing::info!(
251 target: "agent",
252 kind,
253 messages_removed,
254 summary = %summary,
255 "context compaction result"
256 );
257 }
258
259 AgentMessage::ContextUsageUpdate {
260 usage_ratio,
261 context_limit,
262 input_tokens,
263 output_tokens,
264 cache_read_tokens,
265 cache_creation_tokens,
266 reasoning_tokens,
267 total_input_tokens,
268 total_output_tokens,
269 total_cache_read_tokens,
270 total_cache_creation_tokens,
271 total_reasoning_tokens,
272 } => {
273 tracing::info!(
274 target: "agent",
275 kind,
276 usage_ratio = ?usage_ratio,
277 context_limit = ?context_limit,
278 input_tokens,
279 output_tokens,
280 cache_read_tokens = cache_read_tokens.unwrap_or(0),
281 cache_creation_tokens = cache_creation_tokens.unwrap_or(0),
282 reasoning_tokens = reasoning_tokens.unwrap_or(0),
283 total_input_tokens,
284 total_output_tokens,
285 total_cache_read_tokens,
286 total_cache_creation_tokens,
287 total_reasoning_tokens,
288 "context usage"
289 );
290 }
291
292 AgentMessage::ContextCleared => {
293 tracing::info!(target: "agent", kind, "context cleared");
294 }
295
296 AgentMessage::ToolCallUpdate { .. }
297 | AgentMessage::Text { .. }
298 | AgentMessage::Thought { .. }
299 | AgentMessage::Done => {}
300 }
301}
302
303fn setup_tracing(verbose: bool, format: &OutputFormat) {
304 use tracing_subscriber::Layer;
305 use tracing_subscriber::filter::{self, EnvFilter};
306 use tracing_subscriber::fmt;
307 use tracing_subscriber::layer::SubscriberExt;
308 use tracing_subscriber::util::SubscriberInitExt;
309
310 let diag_filter = if verbose { EnvFilter::new("debug,agent=off") } else { EnvFilter::new("error,agent=off") };
311
312 let diag_layer = fmt::layer().with_writer(io::stderr).with_filter(diag_filter);
313
314 let agent_filter = filter::filter_fn(|meta| meta.target().starts_with("agent"));
315
316 match format {
317 OutputFormat::Text => {
318 if verbose {
319 tracing_subscriber::registry().with(diag_layer).init();
320 } else {
321 tracing_subscriber::registry().init();
323 }
324 }
325 OutputFormat::Pretty => {
326 let agent_layer = fmt::layer().with_writer(io::stdout).pretty().with_filter(agent_filter);
327 tracing_subscriber::registry().with(diag_layer).with(agent_layer).init();
328 }
329 OutputFormat::Json => {
330 let agent_layer = fmt::layer().with_writer(io::stdout).json().with_filter(agent_filter);
331 tracing_subscriber::registry().with(diag_layer).with(agent_layer).init();
332 }
333 }
334}
335
336#[cfg(test)]
337mod tests {
338 use super::*;
339 use std::sync::{Arc, Mutex};
340
341 use tracing_subscriber::Layer;
342 use tracing_subscriber::fmt;
343 use tracing_subscriber::layer::SubscriberExt;
344
345 #[test]
346 fn emit_event_emits_complete_text() {
347 let output = with_test_subscriber(|| {
348 emit_event(&AgentMessage::text("id", "hello", true, "model"));
349 });
350 assert!(output.contains("hello"), "expected 'hello' in: {output}");
351 }
352
353 #[test]
354 fn emit_event_skips_incomplete_text() {
355 let output = with_test_subscriber(|| {
356 emit_event(&AgentMessage::text("id", "hello", false, "model"));
357 });
358 assert!(output.is_empty(), "expected empty output, got: {output}");
359 }
360
361 #[test]
362 fn emit_event_emits_complete_thought() {
363 let output = with_test_subscriber(|| {
364 emit_event(&AgentMessage::thought("id", "deep thinking", true, "model"));
365 });
366 assert!(output.contains("deep thinking"), "expected 'deep thinking' in: {output}");
367 }
368
369 #[test]
370 fn emit_event_skips_incomplete_thought() {
371 let output = with_test_subscriber(|| {
372 emit_event(&AgentMessage::thought("id", "partial", false, "model"));
373 });
374 assert!(output.is_empty(), "expected empty output, got: {output}");
375 }
376
377 #[test]
378 fn emit_event_emits_tool_call() {
379 let msg = AgentMessage::ToolCall {
380 request: llm::ToolCallRequest {
381 id: "tc1".to_string(),
382 name: "bash".to_string(),
383 arguments: "{}".to_string(),
384 },
385 model_name: "test".to_string(),
386 };
387 let output = with_test_subscriber(|| {
388 emit_event(&msg);
389 });
390 assert!(output.contains("bash"), "expected 'bash' in: {output}");
391 }
392
393 #[test]
394 fn emit_event_skips_tool_call_updates() {
395 let msg = AgentMessage::ToolCallUpdate {
396 tool_call_id: "tc1".to_string(),
397 chunk: "{\"partial".to_string(),
398 model_name: "test".to_string(),
399 };
400 let output = with_test_subscriber(|| {
401 emit_event(&msg);
402 });
403 assert!(output.is_empty(), "expected empty output, got: {output}");
404 }
405
406 #[test]
407 fn emit_event_emits_tool_result() {
408 let msg = AgentMessage::ToolResult {
409 result: llm::ToolCallResult {
410 id: "tc1".to_string(),
411 name: "bash".to_string(),
412 arguments: "{}".to_string(),
413 result: "ok".to_string(),
414 },
415 result_meta: None,
416 model_name: "test".to_string(),
417 };
418 let output = with_test_subscriber(|| {
419 emit_event(&msg);
420 });
421 assert!(output.contains("bash"), "expected 'bash' in: {output}");
422 assert!(output.contains("ok"), "expected 'ok' in: {output}");
423 }
424
425 #[test]
426 fn emit_event_emits_error() {
427 let msg = AgentMessage::Error { message: "something broke".to_string() };
428 let output = with_test_subscriber(|| {
429 emit_event(&msg);
430 });
431 assert!(output.contains("something broke"), "expected 'something broke' in: {output}");
432 }
433
434 #[test]
435 fn emit_event_skips_done() {
436 let output = with_test_subscriber(|| {
437 emit_event(&AgentMessage::Done);
438 });
439 assert!(output.is_empty(), "expected empty output, got: {output}");
440 }
441
442 #[test]
445 fn emit_text_formats_complete_text() {
446 assert_eq!(format_text(&AgentMessage::text("id", "hello world", true, "m")), Some("hello world".to_string()));
447 }
448
449 #[test]
450 fn emit_text_skips_incomplete_text() {
451 assert_eq!(format_text(&AgentMessage::text("id", "partial", false, "m")), None);
452 }
453
454 #[test]
455 fn emit_text_formats_complete_thought() {
456 assert_eq!(
457 format_text(&AgentMessage::thought("id", "reasoning here", true, "m")),
458 Some("Thought: reasoning here".to_string())
459 );
460 }
461
462 #[test]
463 fn emit_text_skips_incomplete_thought() {
464 assert_eq!(format_text(&AgentMessage::thought("id", "partial", false, "m")), None);
465 }
466
467 #[test]
468 fn emit_text_formats_tool_call() {
469 let msg = AgentMessage::ToolCall {
470 request: llm::ToolCallRequest {
471 id: "tc1".to_string(),
472 name: "bash".to_string(),
473 arguments: r#"{"cmd":"ls"}"#.to_string(),
474 },
475 model_name: "test".to_string(),
476 };
477 assert_eq!(format_text(&msg), Some(r#"Tool call: bash({"cmd":"ls"})"#.to_string()));
478 }
479
480 #[test]
481 fn emit_text_skips_tool_call_updates() {
482 let msg = AgentMessage::ToolCallUpdate {
483 tool_call_id: "tc1".to_string(),
484 chunk: "partial".to_string(),
485 model_name: "test".to_string(),
486 };
487 assert_eq!(format_text(&msg), None);
488 }
489
490 #[test]
491 fn emit_text_formats_tool_result() {
492 let msg = AgentMessage::ToolResult {
493 result: llm::ToolCallResult {
494 id: "tc1".to_string(),
495 name: "bash".to_string(),
496 arguments: "{}".to_string(),
497 result: "output".to_string(),
498 },
499 result_meta: None,
500 model_name: "test".to_string(),
501 };
502 assert_eq!(format_text(&msg), Some("Tool result [bash]: output".to_string()));
503 }
504
505 #[test]
506 fn emit_text_formats_tool_error() {
507 let msg = AgentMessage::ToolError {
508 error: llm::ToolCallError {
509 id: "tc1".to_string(),
510 name: "bash".to_string(),
511 arguments: None,
512 error: "not found".to_string(),
513 },
514 model_name: "test".to_string(),
515 };
516 assert_eq!(format_text(&msg), Some("Tool error [bash]: not found".to_string()));
517 }
518
519 #[test]
520 fn emit_text_formats_error() {
521 let msg = AgentMessage::Error { message: "boom".to_string() };
522 assert_eq!(format_text(&msg), Some("Error: boom".to_string()));
523 }
524
525 #[test]
526 fn emit_text_formats_cancelled() {
527 let msg = AgentMessage::Cancelled { message: "user stopped".to_string() };
528 assert_eq!(format_text(&msg), Some("Cancelled: user stopped".to_string()));
529 }
530
531 #[test]
532 fn emit_text_formats_auto_continue() {
533 let msg = AgentMessage::AutoContinue { attempt: 2, max_attempts: 5 };
534 assert_eq!(format_text(&msg), Some("Continuing (2/5)...".to_string()));
535 }
536
537 #[test]
538 fn emit_text_formats_model_switched() {
539 let msg = AgentMessage::ModelSwitched { previous: "old-model".to_string(), new: "new-model".to_string() };
540 assert_eq!(format_text(&msg), Some("Model switched: old-model -> new-model".to_string()));
541 }
542
543 #[test]
544 fn emit_text_skips_done() {
545 assert_eq!(format_text(&AgentMessage::Done), None);
546 }
547
548 fn tool_progress(progress: f64, total: Option<f64>, message: Option<&str>) -> AgentMessage {
549 AgentMessage::ToolProgress {
550 request: llm::ToolCallRequest {
551 id: "tc1".to_string(),
552 name: "bash".to_string(),
553 arguments: "{}".to_string(),
554 },
555 progress,
556 total,
557 message: message.map(str::to_string),
558 }
559 }
560
561 fn usage_update() -> AgentMessage {
562 AgentMessage::ContextUsageUpdate {
563 usage_ratio: Some(0.25),
564 context_limit: Some(200_000),
565 input_tokens: 1500,
566 output_tokens: 250,
567 cache_read_tokens: Some(400),
568 cache_creation_tokens: Some(100),
569 reasoning_tokens: Some(50),
570 total_input_tokens: 5000,
571 total_output_tokens: 800,
572 total_cache_read_tokens: 1200,
573 total_cache_creation_tokens: 300,
574 total_reasoning_tokens: 150,
575 }
576 }
577
578 #[test]
579 fn emit_text_formats_tool_progress_with_total() {
580 let msg = tool_progress(50.0, Some(100.0), Some("halfway"));
581 assert_eq!(format_text(&msg), Some("Tool progress [bash]: 50/100 - halfway".to_string()));
582 }
583
584 #[test]
585 fn emit_text_formats_tool_progress_without_total() {
586 let msg = tool_progress(42.0, None, None);
587 assert_eq!(format_text(&msg), Some("Tool progress [bash]: 42".to_string()));
588 }
589
590 #[test]
591 fn emit_text_formats_context_compaction_started() {
592 let msg = AgentMessage::ContextCompactionStarted { message_count: 42 };
593 assert_eq!(format_text(&msg), Some("Context compaction started (42 messages)".to_string()));
594 }
595
596 #[test]
597 fn emit_text_formats_context_compaction_result() {
598 let msg = AgentMessage::ContextCompactionResult { summary: "summary here".to_string(), messages_removed: 10 };
599 assert_eq!(format_text(&msg), Some("Context compacted: 10 messages removed. summary here".to_string()));
600 }
601
602 #[test]
603 fn emit_text_formats_context_usage_update() {
604 assert_eq!(
605 format_text(&usage_update()),
606 Some("Tokens: 1500 in, 250 out (total: 5000 in, 800 out)".to_string())
607 );
608 }
609
610 #[test]
611 fn emit_text_formats_context_cleared() {
612 assert_eq!(format_text(&AgentMessage::ContextCleared), Some("Context cleared".to_string()));
613 }
614
615 #[test]
616 fn emit_event_emits_tool_progress() {
617 let output = with_test_subscriber(|| emit_event(&tool_progress(3.0, Some(10.0), Some("step"))));
618 assert!(output.contains("tool_progress"), "missing type: {output}");
619 assert!(output.contains("bash"), "missing tool name: {output}");
620 assert!(output.contains('3'), "missing progress: {output}");
621 }
622
623 #[test]
624 fn emit_event_emits_context_compaction_started() {
625 let msg = AgentMessage::ContextCompactionStarted { message_count: 7 };
626 let output = with_test_subscriber(|| emit_event(&msg));
627 assert!(output.contains("context_compaction_started"), "missing type: {output}");
628 assert!(output.contains('7'), "missing message_count: {output}");
629 }
630
631 #[test]
632 fn emit_event_emits_context_compaction_result() {
633 let msg = AgentMessage::ContextCompactionResult { summary: "done".to_string(), messages_removed: 5 };
634 let output = with_test_subscriber(|| emit_event(&msg));
635 assert!(output.contains("context_compaction_result"), "missing type: {output}");
636 assert!(output.contains("done"), "missing summary: {output}");
637 }
638
639 #[test]
640 fn emit_event_emits_context_usage_update() {
641 let output = with_test_subscriber(|| emit_event(&usage_update()));
642 assert!(output.contains("context_usage"), "missing type: {output}");
643 assert!(output.contains("1500"), "missing input_tokens: {output}");
644 assert!(output.contains("5000"), "missing total_input_tokens: {output}");
645 }
646
647 #[test]
648 fn emit_event_emits_context_cleared() {
649 let output = with_test_subscriber(|| emit_event(&AgentMessage::ContextCleared));
650 assert!(output.contains("context_cleared"), "missing type: {output}");
651 }
652
653 #[test]
654 fn emit_event_includes_type_for_tool_call() {
655 let msg = AgentMessage::ToolCall {
656 request: llm::ToolCallRequest {
657 id: "tc1".to_string(),
658 name: "bash".to_string(),
659 arguments: "{}".to_string(),
660 },
661 model_name: "test".to_string(),
662 };
663 let output = with_test_subscriber(|| emit_event(&msg));
664 assert!(output.contains("tool_call"), "missing type: {output}");
665 }
666
667 fn tool_call_msg() -> AgentMessage {
668 AgentMessage::ToolCall {
669 request: llm::ToolCallRequest {
670 id: "tc1".to_string(),
671 name: "bash".to_string(),
672 arguments: "{}".to_string(),
673 },
674 model_name: "test".to_string(),
675 }
676 }
677
678 fn tool_result_msg() -> AgentMessage {
679 AgentMessage::ToolResult {
680 result: llm::ToolCallResult {
681 id: "tc1".to_string(),
682 name: "bash".to_string(),
683 arguments: "{}".to_string(),
684 result: "ok".to_string(),
685 },
686 result_meta: None,
687 model_name: "test".to_string(),
688 }
689 }
690
691 #[test]
692 fn event_kind_none_for_non_filterable_variants() {
693 assert_eq!(event_kind(&AgentMessage::Done), None);
694 assert_eq!(event_kind(&AgentMessage::text("id", "x", false, "m")), None);
695 assert_eq!(event_kind(&AgentMessage::thought("id", "x", false, "m")), None);
696 assert_eq!(
697 event_kind(&AgentMessage::ToolCallUpdate {
698 tool_call_id: "tc1".to_string(),
699 chunk: "x".to_string(),
700 model_name: "m".to_string(),
701 }),
702 None,
703 );
704 }
705
706 #[test]
707 fn should_emit_empty_filter_allows_everything() {
708 assert!(should_emit(&tool_call_msg(), &[]));
709 assert!(should_emit(&AgentMessage::Error { message: "e".to_string() }, &[]));
710 assert!(should_emit(&AgentMessage::Done, &[]));
711 }
712
713 #[test]
714 fn should_emit_single_type_whitelist() {
715 let filter = &[CliEventKind::ToolCall];
716 assert!(should_emit(&tool_call_msg(), filter));
717 assert!(!should_emit(&tool_result_msg(), filter));
718 assert!(!should_emit(&AgentMessage::Error { message: "e".to_string() }, filter));
719 }
720
721 #[test]
722 fn should_emit_multi_type_whitelist() {
723 let filter = &[CliEventKind::ToolCall, CliEventKind::ToolResult];
724 assert!(should_emit(&tool_call_msg(), filter));
725 assert!(should_emit(&tool_result_msg(), filter));
726 assert!(!should_emit(&AgentMessage::Error { message: "e".to_string() }, filter));
727 }
728
729 #[test]
730 fn should_emit_none_typed_variants_pass_through_even_with_filter() {
731 let filter = &[CliEventKind::ToolCall];
732 assert!(should_emit(&AgentMessage::Done, filter));
733 assert!(should_emit(&AgentMessage::text("id", "x", false, "m"), filter));
734 assert!(should_emit(
735 &AgentMessage::ToolCallUpdate {
736 tool_call_id: "tc1".to_string(),
737 chunk: "x".to_string(),
738 model_name: "m".to_string(),
739 },
740 filter,
741 ));
742 }
743
744 #[tokio::test]
745 async fn stream_output_filter_only_emits_whitelisted_types() {
746 let (tx, rx) = mpsc::channel(16);
747 tx.send(tool_call_msg()).await.unwrap();
748 tx.send(tool_result_msg()).await.unwrap();
749 tx.send(AgentMessage::Error { message: "boom".to_string() }).await.unwrap();
750 tx.send(AgentMessage::Done).await.unwrap();
751 drop(tx);
752
753 let filter = vec![CliEventKind::ToolCall];
754 let (_guard, buf) = test_subscriber_guard();
755
756 let code = stream_output(rx, &OutputFormat::Pretty, &filter).await;
757 assert_eq!(code, ExitCode::SUCCESS);
758
759 let output = String::from_utf8(buf.lock().unwrap().clone()).unwrap();
760 assert!(output.contains("tool_call"), "tool_call missing: {output}");
761 assert!(!output.contains("tool_result"), "tool_result leaked past filter: {output}");
762 assert!(!output.contains("boom"), "error leaked past filter: {output}");
763 }
764
765 #[tokio::test]
766 async fn stream_output_done_breaks_loop_under_filter() {
767 let (tx, rx) = mpsc::channel(4);
768 tx.send(AgentMessage::Done).await.unwrap();
769 let filter = vec![CliEventKind::ToolCall];
770 let code = stream_output(rx, &OutputFormat::Text, &filter).await;
771 assert_eq!(code, ExitCode::SUCCESS);
772 }
773
774 fn test_subscriber_guard() -> (tracing::subscriber::DefaultGuard, Arc<Mutex<Vec<u8>>>) {
775 let buf = Arc::new(Mutex::new(Vec::new()));
776 let buf_clone = Arc::clone(&buf);
777
778 let writer = move || -> TestWriter { TestWriter { buf: Arc::clone(&buf_clone) } };
779
780 let layer = fmt::layer()
781 .with_writer(writer)
782 .with_ansi(false)
783 .with_level(false)
784 .with_target(false)
785 .with_timer(fmt::time::uptime())
786 .with_filter(tracing_subscriber::filter::filter_fn(|meta| meta.target().starts_with("agent")));
787
788 let subscriber = tracing_subscriber::registry().with(layer);
789 let guard = tracing::subscriber::set_default(subscriber);
790 (guard, buf)
791 }
792
793 fn with_test_subscriber<F: FnOnce()>(f: F) -> String {
794 let (_guard, buf) = test_subscriber_guard();
795 f();
796 let bytes = buf.lock().unwrap();
797 String::from_utf8(bytes.clone()).unwrap()
798 }
799
800 #[derive(Clone)]
801 struct TestWriter {
802 buf: Arc<Mutex<Vec<u8>>>,
803 }
804
805 impl io::Write for TestWriter {
806 fn write(&mut self, data: &[u8]) -> io::Result<usize> {
807 self.buf.lock().unwrap().extend_from_slice(data);
808 Ok(data.len())
809 }
810
811 fn flush(&mut self) -> io::Result<()> {
812 Ok(())
813 }
814 }
815}