aether_cli/
runtime.rs

1use crate::error::CliError;
2use aether_core::agent_spec::AgentSpec;
3use aether_core::core::{AgentBuilder, AgentHandle, Prompt};
4use aether_core::events::{AgentMessage, UserMessage};
5use aether_core::mcp::McpBuilder;
6use aether_core::mcp::McpSpawnResult;
7use aether_core::mcp::mcp;
8use aether_core::mcp::run_mcp_task::McpCommand;
9use aether_project::load_agent_catalog;
10use llm::{ChatMessage, LlmModel, ToolDefinition};
11use mcp_servers::McpBuilderExt;
12use mcp_utils::client::oauth::OAuthHandler;
13use mcp_utils::client::{McpClientEvent, McpServerConfig};
14use mcp_utils::status::McpServerStatusEntry;
15use std::path::{Path, PathBuf};
16use tokio::sync::mpsc::{Receiver, Sender};
17use tokio::task::JoinHandle;
18use tracing::debug;
19
20pub struct RuntimeBuilder {
21    cwd: PathBuf,
22    spec: AgentSpec,
23    mcp_configs: Vec<PathBuf>,
24    extra_mcp_servers: Vec<McpServerConfig>,
25    oauth_applicator: Option<Box<dyn FnOnce(McpBuilder) -> McpBuilder + Send>>,
26    prompt_cache_key: Option<String>,
27}
28
29pub struct Runtime {
30    pub agent_tx: Sender<UserMessage>,
31    pub agent_rx: Receiver<AgentMessage>,
32    pub agent_handle: AgentHandle,
33    pub mcp_tx: Sender<McpCommand>,
34    pub event_rx: Receiver<McpClientEvent>,
35    pub server_statuses: Vec<McpServerStatusEntry>,
36    pub mcp_handle: JoinHandle<()>,
37}
38
39pub struct PromptInfo {
40    pub spec: AgentSpec,
41    pub tool_definitions: Vec<ToolDefinition>,
42}
43
44impl RuntimeBuilder {
45    pub fn new(cwd: &Path, model: &str) -> Result<Self, CliError> {
46        let cwd = cwd.canonicalize().map_err(CliError::IoError)?;
47        let parsed_model: LlmModel = model.parse().map_err(|e: String| CliError::ModelError(e))?;
48        let catalog = load_agent_catalog(&cwd).map_err(|e| CliError::AgentError(e.to_string()))?;
49        let spec = catalog.resolve_default(&parsed_model, None, &cwd);
50
51        Ok(Self {
52            cwd,
53            spec,
54            mcp_configs: Vec::new(),
55            extra_mcp_servers: Vec::new(),
56            oauth_applicator: None,
57            prompt_cache_key: None,
58        })
59    }
60
61    pub fn from_spec(cwd: PathBuf, spec: AgentSpec) -> Self {
62        Self {
63            cwd,
64            spec,
65            mcp_configs: Vec::new(),
66            extra_mcp_servers: Vec::new(),
67            oauth_applicator: None,
68            prompt_cache_key: None,
69        }
70    }
71
72    pub fn prompt_cache_key(mut self, key: String) -> Self {
73        self.prompt_cache_key = Some(key);
74        self
75    }
76
77    /// Set the MCP config path overrides. When non-empty, these completely
78    /// replace any paths resolved from the agent's `AgentSpec` (CLI override
79    /// semantics). On collisions across files, the rightmost path wins.
80    pub fn mcp_configs(mut self, paths: Vec<PathBuf>) -> Self {
81        self.mcp_configs = paths;
82        self
83    }
84
85    pub fn extra_servers(mut self, servers: Vec<McpServerConfig>) -> Self {
86        self.extra_mcp_servers = servers;
87        self
88    }
89
90    pub fn oauth_handler<H: OAuthHandler + 'static>(mut self, handler: H) -> Self {
91        self.oauth_applicator = Some(Box::new(|builder| builder.with_oauth_handler(handler)));
92        self
93    }
94
95    pub async fn build(
96        self,
97        custom_prompt: Option<Prompt>,
98        messages: Option<Vec<ChatMessage>>,
99    ) -> Result<Runtime, CliError> {
100        let prompt_cache_key = self.prompt_cache_key.clone();
101        let mcp = self.spawn_mcp().await?;
102
103        let filtered_tools = mcp.spec.tools.apply(mcp.tool_definitions);
104        let mut agent_builder = AgentBuilder::from_spec(&mcp.spec, vec![])
105            .await
106            .map_err(|e| CliError::AgentError(e.to_string()))?
107            .tools(mcp.mcp_tx.clone(), filtered_tools);
108
109        if let Some(key) = prompt_cache_key {
110            agent_builder = agent_builder.prompt_cache_key(key);
111        }
112
113        if let Some(prompt) = custom_prompt {
114            agent_builder = agent_builder.system_prompt(prompt);
115        }
116
117        if let Some(msgs) = messages {
118            agent_builder = agent_builder.messages(msgs);
119        }
120
121        let (agent_tx, agent_rx, agent_handle) =
122            agent_builder.spawn().await.map_err(|e| CliError::AgentError(e.to_string()))?;
123
124        Ok(Runtime {
125            agent_tx,
126            agent_rx,
127            agent_handle,
128            mcp_tx: mcp.mcp_tx,
129            event_rx: mcp.event_rx,
130            server_statuses: mcp.server_statuses,
131            mcp_handle: mcp.mcp_handle,
132        })
133    }
134
135    pub async fn build_prompt_info(self) -> Result<PromptInfo, CliError> {
136        let mcp = self.spawn_mcp().await?;
137        let filtered_tools = mcp.spec.tools.apply(mcp.tool_definitions);
138        Ok(PromptInfo { spec: mcp.spec, tool_definitions: filtered_tools })
139    }
140
141    async fn spawn_mcp(self) -> Result<McpParts, CliError> {
142        let mut builder = mcp().with_builtin_servers(self.cwd.clone(), &self.cwd);
143
144        if !self.extra_mcp_servers.is_empty() {
145            builder = builder.with_servers(self.extra_mcp_servers);
146        }
147
148        if let Some(apply_oauth) = self.oauth_applicator {
149            builder = apply_oauth(builder);
150        }
151
152        let mcp_config_paths: Vec<PathBuf> =
153            if self.mcp_configs.is_empty() { self.spec.mcp_config_paths.clone() } else { self.mcp_configs };
154
155        if !mcp_config_paths.is_empty() {
156            debug!("Loading MCP configs from: {:?}", mcp_config_paths);
157            builder =
158                builder.from_json_files(&mcp_config_paths).await.map_err(|e| CliError::McpError(e.to_string()))?;
159        }
160
161        let McpSpawnResult {
162            tool_definitions,
163            instructions,
164            server_statuses,
165            command_tx: mcp_tx,
166            event_rx,
167            handle: mcp_handle,
168        } = builder.spawn().await.map_err(|e| CliError::McpError(e.to_string()))?;
169
170        let mut spec = self.spec;
171        spec.prompts.push(Prompt::mcp_instructions(instructions));
172
173        Ok(McpParts { spec, tool_definitions, mcp_tx, event_rx, server_statuses, mcp_handle })
174    }
175}
176
177struct McpParts {
178    spec: AgentSpec,
179    tool_definitions: Vec<ToolDefinition>,
180    mcp_tx: Sender<McpCommand>,
181    event_rx: Receiver<McpClientEvent>,
182    server_statuses: Vec<McpServerStatusEntry>,
183    mcp_handle: JoinHandle<()>,
184}
aether_cli/runtime.rs

aether_cli/
runtime.rs
