use dashmap::DashMap;
use parking_lot::RwLock;
use std::collections::HashSet;
use std::future::Future;
use std::path::{Path, PathBuf};
use std::pin::Pin;
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::{Arc, OnceLock};
use std::time::Duration;
use tokio::sync::mpsc;
use tokio::sync::Mutex as AsyncMutex;
use tracing::{error, info, warn};
use zerolaunch_plugin_api::config::{ComponentType, Configurable};
use zerolaunch_plugin_api::plugin::PluginMetadata;
use zerolaunch_plugin_api::PluginContext;
use zerolaunch_plugin_protocol::manifest::Manifest;
use zerolaunch_plugin_protocol::messages::{ComponentKind, KeywordOptimizerInfo};
use zerolaunch_plugin_protocol::ProtocolError;
use crate::adapter::remote_component::{RemoteComponent, RemoteComponentKind};
use crate::host_dispatch::HostCallHandler;
use crate::info::PluginInfo;
use crate::installed::{metadata_from_manifest, InstalledCache};
use crate::process::force_kill_process;
use crate::process::{PluginProcess, ProcessState};
const PROCESS_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5);
pub type RestartCallback =
Arc<dyn Fn(PluginRegistration) -> Pin<Box<dyn Future<Output = ()> + Send>> + Send + Sync>;
pub type CrashCallback = Arc<dyn Fn(PluginRegistration) + Send + Sync>;
#[derive(Clone, Debug)]
pub struct PluginRegistration {
pub plugin_id: String,
pub manifest: Manifest,
pub plugin_dir: PathBuf,
pub metadata: Arc<PluginMetadata>,
pub components: Vec<Arc<RemoteComponent>>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum PluginRuntimeState {
#[serde(rename = "starting")]
Starting,
#[serde(rename = "running")]
Running,
#[serde(rename = "stopped")]
Stopped,
#[serde(rename = "crashed")]
Crashed,
#[serde(rename = "error")]
Error,
#[serde(rename = "unknown")]
#[default]
Unknown,
}
impl From<&ProcessState> for PluginRuntimeState {
fn from(state: &ProcessState) -> Self {
match state {
ProcessState::Starting => Self::Starting,
ProcessState::Running => Self::Running,
ProcessState::Stopped => Self::Stopped,
ProcessState::Crashed { .. } => Self::Crashed,
ProcessState::Error(_) => Self::Error,
}
}
}
struct PluginRestartContext {
manifest: Manifest,
plugin_dir: PathBuf,
host_call_handler: Arc<dyn HostCallHandler>,
crash_tx: mpsc::Sender<String>,
on_restart: RestartCallback,
on_crash: CrashCallback,
restart_count: AtomicU32,
locale: RwLock<String>,
}
impl std::fmt::Debug for PluginRestartContext {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("PluginRestartContext")
.field("manifest", &self.manifest)
.field("plugin_dir", &self.plugin_dir)
.field("host_call_handler", &self.host_call_handler)
.field("crash_tx", &self.crash_tx)
.field("on_restart", &"<RestartCallback>")
.field("restart_count", &self.restart_count)
.finish()
}
}
pub struct PluginHostManager {
pub processes: Arc<DashMap<String, Arc<PluginProcess>>>,
pub plugins: Arc<DashMap<String, PluginRegistration>>,
pub data_dir_root: PathBuf,
pub log_dir_root: PathBuf,
pub(crate) plugins_dir: PathBuf,
pub(crate) installed_cache: Arc<InstalledCache>,
pub(crate) installed_watcher: parking_lot::Mutex<Option<notify::RecommendedWatcher>>,
restart_contexts: Arc<DashMap<String, Arc<PluginRestartContext>>>,
load_lock: Arc<AsyncMutex<()>>,
builtin_component_ids: RwLock<HashSet<String>>,
self_arc: OnceLock<Arc<Self>>,
}
#[derive(Debug, thiserror::Error)]
pub enum PluginLoadError {
#[error("manifest error: {0}")]
Manifest(String),
#[error("protocol error: {0}")]
Protocol(#[from] ProtocolError),
#[error("plugin already loaded: {0}")]
AlreadyLoaded(String),
#[error("plugin not found: {0}")]
NotFound(String),
#[error(
"component id collision for plugin {plugin_id}: '{component_id}' is already registered"
)]
ComponentIdCollision {
plugin_id: String,
component_id: String,
},
}
impl PluginHostManager {
pub fn new(plugins_dir: PathBuf, data_dir_root: PathBuf, log_dir_root: PathBuf) -> Arc<Self> {
let installed_cache = Arc::new(InstalledCache::load(plugins_dir.clone()));
let mgr = Arc::new(Self {
processes: Arc::new(DashMap::new()),
plugins: Arc::new(DashMap::new()),
data_dir_root,
log_dir_root,
plugins_dir,
installed_cache,
installed_watcher: parking_lot::Mutex::new(None),
restart_contexts: Arc::new(DashMap::new()),
load_lock: Arc::new(AsyncMutex::new(())),
builtin_component_ids: RwLock::new(HashSet::new()),
self_arc: OnceLock::new(),
});
if mgr.self_arc.set(Arc::clone(&mgr)).is_err() {
panic!("PluginHostManager self arc already set");
}
mgr
}
pub fn set_builtin_component_ids(&self, ids: HashSet<String>) {
*self.builtin_component_ids.write() = ids;
}
fn spawn_crash_loop(&self, crash_rx: mpsc::Receiver<String>) {
let mgr = Arc::clone(self.self_arc.get().expect("self arc set in load"));
tokio::spawn(async move {
crash_loop(mgr, crash_rx).await;
});
}
pub async fn load(
&self,
plugin_dir: &Path,
host_call_handler: Arc<dyn HostCallHandler>,
on_restart: RestartCallback,
on_crash: CrashCallback,
restart_count: u32,
locale: &str,
) -> Result<PluginRegistration, PluginLoadError> {
let manifest_path = plugin_dir.join("manifest.toml");
let manifest_bytes = std::fs::read_to_string(&manifest_path)
.map_err(|e| PluginLoadError::Manifest(format!("cannot read manifest.toml: {}", e)))?;
let manifest: Manifest = toml::from_str(&manifest_bytes)
.map_err(|e| PluginLoadError::Manifest(format!("invalid manifest: {}", e)))?;
validate_manifest(&manifest, plugin_dir)?;
let plugin_id = manifest.plugin.id.clone();
let _guard = self.load_lock.lock().await;
if self.processes.contains_key(&plugin_id) {
return Err(PluginLoadError::AlreadyLoaded(plugin_id));
}
let data_dir = self.data_dir_root.join(&plugin_id);
let log_dir = self.log_dir_root.clone();
if let Err(e) = std::fs::create_dir_all(&data_dir) {
warn!(
"Failed to create plugin data dir {}: {}",
data_dir.display(),
e
);
}
if let Err(e) = std::fs::create_dir_all(&log_dir) {
warn!(
"Failed to create plugin log dir {}: {}",
log_dir.display(),
e
);
}
info!("Loading plugin {} from {}", plugin_id, plugin_dir.display());
let (crash_tx, crash_rx) = mpsc::channel::<String>(4);
let process = match PluginProcess::spawn(
&manifest,
plugin_dir,
&data_dir,
&log_dir,
host_call_handler.clone(),
crash_tx.clone(),
restart_count,
locale,
)
.await
{
Ok(p) => p,
Err(e) => {
return Err(e.into());
}
};
let client = process.client.clone();
self.restart_contexts.insert(
plugin_id.clone(),
Arc::new(PluginRestartContext {
manifest: manifest.clone(),
plugin_dir: plugin_dir.to_path_buf(),
host_call_handler,
crash_tx,
on_restart,
on_crash,
restart_count: AtomicU32::new(restart_count),
locale: RwLock::new(locale.to_string()),
}),
);
self.spawn_crash_loop(crash_rx);
let process = Arc::new(process);
self.processes.insert(plugin_id.clone(), process.clone());
let init_result = match process.discover_components().await {
Ok(result) => result,
Err(e) => {
self.shutdown_loaded_process(&plugin_id).await;
return Err(PluginLoadError::Protocol(e));
}
};
let adapters = build_components(&plugin_id, plugin_dir, &manifest, client, &init_result);
if let Err(e) = validate_component_contract(&manifest, &adapters.components) {
error!("拒绝加载插件 {}:{}", plugin_id, e);
self.shutdown_loaded_process(&plugin_id).await;
return Err(e);
}
let component_ids = adapters
.components
.iter()
.map(|c| c.component_id().to_string())
.collect::<Vec<_>>();
if let Some(collision) =
find_component_id_collision(&component_ids, |id| self.component_id_is_taken(id))
{
error!(
"拒绝加载插件 {}:组件 id '{}' 已被其他已注册组件占用,清理进程并放弃加载",
plugin_id, collision
);
self.shutdown_loaded_process(&plugin_id).await;
return Err(PluginLoadError::ComponentIdCollision {
plugin_id,
component_id: collision,
});
}
let registered = adapters.clone();
self.plugins.insert(plugin_id.clone(), adapters);
Ok(registered)
}
pub fn plugins_dir(&self) -> &Path {
&self.plugins_dir
}
fn component_id_is_taken(&self, id: &str) -> bool {
let taken_by_plugin = self
.plugins
.iter()
.any(|e| e.value().components.iter().any(|c| c.component_id() == id));
taken_by_plugin || self.builtin_component_ids.read().contains(id)
}
pub async fn unload(&self, plugin_id: &str) -> Result<(), PluginLoadError> {
info!("Unloading plugin {}", plugin_id);
self.teardown(plugin_id).await;
for suffix in ["", ".1"] {
let log_file = self
.log_dir_root
.join(format!("{}.log{}", plugin_id, suffix));
let _ = std::fs::remove_file(&log_file);
}
Ok(())
}
async fn teardown(&self, plugin_id: &str) {
let _guard = self.load_lock.lock().await;
let locale = self
.restart_contexts
.get(plugin_id)
.map(|ctx| ctx.locale.read().clone())
.unwrap_or_default();
let plugin = self.plugins.get(plugin_id).and_then(|reg| {
reg.components
.iter()
.find_map(|comp| Arc::clone(comp).as_plugin())
});
if let Some((_, ctx)) = self.restart_contexts.remove(plugin_id) {
ctx.host_call_handler.detach_event_sink();
}
if let Some(plugin) = plugin {
let mut ctx = PluginContext::new("shutdown");
ctx.plugin_id = Some(plugin_id.to_string());
ctx.locale = locale;
if let Err(e) = plugin.on_shutdown(&ctx).await {
warn!("Plugin {} shutdown hook failed: {}", plugin_id, e);
}
}
self.shutdown_loaded_process(plugin_id).await;
self.plugins.remove(plugin_id);
}
async fn shutdown_loaded_process(&self, plugin_id: &str) {
if let Some((_, proc)) = self.processes.remove(plugin_id) {
match Arc::try_unwrap(proc) {
Ok(process) => process.shutdown(PROCESS_SHUTDOWN_TIMEOUT).await,
Err(arc) => {
warn!(
"Plugin {} process Arc has {} strong references; forcing kill. \n This may indicate a leaked clone of the process handle.",
plugin_id,
Arc::strong_count(&arc)
);
arc.state.write().clone_from(&ProcessState::Stopped);
arc.client.close_transport();
if let Some(pid) = arc.pid {
force_kill_process(pid);
}
}
}
}
}
pub async fn shutdown_all(&self) {
let _guard = self.load_lock.lock().await;
for plugin_id in self
.restart_contexts
.iter()
.map(|e| e.key().clone())
.collect::<Vec<_>>()
{
if let Some((_, ctx)) = self.restart_contexts.remove(&plugin_id) {
ctx.host_call_handler.detach_event_sink();
}
}
self.plugins.clear();
let processes: Vec<Arc<PluginProcess>> = self
.processes
.iter()
.map(|e| (e.key().clone(), Arc::clone(e.value())))
.collect::<Vec<_>>()
.into_iter()
.filter_map(|(plugin_id, proc)| self.processes.remove(&plugin_id).map(|_| proc))
.collect();
futures::future::join_all(processes.into_iter().map(|proc| async move {
match Arc::try_unwrap(proc) {
Ok(process) => process.shutdown(PROCESS_SHUTDOWN_TIMEOUT).await,
Err(arc) => {
warn!(
"Plugin {} process Arc has {} strong references during exit; forcing kill.",
arc.plugin_id,
Arc::strong_count(&arc)
);
arc.state.write().clone_from(&ProcessState::Stopped);
arc.client.close_transport();
if let Some(pid) = arc.pid {
force_kill_process(pid);
}
}
}
}))
.await;
}
pub fn update_locale(&self, plugin_id: &str, locale: &str) {
if let Some(ctx) = self.restart_contexts.get_mut(plugin_id) {
*ctx.locale.write() = locale.to_string();
}
}
pub fn get_plugin_info(
&self,
plugin_id: &str,
is_component_enabled: impl Fn(&str) -> bool,
) -> Option<PluginInfo> {
let entry = self.plugins.get(plugin_id)?;
let registration = entry.value();
let process_state = self
.processes
.get(plugin_id)
.map(|p| PluginRuntimeState::from(&*p.state.read()))
.unwrap_or_default();
let enabled = is_component_enabled(®istration.metadata.id);
Some(PluginInfo::from_metadata(
®istration.metadata,
process_state,
enabled,
))
}
}
fn validate_manifest(manifest: &Manifest, plugin_dir: &Path) -> Result<(), PluginLoadError> {
let id = &manifest.plugin.id;
static PLUGIN_ID_RE: OnceLock<regex::Regex> = OnceLock::new();
let re = PLUGIN_ID_RE.get_or_init(|| {
regex::Regex::new(zerolaunch_plugin_protocol::manifest::PLUGIN_ID_RE).unwrap()
});
if !re.is_match(id) {
return Err(PluginLoadError::Manifest(format!(
"invalid plugin id '{}': must match reverse domain",
id
)));
}
if semver::Version::parse(&manifest.plugin.version).is_err() {
return Err(PluginLoadError::Manifest(format!(
"invalid plugin version '{}'",
manifest.plugin.version
)));
}
if manifest.components.provides.is_empty() {
return Err(PluginLoadError::Manifest(
"components.provides must have at least one entry".into(),
));
}
for p in &manifest.components.provides {
if !zerolaunch_plugin_protocol::manifest::REQUIRED_PROVIDES_VALUES.contains(&p.as_str()) {
return Err(PluginLoadError::Manifest(format!(
"unknown component type '{}'",
p
)));
}
}
let cmd_path = plugin_dir.join(&manifest.runtime.command);
let canonical_cmd = cmd_path
.canonicalize()
.map_err(|e| PluginLoadError::Manifest(format!("command not found: {}", e)))?;
let canonical_plugin_dir = plugin_dir
.canonicalize()
.map_err(|e| PluginLoadError::Manifest(format!("plugin dir canonicalize: {}", e)))?;
if !canonical_cmd.starts_with(&canonical_plugin_dir) {
return Err(PluginLoadError::Manifest(
"command path escapes plugin directory".into(),
));
}
if let Some(icon) = &manifest.icon {
if let Ok(canonical_icon) = plugin_dir.join(&icon.path).canonicalize() {
if !canonical_icon.starts_with(&canonical_plugin_dir) {
return Err(PluginLoadError::Manifest(
"icon path escapes plugin directory".into(),
));
}
}
}
Ok(())
}
fn find_by_id<T: Clone + Default>(map: &[(String, T)], component_id: &str) -> T {
map.iter()
.find(|(id, _)| id == component_id)
.map(|(_, v)| v.clone())
.unwrap_or_default()
}
fn find_settings_value(
values: &[(String, serde_json::Value)],
component_id: &str,
) -> serde_json::Value {
values
.iter()
.find(|(id, _)| id == component_id)
.map(|(_, v)| v.clone())
.unwrap_or(serde_json::Value::Null)
}
fn component_type_of(kind: &ComponentKind) -> ComponentType {
match kind {
ComponentKind::Plugin => ComponentType::Plugin,
ComponentKind::DataSource => ComponentType::DataSource,
ComponentKind::ActionExecutor { .. } => ComponentType::ActionExecutor,
ComponentKind::SearchEngine => ComponentType::SearchEngine,
ComponentKind::ScoreBooster => ComponentType::ScoreBooster,
ComponentKind::KeywordOptimizer => ComponentType::KeywordOptimizer,
ComponentKind::KeywordInjector => ComponentType::KeywordInjector,
}
}
fn validate_component_contract(
manifest: &Manifest,
components: &[Arc<RemoteComponent>],
) -> Result<(), PluginLoadError> {
let plugin_components: Vec<&Arc<RemoteComponent>> =
components.iter().filter(|c| c.is_plugin()).collect();
if plugin_components.len() != 1 {
return Err(PluginLoadError::Manifest(format!(
"components must declare exactly one plugin component, found {}",
plugin_components.len()
)));
}
let plugin_component_id = plugin_components[0].component_id();
if plugin_component_id != manifest.plugin.id {
return Err(PluginLoadError::Manifest(format!(
"plugin component id '{plugin_component_id}' must equal plugin id '{}'",
manifest.plugin.id
)));
}
let mut actual: HashSet<&str> = HashSet::new();
for c in components {
if c.is_plugin() {
actual.insert("plugin");
}
if c.is_data_source() {
actual.insert("data_source");
}
if c.is_action_executor() {
actual.insert("action_executor");
}
if c.is_search_engine() {
actual.insert("search_engine");
}
if c.is_score_booster() {
actual.insert("score_booster");
}
if c.is_keyword_optimizer() {
actual.insert("keyword_optimizer");
}
if c.is_keyword_injector() {
actual.insert("keyword_injector");
}
}
let declared: HashSet<&str> = manifest
.components
.provides
.iter()
.map(String::as_str)
.collect();
if declared != actual {
let mut declared: Vec<&str> = declared.into_iter().collect();
let mut actual: Vec<&str> = actual.into_iter().collect();
declared.sort_unstable();
actual.sort_unstable();
return Err(PluginLoadError::Manifest(format!(
"components.provides {:?} does not match actual components {:?}",
declared, actual
)));
}
Ok(())
}
fn find_component_id_collision(
component_ids: &[String],
component_id_taken: impl Fn(&str) -> bool,
) -> Option<String> {
component_ids
.iter()
.find(|id| component_id_taken(id))
.cloned()
}
fn build_components(
plugin_id: &str,
plugin_dir: &Path,
manifest: &Manifest,
client: Arc<crate::client::JsonRpcClient>,
init_result: &crate::process::InitResult,
) -> PluginRegistration {
let plugin_metadata = Arc::new(metadata_from_manifest(plugin_dir, manifest));
let components: Vec<Arc<RemoteComponent>> = init_result
.components
.iter()
.map(|comp| {
let schema = find_by_id(&init_result.settings_schemas, &comp.component_id);
let settings = find_settings_value(&init_result.settings_values, &comp.component_id);
let config_actions = find_by_id(&init_result.config_actions_map, &comp.component_id);
let default_enabled = find_by_id(&init_result.default_enabled_map, &comp.component_id);
let priority = comp.priority;
let kind = match &comp.kind {
ComponentKind::Plugin => {
let interaction_policy =
find_by_id(&init_result.interaction_policy_map, &comp.component_id);
RemoteComponentKind::Plugin {
metadata: Arc::clone(&plugin_metadata),
interaction_policy: parking_lot::RwLock::new(interaction_policy),
}
}
ComponentKind::DataSource => RemoteComponentKind::DataSource,
ComponentKind::ActionExecutor { target_types } => {
let result_actions =
find_by_id(&init_result.executor_actions_map, &comp.component_id);
RemoteComponentKind::ActionExecutor {
target_types: target_types.clone(),
result_actions,
}
}
ComponentKind::SearchEngine => RemoteComponentKind::SearchEngine,
ComponentKind::ScoreBooster => RemoteComponentKind::ScoreBooster,
ComponentKind::KeywordOptimizer => {
let info = init_result
.keyword_optimizer_info_map
.iter()
.find(|(id, _)| id == &comp.component_id)
.map(|(_, v)| v.clone())
.unwrap_or_else(|| KeywordOptimizerInfo::fallback(priority));
RemoteComponentKind::KeywordOptimizer {
info: parking_lot::RwLock::new(info),
}
}
ComponentKind::KeywordInjector => RemoteComponentKind::KeywordInjector,
};
Arc::new(RemoteComponent::new(
comp.component_id.clone(),
comp.component_name.clone(),
comp.component_description.clone(),
component_type_of(&comp.kind),
priority,
client.clone(),
schema,
settings,
config_actions,
default_enabled,
kind,
))
})
.collect();
let registration = PluginRegistration {
plugin_id: plugin_id.to_string(),
manifest: manifest.clone(),
plugin_dir: plugin_dir.to_path_buf(),
metadata: plugin_metadata,
components,
};
#[cfg(debug_assertions)]
{
for comp in ®istration.components {
if let RemoteComponentKind::Plugin { metadata, .. } = &comp.kind {
debug_assert!(
Arc::ptr_eq(metadata, ®istration.metadata),
"组件元数据与 PluginRegistration.metadata 必须同源共享(Arc::clone)"
);
}
}
}
registration
}
async fn crash_loop(mgr: Arc<PluginHostManager>, mut crash_rx: mpsc::Receiver<String>) {
while let Some(plugin_id) = crash_rx.recv().await {
handle_crash(&mgr, &plugin_id).await;
}
}
async fn handle_crash(mgr: &PluginHostManager, plugin_id: &str) {
warn!("Watchdog triggered restart for plugin: {}", plugin_id);
let prev = mgr.plugins.remove(plugin_id);
mgr.processes.remove(plugin_id);
let Some(ctx) = mgr
.restart_contexts
.get(plugin_id)
.map(|r| Arc::clone(r.value()))
else {
warn!(
"Crash notification for plugin '{}' but no restart context found — \
plugin may have been unloaded concurrently",
plugin_id
);
return;
};
ctx.host_call_handler.detach_event_sink();
if let Some((_, prev)) = prev {
(ctx.on_crash)(prev);
}
let new_count = ctx.restart_count.fetch_add(1, Ordering::SeqCst) + 1;
let max_restart = ctx.manifest.runtime.max_restart;
if new_count > max_restart {
error!(
"Plugin {} exceeded max restarts ({}/{}) — not restarting",
plugin_id, new_count, max_restart
);
mgr.restart_contexts.remove(plugin_id);
return;
}
let locale = ctx.locale.read().clone();
match mgr
.load(
&ctx.plugin_dir,
ctx.host_call_handler.clone(),
ctx.on_restart.clone(),
ctx.on_crash.clone(),
new_count,
&locale,
)
.await
{
Ok(registered) => {
(ctx.on_restart)(registered).await;
info!("Plugin {} successfully restarted", plugin_id);
}
Err(e) => {
error!("Failed to restart plugin {}: {}", plugin_id, e);
mgr.restart_contexts.remove(plugin_id);
}
}
}
#[cfg(test)]
mod tests {
use super::find_component_id_collision;
#[test]
fn collision_check_finds_first_taken_id() {
let ids = vec!["free".to_string(), "taken".to_string(), "other".to_string()];
let taken = |id: &str| id == "taken";
assert_eq!(
find_component_id_collision(&ids, taken),
Some("taken".to_string())
);
}
#[test]
fn collision_check_returns_none_when_all_free() {
let ids = vec!["a".to_string(), "b".to_string()];
assert_eq!(find_component_id_collision(&ids, |_| false), None);
}
}