Skip to main content

aether_cli/headless/
run.rs

1use aether_core::core::{AgentDeps, Prompt};
2use aether_core::events::{
3    AgentEvent, Command, ContextEvent, MessageEvent, ModelEvent, ToolEvent, TurnEvent, TurnOutcome,
4};
5use aether_core::mcp::McpHandle;
6use aether_telemetry::TelemetryRuntime;
7use std::io;
8use std::process::ExitCode;
9use std::sync::Arc;
10use tokio::sync::mpsc;
11use tracing::error;
12
13use crate::telemetry::build_telemetry_runtime;
14
15use super::error::CliError;
16use super::{CliEventKind, RunConfig};
17use crate::output::{OutputFormat, print_message};
18use crate::runtime::RuntimeBuilder;
19use crate::slash_commands::{expand_slash_command, parse_slash_command};
20
21pub async fn run(config: RunConfig) -> Result<ExitCode, CliError> {
22    setup_tracing(config.verbose);
23
24    let telemetry = build_telemetry_runtime(config.telemetry.as_ref(), config.trace_context.clone())?;
25    let result = run_agent(config, telemetry.clone()).await;
26
27    if let Some(telemetry) = telemetry {
28        telemetry.shutdown_or_log();
29    }
30    result
31}
32
33async fn run_agent(config: RunConfig, telemetry: Option<Arc<TelemetryRuntime>>) -> Result<ExitCode, CliError> {
34    let mut spec = config.spec;
35    if let Some(system_prompt) = config.system_prompt {
36        spec.prompts.push(Prompt::text(&system_prompt));
37    }
38
39    let registry = config.agent_catalog.registry().clone();
40    let deps =
41        AgentDeps::new(config.oauth_credential_store, telemetry.as_ref().map(|runtime| runtime.observer_factory()))
42            .with_agent_registry(registry);
43    let (agent, _mcp_snapshot) = RuntimeBuilder::from_spec(config.cwd.clone(), spec)
44        .mcp_sources(config.mcp_config_sources)
45        .agent_deps(deps)
46        .build_ready(vec![])
47        .await?;
48
49    let prompt = expand_prompt(agent.mcp_runtime.handle(), config.prompt).await;
50
51    agent
52        .agent_tx
53        .send(Command::text(&prompt))
54        .await
55        .map_err(|e| CliError::AgentError(format!("Failed to send prompt: {e}")))?;
56
57    let exit_code = stream_output(agent.agent_rx, config.output, &config.events).await;
58
59    drop(agent.agent_tx);
60    agent.agent_handle.await_completion().await;
61
62    Ok(exit_code)
63}
64
65async fn expand_prompt(mcp: &McpHandle, prompt: String) -> String {
66    let Some(slash_command) = parse_slash_command(&prompt) else {
67        return prompt;
68    };
69
70    match expand_slash_command(mcp, slash_command.command_name, slash_command.args_text).await {
71        Ok(expanded) => expanded,
72        Err(error) => {
73            error!("Failed to expand slash command: {error}");
74            prompt
75        }
76    }
77}
78
79async fn stream_output(mut rx: mpsc::Receiver<AgentEvent>, format: OutputFormat, events: &[CliEventKind]) -> ExitCode {
80    while let Some(msg) = rx.recv().await {
81        if should_emit(&msg, events)
82            && let Err(error) = print_message(format, &msg)
83        {
84            eprintln!("Failed to serialize headless event: {error}");
85            return ExitCode::FAILURE;
86        }
87
88        if let Some(outcome) = msg.turn_outcome() {
89            return match outcome {
90                TurnOutcome::Failed { .. } => ExitCode::FAILURE,
91                TurnOutcome::Completed | TurnOutcome::Cancelled => ExitCode::SUCCESS,
92            };
93        }
94    }
95    ExitCode::SUCCESS
96}
97
98fn should_emit(msg: &AgentEvent, include: &[CliEventKind]) -> bool {
99    let Some(kind) = event_kind(msg) else { return false };
100    include.is_empty() || include.contains(&kind)
101}
102
103fn event_kind(msg: &AgentEvent) -> Option<CliEventKind> {
104    match msg {
105        AgentEvent::Message(MessageEvent::Text { is_complete: true, .. }) => Some(CliEventKind::Text),
106        AgentEvent::Message(MessageEvent::Thought { is_complete: true, .. }) => Some(CliEventKind::Thought),
107        AgentEvent::Tool(ToolEvent::Call { .. }) => Some(CliEventKind::ToolCall),
108        AgentEvent::Tool(
109            ToolEvent::Result { .. } | ToolEvent::TaskCreated { .. } | ToolEvent::TaskCompleted { .. },
110        ) => Some(CliEventKind::ToolResult),
111        AgentEvent::Tool(ToolEvent::Error { .. } | ToolEvent::TaskFailed { .. } | ToolEvent::TaskCancelled { .. }) => {
112            Some(CliEventKind::ToolError)
113        }
114        AgentEvent::Turn(TurnEvent::AutoContinue { .. }) => Some(CliEventKind::AutoContinue),
115        AgentEvent::Model(ModelEvent::Switched { .. }) => Some(CliEventKind::ModelSwitched),
116        AgentEvent::Tool(
117            ToolEvent::Progress { .. }
118            | ToolEvent::DisplayUpdate { .. }
119            | ToolEvent::SubAgentProgress { .. }
120            | ToolEvent::TaskStatus { .. },
121        ) => Some(CliEventKind::ToolProgress),
122        AgentEvent::Context(ContextEvent::CompactionStarted { .. }) => Some(CliEventKind::ContextCompactionStarted),
123        AgentEvent::Context(ContextEvent::CompactionEnded { .. }) => Some(CliEventKind::ContextCompactionEnded),
124        AgentEvent::Context(ContextEvent::CompactionResult { .. }) => Some(CliEventKind::ContextCompactionResult),
125        AgentEvent::Context(ContextEvent::UsageUpdated { .. }) => Some(CliEventKind::ContextUsage),
126        AgentEvent::SessionUsage(_) => Some(CliEventKind::SessionUsage),
127        AgentEvent::Context(ContextEvent::Cleared) => Some(CliEventKind::ContextCleared),
128        AgentEvent::Turn(TurnEvent::Started { .. }) => Some(CliEventKind::TurnStarted),
129        AgentEvent::Turn(TurnEvent::Ended { .. }) => Some(CliEventKind::TurnEnded),
130        AgentEvent::Turn(TurnEvent::RetryScheduled { .. }) => Some(CliEventKind::LlmRetryScheduled),
131        AgentEvent::Turn(TurnEvent::LlmCallStarted { .. }) => Some(CliEventKind::LlmCallStarted),
132        AgentEvent::Turn(TurnEvent::LlmCallEnded { .. }) => Some(CliEventKind::LlmCallEnded),
133        AgentEvent::Tool(ToolEvent::DefinitionsUpdated { .. }) => Some(CliEventKind::ToolDefinitionsUpdated),
134        AgentEvent::Message(
135            MessageEvent::Text { is_complete: false, .. } | MessageEvent::Thought { is_complete: false, .. },
136        )
137        | AgentEvent::Tool(ToolEvent::InputStarted { .. } | ToolEvent::InputDelta { .. }) => None,
138    }
139}
140
141fn setup_tracing(verbose: bool) {
142    use tracing_subscriber::Layer;
143    use tracing_subscriber::filter::EnvFilter;
144    use tracing_subscriber::fmt;
145    use tracing_subscriber::layer::SubscriberExt;
146    use tracing_subscriber::util::SubscriberInitExt;
147
148    let filter = if verbose { EnvFilter::new("debug,agent=off") } else { EnvFilter::new("error,agent=off") };
149    let layer = fmt::layer().with_writer(io::stderr).with_filter(filter);
150
151    tracing_subscriber::registry().with(layer).init();
152}
153
154#[cfg(test)]
155mod tests {
156    use aether_core::events::StreamState;
157
158    use super::*;
159    use llm::ContextUsage;
160
161    #[test]
162    fn event_kind_none_for_non_output_fragments() {
163        assert_eq!(event_kind(&AgentEvent::text("id", "x", StreamState::Partial)), None);
164        assert_eq!(event_kind(&AgentEvent::thought("id", "x", StreamState::Partial)), None);
165        assert_eq!(
166            event_kind(&AgentEvent::Tool(ToolEvent::InputStarted { id: "tc1".to_string(), name: "bash".to_string() })),
167            None,
168        );
169        assert_eq!(
170            event_kind(&AgentEvent::Tool(ToolEvent::InputDelta { id: "tc1".to_string(), chunk: "x".to_string() })),
171            None,
172        );
173    }
174
175    #[test]
176    fn event_kind_turn_ended_is_filterable() {
177        assert_eq!(event_kind(&AgentEvent::turn_ended(TurnOutcome::Completed)), Some(CliEventKind::TurnEnded));
178    }
179
180    #[test]
181    fn should_emit_empty_filter_rejects_non_output_events() {
182        assert!(should_emit(&tool_call_msg(), &[]));
183        assert!(should_emit(&AgentEvent::turn_ended(TurnOutcome::failed("e")), &[]));
184        assert!(should_emit(&AgentEvent::turn_ended(TurnOutcome::Completed), &[]));
185        assert!(!should_emit(&AgentEvent::text("id", "x", StreamState::Partial), &[]));
186        assert!(!should_emit(
187            &AgentEvent::Tool(ToolEvent::InputDelta { id: "tc1".to_string(), chunk: "x".to_string() }),
188            &[],
189        ));
190    }
191
192    #[test]
193    fn should_emit_single_type_whitelist() {
194        let filter = &[CliEventKind::ToolCall];
195        assert!(should_emit(&tool_call_msg(), filter));
196        assert!(!should_emit(&tool_result_msg(), filter));
197        assert!(!should_emit(&AgentEvent::turn_ended(TurnOutcome::Completed), filter));
198    }
199
200    #[test]
201    fn should_emit_multi_type_whitelist() {
202        let filter = &[CliEventKind::ToolCall, CliEventKind::ToolResult];
203        assert!(should_emit(&tool_call_msg(), filter));
204        assert!(should_emit(&tool_result_msg(), filter));
205        assert!(!should_emit(&AgentEvent::turn_ended(TurnOutcome::Completed), filter));
206    }
207
208    #[test]
209    fn should_emit_turn_ended_respects_filter() {
210        let msg = AgentEvent::turn_ended(TurnOutcome::Completed);
211        assert!(should_emit(&msg, &[CliEventKind::TurnEnded]));
212        assert!(!should_emit(&msg, &[CliEventKind::ToolCall]));
213    }
214
215    #[test]
216    fn event_kind_covers_every_cli_event_kind() {
217        use clap::ValueEnum;
218
219        let samples = vec![
220            (AgentEvent::text("id", "x", StreamState::Complete), CliEventKind::Text),
221            (AgentEvent::thought("id", "x", StreamState::Complete), CliEventKind::Thought),
222            (tool_call_msg(), CliEventKind::ToolCall),
223            (tool_result_msg(), CliEventKind::ToolResult),
224            (
225                AgentEvent::Tool(ToolEvent::Error {
226                    error: llm::ToolCallError {
227                        id: "tc1".to_string(),
228                        name: "bash".to_string(),
229                        arguments: None,
230                        error: "boom".to_string(),
231                    },
232                }),
233                CliEventKind::ToolError,
234            ),
235            (
236                AgentEvent::Turn(TurnEvent::AutoContinue {
237                    attempt: 1,
238                    max_attempts: 3,
239                    message_id: llm::MessageId::new(),
240                    content: vec![],
241                }),
242                CliEventKind::AutoContinue,
243            ),
244            (
245                AgentEvent::Model(ModelEvent::Switched { previous: "a".to_string(), new: "b".to_string() }),
246                CliEventKind::ModelSwitched,
247            ),
248            (tool_progress(1.0, None, None), CliEventKind::ToolProgress),
249            (
250                AgentEvent::Context(ContextEvent::CompactionStarted {
251                    compaction_id: "compaction".into(),
252                    message_count: 1,
253                }),
254                CliEventKind::ContextCompactionStarted,
255            ),
256            (
257                AgentEvent::Context(ContextEvent::CompactionEnded {
258                    compaction_id: "compaction".into(),
259                    outcome: aether_core::events::CompactionOutcome::Completed,
260                }),
261                CliEventKind::ContextCompactionEnded,
262            ),
263            (
264                AgentEvent::Context(ContextEvent::CompactionResult {
265                    compaction_id: "compaction".into(),
266                    message_id: llm::MessageId::new(),
267                    summary: "s".to_string(),
268                    messages_removed: 1,
269                }),
270                CliEventKind::ContextCompactionResult,
271            ),
272            (usage_update(), CliEventKind::ContextUsage),
273            (
274                AgentEvent::SessionUsage(llm::testing::session_usage_event(1, llm::TokenUsage::new(1, 1))),
275                CliEventKind::SessionUsage,
276            ),
277            (AgentEvent::Context(ContextEvent::Cleared), CliEventKind::ContextCleared),
278            (AgentEvent::Turn(TurnEvent::Started { content: vec![] }), CliEventKind::TurnStarted),
279            (AgentEvent::turn_ended(TurnOutcome::Completed), CliEventKind::TurnEnded),
280            (retry_scheduled(1, 10), CliEventKind::LlmRetryScheduled),
281            (llm_call_started(0), CliEventKind::LlmCallStarted),
282            (
283                AgentEvent::Turn(TurnEvent::LlmCallEnded {
284                    purpose: llm::LlmCallPurpose::Chat,
285                    outcome: aether_core::events::LlmCallOutcome::Cancelled,
286                }),
287                CliEventKind::LlmCallEnded,
288            ),
289            (AgentEvent::Tool(ToolEvent::DefinitionsUpdated { tools: vec![] }), CliEventKind::ToolDefinitionsUpdated),
290        ];
291
292        for kind in CliEventKind::value_variants() {
293            assert!(samples.iter().any(|(_, k)| k == kind), "samples is missing a case for {kind:?}");
294        }
295
296        for (msg, kind) in &samples {
297            assert_eq!(event_kind(msg), Some(*kind), "event_kind disagrees for {kind:?}");
298        }
299    }
300
301    #[tokio::test]
302    async fn stream_output_turn_ended_breaks_loop_under_filter() {
303        let (tx, rx) = mpsc::channel(4);
304        tx.send(AgentEvent::turn_ended(TurnOutcome::Completed)).await.unwrap();
305        let filter = vec![CliEventKind::ToolCall];
306        let code = stream_output(rx, OutputFormat::Text, &filter).await;
307        assert_eq!(code, ExitCode::SUCCESS);
308    }
309
310    #[tokio::test]
311    async fn stream_output_failed_turn_exits_with_failure() {
312        let (tx, rx) = mpsc::channel(4);
313        tx.send(AgentEvent::turn_ended(TurnOutcome::failed("boom"))).await.unwrap();
314        let code = stream_output(rx, OutputFormat::Text, &[]).await;
315        assert_eq!(code, ExitCode::FAILURE);
316    }
317
318    fn tool_call_msg() -> AgentEvent {
319        AgentEvent::Tool(ToolEvent::Call {
320            request: llm::ToolCallRequest {
321                id: "tc1".to_string(),
322                name: "bash".to_string(),
323                arguments: "{}".to_string(),
324            },
325        })
326    }
327
328    fn tool_result_msg() -> AgentEvent {
329        AgentEvent::Tool(ToolEvent::Result {
330            result: llm::ToolCallResult {
331                id: "tc1".to_string(),
332                name: "bash".to_string(),
333                arguments: "{}".to_string(),
334                result: "ok".to_string(),
335            },
336            result_meta: None,
337        })
338    }
339
340    fn tool_progress(progress: f64, total: Option<f64>, message: Option<&str>) -> AgentEvent {
341        AgentEvent::Tool(ToolEvent::Progress {
342            request: llm::ToolCallRequest {
343                id: "tc1".to_string(),
344                name: "bash".to_string(),
345                arguments: "{}".to_string(),
346            },
347            progress,
348            total,
349            message: message.map(str::to_string),
350        })
351    }
352
353    fn retry_scheduled(attempt: u32, delay_ms: u64) -> AgentEvent {
354        AgentEvent::Turn(TurnEvent::RetryScheduled {
355            purpose: llm::LlmCallPurpose::Chat,
356            attempt,
357            max_attempts: 3,
358            delay_ms,
359        })
360    }
361
362    fn llm_call_started(attempt: u32) -> AgentEvent {
363        AgentEvent::Turn(TurnEvent::LlmCallStarted {
364            purpose: llm::LlmCallPurpose::Chat,
365            model: llm::ModelIdentity::default(),
366            display_name: "test".to_string(),
367            attempt,
368            max_attempts: 3,
369        })
370    }
371
372    fn usage_update() -> AgentEvent {
373        AgentEvent::Context(ContextEvent::UsageUpdated {
374            usage: ContextUsage {
375                input_tokens: 100_000.into(),
376                context_limit: Some(200_000.into()),
377                usage_ratio: Some(0.5),
378            },
379        })
380    }
381}