aether_cli/
runtime.rs

1use crate::error::CliError;
2use aether_auth::OAuthCredentialStorage;
3use aether_core::agent_spec::{AgentSpec, McpConfigSource};
4use aether_core::core::{AgentBuilder, AgentHandle, Prompt};
5use aether_core::events::{AgentMessage, Command};
6use aether_core::mcp::McpBuilder;
7use aether_core::mcp::McpSpawnResult;
8use aether_core::mcp::mcp;
9use aether_core::mcp::run_mcp_task::McpCommand;
10use llm::{ChatMessage, LlmModel, ToolDefinition};
11use mcp_servers::McpBuilderExt;
12use mcp_utils::client::{McpClientEvent, McpConnectionDetails, McpServer, OAuthHandlerFactory};
13use std::path::{Path, PathBuf};
14use std::sync::Arc;
15use tokio::sync::mpsc::{Receiver, Sender};
16use tokio::task::JoinHandle;
17use tracing::debug;
18
19pub struct RuntimeBuilder {
20    cwd: PathBuf,
21    spec: AgentSpec,
22    mcp_config_sources: Vec<McpConfigSource>,
23    extra_mcp_servers: Vec<McpServer>,
24    oauth_applicator: Option<Box<dyn FnOnce(McpBuilder) -> McpBuilder + Send>>,
25    oauth_credential_store: Option<Arc<dyn OAuthCredentialStorage>>,
26    prompt_cache_key: Option<String>,
27}
28
29pub struct Runtime {
30    pub agent_tx: Sender<Command>,
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 mcp_handle: JoinHandle<()>,
36}
37
38pub struct PromptInfo {
39    pub spec: AgentSpec,
40    pub tool_definitions: Vec<ToolDefinition>,
41}
42
43impl RuntimeBuilder {
44    pub fn new(cwd: &Path, model: &str) -> Result<Self, CliError> {
45        let cwd = cwd.canonicalize().map_err(CliError::IoError)?;
46        let parsed_model: LlmModel = model.parse().map_err(CliError::ModelError)?;
47        let spec = AgentSpec::default_spec(&parsed_model, None, Vec::new());
48
49        Ok(Self {
50            cwd,
51            spec,
52            mcp_config_sources: Vec::new(),
53            extra_mcp_servers: Vec::new(),
54            oauth_applicator: None,
55            oauth_credential_store: None,
56            prompt_cache_key: None,
57        })
58    }
59
60    pub fn from_spec(cwd: PathBuf, spec: AgentSpec) -> Self {
61        Self {
62            cwd,
63            spec,
64            mcp_config_sources: Vec::new(),
65            extra_mcp_servers: Vec::new(),
66            oauth_applicator: None,
67            oauth_credential_store: 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 MCP config source overrides. When non-empty, these completely
78    /// replace any sources resolved from the agent's `AgentSpec`.
79    pub fn mcp_sources(mut self, sources: Vec<McpConfigSource>) -> Self {
80        self.mcp_config_sources = sources;
81        self
82    }
83
84    pub fn extra_servers(mut self, servers: Vec<McpServer>) -> Self {
85        self.extra_mcp_servers = servers;
86        self
87    }
88
89    pub fn oauth_handler_factory(mut self, factory: OAuthHandlerFactory) -> Self {
90        self.oauth_applicator = Some(Box::new(|builder| builder.with_oauth_handler_factory(factory)));
91        self
92    }
93
94    pub fn oauth_credential_store(mut self, store: Arc<dyn OAuthCredentialStorage>) -> Self {
95        self.oauth_credential_store = Some(store);
96        self
97    }
98
99    pub async fn build(
100        self,
101        custom_prompt: Option<Prompt>,
102        messages: Option<Vec<ChatMessage>>,
103    ) -> Result<Runtime, CliError> {
104        let prompt_cache_key = self.prompt_cache_key.clone();
105        let oauth_credential_store = self.oauth_credential_store.clone();
106        let (spec, spawn) = self.spawn_mcp().await?;
107        let McpSpawnResult { command_tx: mcp_tx, event_rx, handle: mcp_handle } = spawn;
108
109        let mut agent_builder = AgentBuilder::from_spec(&spec, vec![], oauth_credential_store)
110            .await
111            .map_err(|e| CliError::AgentError(e.to_string()))?
112            .tools(mcp_tx.clone(), Vec::new());
113
114        if let Some(key) = prompt_cache_key {
115            agent_builder = agent_builder.prompt_cache_key(key);
116        }
117
118        if let Some(prompt) = custom_prompt {
119            agent_builder = agent_builder.system_prompt(prompt);
120        }
121
122        if let Some(msgs) = messages {
123            agent_builder = agent_builder.messages(msgs);
124        }
125
126        let (agent_tx, agent_rx, agent_handle) =
127            agent_builder.spawn().await.map_err(|e| CliError::AgentError(e.to_string()))?;
128
129        Ok(Runtime { agent_tx, agent_rx, agent_handle, mcp_tx, event_rx, mcp_handle })
130    }
131
132    /// Spawn MCP, block until every server finishes its initial connection,
133    /// then build an agent pre-seeded with the connection's filtered tools and
134    /// MCP instructions. Returns the live [`Runtime`] plus the bootstrap
135    /// snapshot. Unlike [`build`](Self::build) (which starts with no tools and
136    /// relies on the MCP event stream to populate them later), this is for
137    /// callers that need the agent ready to use tools on its first turn.
138    pub async fn build_ready(self, messages: Vec<ChatMessage>) -> Result<(Runtime, McpConnectionDetails), CliError> {
139        let prompt_cache_key = self.prompt_cache_key.clone();
140        let oauth_credential_store = self.oauth_credential_store.clone();
141        let (spec, mut spawn) = self.spawn_mcp().await?;
142        let snapshot = spawn
143            .block_until_ready()
144            .await
145            .ok_or_else(|| CliError::McpError("MCP bootstrap aborted before completion".to_string()))?;
146        let McpSpawnResult { command_tx: mcp_tx, event_rx, handle: mcp_handle } = spawn;
147
148        let filtered_tools = spec.tools.apply(snapshot.tool_definitions.clone());
149        let mut runtime_spec = spec;
150        runtime_spec.prompts.push(Prompt::McpInstructions(snapshot.instructions.clone()));
151
152        let mut agent_builder = AgentBuilder::from_spec(&runtime_spec, vec![], oauth_credential_store)
153            .await
154            .map_err(|e| CliError::AgentError(e.to_string()))?
155            .tools(mcp_tx.clone(), filtered_tools)
156            .messages(messages);
157
158        if let Some(key) = prompt_cache_key {
159            agent_builder = agent_builder.prompt_cache_key(key);
160        }
161
162        let (agent_tx, agent_rx, agent_handle) =
163            agent_builder.spawn().await.map_err(|e| CliError::AgentError(e.to_string()))?;
164
165        Ok((Runtime { agent_tx, agent_rx, agent_handle, mcp_tx, event_rx, mcp_handle }, snapshot))
166    }
167
168    pub async fn build_prompt_info(self) -> Result<PromptInfo, CliError> {
169        let (spec, mut spawn) = self.spawn_mcp().await?;
170        let details = spawn
171            .block_until_ready()
172            .await
173            .ok_or_else(|| CliError::McpError("MCP bootstrap aborted before completion".to_string()))?;
174        let filtered_tools = spec.tools.apply(details.tool_definitions);
175        Ok(PromptInfo { spec, tool_definitions: filtered_tools })
176    }
177
178    async fn spawn_mcp(self) -> Result<(AgentSpec, McpSpawnResult), CliError> {
179        let mut builder = mcp(&self.cwd).with_builtin_servers(self.cwd.clone(), &self.cwd);
180
181        if !self.extra_mcp_servers.is_empty() {
182            builder = builder.with_servers(self.extra_mcp_servers);
183        }
184
185        if let Some(apply_oauth) = self.oauth_applicator {
186            builder = apply_oauth(builder);
187        }
188
189        if let Some(store) = self.oauth_credential_store {
190            builder = builder.with_oauth_credential_store(store);
191        }
192
193        let mcp_config_sources: Vec<McpConfigSource> = if self.mcp_config_sources.is_empty() {
194            self.spec.mcp_config_sources.clone()
195        } else {
196            self.mcp_config_sources
197        };
198
199        if !mcp_config_sources.is_empty() {
200            debug!("Loading MCP configs from: {:?}", mcp_config_sources);
201            builder = builder
202                .from_mcp_config_sources(&mcp_config_sources)
203                .await
204                .map_err(|e| CliError::McpError(e.to_string()))?;
205        }
206
207        let spawn = builder.spawn().await.map_err(|e| CliError::McpError(e.to_string()))?;
208        Ok((self.spec, spawn))
209    }
210}
aether_cli/runtime.rs

aether_cli/
runtime.rs
