aether_cli/
runtime.rs

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

aether_cli/
runtime.rs
