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

aether_cli/
runtime.rs
