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::mcp;
7use aether_core::mcp::{McpHandle, McpRuntime, McpSession};
8use llm::{ChatMessage, ToolDefinition};
9use mcp_servers::McpBuilderExt;
10use mcp_utils::client::{McpClientEvent, McpConnectionDetails, McpServer, OAuthHandlerFactory};
11use std::path::PathBuf;
12use tokio::sync::mpsc::{Receiver, Sender};
13use tracing::debug;
14
15pub struct RuntimeBuilder {
16    cwd: PathBuf,
17    spec: AgentSpec,
18    mcp_config_sources: Vec<McpConfigSource>,
19    extra_mcp_servers: Vec<McpServer>,
20    oauth_applicator: Option<Box<dyn FnOnce(McpBuilder) -> McpBuilder + Send>>,
21    agent_deps: AgentDeps,
22}
23
24pub struct Runtime {
25    pub agent_tx: Sender<Command>,
26    pub agent_rx: Receiver<AgentEvent>,
27    pub agent_handle: AgentHandle,
28    pub event_rx: Receiver<McpClientEvent>,
29    pub mcp_runtime: McpRuntime,
30}
31
32pub struct PromptInfo {
33    pub spec: AgentSpec,
34    pub tool_definitions: Vec<ToolDefinition>,
35}
36
37impl RuntimeBuilder {
38    pub fn from_spec(cwd: PathBuf, spec: AgentSpec) -> Self {
39        Self {
40            cwd,
41            spec,
42            mcp_config_sources: Vec::new(),
43            extra_mcp_servers: Vec::new(),
44            oauth_applicator: None,
45            agent_deps: AgentDeps::default(),
46        }
47    }
48
49    pub fn agent_deps(mut self, deps: AgentDeps) -> Self {
50        self.agent_deps = deps;
51        self
52    }
53
54    /// Set MCP config source overrides. When non-empty, these completely
55    /// replace any sources resolved from the agent's `AgentSpec`.
56    pub fn mcp_sources(mut self, sources: Vec<McpConfigSource>) -> Self {
57        self.mcp_config_sources = sources;
58        self
59    }
60
61    pub fn extra_servers(mut self, servers: Vec<McpServer>) -> Self {
62        self.extra_mcp_servers = servers;
63        self
64    }
65
66    pub fn oauth_handler_factory(mut self, factory: OAuthHandlerFactory) -> Self {
67        self.oauth_applicator = Some(Box::new(|builder| builder.with_oauth_handler_factory(factory)));
68        self
69    }
70
71    pub async fn build(
72        self,
73        custom_prompt: Option<Prompt>,
74        messages: Option<Vec<ChatMessage>>,
75    ) -> Result<Runtime, CliError> {
76        let deps = self.agent_deps.clone();
77        let (spec, session) = self.spawn_mcp().await?;
78        let mcp = session.handle().clone();
79
80        let (agent_tx, agent_rx, agent_handle) = spawn_agent(&spec, &deps, mcp, Vec::new(), |mut agent_builder| {
81            if let Some(prompt) = custom_prompt {
82                agent_builder = agent_builder.system_prompt(prompt);
83            }
84            if let Some(msgs) = messages {
85                agent_builder = agent_builder.messages(msgs);
86            }
87            agent_builder
88        })
89        .await?;
90        let (mcp_runtime, event_rx) = session.connect_agent(agent_tx.clone()).await.split();
91
92        Ok(Runtime { agent_tx, agent_rx, agent_handle, event_rx, mcp_runtime })
93    }
94
95    /// Spawn MCP, block until every server finishes its initial connection,
96    /// then connect the agent to the session's filtered tools and MCP
97    /// instructions before returning. Returns the live [`Runtime`] plus the
98    /// bootstrap snapshot for callers that need the agent ready to use tools on
99    /// its first turn.
100    pub async fn build_ready(self, messages: Vec<ChatMessage>) -> Result<(Runtime, McpConnectionDetails), CliError> {
101        let deps = self.agent_deps.clone();
102        let (spec, mut session) = self.spawn_mcp().await?;
103        let snapshot = session
104            .block_until_ready()
105            .await
106            .ok_or_else(|| CliError::McpError("MCP bootstrap aborted before completion".to_string()))?;
107        let mcp = session.handle().clone();
108
109        let (agent_tx, agent_rx, agent_handle) =
110            spawn_agent(&spec, &deps, mcp, Vec::new(), |agent_builder| agent_builder.messages(messages)).await?;
111        let (mcp_runtime, event_rx) = session.connect_agent(agent_tx.clone()).await.split();
112
113        Ok((Runtime { agent_tx, agent_rx, agent_handle, event_rx, mcp_runtime }, snapshot))
114    }
115
116    pub async fn build_prompt_info(self) -> Result<PromptInfo, CliError> {
117        let (spec, mut session) = self.spawn_mcp().await?;
118        let details = session
119            .block_until_ready()
120            .await
121            .ok_or_else(|| CliError::McpError("MCP bootstrap aborted before completion".to_string()))?;
122        let filtered_tools = details.tool_definitions();
123        Ok(PromptInfo { spec, tool_definitions: filtered_tools })
124    }
125
126    async fn spawn_mcp(self) -> Result<(AgentSpec, McpSession), CliError> {
127        let deps = self.agent_deps.clone();
128        let mut builder = mcp(&self.cwd).with_tool_filter(self.spec.tools.clone());
129
130        if let Some(apply_oauth) = self.oauth_applicator {
131            builder = apply_oauth(builder);
132        }
133
134        builder = builder.with_agent_deps(deps).with_builtin_servers();
135
136        if !self.extra_mcp_servers.is_empty() {
137            builder = builder.with_servers(self.extra_mcp_servers);
138        }
139
140        let mcp_config_sources: Vec<McpConfigSource> = if self.mcp_config_sources.is_empty() {
141            self.spec.mcp_config_sources.clone()
142        } else {
143            self.mcp_config_sources
144        };
145
146        if !mcp_config_sources.is_empty() {
147            debug!("Loading MCP configs from: {:?}", mcp_config_sources);
148            builder =
149                builder.from_mcp_config_sources(&mcp_config_sources).map_err(|e| CliError::McpError(e.to_string()))?;
150        }
151
152        let spawn = builder.spawn().await.map_err(|e| CliError::McpError(e.to_string()))?;
153        Ok((self.spec, spawn))
154    }
155}
156
157async fn spawn_agent(
158    spec: &AgentSpec,
159    deps: &AgentDeps,
160    mcp: McpHandle,
161    tool_definitions: Vec<ToolDefinition>,
162    configure: impl FnOnce(AgentBuilder) -> AgentBuilder,
163) -> Result<(Sender<Command>, Receiver<AgentEvent>, AgentHandle), CliError> {
164    let builder = AgentBuilder::from_spec(spec, vec![], deps)
165        .await
166        .map_err(|error| CliError::AgentError(error.to_string()))?
167        .tools(mcp, tool_definitions);
168
169    configure(builder).spawn().await.map_err(|error| CliError::AgentError(error.to_string()))
170}
aether_cli/runtime.rs

aether_cli/
runtime.rs
