Skip to main content

aether_cli/headless/
run.rs

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, RunConfig};
9use crate::output::OutputFormat;
10use crate::runtime::RuntimeBuilder;
11
12pub async fn run(config: RunConfig) -> Result<ExitCode, CliError> {
13    setup_tracing(config.verbose);
14
15    let mut spec = config.spec;
16    if let Some(system_prompt) = config.system_prompt {
17        spec.prompts.push(Prompt::text(&system_prompt));
18    }
19
20    let (agent, _mcp_snapshot) = RuntimeBuilder::from_spec(config.cwd.clone(), spec)
21        .mcp_sources(config.mcp_config_sources)
22        .oauth_credential_store(config.oauth_credential_store)
23        .build_ready(vec![])
24        .await?;
25
26    agent
27        .agent_tx
28        .send(Command::text(&config.prompt))
29        .await
30        .map_err(|e| CliError::AgentError(format!("Failed to send prompt: {e}")))?;
31
32    Ok(stream_output(agent.agent_rx, config.output, &config.events).await)
33}
34
35async fn stream_output(
36    mut rx: mpsc::Receiver<AgentMessage>,
37    format: OutputFormat,
38    events: &[CliEventKind],
39) -> ExitCode {
40    while let Some(msg) = rx.recv().await {
41        if should_emit(&msg, events)
42            && let Err(error) = print_message(format, &msg)
43        {
44            eprintln!("Failed to serialize headless event: {error}");
45            return ExitCode::FAILURE;
46        }
47
48        if matches!(msg, AgentMessage::Done) {
49            break;
50        }
51    }
52    ExitCode::SUCCESS
53}
54
55fn print_message(format: OutputFormat, msg: &AgentMessage) -> Result<(), serde_json::Error> {
56    match format {
57        OutputFormat::Text => {
58            if let Some(text) = format_text(msg) {
59                if matches!(msg, AgentMessage::Error { .. }) {
60                    eprintln!("{text}");
61                } else {
62                    println!("{text}");
63                }
64            }
65        }
66        OutputFormat::Pretty => println!("{}", serde_json::to_string_pretty(msg)?),
67        OutputFormat::Json => println!("{}", serde_json::to_string(msg)?),
68    }
69
70    Ok(())
71}
72
73fn should_emit(msg: &AgentMessage, include: &[CliEventKind]) -> bool {
74    let Some(kind) = event_kind(msg) else { return false };
75    include.is_empty() || include.contains(&kind)
76}
77
78fn event_kind(msg: &AgentMessage) -> Option<CliEventKind> {
79    match msg {
80        AgentMessage::Text { is_complete: true, .. } => Some(CliEventKind::Text),
81        AgentMessage::Thought { is_complete: true, .. } => Some(CliEventKind::Thought),
82        AgentMessage::ToolCall { .. } => Some(CliEventKind::ToolCall),
83        AgentMessage::ToolResult { .. } => Some(CliEventKind::ToolResult),
84        AgentMessage::ToolError { .. } => Some(CliEventKind::ToolError),
85        AgentMessage::Error { .. } => Some(CliEventKind::Error),
86        AgentMessage::Cancelled { .. } => Some(CliEventKind::Cancelled),
87        AgentMessage::AutoContinue { .. } => Some(CliEventKind::AutoContinue),
88        AgentMessage::Retrying { .. } => Some(CliEventKind::Retrying),
89        AgentMessage::ModelSwitched { .. } => Some(CliEventKind::ModelSwitched),
90        AgentMessage::ToolProgress { .. } => Some(CliEventKind::ToolProgress),
91        AgentMessage::ContextCompactionStarted { .. } => Some(CliEventKind::ContextCompactionStarted),
92        AgentMessage::ContextCompactionResult { .. } => Some(CliEventKind::ContextCompactionResult),
93        AgentMessage::ContextUsageUpdate { .. } => Some(CliEventKind::ContextUsage),
94        AgentMessage::ContextCleared => Some(CliEventKind::ContextCleared),
95        AgentMessage::Done => Some(CliEventKind::Done),
96        AgentMessage::Text { is_complete: false, .. }
97        | AgentMessage::Thought { is_complete: false, .. }
98        | AgentMessage::ToolCallUpdate { .. } => None,
99    }
100}
101
102fn format_text(msg: &AgentMessage) -> Option<String> {
103    match msg {
104        AgentMessage::Text { chunk, is_complete: true, .. } => Some(chunk.clone()),
105
106        AgentMessage::Thought { chunk, is_complete: true, .. } => Some(format!("Thought: {chunk}")),
107
108        AgentMessage::ToolCall { request, .. } => Some(format!("Tool call: {}({})", request.name, request.arguments)),
109
110        AgentMessage::ToolResult { result, .. } => Some(format!("Tool result [{}]: {}", result.name, result.result)),
111
112        AgentMessage::ToolError { error, .. } => Some(format!("Tool error [{}]: {}", error.name, error.error)),
113
114        AgentMessage::Error { message } => Some(format!("Error: {message}")),
115
116        AgentMessage::Cancelled { message } => Some(format!("Cancelled: {message}")),
117
118        AgentMessage::AutoContinue { attempt, max_attempts } => {
119            Some(format!("Continuing ({attempt}/{max_attempts})..."))
120        }
121
122        AgentMessage::Retrying { attempt, max_attempts, delay_ms, error } => {
123            Some(format!("Retrying ({attempt}/{max_attempts}) in {delay_ms}ms: {error}"))
124        }
125
126        AgentMessage::ModelSwitched { previous, new } => Some(format!("Model switched: {previous} -> {new}")),
127
128        AgentMessage::ToolProgress { request, progress, total, message } => {
129            let bar = match total {
130                Some(t) => format!("{progress}/{t}"),
131                None => format!("{progress}"),
132            };
133            let suffix = message.as_deref().map(|m| format!(" - {m}")).unwrap_or_default();
134            Some(format!("Tool progress [{}]: {bar}{suffix}", request.name))
135        }
136
137        AgentMessage::ContextCompactionStarted { message_count } => {
138            Some(format!("Context compaction started ({message_count} messages)"))
139        }
140
141        AgentMessage::ContextCompactionResult { summary, messages_removed } => {
142            Some(format!("Context compacted: {messages_removed} messages removed. {summary}"))
143        }
144
145        AgentMessage::ContextUsageUpdate {
146            input_tokens,
147            output_tokens,
148            total_input_tokens,
149            total_output_tokens,
150            ..
151        } => Some(format!(
152            "Tokens: {input_tokens} in, {output_tokens} out (total: {total_input_tokens} in, {total_output_tokens} out)"
153        )),
154
155        AgentMessage::ContextCleared => Some("Context cleared".to_string()),
156
157        AgentMessage::Done => Some("Done".to_string()),
158
159        AgentMessage::ToolCallUpdate { .. } | AgentMessage::Text { .. } | AgentMessage::Thought { .. } => None,
160    }
161}
162
163fn setup_tracing(verbose: bool) {
164    use tracing_subscriber::Layer;
165    use tracing_subscriber::filter::EnvFilter;
166    use tracing_subscriber::fmt;
167    use tracing_subscriber::layer::SubscriberExt;
168    use tracing_subscriber::util::SubscriberInitExt;
169
170    let filter = if verbose { EnvFilter::new("debug,agent=off") } else { EnvFilter::new("error,agent=off") };
171    let layer = fmt::layer().with_writer(io::stderr).with_filter(filter);
172
173    tracing_subscriber::registry().with(layer).init();
174}
175
176#[cfg(test)]
177mod tests {
178    use super::*;
179
180    #[test]
181    fn format_text_formats_complete_text() {
182        assert_eq!(format_text(&AgentMessage::text("id", "hello world", true, "m")), Some("hello world".to_string()));
183    }
184
185    #[test]
186    fn format_text_skips_incomplete_text() {
187        assert_eq!(format_text(&AgentMessage::text("id", "partial", false, "m")), None);
188    }
189
190    #[test]
191    fn format_text_formats_complete_thought() {
192        assert_eq!(
193            format_text(&AgentMessage::thought("id", "reasoning here", true, "m")),
194            Some("Thought: reasoning here".to_string())
195        );
196    }
197
198    #[test]
199    fn format_text_skips_incomplete_thought() {
200        assert_eq!(format_text(&AgentMessage::thought("id", "partial", false, "m")), None);
201    }
202
203    #[test]
204    fn format_text_formats_tool_call() {
205        let msg = AgentMessage::ToolCall {
206            request: llm::ToolCallRequest {
207                id: "tc1".to_string(),
208                name: "bash".to_string(),
209                arguments: r#"{"cmd":"ls"}"#.to_string(),
210            },
211            model_name: "test".to_string(),
212        };
213        assert_eq!(format_text(&msg), Some(r#"Tool call: bash({"cmd":"ls"})"#.to_string()));
214    }
215
216    #[test]
217    fn format_text_skips_tool_call_updates() {
218        let msg = AgentMessage::ToolCallUpdate {
219            tool_call_id: "tc1".to_string(),
220            chunk: "partial".to_string(),
221            model_name: "test".to_string(),
222        };
223        assert_eq!(format_text(&msg), None);
224    }
225
226    #[test]
227    fn format_text_formats_tool_result() {
228        assert_eq!(format_text(&tool_result_msg()), Some("Tool result [bash]: ok".to_string()));
229    }
230
231    #[test]
232    fn format_text_formats_tool_error() {
233        let msg = AgentMessage::ToolError {
234            error: llm::ToolCallError {
235                id: "tc1".to_string(),
236                name: "bash".to_string(),
237                arguments: None,
238                error: "not found".to_string(),
239            },
240            model_name: "test".to_string(),
241        };
242        assert_eq!(format_text(&msg), Some("Tool error [bash]: not found".to_string()));
243    }
244
245    #[test]
246    fn format_text_formats_error() {
247        let msg = AgentMessage::Error { message: "boom".to_string() };
248        assert_eq!(format_text(&msg), Some("Error: boom".to_string()));
249    }
250
251    #[test]
252    fn format_text_formats_cancelled() {
253        let msg = AgentMessage::Cancelled { message: "user stopped".to_string() };
254        assert_eq!(format_text(&msg), Some("Cancelled: user stopped".to_string()));
255    }
256
257    #[test]
258    fn format_text_formats_auto_continue() {
259        let msg = AgentMessage::AutoContinue { attempt: 2, max_attempts: 5 };
260        assert_eq!(format_text(&msg), Some("Continuing (2/5)...".to_string()));
261    }
262
263    #[test]
264    fn format_text_formats_model_switched() {
265        let msg = AgentMessage::ModelSwitched { previous: "old-model".to_string(), new: "new-model".to_string() };
266        assert_eq!(format_text(&msg), Some("Model switched: old-model -> new-model".to_string()));
267    }
268
269    #[test]
270    fn format_text_renders_done() {
271        assert_eq!(format_text(&AgentMessage::Done), Some("Done".to_string()));
272    }
273
274    #[test]
275    fn format_text_formats_tool_progress_with_total() {
276        let msg = tool_progress(50.0, Some(100.0), Some("halfway"));
277        assert_eq!(format_text(&msg), Some("Tool progress [bash]: 50/100 - halfway".to_string()));
278    }
279
280    #[test]
281    fn format_text_formats_tool_progress_without_total() {
282        let msg = tool_progress(42.0, None, None);
283        assert_eq!(format_text(&msg), Some("Tool progress [bash]: 42".to_string()));
284    }
285
286    #[test]
287    fn format_text_formats_context_compaction_started() {
288        let msg = AgentMessage::ContextCompactionStarted { message_count: 42 };
289        assert_eq!(format_text(&msg), Some("Context compaction started (42 messages)".to_string()));
290    }
291
292    #[test]
293    fn format_text_formats_context_compaction_result() {
294        let msg = AgentMessage::ContextCompactionResult { summary: "summary here".to_string(), messages_removed: 10 };
295        assert_eq!(format_text(&msg), Some("Context compacted: 10 messages removed. summary here".to_string()));
296    }
297
298    #[test]
299    fn format_text_formats_context_usage_update() {
300        assert_eq!(
301            format_text(&usage_update()),
302            Some("Tokens: 1500 in, 250 out (total: 5000 in, 800 out)".to_string())
303        );
304    }
305
306    #[test]
307    fn format_text_formats_context_cleared() {
308        assert_eq!(format_text(&AgentMessage::ContextCleared), Some("Context cleared".to_string()));
309    }
310
311    #[test]
312    fn event_kind_none_for_non_output_fragments() {
313        assert_eq!(event_kind(&AgentMessage::text("id", "x", false, "m")), None);
314        assert_eq!(event_kind(&AgentMessage::thought("id", "x", false, "m")), None);
315        assert_eq!(
316            event_kind(&AgentMessage::ToolCallUpdate {
317                tool_call_id: "tc1".to_string(),
318                chunk: "x".to_string(),
319                model_name: "m".to_string(),
320            }),
321            None,
322        );
323    }
324
325    #[test]
326    fn event_kind_done_is_filterable() {
327        assert_eq!(event_kind(&AgentMessage::Done), Some(CliEventKind::Done));
328    }
329
330    #[test]
331    fn should_emit_empty_filter_rejects_non_output_events() {
332        assert!(should_emit(&tool_call_msg(), &[]));
333        assert!(should_emit(&AgentMessage::Error { message: "e".to_string() }, &[]));
334        assert!(should_emit(&AgentMessage::Done, &[]));
335        assert!(!should_emit(&AgentMessage::text("id", "x", false, "m"), &[]));
336        assert!(!should_emit(
337            &AgentMessage::ToolCallUpdate {
338                tool_call_id: "tc1".to_string(),
339                chunk: "x".to_string(),
340                model_name: "m".to_string(),
341            },
342            &[],
343        ));
344    }
345
346    #[test]
347    fn should_emit_single_type_whitelist() {
348        let filter = &[CliEventKind::ToolCall];
349        assert!(should_emit(&tool_call_msg(), filter));
350        assert!(!should_emit(&tool_result_msg(), filter));
351        assert!(!should_emit(&AgentMessage::Error { message: "e".to_string() }, filter));
352    }
353
354    #[test]
355    fn should_emit_multi_type_whitelist() {
356        let filter = &[CliEventKind::ToolCall, CliEventKind::ToolResult];
357        assert!(should_emit(&tool_call_msg(), filter));
358        assert!(should_emit(&tool_result_msg(), filter));
359        assert!(!should_emit(&AgentMessage::Error { message: "e".to_string() }, filter));
360    }
361
362    #[test]
363    fn should_emit_done_respects_filter() {
364        assert!(should_emit(&AgentMessage::Done, &[CliEventKind::Done]));
365        assert!(!should_emit(&AgentMessage::Done, &[CliEventKind::ToolCall]));
366    }
367
368    #[test]
369    fn cli_event_kind_names_match_agent_message_type_tags() {
370        use clap::ValueEnum;
371
372        let samples = vec![
373            (AgentMessage::text("id", "x", true, "m"), CliEventKind::Text),
374            (AgentMessage::thought("id", "x", true, "m"), CliEventKind::Thought),
375            (tool_call_msg(), CliEventKind::ToolCall),
376            (tool_result_msg(), CliEventKind::ToolResult),
377            (
378                AgentMessage::ToolError {
379                    error: llm::ToolCallError {
380                        id: "tc1".to_string(),
381                        name: "bash".to_string(),
382                        arguments: None,
383                        error: "boom".to_string(),
384                    },
385                    model_name: "test".to_string(),
386                },
387                CliEventKind::ToolError,
388            ),
389            (AgentMessage::Error { message: "e".to_string() }, CliEventKind::Error),
390            (AgentMessage::Cancelled { message: "c".to_string() }, CliEventKind::Cancelled),
391            (AgentMessage::AutoContinue { attempt: 1, max_attempts: 3 }, CliEventKind::AutoContinue),
392            (
393                AgentMessage::Retrying { attempt: 1, max_attempts: 3, delay_ms: 10, error: "e".to_string() },
394                CliEventKind::Retrying,
395            ),
396            (
397                AgentMessage::ModelSwitched { previous: "a".to_string(), new: "b".to_string() },
398                CliEventKind::ModelSwitched,
399            ),
400            (tool_progress(1.0, None, None), CliEventKind::ToolProgress),
401            (AgentMessage::ContextCompactionStarted { message_count: 1 }, CliEventKind::ContextCompactionStarted),
402            (
403                AgentMessage::ContextCompactionResult { summary: "s".to_string(), messages_removed: 1 },
404                CliEventKind::ContextCompactionResult,
405            ),
406            (usage_update(), CliEventKind::ContextUsage),
407            (AgentMessage::ContextCleared, CliEventKind::ContextCleared),
408            (AgentMessage::Done, CliEventKind::Done),
409        ];
410
411        for kind in CliEventKind::value_variants() {
412            assert!(samples.iter().any(|(_, k)| k == kind), "samples is missing a case for {kind:?}");
413        }
414
415        for (msg, kind) in &samples {
416            assert_eq!(event_kind(msg), Some(*kind), "event_kind disagrees for {kind:?}");
417            let tag = serde_json::to_value(msg).unwrap();
418            let tag = tag["type"].as_str().unwrap().to_string();
419            let clap_name = kind.to_possible_value().unwrap();
420            assert_eq!(tag, clap_name.get_name(), "`--events` value and serialized `type` tag diverged for {kind:?}");
421        }
422    }
423
424    #[tokio::test]
425    async fn stream_output_done_breaks_loop_under_filter() {
426        let (tx, rx) = mpsc::channel(4);
427        tx.send(AgentMessage::Done).await.unwrap();
428        let filter = vec![CliEventKind::ToolCall];
429        let code = stream_output(rx, OutputFormat::Text, &filter).await;
430        assert_eq!(code, ExitCode::SUCCESS);
431    }
432
433    fn tool_call_msg() -> AgentMessage {
434        AgentMessage::ToolCall {
435            request: llm::ToolCallRequest {
436                id: "tc1".to_string(),
437                name: "bash".to_string(),
438                arguments: "{}".to_string(),
439            },
440            model_name: "test".to_string(),
441        }
442    }
443
444    fn tool_result_msg() -> AgentMessage {
445        AgentMessage::ToolResult {
446            result: llm::ToolCallResult {
447                id: "tc1".to_string(),
448                name: "bash".to_string(),
449                arguments: "{}".to_string(),
450                result: "ok".to_string(),
451            },
452            result_meta: None,
453            model_name: "test".to_string(),
454        }
455    }
456
457    fn tool_progress(progress: f64, total: Option<f64>, message: Option<&str>) -> AgentMessage {
458        AgentMessage::ToolProgress {
459            request: llm::ToolCallRequest {
460                id: "tc1".to_string(),
461                name: "bash".to_string(),
462                arguments: "{}".to_string(),
463            },
464            progress,
465            total,
466            message: message.map(str::to_string),
467        }
468    }
469
470    fn usage_update() -> AgentMessage {
471        AgentMessage::ContextUsageUpdate {
472            usage_ratio: Some(0.25),
473            context_limit: Some(200_000),
474            input_tokens: 1500,
475            output_tokens: 250,
476            cache_read_tokens: Some(400),
477            cache_creation_tokens: Some(100),
478            reasoning_tokens: Some(50),
479            total_input_tokens: 5000,
480            total_output_tokens: 800,
481            total_cache_read_tokens: 1200,
482            total_cache_creation_tokens: 300,
483            total_reasoning_tokens: 150,
484        }
485    }
486}