1use std::sync::Arc;
10
11use envconfig::Envconfig;
12use futures::StreamExt;
13use objectiveai_sdk::cli::command::CommandExecutor;
14use objectiveai_sdk::cli::command::plugins;
15use objectiveai_sdk::cli::command::tools;
16use rmcp::transport::streamable_http_server::{StreamableHttpServerConfig, StreamableHttpService};
17use tokio_util::sync::CancellationToken;
18
19use crate::agent_args_registry::AgentArgumentsRegistry;
20use crate::header_session_manager::HeaderSessionManager;
21use crate::objectiveai::ObjectiveAiMcpCli;
22
23#[derive(Envconfig)]
24struct EnvConfigBuilder {
25 #[envconfig(from = "ADDRESS")]
26 address: Option<String>,
27 #[envconfig(from = "PORT")]
28 port: Option<u16>,
29 #[envconfig(from = "SUPPRESS_OUTPUT")]
30 suppress_output: Option<String>,
31 #[envconfig(from = "OBJECTIVEAI_DIR")]
32 objectiveai_dir: Option<String>,
33 #[envconfig(from = "OBJECTIVEAI_STATE")]
34 objectiveai_state: Option<String>,
35}
36
37impl EnvConfigBuilder {
38 fn build(self) -> ConfigBuilder {
39 ConfigBuilder {
40 address: self.address,
41 port: self.port,
42 suppress_output: self
43 .suppress_output
44 .map(|v| matches!(v.to_ascii_lowercase().as_str(), "1" | "true" | "yes" | "on")),
45 objectiveai_dir: self.objectiveai_dir,
46 objectiveai_state: self.objectiveai_state,
47 }
48 }
49}
50
51#[derive(Default)]
52pub struct ConfigBuilder {
53 pub address: Option<String>,
54 pub port: Option<u16>,
55 pub suppress_output: Option<bool>,
56 pub objectiveai_dir: Option<String>,
57 pub objectiveai_state: Option<String>,
58}
59
60impl Envconfig for ConfigBuilder {
61 #[allow(deprecated)]
62 fn init() -> Result<Self, envconfig::Error> {
63 EnvConfigBuilder::init().map(|e| e.build())
64 }
65
66 fn init_from_env() -> Result<Self, envconfig::Error> {
67 EnvConfigBuilder::init_from_env().map(|e| e.build())
68 }
69
70 fn init_from_hashmap(
71 hashmap: &std::collections::HashMap<String, String>,
72 ) -> Result<Self, envconfig::Error> {
73 EnvConfigBuilder::init_from_hashmap(hashmap).map(|e| e.build())
74 }
75}
76
77impl ConfigBuilder {
78 pub fn build(self) -> Config {
79 Config {
80 address: self.address.unwrap_or_else(|| "127.0.0.1".to_string()),
81 port: self.port.unwrap_or(0),
82 suppress_output: self.suppress_output.unwrap_or(false),
83 objectiveai_dir: match self.objectiveai_dir {
86 Some(dir) => std::path::PathBuf::from(dir),
87 None => dirs::home_dir()
88 .unwrap_or_else(|| std::path::PathBuf::from("."))
89 .join(".objectiveai"),
90 },
91 objectiveai_state: self
92 .objectiveai_state
93 .unwrap_or_else(|| "default".to_string()),
94 }
95 }
96}
97
98pub struct Config {
99 pub address: String,
100 pub port: u16,
101 pub suppress_output: bool,
102 pub objectiveai_dir: std::path::PathBuf,
104 pub objectiveai_state: String,
106}
107
108pub async fn setup<E>(
113 config: Config,
114 executor: E,
115) -> std::io::Result<(tokio::net::TcpListener, axum::Router)>
116where
117 E: CommandExecutor + Send + Sync + 'static,
118 E::Error: std::fmt::Display + Send + 'static,
119{
120 let Config {
121 address,
122 port,
123 suppress_output: _,
124 objectiveai_dir: _,
125 objectiveai_state: _,
126 } = config;
127
128 let executor = Arc::new(executor);
129
130 let (plugins_list, tools_list) =
138 tokio::join!(list_plugins(&*executor), list_tools(&*executor));
139 let plugins_list = Arc::new(plugins_list);
140 let tools_list = Arc::new(tools_list);
141
142 let registry = Arc::new(AgentArgumentsRegistry::new());
149
150 let server = ObjectiveAiMcpCli::with_plugins_and_tools(
151 executor.clone(),
152 (*plugins_list).clone(),
153 (*tools_list).clone(),
154 registry.clone(),
155 );
156 let session_manager = Arc::new(HeaderSessionManager::new(
157 registry.clone(),
158 server.clone(),
159 tools_list.clone(),
160 plugins_list.clone(),
161 ));
162 let ct = CancellationToken::new();
163
164 let service: StreamableHttpService<ObjectiveAiMcpCli<E>, HeaderSessionManager<E>> =
165 StreamableHttpService::new(move || Ok(server.clone()), session_manager, {
166 let mut cfg = StreamableHttpServerConfig::default();
168 cfg.stateful_mode = true;
169 cfg.sse_keep_alive = None;
170 cfg.cancellation_token = ct.child_token();
171 cfg
172 });
173
174 let router = axum::Router::new().fallback_service(service);
175 let listener = tokio::net::TcpListener::bind(format!("{address}:{port}")).await?;
176 Ok((listener, router))
177}
178
179pub async fn serve(
180 listener: tokio::net::TcpListener,
181 app: axum::Router,
182) -> std::io::Result<()> {
183 axum::serve(listener, app).await
184}
185
186pub async fn run<E>(config: Config, executor: E) -> std::io::Result<()>
187where
188 E: CommandExecutor + Send + Sync + 'static,
189 E::Error: std::fmt::Display + Send + 'static,
190{
191 let suppress_output = config.suppress_output;
192 let (listener, app) = setup(config, executor).await?;
193
194 let addr = listener.local_addr()?;
199 let connect_ip = match addr.ip() {
200 std::net::IpAddr::V4(v4) if v4.is_unspecified() => {
201 std::net::IpAddr::V4(std::net::Ipv4Addr::LOCALHOST)
202 }
203 std::net::IpAddr::V6(v6) if v6.is_unspecified() => {
204 std::net::IpAddr::V6(std::net::Ipv6Addr::LOCALHOST)
205 }
206 ip => ip,
207 };
208 let connect_url =
209 format!("http://{}", std::net::SocketAddr::new(connect_ip, addr.port()));
210 objectiveai_sdk::process::print_ready(Some(&connect_url));
214
215 if !suppress_output {
216 eprintln!("listening on {addr}");
217 }
218 serve(listener, app).await
219}
220
221async fn list_plugins<E>(executor: &E) -> Vec<plugins::list::ResponseItem>
225where
226 E: CommandExecutor + Send + Sync + 'static,
227 E::Error: std::fmt::Display + Send + 'static,
228{
229 let request = plugins::list::Request {
230 path_type: plugins::list::Path::PluginsList,
231 offset: None,
232 limit: None,
233 base: Default::default(),
234 };
235 match plugins::list::execute(executor, request, None).await {
236 Ok(stream) => stream.filter_map(|r| async move { r.ok() }).collect().await,
237 Err(_) => Vec::new(),
238 }
239}
240
241async fn list_tools<E>(executor: &E) -> Vec<tools::list::ResponseItem>
243where
244 E: CommandExecutor + Send + Sync + 'static,
245 E::Error: std::fmt::Display + Send + 'static,
246{
247 let request = tools::list::Request {
248 path_type: tools::list::Path::ToolsList,
249 offset: None,
250 limit: None,
251 base: Default::default(),
252 };
253 match tools::list::execute(executor, request, None).await {
254 Ok(stream) => stream.filter_map(|r| async move { r.ok() }).collect().await,
255 Err(_) => Vec::new(),
256 }
257}