Skip to main content

objectiveai_mcp/
run.rs

1//! ObjectiveAI MCP server.
2//!
3//! Other crates can `use objectiveai_mcp::{ConfigBuilder, run}` and
4//! spawn the server in-process without going through the binary. The
5//! executor is a required argument — `main.rs` builds a
6//! `BinaryExecutor` from `Config.objectiveai_dir`; other callers can
7//! pass any `CommandExecutor` impl.
8
9use 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            // Layout root (OBJECTIVEAI_DIR). Same default as the api
84            // and the viewer.
85            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    /// Layout root (`OBJECTIVEAI_DIR`); default `~/.objectiveai`.
103    pub objectiveai_dir: std::path::PathBuf,
104    /// State name (`OBJECTIVEAI_STATE`); default `"default"`.
105    pub objectiveai_state: String,
106}
107
108/// Build the rmcp `(TcpListener, axum::Router)` pair. The executor is
109/// `Arc`-wrapped internally so the dynamic plugin / tool route
110/// closures and the static `ObjectiveAI` handler share a single
111/// instance.
112pub 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    // Discover registered plugins + tools concurrently — both are
131    // independent SDK calls that each spawn one cli subprocess.
132    // The lists are shared (via Arc) between `with_plugins_and_tools`
133    // (which registers one dynamic tool per entry) and
134    // `HeaderSessionManager::new` (which validates the optional
135    // `X-OBJECTIVEAI-MCP-{TOOLS,PLUGINS}` filter sets against the
136    // same installed manifest at connect time).
137    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    // Shared per-rmcp-session bag of SessionState (wraps the
143    // legacy AgentArguments identity bag alongside the three
144    // optional X-OBJECTIVEAI-MCP-* filter values). Populated by
145    // the HeaderSessionManager on every initialize (fresh + lazy
146    // reconnect); consumed by every tool dispatcher and the
147    // hand-written `list_tools` handler.
148    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            // rmcp 1.7 marks `StreamableHttpServerConfig` `#[non_exhaustive]`.
167            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    // The daemon spawns exactly one mcp per state and holds it as a
195    // leashed child; the URL clients connect with (wildcard binds map
196    // to loopback) is announced via the stdout readiness handshake
197    // below.
198    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    // Readiness handshake: the daemon (this server's sole spawner and
211    // leash-holder) blocks on this stdout line and caches the address.
212    // No lockfile — the daemon owns this process's lifetime outright.
213    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
221/// Best-effort plugin discovery: drain `plugins list`, drop per-item
222/// errors. A discovery failure (cli binary missing, etc.) returns an
223/// empty list so the server still starts.
224async 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
241/// Best-effort tool discovery; same shape as `list_plugins`.
242async 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}