use std::collections::BTreeMap;
use std::process::Stdio;
use anyhow::{Context, Result};
use mcpmesh_net::errors::synthesized_limited;
use mcpmesh_net::transport::NdjsonTransport;
use mcpmesh_net::{PeerIdentity, SessionBackend, SessionTransport};
use serde_json::Value;
use tokio::io::{AsyncRead, AsyncWrite};
use tokio::process::Command;
use tokio::sync::Semaphore;
use crate::audit::RequestAuditor;
const CONCURRENCY_RETRY_MS: u64 = 1000;
pub struct SpawnBackend {
pub cmd: Vec<String>,
pub env: BTreeMap<String, String>,
pub cwd: Option<String>,
pub concurrency: std::sync::Arc<Semaphore>,
pub service: String,
pub audit: crate::audit::AuditSink,
pub limiter: std::sync::Arc<crate::limits::RateLimiter>,
}
#[async_trait::async_trait]
impl SessionBackend for SpawnBackend {
async fn run(
&self,
identity: Option<PeerIdentity>,
initialize: Value,
transport: SessionTransport,
) -> anyhow::Result<()> {
self.run_over(identity, initialize, transport).await
}
}
impl SpawnBackend {
pub async fn run_over<R, W>(
&self,
identity: Option<PeerIdentity>,
initialize: Value,
mut transport: NdjsonTransport<R, W>,
) -> Result<()>
where
R: AsyncRead + Send + Unpin,
W: AsyncWrite + Send + Unpin,
{
let _permit = match self.concurrency.clone().try_acquire_owned() {
Ok(permit) => permit,
Err(_) => {
let id = initialize.get("id").cloned().unwrap_or(Value::Null);
let _ = transport
.send_value(synthesized_limited(id, CONCURRENCY_RETRY_MS))
.await;
let _ = transport.shutdown().await;
return Ok(());
}
};
let (program, args) = self.cmd.split_first().context("run backend cmd is empty")?;
let mut command = Command::new(program);
command
.args(args)
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::inherit()) .kill_on_drop(true); command.envs(
self.env
.iter()
.filter(|(k, _)| !k.to_ascii_uppercase().starts_with("MCPMESH_")),
);
if let Some(cwd) = &self.cwd {
command.current_dir(cwd);
}
match &identity {
Some(id) => {
command.env("MCPMESH_PEER_NAME", &id.name);
match &id.user_id {
Some(user) => command.env("MCPMESH_PEER_USER", user),
None => command.env_remove("MCPMESH_PEER_USER"),
};
command.env("MCPMESH_PEER_GROUPS", id.groups.join(","));
command.env("MCPMESH_PEER_EID", id.endpoint.principal());
}
None => {
command.env_remove("MCPMESH_PEER_NAME");
command.env_remove("MCPMESH_PEER_USER");
command.env_remove("MCPMESH_PEER_GROUPS");
command.env_remove("MCPMESH_PEER_EID");
}
}
let mut child = command
.spawn()
.with_context(|| format!("spawn run backend child `{program}`"))?;
let child_stdin = child.stdin.take().expect("stdin piped above");
let child_stdout = child.stdout.take().expect("stdout piped above");
let peer = identity
.as_ref()
.map(|id| id.user_id.clone().unwrap_or_else(|| id.name.clone()));
let _session = self
.audit
.session(peer.clone().unwrap_or_default(), self.service.clone());
let auditor = RequestAuditor::new(self.audit.clone(), peer.clone(), self.service.clone());
let outcome = super::pump(
initialize,
&mut transport,
child_stdout,
child_stdin,
auditor,
crate::limits::RateGate::new(
self.limiter.clone(),
identity.as_ref().map(|i| i.endpoint),
),
)
.await;
outcome
}
}