use crate::cache::Cache;
use crate::config::{Config, ServerConfig};
use crate::model::ToolInfo;
use crate::oauth::{OAuth, TokenStore};
use anyhow::{anyhow, bail};
use rmcp::ServiceExt;
use rmcp::model::CallToolResult;
use std::collections::BTreeMap;
use std::sync::Arc;
use tokio::sync::{Mutex, OnceCell, RwLock};
struct Entry {
cfg: ServerConfig,
client: RwLock<OnceCell<Arc<rmcp::service::RunningService<rmcp::service::RoleClient, ()>>>>,
failed: Mutex<Option<String>>,
tools: Mutex<Vec<ToolInfo>>,
instructions: Mutex<Option<String>>,
oauth: Option<Arc<OAuth>>,
token_store: Option<TokenStore>,
}
pub struct Upstreams {
entries: BTreeMap<String, Arc<Entry>>,
cache: Mutex<Cache>,
config_hash: u64,
}
const CONNECT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);
async fn connect(
e: &Entry,
) -> anyhow::Result<rmcp::service::RunningService<rmcp::service::RoleClient, ()>> {
let cfg = &e.cfg;
let fut = async {
if let Some(cmd) = &cfg.command {
let mut c = tokio::process::Command::new(cmd);
c.args(&cfg.args).envs(&cfg.env);
let (transport, stderr) = rmcp::transport::TokioChildProcess::builder(c)
.stderr(std::process::Stdio::piped())
.spawn()?;
if let Some(mut err) = stderr {
use tokio::io::AsyncBufReadExt;
tokio::spawn(async move {
let mut lines = tokio::io::BufReader::new(&mut err).lines();
while let Ok(Some(line)) = lines.next_line().await {
tracing::debug!(target: "upstream-stderr", "{line}");
}
});
}
Ok(().serve(transport).await?)
} else if let Some(url) = &cfg.url {
let mut config = rmcp::transport::streamable_http_client::StreamableHttpClientTransportConfig::with_uri(url.clone());
for (k, v) in &cfg.headers {
config.custom_headers.insert(
http::HeaderName::from_bytes(k.as_bytes())?,
http::HeaderValue::from_str(v)?,
);
}
if let Some(oauth) = &e.oauth {
let store = e.token_store.as_ref().unwrap();
match oauth.access_token(cfg, store).await? {
Some(token) => {
config.custom_headers.insert(
http::header::AUTHORIZATION,
http::HeaderValue::from_str(&format!("Bearer {token}"))?,
);
}
None => {
let auth_url = oauth.begin_flow(cfg, store).await?;
bail!(
"server requires OAuth authorization. Open in a browser: {auth_url} — then retry. Headless? Call authorize_server with the final redirect URL as pasted_url."
);
}
}
}
let transport = rmcp::transport::StreamableHttpClientTransport::from_config(config);
Ok(().serve(transport).await?)
} else {
bail!("server has neither command nor url")
}
};
match tokio::time::timeout(CONNECT_TIMEOUT, fut).await {
Ok(r) => r,
Err(_) => Err(anyhow!(
"connect timed out after {}s",
CONNECT_TIMEOUT.as_secs()
)),
}
}
impl Upstreams {
pub fn new(cfg: Config, cache: Cache, config_hash: u64) -> Upstreams {
let token_lock = Arc::new(Mutex::new(()));
let entries = cfg
.mcp_servers
.iter()
.map(|(name, sc)| {
let tools = cache.servers.get(name).cloned().unwrap_or_default();
let instr = cache.instructions.get(name).cloned();
(
name.clone(),
Arc::new(Entry {
cfg: sc.clone(),
client: RwLock::new(OnceCell::new()),
failed: Mutex::new(None),
tools: Mutex::new(tools),
instructions: Mutex::new(instr),
oauth: sc.oauth.then(OAuth::new),
token_store: sc.oauth.then(|| TokenStore::new(name, token_lock.clone())),
}),
)
})
.collect();
Upstreams {
entries,
cache: Mutex::new(cache),
config_hash,
}
}
pub fn server_names(&self) -> Vec<String> {
self.entries.keys().cloned().collect()
}
pub fn status(&self, name: &str) -> &'static str {
match self.entries.get(name) {
None => "unknown",
Some(e)
if e.client
.try_read()
.map(|c| c.initialized())
.unwrap_or(false) =>
{
"connected"
}
Some(e) => match e.failed.try_lock() {
Ok(g) if g.is_some() => "unavailable",
_ => "cold",
},
}
}
pub fn instructions(&self, name: &str) -> Option<String> {
self.entries
.get(name)?
.instructions
.try_lock()
.ok()?
.clone()
}
fn entry(&self, name: &str) -> anyhow::Result<&Arc<Entry>> {
self.entries.get(name).ok_or_else(|| {
anyhow!(
"unknown server {name:?}; available: {}",
self.server_names().join(", ")
)
})
}
async fn ensure(
&self,
name: &str,
) -> anyhow::Result<Arc<rmcp::service::RunningService<rmcp::service::RoleClient, ()>>> {
let e = self.entry(name)?;
let cell = e.client.read().await;
let r = cell
.get_or_try_init(|| async {
match connect(e).await {
Ok(c) => {
*e.failed.lock().await = None;
Ok(Arc::new(c))
}
Err(err) => {
*e.failed.lock().await = Some(err.to_string());
Err(err)
}
}
})
.await?
.clone();
drop(cell);
if e.tools.lock().await.is_empty() {
self.refresh_inner(name, &r).await.ok();
}
Ok(r)
}
pub async fn refresh(&self, name: &str) -> anyhow::Result<()> {
let client = self.ensure(name).await?;
self.refresh_inner(name, &client).await
}
async fn refresh_inner(
&self,
name: &str,
client: &rmcp::service::RunningService<rmcp::service::RoleClient, ()>,
) -> anyhow::Result<()> {
let e = self.entry(name)?;
let listed = client.list_all_tools().await?;
let mut out = Vec::new();
for t in listed {
if !e.cfg.is_allowed(&t.name) {
continue;
}
out.push(ToolInfo {
name: t.name.to_string(),
description: t.description.map(|d| d.to_string()),
schema: serde_json::to_value(&t.input_schema)?,
annotations: t.annotations.map(serde_json::to_value).transpose()?,
});
}
*e.tools.lock().await = out.clone();
if let Some(info) = client.peer_info()
&& let Some(instr) = info.instructions.clone()
{
*e.instructions.lock().await = Some(instr.clone());
self.cache
.lock()
.await
.instructions
.insert(name.to_string(), instr);
}
self.cache
.lock()
.await
.servers
.insert(name.to_string(), out);
self.save_cache().await;
Ok(())
}
pub fn cached_tools(&self, name: &str) -> Vec<ToolInfo> {
self.entries
.get(name)
.and_then(|e| e.tools.try_lock().ok().map(|g| g.clone()))
.unwrap_or_default()
}
pub async fn tools(&self, name: &str) -> anyhow::Result<Vec<ToolInfo>> {
let e = self.entry(name)?;
let cached = e.tools.lock().await.clone();
if !cached.is_empty() {
return Ok(cached);
}
self.refresh(name).await?;
Ok(e.tools.lock().await.clone())
}
pub async fn call(
&self,
name: &str,
tool: &str,
args: Option<serde_json::Map<String, serde_json::Value>>,
) -> anyhow::Result<CallToolResult> {
let e = self.entry(name)?;
if !e.cfg.is_allowed(tool) {
bail!("tool {tool:?} on server {name:?} is blocked by config");
}
let client = self.ensure(name).await?;
let mut params = rmcp::model::CallToolRequestParams::new(tool.to_string());
if let Some(a) = args {
params = params.with_arguments(a);
}
match client.call_tool(params.clone()).await {
Ok(r) => Ok(r),
Err(err) => {
tracing::debug!(%err, "call failed, reconnecting and retrying once");
e.client.write().await.take();
let client = self.ensure(name).await?;
self.refresh_inner(name, &client).await.ok();
Ok(client.call_tool(params).await?)
}
}
}
pub async fn oauth_begin(&self, name: &str) -> anyhow::Result<Option<String>> {
let e = self.entry(name)?;
let (Some(oauth), Some(store)) = (&e.oauth, &e.token_store) else {
bail!("server {name:?} does not have \"oauth\": true in the config");
};
match oauth.access_token(&e.cfg, store).await? {
Some(_) => Ok(None),
None => Ok(Some(oauth.begin_flow(&e.cfg, store).await?)),
}
}
pub async fn oauth_complete(&self, name: &str, pasted: &str) -> anyhow::Result<()> {
let e = self.entry(name)?;
let Some(oauth) = &e.oauth else {
bail!("server {name:?} does not have \"oauth\": true in the config");
};
oauth.complete_with_url(pasted).await
}
pub async fn save_cache(&self) {
if let Err(e) = self.cache.lock().await.save(self.config_hash) {
tracing::warn!(%e, "failed to save index cache");
}
}
}