Skip to main content

aether_cli/headless/
run.rs

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