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