use std::path::PathBuf;
use std::sync::Arc;
use std::time::{Duration, Instant};
use arc_swap::ArcSwapOption;
use dashmap::DashMap;
use nexo_broker::{AnyBroker, BrokerHandle};
use nexo_config::AppConfig;
use nexo_llm::LlmRegistry;
use tokio::sync::{mpsc, Mutex};
use tokio_util::sync::CancellationToken;
use crate::agent::runtime::ReloadCommand;
use crate::agent::spawn::{AgentSpawnerFn, SharedRuntimeContext};
use crate::runtime_snapshot::RuntimeSnapshot;
use crate::telemetry;
pub struct AgentReloadHandle {
pub reload_tx: mpsc::Sender<ReloadCommand>,
pub known_tools: Arc<Vec<String>>,
}
pub type PostReloadHook = Box<dyn Fn() + Send + Sync>;
pub struct ConfigReloadCoordinator {
config_dir: PathBuf,
runtimes: DashMap<String, AgentReloadHandle>,
llm_registry: Arc<LlmRegistry>,
version: Mutex<u64>,
gate: Mutex<()>,
broker: ArcSwapOption<AnyBroker>,
post_hooks: Mutex<Vec<PostReloadHook>>,
shared_ctx: ArcSwapOption<SharedRuntimeContext>,
spawner: ArcSwapOption<AgentSpawnerFn>,
shutdown: CancellationToken,
}
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct ReloadOutcome {
pub version: u64,
pub applied: Vec<String>,
pub rejected: Vec<ReloadRejection>,
pub elapsed_ms: u64,
}
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct ReloadRejection {
pub agent_id: Option<String>,
pub reason: String,
}
impl ConfigReloadCoordinator {
pub fn new(
config_dir: PathBuf,
llm_registry: Arc<LlmRegistry>,
shutdown: CancellationToken,
) -> Self {
Self {
config_dir,
runtimes: DashMap::new(),
llm_registry,
version: Mutex::new(0),
gate: Mutex::new(()),
broker: ArcSwapOption::from(None),
post_hooks: Mutex::new(Vec::new()),
shared_ctx: ArcSwapOption::from(None),
spawner: ArcSwapOption::from(None),
shutdown,
}
}
pub fn with_shared_context(self, shared: Arc<SharedRuntimeContext>) -> Self {
self.shared_ctx.store(Some(shared));
self
}
pub fn shared_context(&self) -> Option<Arc<SharedRuntimeContext>> {
self.shared_ctx.load_full()
}
pub fn set_spawner(&self, spawner: Arc<AgentSpawnerFn>) {
self.spawner.store(Some(spawner));
}
pub fn spawner(&self) -> Option<Arc<AgentSpawnerFn>> {
self.spawner.load_full()
}
pub fn unregister(&self, agent_id: &str) -> Option<AgentReloadHandle> {
self.runtimes.remove(agent_id).map(|(_, handle)| handle)
}
pub async fn register_post_hook(&self, hook: PostReloadHook) {
self.post_hooks.lock().await.push(hook);
}
pub fn register(
&self,
agent_id: impl Into<String>,
reload_tx: mpsc::Sender<ReloadCommand>,
known_tools: Arc<Vec<String>>,
) {
self.runtimes.insert(
agent_id.into(),
AgentReloadHandle {
reload_tx,
known_tools,
},
);
}
pub async fn reload(&self) -> ReloadOutcome {
let _gate = self.gate.lock().await;
let started = Instant::now();
let mut applied: Vec<String> = Vec::new();
let mut rejected: Vec<ReloadRejection> = Vec::new();
let cfg = match AppConfig::load(&self.config_dir) {
Ok(c) => c,
Err(e) => {
telemetry::inc_config_reload_rejected();
tracing::warn!(error = %e, "config reload: load failed, keeping previous snapshot");
rejected.push(ReloadRejection {
agent_id: None,
reason: format!("AppConfig::load: {e}"),
});
let current = *self.version.lock().await;
return ReloadOutcome {
version: current,
applied,
rejected,
elapsed_ms: started.elapsed().as_millis() as u64,
};
}
};
let known_providers =
crate::agent::KnownProviders::new(cfg.llm.providers.keys().map(String::as_str));
if let Err(e) = crate::agent::validate_agents_with_providers(
&cfg.agents.agents,
&cfg.plugins,
&crate::agent::KnownTools::default(),
&known_providers,
) {
telemetry::inc_config_reload_rejected();
tracing::warn!(error = %e, "config reload: validation failed, keeping previous snapshot");
rejected.push(ReloadRejection {
agent_id: None,
reason: format!("validation: {e}"),
});
let current = *self.version.lock().await;
return ReloadOutcome {
version: current,
applied,
rejected,
elapsed_ms: started.elapsed().as_millis() as u64,
};
}
let mut version_guard = self.version.lock().await;
*version_guard += 1;
let new_version = *version_guard;
drop(version_guard);
for agent_cfg in &cfg.agents.agents {
if !self.runtimes.contains_key(&agent_cfg.id) {
let Some(spawner) = self.spawner() else {
rejected.push(ReloadRejection {
agent_id: Some(agent_cfg.id.clone()),
reason: "adding a new agent at runtime is not supported; \
set a spawner via ConfigReloadCoordinator::set_spawner"
.into(),
});
continue;
};
match spawner.call(agent_cfg.clone()).await {
Ok(spawned) => {
self.register(spawned.agent_id.clone(), spawned.reload_tx, spawned.known_tools);
applied.push(spawned.agent_id.clone());
if let Some(b) = self.broker.load_full() {
let evt = nexo_broker::Event::new(
"events.runtime.agent.spawned",
"config_reload",
serde_json::json!({
"agent_id": spawned.agent_id,
"version": new_version,
}),
);
let _ = b.publish("events.runtime.agent.spawned", evt).await;
}
tracing::info!(
agent = %agent_cfg.id,
"hot-spawned agent via ConfigReloadCoordinator",
);
}
Err(e) => {
rejected.push(ReloadRejection {
agent_id: Some(agent_cfg.id.clone()),
reason: format!("spawn: {e}"),
});
}
}
continue;
}
let Some(handle) = self.runtimes.get(&agent_cfg.id) else {
rejected.push(ReloadRejection {
agent_id: Some(agent_cfg.id.clone()),
reason: "agent vanished between hot-spawn check and snapshot build".into(),
});
continue;
};
let known_strs: Vec<&str> = handle.known_tools.iter().map(|s| s.as_str()).collect();
let catalog = crate::agent::KnownTools::new(known_strs);
if let Err(e) = crate::agent::validate_agent(agent_cfg, &cfg.plugins, &catalog)
{
rejected.push(ReloadRejection {
agent_id: Some(agent_cfg.id.clone()),
reason: format!("post-assembly validation: {e}"),
});
continue;
}
let snap = match RuntimeSnapshot::build(
Arc::new(agent_cfg.clone()),
&self.llm_registry,
&cfg.llm,
new_version,
) {
Ok(s) => Arc::new(s),
Err(e) => {
rejected.push(ReloadRejection {
agent_id: Some(agent_cfg.id.clone()),
reason: format!("snapshot build: {e}"),
});
continue;
}
};
match handle.reload_tx.send(ReloadCommand::Apply(snap)).await {
Ok(()) => applied.push(agent_cfg.id.clone()),
Err(e) => rejected.push(ReloadRejection {
agent_id: Some(agent_cfg.id.clone()),
reason: format!("dispatch: {e}"),
}),
}
}
let removed_ids: Vec<String> = self
.runtimes
.iter()
.filter_map(|entry| {
let id = entry.key();
if !cfg.agents.agents.iter().any(|a| &a.id == id) {
Some(id.clone())
} else {
None
}
})
.collect();
for id in removed_ids {
let Some(handle) = self.unregister(&id) else {
continue;
};
if let Err(e) = handle.reload_tx.send(ReloadCommand::Shutdown).await {
tracing::warn!(
agent = %id,
error = %e,
"hot-remove: runtime mailbox dispatch failed (task may have exited already)",
);
}
if let Some(b) = self.broker.load_full() {
let evt = nexo_broker::Event::new(
"events.runtime.agent.removed",
"config_reload",
serde_json::json!({
"agent_id": id,
"version": new_version,
}),
);
let _ = b.publish("events.runtime.agent.removed", evt).await;
}
applied.push(id.clone());
tracing::info!(agent = %id, "hot-removed agent via ConfigReloadCoordinator");
}
let elapsed_ms = started.elapsed().as_millis() as u64;
telemetry::observe_config_reload_latency_ms(elapsed_ms);
if !applied.is_empty() {
telemetry::inc_config_reload_applied();
tracing::info!(
version = new_version,
applied = ?applied,
rejected_count = rejected.len(),
elapsed_ms,
"config reload applied",
);
let hooks = self.post_hooks.lock().await;
for hook in hooks.iter() {
hook();
}
drop(hooks);
if let Some(broker) = self.broker.load_full() {
let payload = serde_json::json!({
"version": new_version,
"applied": &applied,
"rejected": &rejected,
"elapsed_ms": elapsed_ms,
});
let topic = "events.runtime.config.reloaded";
let evt = nexo_broker::Event::new(topic, "config-reload", payload);
if let Err(e) = broker.publish(topic, evt).await {
tracing::warn!(error = %e, "failed to publish events.runtime.config.reloaded");
}
}
}
if !rejected.is_empty() {
tracing::warn!(
version = new_version,
rejected = ?rejected,
"config reload: partial rejects",
);
}
ReloadOutcome {
version: new_version,
applied,
rejected,
elapsed_ms,
}
}
pub async fn start(
self: Arc<Self>,
broker: AnyBroker,
reload: nexo_config::RuntimeReloadConfig,
) -> anyhow::Result<()> {
if !reload.enabled {
tracing::info!("config hot-reload disabled via runtime.yaml");
return Ok(());
}
self.broker.store(Some(Arc::new(broker.clone())));
let watcher_rx = crate::config_watch::spawn_config_watcher(
self.config_dir.clone(),
reload.extra_watch_paths.clone(),
Duration::from_millis(reload.debounce_ms),
self.shutdown.clone(),
)?;
let coord_watcher = Arc::clone(&self);
tokio::spawn(async move {
let mut rx = watcher_rx;
while let Some(()) = rx.recv().await {
if coord_watcher.shutdown.is_cancelled() {
break;
}
let _ = coord_watcher.reload().await;
}
});
let mut sub = broker.subscribe("control.reload").await?;
let coord_broker = Arc::clone(&self);
let broker_clone = broker.clone();
tokio::spawn(async move {
loop {
if coord_broker.shutdown.is_cancelled() {
break;
}
let Some(_event) = sub.next().await else {
break;
};
let outcome = coord_broker.reload().await;
let ack_topic = "control.reload.ack";
let payload = serde_json::to_value(&outcome)
.unwrap_or_else(|e| serde_json::json!({ "error": e.to_string() }));
let evt = nexo_broker::Event::new(ack_topic, "config-reload", payload);
if let Err(e) = broker_clone.publish(ack_topic, evt).await {
tracing::warn!(error = %e, "failed to publish control.reload.ack");
}
}
});
Ok(())
}
pub async fn version(&self) -> u64 {
*self.version.lock().await
}
}
#[cfg(test)]
impl ConfigReloadCoordinator {
pub async fn post_hooks_len_for_test(&self) -> usize {
self.post_hooks.lock().await.len()
}
pub async fn fire_post_hooks_for_test(&self) {
let hooks = self.post_hooks.lock().await;
for hook in hooks.iter() {
hook();
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn reload_with_no_config_dir_falls_back_to_defaults() {
let coord = Arc::new(ConfigReloadCoordinator::new(
PathBuf::from("/nonexistent-config-dir-xyz"),
Arc::new(LlmRegistry::with_builtins()),
CancellationToken::new(),
));
let outcome = coord.reload().await;
assert!(
outcome.rejected.is_empty(),
"a missing config dir is tolerated, not rejected: {:?}",
outcome.rejected
);
assert!(
outcome.applied.is_empty(),
"default config has no agents to apply: {:?}",
outcome.applied
);
}
#[tokio::test]
async fn version_starts_at_zero() {
let coord = ConfigReloadCoordinator::new(
PathBuf::from("."),
Arc::new(LlmRegistry::with_builtins()),
CancellationToken::new(),
);
assert_eq!(coord.version().await, 0);
}
#[tokio::test]
async fn post_hooks_register_and_can_be_invoked_in_order() {
use std::sync::atomic::{AtomicUsize, Ordering};
let coord = ConfigReloadCoordinator::new(
PathBuf::from("."),
Arc::new(LlmRegistry::with_builtins()),
CancellationToken::new(),
);
let order = Arc::new(AtomicUsize::new(0));
let a_witness = Arc::new(AtomicUsize::new(0));
let b_witness = Arc::new(AtomicUsize::new(0));
{
let order = Arc::clone(&order);
let w = Arc::clone(&a_witness);
coord
.register_post_hook(Box::new(move || {
w.store(order.fetch_add(1, Ordering::SeqCst) + 1, Ordering::SeqCst);
}))
.await;
}
{
let order = Arc::clone(&order);
let w = Arc::clone(&b_witness);
coord
.register_post_hook(Box::new(move || {
w.store(order.fetch_add(1, Ordering::SeqCst) + 1, Ordering::SeqCst);
}))
.await;
}
let hooks = coord.post_hooks.lock().await;
assert_eq!(hooks.len(), 2);
for hook in hooks.iter() {
hook();
}
drop(hooks);
assert_eq!(a_witness.load(Ordering::SeqCst), 1);
assert_eq!(b_witness.load(Ordering::SeqCst), 2);
}
}