use std::{
path::{Path, PathBuf},
sync::{
Arc,
atomic::{AtomicUsize, Ordering},
},
time::Duration,
};
use anyhow::Context as _;
use rmcp::{
ErrorData, ServerHandler, ServiceExt,
model::{
CallToolRequestParams, CallToolResponse, CallToolResult, ErrorCode, ListToolsResult,
PaginatedRequestParams, ServerCapabilities, ServerConfig, Tool,
},
service::{ClientInitializeError, RequestContext, RoleServer, RunningService, ServiceError},
transport::{
StreamableHttpClientTransport,
streamable_http_client::{StreamableHttpClientTransportConfig, StreamableHttpError},
},
};
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
use tokio::sync::{Mutex, RwLock};
use tokio_util::sync::CancellationToken;
use crate::context::{self, Inputs, Resolved};
pub const SESSION_PREFIX: &str = "s-";
const SESSION_HEX: usize = 32;
const PRESENCE_TTL_SECS: i64 = 900;
const KEEPALIVE_EVERY: Duration = Duration::from_secs(300);
const SESSION_TTL_ENV: &str = "BUS_SESSION_TTL_SECS";
const RENEW_LEAD_ENV: &str = "BUS_SESSION_RENEW_LEAD_SECS";
const DRAIN_TIMEOUT: Duration = Duration::from_secs(5);
const EXIT_TIMEOUT: Duration = Duration::from_secs(3);
pub const CONFIGURE_TOOL: &str = "configure_session";
pub const STATUS_TOOL: &str = "session_status";
type Remote = RunningService<rmcp::RoleClient, rmcp::model::ClientConfig>;
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, JsonSchema)]
#[serde(rename_all = "kebab-case")]
pub enum Binding {
Explicit,
ClaudeCode,
RequestMeta,
Instance,
}
#[derive(Clone, Debug, Default)]
pub struct ProxyOptions {
pub inputs: Inputs,
pub project: Option<String>,
pub role: Option<String>,
pub channel: Option<String>,
pub host_session: Option<String>,
pub state_dir: PathBuf,
}
pub use crate::context::session_for_host as session_for;
fn check_label(field: &str, raw: &str) -> Result<Option<String>, ErrorData> {
crate::store::presence::normalize_label(field, raw)
.map(|v| (!v.is_empty()).then_some(v))
.map_err(|e| ErrorData::invalid_params(e.to_string(), None))
}
fn random_session() -> String {
let raw = crate::auth::generate_token();
format!(
"{SESSION_PREFIX}{}",
&raw[crate::auth::TOKEN_PREFIX.len()..crate::auth::TOKEN_PREFIX.len() + SESSION_HEX]
)
}
fn env_secs(name: &str) -> Option<i64> {
std::env::var(name)
.ok()
.and_then(|v| v.trim().parse::<i64>().ok())
.filter(|v| *v > 0)
}
fn requested_session_ttl() -> Option<i64> {
env_secs(SESSION_TTL_ENV).map(|v| v.clamp(60, crate::auth::MAX_SESSION_TTL_SECS))
}
fn renewal_lead_secs(lifetime: i64, override_secs: Option<i64>) -> i64 {
let lifetime = lifetime.max(2);
override_secs.unwrap_or(lifetime / 2).clamp(1, lifetime - 1)
}
#[derive(Clone, Debug)]
pub struct SessionProof {
pub token: String,
pub session_id: String,
pub epoch: i64,
pub expires_at: String,
}
enum Renewal {
Renewed,
Refused,
Retry,
Nothing,
}
struct Connected {
resolved: Resolved,
agent: String,
team: String,
remote: Arc<Remote>,
tools: Vec<Tool>,
remote_instructions: Option<String>,
proof: Option<SessionProof>,
ct: CancellationToken,
}
struct State {
connected: Option<Connected>,
disconnected_reason: Option<String>,
session: String,
binding: Binding,
host_id: Option<String>,
project: Option<String>,
role: Option<String>,
channel: Option<String>,
generation: u64,
}
struct InFlight(Arc<AtomicUsize>);
impl InFlight {
fn enter(counter: &Arc<AtomicUsize>) -> Self {
counter.fetch_add(1, Ordering::SeqCst);
Self(counter.clone())
}
}
impl Drop for InFlight {
fn drop(&mut self) {
self.0.fetch_sub(1, Ordering::SeqCst);
}
}
#[derive(Clone)]
pub struct Proxy {
state: Arc<RwLock<State>>,
in_flight: Arc<AtomicUsize>,
switch: Arc<Mutex<()>>,
wake: Arc<tokio::sync::Notify>,
opts: Arc<ProxyOptions>,
project_dir: PathBuf,
}
#[derive(Debug, Default, Deserialize, JsonSchema)]
pub struct ConfigureArgs {
#[serde(default)]
pub role: Option<String>,
#[serde(default)]
pub project: Option<String>,
#[serde(default)]
pub channel: Option<String>,
#[serde(default)]
pub profile: Option<String>,
}
#[derive(Debug, Serialize, JsonSchema)]
pub struct Status {
pub connected: bool,
pub error: Option<String>,
pub agent: Option<String>,
pub team: Option<String>,
pub session: String,
pub address: Option<String>,
pub project: Option<String>,
pub role: Option<String>,
pub channel: Option<String>,
pub profile: Option<String>,
pub binding: Binding,
pub project_root: Option<String>,
pub bus: Option<String>,
}
#[derive(Debug, Serialize, JsonSchema)]
pub struct ConfigureResult {
pub status: Status,
pub previous: Option<PreviousIdentity>,
}
#[derive(Debug, Serialize, JsonSchema)]
pub struct PreviousIdentity {
pub agent: String,
pub team: String,
pub session: String,
pub open_claims: Vec<String>,
pub held_locks: Vec<String>,
}
fn schema_of<T: JsonSchema>() -> Arc<rmcp::model::JsonObject> {
let schema = schemars::schema_for!(T);
match serde_json::to_value(schema) {
Ok(Value::Object(map)) => Arc::new(map),
_ => Arc::new(rmcp::model::JsonObject::new()),
}
}
fn local_tools() -> Vec<Tool> {
vec![
Tool::new(
CONFIGURE_TOOL,
"Set how THIS window presents itself on the bus: role (implementation, design, \
review, …), project and default channel. Metadata only — it never changes who \
you are or your session id, so cursors, claims and locks stay yours. `profile` \
switches to another locally approved credential of the same team after \
verifying it; a different team needs a new conversation. Affects this window \
only.",
schema_of::<ConfigureArgs>(),
)
.with_title("Configure this session")
.with_output_schema::<ConfigureResult>(),
Tool::new(
STATUS_TOOL,
"Who this window is on the bus (verified agent and team), its session id and \
address (`agent/session`, what teammates use to reach exactly this window), \
project, role and default channel. Never returns credentials.",
schema_of::<EmptyArgs>(),
)
.with_title("Session status")
.with_output_schema::<Status>(),
]
}
#[derive(Debug, Default, Deserialize, JsonSchema)]
pub struct EmptyArgs {}
async fn connect_remote(
url: &str,
credential: &str,
session: &str,
epoch: Option<i64>,
) -> anyhow::Result<Remote> {
let mut config = StreamableHttpClientTransportConfig::with_uri(url.to_owned());
config.auth_header = Some(credential.to_owned());
config.allow_stateless = true;
config.custom_headers.insert(
crate::auth::SESSION_HEADER.parse()?,
session
.parse()
.context("session label is not a valid header value")?,
);
if let Some(epoch) = epoch {
config.custom_headers.insert(
crate::auth::EPOCH_HEADER.parse()?,
epoch
.to_string()
.parse()
.context("epoch is not a valid header value")?,
);
}
let transport = StreamableHttpClientTransport::from_config(config);
let remote = rmcp::model::ClientConfig::default()
.serve(transport)
.await
.map_err(|e| {
tracing::warn!(error = %e, url, "could not open a connection to the bus");
let lost = || "the connection failed before an answer came back".to_owned();
let (why, rejected) = match &e {
ClientInitializeError::JsonRpcError(data) => (data.message.to_string(), false),
ClientInitializeError::TransportError { error, .. } => {
let rejected = matches!(
http_error_in(&*error.error),
Some(StreamableHttpError::AuthRequired(_))
);
let why = match refusal_in(&*error.error) {
Some(r) => r.text(),
None if rejected => "the bus rejected the credential".to_owned(),
None => lost(),
};
(why, rejected)
}
_ => (lost(), false),
};
let text = format!("could not connect to the bus: {why}");
if rejected {
anyhow::Error::new(Verdict::Unauthorized).context(text)
} else {
anyhow::anyhow!(text)
}
})?;
let peer_info = remote.peer_info();
if let Some(si) = peer_info.as_ref().and_then(|i| i.server_info.as_ref()) {
let ours = env!("CARGO_PKG_VERSION");
if si.name == "ai-crew-sync" && si.version != ours {
tracing::warn!(
binary = ours,
bus = %si.version,
"this binary and the bus run different ai-crew-sync versions; \
if tools fail to load or calls are refused, align the two \
before debugging anything else"
);
} else if si.name != "ai-crew-sync" {
tracing::debug!(
server = %si.name,
version = %si.version,
"the bus did not identify an ai-crew-sync version (0.7.0 or older)"
);
}
}
Ok(remote)
}
fn unauthorized(e: &ServiceError) -> bool {
let ServiceError::TransportSend(sent) = e else {
return false;
};
match http_error_in(&*sent.error) {
Some(http) => matches!(http, StreamableHttpError::AuthRequired(_)),
None => sent.error.to_string().contains("Auth required"),
}
}
fn http_error_in<'a>(
root: &'a (dyn std::error::Error + 'static),
) -> Option<&'a StreamableHttpError<reqwest::Error>> {
let mut cause = Some(root);
while let Some(err) = cause {
if let Some(http) = err.downcast_ref::<StreamableHttpError<reqwest::Error>>() {
return Some(http);
}
cause = err.source();
}
None
}
fn no_such_tool(e: &ServiceError) -> bool {
match e {
ServiceError::McpError(err) => {
err.code == ErrorCode::METHOD_NOT_FOUND
|| (err.code == ErrorCode::INVALID_PARAMS && err.message.trim() == "tool not found")
}
_ => false,
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
enum Verdict {
#[error("the bus rejected the credential")]
Unauthorized,
#[error("the bus has no such tool")]
NoSuchTool,
}
fn verdict(e: &ServiceError) -> Option<Verdict> {
if unauthorized(e) {
Some(Verdict::Unauthorized)
} else if no_such_tool(e) {
Some(Verdict::NoSuchTool)
} else {
None
}
}
fn verdict_of(e: &anyhow::Error) -> Option<Verdict> {
e.chain().find_map(|c| c.downcast_ref::<Verdict>().copied())
}
fn settled_ids(reply: Option<&Value>, sent: &[String]) -> std::collections::HashSet<String> {
let mut settled = std::collections::HashSet::new();
let Some(reply) = reply else {
return settled;
};
let listed = reply.get("confirmed_ids").is_some() || reply.get("already_confirmed").is_some();
if listed {
for key in ["confirmed_ids", "already_confirmed"] {
if let Some(ids) = reply.get(key).and_then(|v| v.as_array()) {
settled.extend(ids.iter().filter_map(|v| v.as_str()).map(str::to_owned));
}
}
} else if reply.get("confirmed").and_then(|v| v.as_i64()) == Some(sent.len() as i64) {
settled.extend(sent.iter().cloned());
}
settled
}
#[derive(Debug, PartialEq, Eq)]
struct Refusal {
status: u16,
said: Option<String>,
}
impl Refusal {
fn text(&self) -> String {
match &self.said {
Some(said) => format!("the bus refused it before running it: {said}"),
None => format!(
"the bus refused it with HTTP {} before running it",
self.status
),
}
}
}
fn refusal_in(root: &(dyn std::error::Error + 'static)) -> Option<Refusal> {
let StreamableHttpError::UnexpectedServerResponse(msg) = http_error_in(root)? else {
return None;
};
let rest = msg.strip_prefix("HTTP ")?;
let (head, body) = rest.split_once(": ").unwrap_or((rest, ""));
let status: u16 = head.split_whitespace().next()?.parse().ok()?;
let said = serde_json::from_str::<Value>(body)
.ok()
.and_then(|v| v.get("error")?.as_str().map(str::to_owned));
(said.is_some() || (400..500).contains(&status)).then_some(Refusal { status, said })
}
fn refusal(e: &ServiceError) -> Option<Refusal> {
match e {
ServiceError::TransportSend(sent) => refusal_in(&*sent.error),
_ => None,
}
}
fn remote_error_text(e: &ServiceError) -> String {
if let ServiceError::McpError(data) = e {
return data.message.to_string();
}
tracing::warn!(error = %e, "the call to the bus failed in transport");
match refusal(e) {
Some(r) => r.text(),
None => "the connection to the bus failed before an answer came back".to_owned(),
}
}
async fn call_remote(remote: &Remote, name: &str, args: Value) -> anyhow::Result<Value> {
let arguments: rmcp::model::JsonObject =
serde_json::from_value(args).context("arguments must be an object")?;
let result = remote
.call_tool(CallToolRequestParams::new(name.to_owned()).with_arguments(arguments))
.await
.map_err(|e| {
let text = format!("{name}: {}", remote_error_text(&e));
match verdict(&e) {
Some(v) => anyhow::Error::new(v).context(text),
None => anyhow::anyhow!(text),
}
})?;
if result.is_error == Some(true) {
let said: Vec<String> = result
.content
.iter()
.filter_map(|c| c.as_text().map(|t| t.text.clone()))
.collect();
anyhow::bail!("{name}: {}", said.join(" "));
}
Ok(result.structured_content.unwrap_or(Value::Null))
}
async fn register_new(remote: &Remote, session: &str) -> anyhow::Result<Option<SessionProof>> {
let mut args = json!({ "session": session });
if let Some(ttl) = requested_session_ttl() {
args["ttl_seconds"] = json!(ttl);
}
match call_remote(remote, "register_session", args).await {
Ok(v) => match v["session_token"].as_str() {
Some(token) => Ok(Some(SessionProof {
token: token.to_owned(),
session_id: v["session_id"].as_str().unwrap_or_default().to_owned(),
epoch: v["epoch"].as_i64().unwrap_or(1),
expires_at: v["expires_at"].as_str().unwrap_or_default().to_owned(),
})),
None => anyhow::bail!(
"the bus accepted register_session but returned no credential; refusing to \
continue with an asserted label while reporting a proven identity"
),
},
Err(e) => {
let text = e.to_string();
if verdict_of(&e) == Some(Verdict::NoSuchTool) {
tracing::warn!(
"this bus does not issue session credentials; continuing with the label only"
);
Ok(None)
} else {
Err(anyhow::anyhow!(
"could not register this window's session: {text}"
))
}
}
}
}
async fn resume_with(
url: &str,
prior: &SessionProof,
session: &str,
) -> anyhow::Result<Option<SessionProof>> {
let remote = connect_remote(url, &prior.token, session, Some(prior.epoch))
.await
.context("the stored session credential could not open a connection")?;
let mut args = json!({});
if let Some(ttl) = requested_session_ttl() {
args["ttl_seconds"] = json!(ttl);
}
let outcome = call_remote(&remote, "resume_session", args).await;
let _ = remote.cancel().await;
let v =
outcome.context("this window's session could not be resumed; it may have been revoked")?;
let token = v["session_token"]
.as_str()
.context("resume_session returned no credential")?;
Ok(Some(SessionProof {
token: token.to_owned(),
session_id: v["session_id"].as_str().unwrap_or_default().to_owned(),
epoch: v["epoch"].as_i64().unwrap_or(1),
expires_at: v["expires_at"].as_str().unwrap_or_default().to_owned(),
}))
}
async fn establish(
inputs: &Inputs,
session: &str,
existing_proof: Option<SessionProof>,
) -> anyhow::Result<(
Resolved,
String,
String,
Remote,
Vec<Tool>,
Option<String>,
Option<SessionProof>,
)> {
let resolved = context::resolve(inputs)?;
for w in &resolved.warnings {
tracing::warn!("{w}");
}
let rejected = || {
anyhow::anyhow!(
"the bus rejected this window's credential — it has been revoked or \
rotated{}. The credential came from {}. Issue a new token \
(`ai-crew-sync admin token issue --save`) or select another \
approved profile",
resolved
.profile
.as_deref()
.map(|p| format!(" (profile '{p}')"))
.unwrap_or_default(),
resolved.credential_provenance()
)
};
let remote = match connect_remote(&resolved.mcp_url, &resolved.token, session, None).await {
Ok(remote) => remote,
Err(e) if verdict_of(&e) == Some(Verdict::Unauthorized) => return Err(rejected()),
Err(e) => return Err(e),
};
let me = match call_remote(&remote, "whoami", json!({})).await {
Ok(me) => me,
Err(e) => {
let raw = e.to_string();
let _ = remote.cancel().await;
if verdict_of(&e) == Some(Verdict::Unauthorized) {
return Err(rejected());
}
anyhow::bail!("the bus did not accept the credential: {raw}");
}
};
let agent = me["agent"].as_str().unwrap_or_default().to_owned();
let team = me["team"].as_str().unwrap_or_default().to_owned();
if let Some((exp_team, exp_agent)) = &resolved.expected
&& (&agent != exp_agent || &team != exp_team)
{
let _ = remote.cancel().await;
anyhow::bail!(
"profile '{}' expects {exp_agent}@{exp_team} but the token authenticates as \
{agent}@{team}; fix the profile or its token entry",
resolved.profile.as_deref().unwrap_or("?")
);
}
let stored = existing_proof;
let proof = match &stored {
Some(prior) => resume_with(&resolved.mcp_url, prior, session).await?,
None => register_new(&remote, session).await?,
};
let (remote, tools, instructions) = match &proof {
Some(proof) => {
let _ = remote.cancel().await;
let remote =
connect_remote(&resolved.mcp_url, &proof.token, session, Some(proof.epoch))
.await
.context("the session credential could not open a connection")?;
let tools = remote.list_all_tools().await.map_err(|e| {
anyhow::anyhow!("could not list the bus's tools: {}", remote_error_text(&e))
})?;
let instructions = remote.peer_info().and_then(|i| i.instructions.clone());
(remote, tools, instructions)
}
None => {
let tools = remote.list_all_tools().await.map_err(|e| {
anyhow::anyhow!("could not list the bus's tools: {}", remote_error_text(&e))
})?;
let instructions = remote.peer_info().and_then(|i| i.instructions.clone());
(remote, tools, instructions)
}
};
Ok((resolved, agent, team, remote, tools, instructions, proof))
}
fn git_place(dir: &Path) -> (Option<String>, Option<String>) {
let run = |args: &[&str]| {
std::process::Command::new("git")
.args(args)
.current_dir(dir)
.output()
.ok()
.filter(|o| o.status.success())
.map(|o| String::from_utf8_lossy(&o.stdout).trim().to_owned())
.filter(|s| !s.is_empty())
};
let repo = run(&["config", "--get", "remote.origin.url"]).map(|url| {
let trimmed = url.trim_end_matches(".git");
let tail: Vec<&str> = trimmed.rsplit(['/', ':']).take(2).collect();
if tail.len() == 2 {
format!("{}/{}", tail[1], tail[0])
} else {
trimmed.to_owned()
}
});
let branch = run(&["branch", "--show-current"]);
(repo, branch)
}
impl Proxy {
pub async fn start(opts: ProxyOptions) -> Self {
let project_dir = opts
.inputs
.project_dir
.clone()
.or_else(|| std::env::var_os("CLAUDE_PROJECT_DIR").map(PathBuf::from))
.or_else(|| std::env::current_dir().ok())
.unwrap_or_else(|| PathBuf::from("."));
let (host_id, binding) = if let Some(id) = opts
.host_session
.as_deref()
.map(str::trim)
.filter(|s| !s.is_empty())
{
(Some(id.to_owned()), Binding::Explicit)
} else if let Some(id) = std::env::var("CLAUDE_CODE_SESSION_ID")
.ok()
.map(|s| s.trim().to_owned())
.filter(|s| !s.is_empty())
{
(Some(id), Binding::ClaudeCode)
} else {
(None, Binding::Instance)
};
let session = match &host_id {
Some(id) => session_for(id),
None => random_session(),
};
let proxy = Self {
state: Arc::new(RwLock::new(State {
connected: None,
disconnected_reason: None,
session,
binding,
host_id,
project: opts.project.clone(),
role: opts.role.clone(),
channel: opts.channel.clone(),
generation: 0,
})),
in_flight: Arc::new(AtomicUsize::new(0)),
switch: Arc::new(Mutex::new(())),
wake: Arc::new(tokio::sync::Notify::new()),
opts: Arc::new(opts),
project_dir,
};
let inputs = proxy.opts.inputs.clone();
if let Err(e) = proxy.connect_with(&inputs, true).await {
tracing::warn!(error = %e, "proxy started without a bus connection");
proxy.state.write().await.disconnected_reason = Some(format!("{e:#}"));
}
proxy
}
async fn connect_with(
&self,
inputs: &Inputs,
reuse_proof: bool,
) -> anyhow::Result<Option<PreviousIdentity>> {
let _guard = self.switch.lock().await;
let (session, binding_key) = {
let st = self.state.read().await;
(
st.session.clone(),
st.host_id.clone().unwrap_or_else(|| st.session.clone()),
)
};
let existing = reuse_proof
.then(|| context::read_binding(&self.opts.state_dir, &binding_key))
.flatten()
.and_then(|b| match (b.session_token, b.session_id, b.epoch) {
(Some(token), Some(session_id), Some(epoch)) if !token.is_empty() => {
Some(SessionProof {
token,
session_id,
epoch,
expires_at: b.expires_at.unwrap_or_default(),
})
}
_ => None,
});
let (resolved, agent, team, remote, tools, instructions, proof) =
establish(inputs, &session, existing).await?;
{
let st = self.state.read().await;
if let Some(old) = &st.connected
&& old.team != team
{
let _ = remote.cancel().await;
anyhow::bail!(
"this conversation is bound to team '{}'; the profile '{}' belongs to team \
'{team}'. Switching teams inside a conversation is not allowed — the \
transcript already holds '{}' material. Start a new conversation with \
that profile instead",
old.team,
resolved.profile.as_deref().unwrap_or("?"),
old.team
);
}
}
{
let mut st = self.state.write().await;
if st.project.is_none() {
st.project = resolved.project.clone();
}
if st.channel.is_none() {
st.channel = resolved.channel.clone();
}
}
let previous = {
let old = {
let mut st = self.state.write().await;
st.connected.take()
};
match old {
Some(old) => {
old.ct.cancel();
let started = std::time::Instant::now();
while self.in_flight.load(Ordering::SeqCst) > 0
&& started.elapsed() < DRAIN_TIMEOUT
{
tokio::time::sleep(Duration::from_millis(20)).await;
}
let held = report_holdings(&old.remote, &old.agent, &old.team, &session).await;
if old.proof.is_some() {
let _ = tokio::time::timeout(
EXIT_TIMEOUT,
call_remote(&old.remote, "revoke_session", json!({})),
)
.await;
}
let _ = tokio::time::timeout(
EXIT_TIMEOUT,
call_remote(
&old.remote,
"heartbeat",
json!({"status": "idle", "ttl_seconds": 30}),
),
)
.await;
close_remote(old.remote).await;
Some(held)
}
None => None,
}
};
let connected = Connected {
resolved,
agent,
team,
remote: Arc::new(remote),
tools,
remote_instructions: instructions,
proof,
ct: CancellationToken::new(),
};
{
let mut st = self.state.write().await;
st.connected = Some(connected);
st.disconnected_reason = None;
st.generation += 1;
}
self.wake.notify_one();
self.heartbeat("active").await;
self.write_binding().await;
Ok(previous)
}
async fn heartbeat(&self, status: &str) {
let (remote, project, role) = {
let st = self.state.read().await;
let Some(c) = &st.connected else { return };
(c.remote.clone(), st.project.clone(), st.role.clone())
};
let (repo, branch) = git_place(&self.project_dir);
let mut args = json!({"status": status, "ttl_seconds": PRESENCE_TTL_SECS});
if let Some(r) = repo {
args["repo"] = Value::String(r);
}
if let Some(b) = branch {
args["branch"] = Value::String(b);
}
args["project"] = Value::String(project.unwrap_or_default());
args["role"] = Value::String(role.unwrap_or_default());
if let Err(e) = call_remote(&remote, "heartbeat", args).await {
tracing::warn!(error = %e, "heartbeat failed");
}
}
async fn write_binding(&self) {
let st = self.state.read().await;
let key = st.host_id.clone().unwrap_or_else(|| st.session.clone());
let record = json!({
"host_id_present": st.host_id.is_some(),
"binding": st.binding,
"session": st.session,
"project": st.project,
"role": st.role,
"channel": st.channel,
"profile": st.connected.as_ref().and_then(|c| c.resolved.profile.clone()),
"agent": st.connected.as_ref().map(|c| c.agent.clone()),
"team": st.connected.as_ref().map(|c| c.team.clone()),
"mcp_url": st.connected.as_ref().map(|c| c.resolved.mcp_url.clone()),
"session_token": st.connected.as_ref().and_then(|c| c.proof.as_ref().map(|p| p.token.clone())),
"session_id": st.connected.as_ref().and_then(|c| c.proof.as_ref().map(|p| p.session_id.clone())),
"epoch": st.connected.as_ref().and_then(|c| c.proof.as_ref().map(|p| p.epoch)),
"expires_at": st.connected.as_ref().and_then(|c| c.proof.as_ref().map(|p| p.expires_at.clone())),
"proxy_pid": std::process::id(),
"updated_at": chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Secs, true),
});
drop(st);
let path = context::binding_path(&self.opts.state_dir, &key);
let text = record.to_string();
if let Err(e) = under_config_lock(self.opts.state_dir.clone(), move || {
context::write_binding_file(&path, &text)
})
.await
{
tracing::warn!(error = %e, "could not write the session binding");
}
}
async fn status(&self) -> Status {
let st = self.state.read().await;
let c = st.connected.as_ref();
Status {
connected: c.is_some(),
error: st.disconnected_reason.clone(),
agent: c.map(|c| c.agent.clone()),
team: c.map(|c| c.team.clone()),
session: st.session.clone(),
address: c.map(|c| format!("{}/{}", c.agent, st.session)),
project: st.project.clone(),
role: st.role.clone(),
channel: st.channel.clone(),
profile: c.and_then(|c| c.resolved.profile.clone()),
binding: st.binding,
project_root: c
.and_then(|c| c.resolved.project_root.as_ref())
.map(|p| p.display().to_string()),
bus: c.map(|c| c.resolved.mcp_url.clone()),
}
}
async fn configure(&self, args: ConfigureArgs) -> anyhow::Result<ConfigureResult> {
let mut previous = None;
let staged_role = match args.role {
Some(v) => Some(check_label("role", &v).map_err(|e| anyhow::anyhow!("{}", e.message))?),
None => None,
};
let staged_project = match args.project {
Some(v) => {
Some(check_label("project", &v).map_err(|e| anyhow::anyhow!("{}", e.message))?)
}
None => None,
};
let staged_channel = match args.channel {
Some(v) => {
Some(check_label("channel", &v).map_err(|e| anyhow::anyhow!("{}", e.message))?)
}
None => None,
};
if let Some(profile) = args
.profile
.map(|p| p.trim().to_owned())
.filter(|p| !p.is_empty())
{
let mut inputs = self.opts.inputs.clone();
inputs.profile = Some(profile);
inputs.explicit_token = None;
inputs.explicit_url = None;
previous = self.connect_with(&inputs, false).await?;
}
{
let mut st = self.state.write().await;
if let Some(role) = staged_role {
st.role = role;
}
if let Some(project) = staged_project {
st.project = project;
}
if let Some(channel) = staged_channel {
st.channel = channel;
}
}
self.heartbeat("active").await;
self.write_binding().await;
Ok(ConfigureResult {
status: self.status().await,
previous,
})
}
async fn observe_meta(&self, meta: &rmcp::model::RequestMetaObject) -> Result<(), ErrorData> {
let thread = meta
.0
.0
.get("threadId")
.or_else(|| meta.0.0.get("sessionId"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|s| !s.is_empty());
let Some(thread) = thread else { return Ok(()) };
let (bound, current) = {
let st = self.state.read().await;
(st.host_id.clone(), st.binding)
};
match bound {
Some(id) if id == thread => Ok(()),
Some(_) if current == Binding::RequestMeta => Err(ErrorData::invalid_request(
"this proxy instance is bound to another conversation; a second one is using \
the same MCP process, which is not supported. Configure the host to start \
one `ai-crew-sync mcp proxy` per conversation",
None,
)),
Some(_) => Ok(()),
None => self.rebind(thread.to_owned()).await.map_err(|e| {
ErrorData::internal_error(format!("could not bind the conversation: {e:#}"), None)
}),
}
}
async fn rebind(&self, host_id: String) -> anyhow::Result<()> {
let previous = {
let st = self.state.read().await;
(st.host_id.clone(), st.binding, st.session.clone())
};
{
let mut st = self.state.write().await;
st.host_id = Some(host_id.clone());
st.binding = Binding::RequestMeta;
st.session = session_for(&host_id);
}
let inputs = {
let st = self.state.read().await;
match &st.connected {
Some(c) => {
let mut i = self.opts.inputs.clone();
i.profile = c.resolved.profile.clone();
i
}
None => self.opts.inputs.clone(),
}
};
match self.connect_with(&inputs, true).await {
Ok(_) => Ok(()),
Err(e) => {
let mut st = self.state.write().await;
(st.host_id, st.binding, st.session) = previous;
st.disconnected_reason = Some(format!("{e:#}"));
Err(e)
}
}
}
async fn spool_and_confirm(
&self,
result: CallToolResult,
host_ct: CancellationToken,
) -> CallToolResult {
let Some(structured) = result.structured_content.clone() else {
return result;
};
let session = self.state.read().await.session.clone();
let path = crate::spool::spool_path(&self.opts.state_dir, &session);
let mut entries: Vec<crate::spool::Entry> = structured
.get("references")
.and_then(|v| v.as_array())
.map(|refs| {
refs.iter()
.filter_map(|r| {
Some(crate::spool::Entry {
delivery_id: r.get("delivery_id")?.as_str()?.to_owned(),
message_id: r.get("message_id")?.as_str()?.to_owned(),
conversation_id: r
.get("conversation_id")
.and_then(|v| v.as_str())
.unwrap_or_default()
.to_owned(),
seq: r.get("seq").and_then(|v| v.as_i64()).unwrap_or(0),
from_address: r
.get("from_address")
.and_then(|v| v.as_str())
.unwrap_or_default()
.to_owned(),
created_at: r
.get("created_at")
.and_then(|v| v.as_str())
.unwrap_or_default()
.to_owned(),
confirmed: false,
spooled_at: chrono::Utc::now().to_rfc3339(),
})
})
.collect::<Vec<_>>()
})
.unwrap_or_default();
let mut held = crate::spool::read(&path);
let spooled = match crate::spool::append(&path, &entries) {
Ok(written) => written,
Err(e) => {
tracing::warn!(error = %e, "could not spool inbox references; not confirming");
return result;
}
};
held.extend(spooled.iter().cloned());
entries.clear();
let to_confirm = crate::spool::unconfirmed(&held);
if to_confirm.is_empty() {
return result;
}
let mut params = rmcp::model::JsonObject::new();
params.insert(
"delivery_ids".into(),
Value::Array(
to_confirm
.iter()
.map(|id| Value::String(id.clone()))
.collect(),
),
);
let confirm =
CallToolRequestParams::new("confirm_inbox_delivery".to_string()).with_arguments(params);
match self.forward(confirm, host_ct).await {
Ok(confirmation) => {
let settled = settled_ids(confirmation.structured_content.as_ref(), &to_confirm);
let sent: std::collections::HashSet<&String> = to_confirm.iter().collect();
let mut left = 0usize;
for entry in held.iter_mut() {
if sent.contains(&entry.delivery_id) {
if settled.contains(&entry.delivery_id) {
entry.confirmed = true;
} else {
left += 1;
}
}
}
if left > 0 {
tracing::warn!(
sent = to_confirm.len(),
left,
"the bus did not settle every reference sent; keeping the rest in \
the spool"
);
}
if let Err(e) = crate::spool::rewrite(&path, &held) {
tracing::warn!(error = %e, "could not compact the inbox spool");
}
}
Err(e) => tracing::warn!(error = %e, "could not confirm inbox delivery"),
}
result
}
async fn forward(
&self,
request: CallToolRequestParams,
host_ct: CancellationToken,
) -> Result<CallToolResult, ErrorData> {
let (remote, ct, generation, _guard) = {
let st = self.state.read().await;
let Some(c) = &st.connected else {
return Err(ErrorData::invalid_request(
format!(
"not connected to the bus: {}. Call {CONFIGURE_TOOL} with an approved \
profile, or fix the local configuration and start a new conversation",
st.disconnected_reason
.as_deref()
.unwrap_or("no profile resolved")
),
None,
));
};
(
c.remote.clone(),
c.ct.clone(),
st.generation,
InFlight::enter(&self.in_flight),
)
};
let name = request.name.to_string();
let profile = {
let st = self.state.read().await;
st.connected
.as_ref()
.and_then(|c| c.resolved.profile.clone())
};
let outcome = tokio::select! {
r = remote.call_tool(request) => r.map_err(|e| {
if unauthorized(&e) {
self.mark_unauthorized(&profile);
ErrorData::invalid_request(
format!(
"{name}: the bus rejected this window's credential — it has been \
revoked or rotated{}. Issue a new token (`ai-crew-sync admin \
token issue --save`) and call {CONFIGURE_TOOL} with an approved \
profile; nothing was sent",
profile
.as_deref()
.map(|p| format!(" (profile '{p}')"))
.unwrap_or_default()
),
None,
)
} else if let ServiceError::McpError(data) = e {
data
} else if let Some(r) = refusal(&e) {
tracing::warn!(error = %e, tool = %name, "the bus refused a forwarded call");
ErrorData::invalid_request(format!("{name}: {}", r.text()), None)
} else {
tracing::warn!(error = %e, tool = %name, "a forwarded call failed in transport");
ErrorData::internal_error(
format!(
"{name} could not reach the bus: the connection failed before an \
answer came back, so the call may or may not have run. Check \
before repeating anything that is not safe to repeat."
),
None,
)
}
}),
_ = ct.cancelled() => Err(ErrorData::invalid_request(
format!(
"{name} was cancelled: this window switched credentials while the call was \
in flight (generation {generation}). Nothing was replayed; call again if \
it is still wanted, as the new identity"
),
None,
)),
_ = host_ct.cancelled() => Err(ErrorData::invalid_request(
format!("{name} was cancelled by the client"),
None,
)),
};
outcome
}
fn mark_unauthorized(&self, profile: &Option<String>) {
if let Ok(mut st) = self.state.try_write() {
st.disconnected_reason = Some(format!(
"the bus rejected the credential{} (revoked or rotated)",
profile
.as_deref()
.map(|p| format!(" of profile '{p}'"))
.unwrap_or_default()
));
}
}
async fn default_channel(&self) -> Option<String> {
let st = self.state.read().await;
st.channel.clone().or_else(|| st.project.clone())
}
fn instructions(&self, st: &State) -> String {
let mut lines = Vec::new();
match &st.connected {
Some(c) => {
lines.push(format!(
"[ai-crew-sync] You are agent '{}' on team '{}', in session '{}'. Teammates \
reach exactly this window at '{}/{}'.",
c.agent, c.team, st.session, c.agent, st.session
));
lines.push(format!(
"- project: {}, role: {}, default channel: {}. Change them with \
{CONFIGURE_TOOL}; see them with {STATUS_TOOL}. Find teammates' windows \
with list_sessions.",
st.project.as_deref().unwrap_or("(none — set it)"),
st.role.as_deref().unwrap_or("(none — set it)"),
st.channel
.as_deref()
.or(st.project.as_deref())
.unwrap_or("(none)"),
));
lines.push(
"- Nothing is pushed into an idle turn: call read_messages or wait_for_updates \
to receive what teammates sent."
.to_owned(),
);
if let Some(remote) = &c.remote_instructions {
lines.push(String::new());
lines.push(remote.clone());
}
}
None => {
lines.push(format!(
"[ai-crew-sync] Not connected to the team bus: {}. Only {CONFIGURE_TOOL} and \
{STATUS_TOOL} are available until a locally approved profile connects.",
st.disconnected_reason
.as_deref()
.unwrap_or("no profile resolved")
));
}
}
lines.join("\n")
}
pub async fn keepalive(self, ct: CancellationToken) {
let mut next_heartbeat = tokio::time::Instant::now() + KEEPALIVE_EVERY;
let mut not_before: Option<(u64, tokio::time::Instant)> = None;
loop {
let deadline = self.renewal_deadline().await;
let renew_at = deadline.map(|(generation, at)| match not_before {
Some((for_generation, nb)) if for_generation == generation => at.max(nb),
_ => at,
});
let renew_sleep = tokio::time::sleep_until(
renew_at.unwrap_or_else(|| tokio::time::Instant::now() + KEEPALIVE_EVERY),
);
tokio::select! {
_ = ct.cancelled() => return,
_ = self.wake.notified() => {}
_ = tokio::time::sleep_until(next_heartbeat) => {
self.heartbeat("active").await;
next_heartbeat = tokio::time::Instant::now() + KEEPALIVE_EVERY;
}
_ = renew_sleep, if renew_at.is_some() => {
let pause = match self.renew_credential().await {
Renewal::Refused => KEEPALIVE_EVERY,
Renewal::Renewed | Renewal::Retry | Renewal::Nothing => self.renewal_retry().await,
};
if let Some((generation, _)) = deadline {
not_before = Some((generation, tokio::time::Instant::now() + pause));
}
}
}
}
}
async fn renewal_deadline(&self) -> Option<(u64, tokio::time::Instant)> {
let (generation, expires_at) = {
let st = self.state.read().await;
let expires_at = st
.connected
.as_ref()
.and_then(|c| c.proof.as_ref())
.map(|p| p.expires_at.clone())?;
(st.generation, expires_at)
};
let expires_at = chrono::DateTime::parse_from_rfc3339(&expires_at).ok()?;
let remaining = (expires_at.with_timezone(&chrono::Utc) - chrono::Utc::now())
.num_seconds()
.max(0);
let due_in = (remaining - self.renewal_lead().await).max(0) as u64;
Some((
generation,
tokio::time::Instant::now() + Duration::from_secs(due_in),
))
}
async fn renewal_lead(&self) -> i64 {
let lifetime = requested_session_ttl().unwrap_or(crate::auth::SESSION_TTL_SECS);
renewal_lead_secs(lifetime, env_secs(RENEW_LEAD_ENV))
}
async fn renewal_retry(&self) -> Duration {
Duration::from_secs((self.renewal_lead().await / 4).clamp(2, 60) as u64)
}
async fn renew_credential(&self) -> Renewal {
let (remote, proof, generation, ct, profile) = {
let st = self.state.read().await;
let Some(c) = &st.connected else {
return Renewal::Nothing;
};
let Some(p) = &c.proof else {
return Renewal::Nothing;
};
(
c.remote.clone(),
p.clone(),
st.generation,
c.ct.clone(),
c.resolved.profile.clone(),
)
};
let mut args = json!({});
if let Some(ttl) = requested_session_ttl() {
args["ttl_seconds"] = json!(ttl);
}
let outcome = tokio::select! {
_ = ct.cancelled() => return Renewal::Nothing,
r = call_remote(&remote, "renew_session", args) => r,
};
match outcome {
Ok(v) => {
let Some(expires_at) = v["expires_at"].as_str().map(str::to_owned) else {
tracing::warn!("renew_session answered without an expiry; keeping the old one");
return Renewal::Retry;
};
if v["epoch"].as_i64().is_some_and(|e| e != proof.epoch) {
tracing::warn!(
"renew_session answered for another epoch; keeping the credential this \
window holds"
);
return Renewal::Retry;
}
{
let mut st = self.state.write().await;
if st.generation != generation {
return Renewal::Nothing;
}
let Some(current) = st.connected.as_mut().and_then(|c| c.proof.as_mut()) else {
return Renewal::Nothing;
};
if current.session_id != proof.session_id || current.epoch != proof.epoch {
return Renewal::Nothing;
}
current.expires_at = expires_at.clone();
}
self.stamp_binding_expiry(&proof, &expires_at).await;
tracing::debug!(expires_at = %expires_at, "session credential renewed");
Renewal::Renewed
}
Err(e) => match verdict_of(&e) {
Some(Verdict::Unauthorized) => {
tracing::warn!(error = %e, "the bus refused to renew this window's credential");
self.mark_unauthorized(&profile);
Renewal::Refused
}
Some(Verdict::NoSuchTool) => {
tracing::debug!("this bus does not renew credentials");
Renewal::Nothing
}
None => {
tracing::warn!(error = %e, "could not renew this window's credential; retrying");
Renewal::Retry
}
},
}
}
async fn stamp_binding_expiry(&self, proof: &SessionProof, expires_at: &str) {
let key = {
let st = self.state.read().await;
st.host_id.clone().unwrap_or_else(|| st.session.clone())
};
let path = context::binding_path(&self.opts.state_dir, &key);
let (session_id, epoch, expires_at) =
(proof.session_id.clone(), proof.epoch, expires_at.to_owned());
let stamped = under_config_lock(self.opts.state_dir.clone(), move || {
let Ok(text) = std::fs::read_to_string(&path) else {
return Ok(false);
};
let Ok(mut value) = serde_json::from_str::<Value>(&text) else {
return Ok(false);
};
let same = value["session_id"].as_str() == Some(session_id.as_str())
&& value["epoch"].as_i64() == Some(epoch);
if !same {
return Ok(false);
}
if let Some(map) = value.as_object_mut() {
map.insert("expires_at".into(), json!(expires_at));
map.insert(
"updated_at".into(),
json!(chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Secs, true)),
);
}
context::write_binding_file(&path, &value.to_string())?;
Ok(true)
})
.await;
match stamped {
Ok(true) => {}
Ok(false) => tracing::debug!(
binding = %key,
"another instance owns this binding now, or it is gone; not stamping"
),
Err(e) => {
tracing::warn!(error = %e, binding = %key, "could not record the renewed expiry")
}
}
}
pub async fn shutdown(&self) {
let remote = {
let st = self.state.read().await;
st.connected.as_ref().map(|c| c.remote.clone())
};
if let Some(remote) = remote {
let _ = tokio::time::timeout(
EXIT_TIMEOUT,
call_remote(
&remote,
"heartbeat",
json!({"status": "idle", "ttl_seconds": 120}),
),
)
.await;
close_remote(remote).await;
}
self.mark_closed().await;
}
async fn mark_closed(&self) {
let (key, mine) = {
let st = self.state.read().await;
let key = st.host_id.clone().unwrap_or_else(|| st.session.clone());
let mine = st
.connected
.as_ref()
.and_then(|c| c.proof.as_ref().map(|p| (p.session_id.clone(), p.epoch)));
(key, mine)
};
let path = context::binding_path(&self.opts.state_dir, &key);
let shown = key.clone();
let outcome = under_config_lock(self.opts.state_dir.clone(), move || {
let Ok(text) = std::fs::read_to_string(&path) else {
return Ok(false);
};
let Ok(mut value) = serde_json::from_str::<Value>(&text) else {
return Ok(false);
};
if let Some((session_id, epoch)) = mine {
let same = value["session_id"].as_str() == Some(session_id.as_str())
&& value["epoch"].as_i64() == Some(epoch);
if !same {
return Ok(false);
}
}
if let Some(map) = value.as_object_mut() {
map.insert("closed_at".into(), json!(chrono::Utc::now().to_rfc3339()));
}
context::write_binding_file(&path, &value.to_string())?;
Ok(true)
})
.await;
match outcome {
Ok(true) => {}
Ok(false) => tracing::debug!(
binding = %shown,
"another instance owns this binding now, or it is gone; leaving it alone"
),
Err(e) => {
tracing::warn!(error = %e, binding = %shown, "could not mark the binding closed")
}
}
}
}
async fn under_config_lock<T: Send + 'static>(
dir: PathBuf,
f: impl FnOnce() -> anyhow::Result<T> + Send + 'static,
) -> anyhow::Result<T> {
tokio::task::spawn_blocking(move || context::with_config_lock(&dir, f))
.await
.map_err(|e| anyhow::anyhow!("the binding writer task failed: {e}"))?
}
async fn close_remote(remote: Arc<Remote>) {
if let Ok(owned) = Arc::try_unwrap(remote) {
let _ = owned.cancel().await;
}
}
async fn report_holdings(
remote: &Remote,
agent: &str,
team: &str,
session: &str,
) -> PreviousIdentity {
let open_claims = call_remote(remote, "list_tasks", json!({"mine_only": true}))
.await
.ok()
.and_then(|v| v["tasks"].as_array().cloned())
.unwrap_or_default()
.iter()
.filter(|t| t["status"] == "claimed")
.filter_map(|t| t["key"].as_str().map(str::to_owned))
.collect();
let held_locks = call_remote(remote, "list_locks", json!({}))
.await
.ok()
.and_then(|v| v["locks"].as_array().cloned())
.unwrap_or_default()
.iter()
.filter(|l| {
l["holder"] == agent && l["holder_session"].as_str().unwrap_or_default() == session
})
.filter_map(|l| l["name"].as_str().map(str::to_owned))
.collect();
PreviousIdentity {
agent: agent.to_owned(),
team: team.to_owned(),
session: session.to_owned(),
open_claims,
held_locks,
}
}
fn tool_error(msg: String) -> CallToolResult {
CallToolResult::error(vec![rmcp::model::ContentBlock::text(msg)])
}
impl ServerHandler for Proxy {
fn get_info(&self) -> ServerConfig {
let mut info = ServerConfig::new(ServerCapabilities::builder().enable_tools().build());
let text = match self.state.try_read() {
Ok(st) => self.instructions(&st),
Err(_) => format!("[ai-crew-sync] initialising; call {STATUS_TOOL} for details."),
};
info.instructions = Some(text);
info
}
async fn list_tools(
&self,
_request: Option<PaginatedRequestParams>,
_context: RequestContext<RoleServer>,
) -> Result<ListToolsResult, ErrorData> {
let mut tools = local_tools();
if let Some(c) = &self.state.read().await.connected {
tools.extend(c.tools.iter().cloned());
}
Ok(ListToolsResult::with_all_items(tools))
}
async fn call_tool(
&self,
request: CallToolRequestParams,
context: RequestContext<RoleServer>,
) -> Result<CallToolResponse, ErrorData> {
self.observe_meta(&context.meta).await?;
match request.name.as_ref() {
STATUS_TOOL => {
let status = self.status().await;
let value = serde_json::to_value(status)
.map_err(|e| ErrorData::internal_error(e.to_string(), None))?;
Ok(CallToolResult::structured(value).into())
}
CONFIGURE_TOOL => {
let args: ConfigureArgs = match request.arguments {
Some(map) => serde_json::from_value(Value::Object(map))
.map_err(|e| ErrorData::invalid_params(e.to_string(), None))?,
None => ConfigureArgs::default(),
};
match self.configure(args).await {
Ok(result) => {
let value = serde_json::to_value(result)
.map_err(|e| ErrorData::internal_error(e.to_string(), None))?;
Ok(CallToolResult::structured(value).into())
}
Err(e) => Ok(tool_error(format!("{e:#}")).into()),
}
}
"fetch_conversation_inbox" => {
let result = self.forward(request, context.ct.clone()).await?;
Ok(self.spool_and_confirm(result, context.ct).await.into())
}
"post_message" => {
let mut request = request;
if let Some(channel) = self.default_channel().await {
let args = request.arguments.get_or_insert_with(Default::default);
let addressed = args.contains_key("channel") || args.contains_key("to");
if !addressed {
args.insert("channel".into(), Value::String(channel));
}
}
Ok(self.forward(request, context.ct).await?.into())
}
_ => Ok(self.forward(request, context.ct).await?.into()),
}
}
}
pub async fn run(opts: ProxyOptions) -> anyhow::Result<()> {
let proxy = Proxy::start(opts).await;
let ct = CancellationToken::new();
let keepalive = tokio::spawn(proxy.clone().keepalive(ct.child_token()));
let running = proxy
.clone()
.serve(rmcp::transport::stdio())
.await
.context("MCP initialize over stdio failed")?;
let quit = running.waiting().await;
tracing::debug!(?quit, "host closed the connection");
ct.cancel();
let _ = keepalive.await;
proxy.shutdown().await;
Ok(())
}
#[cfg(test)]
mod renewal_tests {
#[test]
fn renewal_lead_is_half_the_lifetime_unless_overridden_and_inside_it() {
assert_eq!(super::renewal_lead_secs(24 * 3600, None), 12 * 3600);
assert_eq!(super::renewal_lead_secs(60, None), 30);
assert_eq!(super::renewal_lead_secs(60, Some(50)), 50);
assert_eq!(
super::renewal_lead_secs(60, Some(600)),
59,
"never past the lifetime"
);
assert_eq!(
super::renewal_lead_secs(1, None),
1,
"a degenerate lifetime still yields a lead"
);
}
}
#[cfg(test)]
mod unauthorized_tests {
use rmcp::{
RoleClient,
model::{ErrorCode, ErrorData},
transport::{
DynamicTransportError, StreamableHttpClientTransport,
streamable_http_client::{AuthRequiredError, StreamableHttpError},
},
};
use super::*;
fn transport(e: StreamableHttpError<reqwest::Error>) -> ServiceError {
ServiceError::TransportSend(DynamicTransportError::new::<
StreamableHttpClientTransport<reqwest::Client>,
RoleClient,
>(e))
}
#[test]
fn a_rejected_bearer_is_the_transport_saying_so() {
let e = transport(StreamableHttpError::AuthRequired(AuthRequiredError::new(
"Bearer".into(),
)));
assert!(unauthorized(&e));
}
#[test]
fn a_refusal_that_spells_401_is_still_a_refusal() {
let e = ServiceError::McpError(ErrorData::invalid_request(
"you do not hold the claim on 'api#1': it is held by joaquin (session \
's-a7da401d8d70'), the lease expires in 401s",
None,
));
assert!(!unauthorized(&e));
let e = ServiceError::McpError(ErrorData::internal_error("Auth required", None));
assert!(
!unauthorized(&e),
"not even when it borrows the transport's words"
);
}
#[test]
fn a_missing_tool_is_the_code_saying_so() {
let e = ServiceError::McpError(ErrorData::invalid_params("tool not found", None));
assert!(no_such_tool(&e));
assert_eq!(verdict(&e), Some(Verdict::NoSuchTool));
let e = ServiceError::McpError(ErrorData::new(
ErrorCode::METHOD_NOT_FOUND,
"Method not found",
None,
));
assert!(no_such_tool(&e));
}
#[test]
fn a_refusal_that_spells_a_missing_method_is_still_a_refusal() {
let e = ServiceError::McpError(ErrorData::invalid_params(
"conflict: session 's-32601f03e877' is already registered and still live. \
Holding the agent token does not make you that window: Method aside, \
reconnect it with resume_session",
None,
));
assert!(!no_such_tool(&e));
assert_eq!(verdict(&e), None);
let e = ServiceError::McpError(ErrorData::invalid_params("not found: message 32601", None));
assert!(!no_such_tool(&e));
let e = ServiceError::McpError(ErrorData::invalid_params(
"tool not found: the deploy tool named in `depends_on` does not exist",
None,
));
assert!(!no_such_tool(&e));
assert!(!no_such_tool(&ServiceError::TransportClosed));
}
#[test]
fn a_verdict_survives_the_anyhow_chain() {
let e = anyhow::Error::new(Verdict::NoSuchTool).context("register_session failed");
assert_eq!(verdict_of(&e), Some(Verdict::NoSuchTool));
let e = anyhow::anyhow!("register_session failed: tool not found -32601 Method");
assert_eq!(verdict_of(&e), None, "words are not a verdict");
}
#[test]
fn an_http_refusal_is_read_for_its_shape() {
let answered = |msg: &str| {
transport(StreamableHttpError::UnexpectedServerResponse(
msg.to_owned().into(),
))
};
let e = answered(r#"HTTP 429 Too Many Requests: {"error":"rate limit exceeded"}"#);
assert_eq!(
refusal(&e),
Some(Refusal {
status: 429,
said: Some("rate limit exceeded".into())
})
);
assert!(remote_error_text(&e).contains("rate limit exceeded"));
assert!(!unauthorized(&e));
let e = answered("HTTP 404 Not Found: <html>nope</html>");
assert_eq!(
refusal(&e),
Some(Refusal {
status: 404,
said: None
})
);
assert_eq!(refusal(&answered("HTTP 504 Gateway Timeout: ")), None);
assert_eq!(
refusal(&answered("invalid www-authenticate header value")),
None
);
assert_eq!(refusal(&ServiceError::TransportClosed), None);
}
#[test]
fn another_transport_failure_is_not_a_rejected_bearer() {
let e = transport(StreamableHttpError::UnexpectedContentType(Some(
"text/html; 401".into(),
)));
assert!(!unauthorized(&e));
assert!(!unauthorized(&ServiceError::TransportClosed));
}
}