1use crate::error::CliError;
2use aether_core::agent_spec::AgentSpec;
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 aether_project::load_agent_catalog;
10use llm::{ChatMessage, LlmModel, ToolDefinition};
11use mcp_servers::McpBuilderExt;
12use mcp_utils::client::oauth::OAuthHandler;
13use mcp_utils::client::{ElicitationRequest, McpServerConfig};
14use mcp_utils::status::McpServerStatusEntry;
15use std::path::{Path, PathBuf};
16use tokio::sync::mpsc::{Receiver, Sender};
17use tokio::task::JoinHandle;
18use tracing::debug;
19
20pub struct RuntimeBuilder {
21 cwd: PathBuf,
22 spec: AgentSpec,
23 mcp_config: Option<PathBuf>,
24 extra_mcp_servers: Vec<McpServerConfig>,
25 oauth_applicator: Option<Box<dyn FnOnce(McpBuilder) -> McpBuilder + Send>>,
26 prompt_cache_key: Option<String>,
27}
28
29pub struct Runtime {
30 pub agent_tx: Sender<UserMessage>,
31 pub agent_rx: Receiver<AgentMessage>,
32 pub agent_handle: AgentHandle,
33 pub mcp_tx: Sender<McpCommand>,
34 pub elicitation_rx: Receiver<ElicitationRequest>,
35 pub server_statuses: Vec<McpServerStatusEntry>,
36 pub mcp_handle: JoinHandle<()>,
37}
38
39pub struct PromptInfo {
40 pub spec: AgentSpec,
41 pub tool_definitions: Vec<ToolDefinition>,
42}
43
44impl RuntimeBuilder {
45 pub fn new(cwd: &Path, model: &str) -> Result<Self, CliError> {
46 let cwd = cwd.canonicalize().map_err(CliError::IoError)?;
47 let parsed_model: LlmModel = model.parse().map_err(|e: String| CliError::ModelError(e))?;
48 let catalog = load_agent_catalog(&cwd).map_err(|e| CliError::AgentError(e.to_string()))?;
49 let spec = catalog.resolve_default(&parsed_model, None, &cwd);
50
51 Ok(Self {
52 cwd,
53 spec,
54 mcp_config: None,
55 extra_mcp_servers: Vec::new(),
56 oauth_applicator: None,
57 prompt_cache_key: None,
58 })
59 }
60
61 pub fn from_spec(cwd: PathBuf, spec: AgentSpec) -> Self {
62 Self {
63 cwd,
64 spec,
65 mcp_config: None,
66 extra_mcp_servers: Vec::new(),
67 oauth_applicator: 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 pub fn mcp_config(mut self, path: PathBuf) -> Self {
78 self.mcp_config = Some(path);
79 self
80 }
81
82 pub fn mcp_config_opt(self, path: Option<PathBuf>) -> Self {
83 match path {
84 Some(p) => self.mcp_config(p),
85 None => self,
86 }
87 }
88
89 pub fn extra_servers(mut self, servers: Vec<McpServerConfig>) -> Self {
90 self.extra_mcp_servers = servers;
91 self
92 }
93
94 pub fn oauth_handler<H: OAuthHandler + 'static>(mut self, handler: H) -> Self {
95 self.oauth_applicator = Some(Box::new(|builder| builder.with_oauth_handler(handler)));
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 mcp = self.spawn_mcp().await?;
106
107 let filtered_tools = mcp.spec.tools.apply(mcp.tool_definitions);
108 let mut agent_builder = AgentBuilder::from_spec(&mcp.spec, vec![])
109 .map_err(|e| CliError::AgentError(e.to_string()))?
110 .tools(mcp.mcp_tx.clone(), filtered_tools);
111
112 if let Some(key) = prompt_cache_key {
113 agent_builder = agent_builder.prompt_cache_key(key);
114 }
115
116 if let Some(prompt) = custom_prompt {
117 agent_builder = agent_builder.system_prompt(prompt);
118 }
119
120 if let Some(msgs) = messages {
121 agent_builder = agent_builder.messages(msgs);
122 }
123
124 let (agent_tx, agent_rx, agent_handle) = agent_builder
125 .spawn()
126 .await
127 .map_err(|e| CliError::AgentError(e.to_string()))?;
128
129 Ok(Runtime {
130 agent_tx,
131 agent_rx,
132 agent_handle,
133 mcp_tx: mcp.mcp_tx,
134 elicitation_rx: mcp.elicitation_rx,
135 server_statuses: mcp.server_statuses,
136 mcp_handle: mcp.mcp_handle,
137 })
138 }
139
140 pub async fn build_prompt_info(self) -> Result<PromptInfo, CliError> {
141 let mcp = self.spawn_mcp().await?;
142 let filtered_tools = mcp.spec.tools.apply(mcp.tool_definitions);
143 Ok(PromptInfo {
144 spec: mcp.spec,
145 tool_definitions: filtered_tools,
146 })
147 }
148
149 async fn spawn_mcp(self) -> Result<McpParts, CliError> {
150 let mut builder = mcp().with_builtin_servers(self.cwd.clone(), &self.cwd);
151
152 if !self.extra_mcp_servers.is_empty() {
153 builder = builder.with_servers(self.extra_mcp_servers);
154 }
155
156 if let Some(apply_oauth) = self.oauth_applicator {
157 builder = apply_oauth(builder);
158 }
159
160 let mcp_config_path = self.mcp_config.or(self.spec.mcp_config_path.clone());
161
162 if let Some(ref config_path) = mcp_config_path {
163 debug!("Loading MCP config from: {}", config_path.display());
164 let config_str = config_path
165 .to_str()
166 .ok_or_else(|| CliError::McpError("Invalid MCP config path".to_string()))?;
167
168 builder = builder
169 .from_json_file(config_str)
170 .await
171 .map_err(|e| CliError::McpError(e.to_string()))?;
172 }
173
174 let McpSpawnResult {
175 tool_definitions,
176 instructions,
177 server_statuses,
178 command_tx: mcp_tx,
179 elicitation_rx,
180 handle: mcp_handle,
181 } = builder
182 .spawn()
183 .await
184 .map_err(|e| CliError::McpError(e.to_string()))?;
185
186 let mut spec = self.spec;
187 spec.prompts.push(Prompt::mcp_instructions(instructions));
188
189 Ok(McpParts {
190 spec,
191 tool_definitions,
192 mcp_tx,
193 elicitation_rx,
194 server_statuses,
195 mcp_handle,
196 })
197 }
198}
199
200struct McpParts {
201 spec: AgentSpec,
202 tool_definitions: Vec<ToolDefinition>,
203 mcp_tx: Sender<McpCommand>,
204 elicitation_rx: Receiver<ElicitationRequest>,
205 server_statuses: Vec<McpServerStatusEntry>,
206 mcp_handle: JoinHandle<()>,
207}