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