use std::collections::BTreeSet;
use std::time::Duration;
use rmcp::service::{Peer, RoleServer};
use serde_json::{Value, json};
use tracing::{debug, info, warn};
use uuid::Uuid;
use crate::socket_client::SocketClient;
use super::server::{RESPONSE_PREFIX, TOOLS_LIST_TOPIC, new_req_id, unwrap_reply_payload};
const CAPSULES_LOADED_TOPIC: &str = "astrid.v1.capsules_loaded";
const ENUMERATE_DEADLINE: Duration = Duration::from_secs(55);
pub(super) async fn run(peer: Peer<RoleServer>, principal: String) {
let session = astrid_core::SessionId::from_uuid(Uuid::new_v4());
let mut watch_client = match SocketClient::connect(session).await {
Ok(c) => c,
Err(e) => {
warn!(error = %e, "MCP hot-reload watcher: failed to open watch uplink; live tool-reload pushes disabled");
return;
},
};
info!("MCP hot-reload watcher: watching for {CAPSULES_LOADED_TOPIC} broadcasts");
let mut last_known: BTreeSet<String> = match enumerate_tool_names(&principal).await {
Ok(names) => names,
Err(e) => {
warn!(error = %e, "MCP hot-reload watcher: baseline seed failed; starting from empty set");
BTreeSet::new()
},
};
loop {
let frame = match watch_client.read_raw_frame().await {
Ok(Some(bytes)) => bytes,
Ok(None) => {
debug!("MCP hot-reload watcher: watch uplink closed; stopping");
return;
},
Err(e) => {
warn!(error = %e, "MCP hot-reload watcher: watch uplink read failed; stopping");
return;
},
};
let Ok(raw) = serde_json::from_slice::<Value>(&frame) else {
continue;
};
if raw.get("topic").and_then(Value::as_str) != Some(CAPSULES_LOADED_TOPIC) {
continue;
}
debug!("MCP hot-reload watcher: capsules_loaded received; re-enumerating tools");
let names = match enumerate_tool_names(&principal).await {
Ok(names) => names,
Err(e) => {
warn!(error = %e, "MCP hot-reload watcher: tool re-enumeration failed; keeping last-known set");
continue;
},
};
if names == last_known {
debug!("MCP hot-reload watcher: tool set unchanged; suppressing notification");
continue;
}
if let Err(e) = peer.notify_tool_list_changed().await {
warn!(error = %e, "MCP hot-reload watcher: notify failed (peer closed); stopping");
return;
}
info!(
tools = names.len(),
"MCP hot-reload watcher: tool set changed; pushed tools/list_changed"
);
last_known = names;
}
}
async fn enumerate_tool_names(principal: &str) -> anyhow::Result<BTreeSet<String>> {
let session = astrid_core::SessionId::from_uuid(Uuid::new_v4());
let mut client = SocketClient::connect(session).await?;
let req_id = new_req_id();
let reply_topic = format!("{RESPONSE_PREFIX}{req_id}");
let body = json!({ "req_id": req_id });
let msg = astrid_types::ipc::IpcMessage::new(
TOOLS_LIST_TOPIC,
astrid_types::ipc::IpcPayload::RawJson(body),
Uuid::nil(),
)
.with_principal(principal.to_string());
client.send_message(msg).await?;
let raw = client
.read_until_topic(&reply_topic, ENUMERATE_DEADLINE)
.await?;
let reply = unwrap_reply_payload(&raw);
let names = reply
.get("tools")
.and_then(Value::as_array)
.map(|arr| {
arr.iter()
.filter_map(|t| t.get("name").and_then(Value::as_str))
.map(str::to_string)
.collect()
})
.unwrap_or_default();
Ok(names)
}