aether_cli/
runtime.rs

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

aether_cli/
runtime.rs
