af_agent/

plugin.rs

use af_agent_session::ContentBlock;
use af_context::{ProfileRevisionId, RunId, SessionId, SubjectId, TenantId};
use std::collections::{BTreeMap, BTreeSet};
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;

use async_trait::async_trait;
use semver::{Version, VersionReq};
use serde::{Deserialize, Serialize};
use serde_json::Value;

use crate::{
    validate_json_schema, ContextContributor, PromptAuthority, PromptSection, Tool, ToolRegistry,
};

/// Deployment-time declaration of a trusted plugin.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct PluginManifest {
    /// Stable identifier of this record.
    pub id: String,
    /// Semantic version string.
    pub version: String,
    /// Required plugins and the semver ranges they must satisfy.
    #[serde(default)]
    pub dependencies: BTreeMap<String, String>,
    /// JSON Schema for the plugin's Profile configuration.
    #[serde(default)]
    pub config_schema: Value,
    /// Permissions granted or required.
    #[serde(default)]
    pub permissions: BTreeSet<PluginPermission>,
}

impl PluginManifest {
    /// Reject invalid ids and versions.
    pub fn validate(&self) -> Result<(), PluginError> {
        if self.id.trim().is_empty()
            || !self.id.chars().all(|ch| {
                ch.is_ascii_lowercase() || ch.is_ascii_digit() || matches!(ch, '.' | '-' | '_')
            })
        {
            return Err(PluginError::InvalidManifest(format!(
                "invalid plugin id {}",
                self.id
            )));
        }
        Version::parse(&self.version)
            .map_err(|error| PluginError::InvalidManifest(format!("{}: {error}", self.id)))?;
        Ok(())
    }
}

/// Contribution kinds a plugin may register.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum PluginPermission {
    /// Prompt sections.
    Prompt,
    /// Context contributors.
    Context,
    /// Tools.
    Tool,
    /// Hooks.
    Hook,
}

/// Redacted description of a registered plugin.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct PluginDescriptor {
    /// Stable identifier of this record.
    pub id: String,
    /// Semantic version string.
    pub version: String,
    /// Declared dependency ranges.
    pub dependencies: BTreeMap<String, String>,
    /// Permissions granted or required.
    pub permissions: BTreeSet<PluginPermission>,
}

/// Identity a plugin is mounted under for one Session.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PluginMountContext {
    /// Tenant that owns this record.
    pub tenant_id: TenantId,
    /// Subject (user or service principal) acting on or owning this record.
    pub subject_id: SubjectId,
    /// Session this record belongs to.
    pub session_id: SessionId,
    /// Immutable Agent Profile revision pinned by the Session.
    pub profile_revision_id: ProfileRevisionId,
}

/// Result of a pre-tool hook.
#[derive(Debug, Clone, PartialEq)]
pub enum HookDecision {
    /// Let the tool run.
    Continue,
    /// Refuse the call; the model sees `reason`.
    Deny {
        /// Why the call was refused.
        reason: String,
    },
    /// Park the Run until the user answers.
    WaitForInput {
        /// `action` (approval) or `question`.
        kind: String,
        /// Interaction payload shown to the user.
        payload: Value,
    },
}

/// Boxed tool execution passed through `around_tool`.
pub type ToolExecutionFuture<'a> = Pin<Box<dyn Future<Output = Result<Value, String>> + Send + 'a>>;

/// Position of the loop when a step-level hook runs.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct StepContext {
    /// The originating request.
    pub request: af_context::RequestContext,
    /// Session this record belongs to.
    pub session_id: SessionId,
    /// Run this record belongs to.
    pub run_id: RunId,
    /// 1-based Turn number inside the Run.
    pub turn: u32,
    /// 1-based step number inside the Turn.
    pub step: u32,
}

/// Whether the loop may open the proposed step. A rejection closes the Run
/// without a model request and is recorded as a `step_rejected` extension event.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PreStepDecision {
    /// Open the step.
    Enter,
    /// Do not open the step; the Run is cancelled with `reason`.
    Reject {
        /// Why the step was refused.
        reason: String,
    },
}

/// Mutable per-request model policy. Hooks may narrow or redirect it before the
/// request is frozen into `ModelRequestPrepared`.
#[derive(Debug, Clone, PartialEq)]
pub struct ModelRequestPolicy {
    /// Model identifier as registered in the model registry.
    pub model: String,
    /// Reasoning budget for this request.
    pub reasoning_effort: Option<af_llm::ReasoningEffort>,
    /// Sampling temperature.
    pub temperature: f32,
    /// Upper bound on prompt plus completion tokens.
    pub max_tokens: Option<u32>,
}

/// One failed model attempt as seen by [`Hook::request_error`].
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ModelFailure {
    /// 1-based attempt counter.
    pub attempt: u32,
    /// Error message when the operation failed.
    pub error: String,
    /// Whether the failure is safe to retry.
    pub retryable: bool,
    /// Whether the runtime still has attempts left.
    pub attempts_remaining: bool,
}

/// Recovery ownership for a failed model attempt. `Default` keeps the
/// runtime's retry policy; `Retry` forces one more attempt while attempts
/// remain; `Terminal` stops retrying even for a retryable failure.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RequestErrorAction {
    /// Keep the runtime's retry policy.
    Default,
    /// Force one more attempt while attempts remain.
    Retry,
    /// Stop retrying now.
    Terminal,
}

/// Whether a turn that owes no tool result may close. `Steer` appends the
/// content as a logged user message and runs another step instead.
#[derive(Debug, Clone, PartialEq)]
pub enum TurnStopDecision {
    /// Close the turn.
    Stop,
    /// Append `content` as a logged user message and run one more step.
    Steer {
        /// Steering input.
        content: Vec<ContentBlock>,
    },
}

/// Interception points around the loop; every method has a no-op default.
#[async_trait]
pub trait Hook: Send + Sync {
    /// Runs before a step opens. The first `Reject` across hooks wins.
    async fn pre_step(&self, _context: &StepContext) -> PreStepDecision {
        PreStepDecision::Enter
    }
    /// Runs before each model request is frozen; hooks apply in mount order.
    async fn request(&self, _context: &StepContext, _policy: &mut ModelRequestPolicy) {}
    /// Runs after each failed model attempt. The first non-`Default` action wins.
    async fn request_error(
        &self,
        _context: &StepContext,
        _failure: &ModelFailure,
    ) -> RequestErrorAction {
        RequestErrorAction::Default
    }
    /// Runs when the model owes nothing more. The first `Steer` across hooks wins.
    async fn turn_stopping(&self, _context: &StepContext) -> TurnStopDecision {
        TurnStopDecision::Stop
    }
    /// Runs before a tool executes; may deny it or park the Run on an interaction.
    async fn before_tool(
        &self,
        _context: &crate::ToolExecutionContext,
        _tool: &str,
        _arguments: &Value,
    ) -> HookDecision {
        HookDecision::Continue
    }
    /// Wrap the actual tool execution; call `next` to run it.
    fn around_tool<'a>(
        &'a self,
        _context: &'a crate::ToolExecutionContext,
        _tool: &'a str,
        _arguments: &'a Value,
        next: ToolExecutionFuture<'a>,
    ) -> ToolExecutionFuture<'a> {
        next
    }
    /// Runs after a tool result is committed; an error is recorded but cannot replace the result.
    async fn after_tool(
        &self,
        _context: &crate::ToolExecutionContext,
        _tool: &str,
        _arguments: &Value,
        _result: &Value,
    ) -> Result<(), PluginError> {
        Ok(())
    }
}

#[derive(Default, Clone)]
struct PluginContributions {
    prompts: BTreeMap<String, PromptSection>,
    contexts: Vec<Arc<dyn ContextContributor>>,
    tools: ToolRegistry,
    hooks: Vec<Arc<dyn Hook>>,
}

impl PluginContributions {
    fn merge(&mut self, additions: Self) -> Result<(), PluginError> {
        let mut merged = self.clone();
        for (id, prompt) in additions.prompts {
            insert_once(&mut merged.prompts, id, prompt)?;
        }
        merged
            .tools
            .extend(&additions.tools)
            .map_err(PluginError::Conflict)?;
        merged.contexts.extend(additions.contexts);
        merged.hooks.extend(additions.hooks);
        *self = merged;
        Ok(())
    }
}

/// Scoped registrar a plugin contributes through during activation; permissions are checked per call.
pub struct AgentRegistrar<'a> {
    owner: &'a PluginManifest,
    contributions: &'a mut PluginContributions,
}

impl AgentRegistrar<'_> {
    /// Register a prompt section under the plugin's authority.
    pub fn prompt(
        &mut self,
        id: impl Into<String>,
        prompt: impl Into<String>,
    ) -> Result<(), PluginError> {
        require(self.owner, PluginPermission::Prompt)?;
        let id = format!("plugin.{}.{}", self.owner.id, id.into());
        insert_once(
            &mut self.contributions.prompts,
            id.clone(),
            PromptSection {
                id,
                order: 400,
                authority: PromptAuthority::Plugin,
                source: self.owner.id.clone(),
                version: self.owner.version.clone(),
                content: prompt.into(),
            },
        )
    }
    /// Register a tool.
    pub fn tool(&mut self, tool: Arc<dyn Tool>) -> Result<(), PluginError> {
        require(self.owner, PluginPermission::Tool)?;
        if self.contributions.tools.contains(tool.name()) {
            return Err(PluginError::Conflict(format!("tool {}", tool.name())));
        }
        self.contributions
            .tools
            .register(tool)
            .map(|_| ())
            .map_err(PluginError::Conflict)
    }
    /// Register a context contributor.
    pub fn context(&mut self, contributor: Arc<dyn ContextContributor>) -> Result<(), PluginError> {
        require(self.owner, PluginPermission::Context)?;
        self.contributions.contexts.push(contributor);
        Ok(())
    }
    /// Register a hook.
    pub fn hook(&mut self, hook: Arc<dyn Hook>) -> Result<(), PluginError> {
        require(self.owner, PluginPermission::Hook)?;
        self.contributions.hooks.push(hook);
        Ok(())
    }
}

/// Per-Session resource a plugin holds; released on unmount.
#[async_trait]
pub trait PluginLease: Send + Sync {
    /// Cleanup must be idempotent because a failed release is retried.
    async fn unmount(&mut self) -> Result<(), PluginError>;
}

/// Trusted deployment-time extension compiled into the host.
#[async_trait]
pub trait AgentPlugin: Send + Sync {
    /// Static manifest.
    fn manifest(&self) -> &PluginManifest;

    /// Creates the cleanup lease before activation performs any side effects.
    /// The lease must be safe to unmount before, during, or after activation.
    fn lease(&self, context: &PluginMountContext, config: &Value) -> Box<dyn PluginLease>;

    /// Contribute prompts, tools, contexts and hooks for one mount.
    async fn activate(
        &self,
        context: &PluginMountContext,
        config: &Value,
        registrar: &mut AgentRegistrar<'_>,
    ) -> Result<(), PluginError>;
}

/// Result of an atomic mount: merged contributions plus the leases to release.
pub struct MountedPlugins {
    contributions: PluginContributions,
    leases: Arc<LeaseSet>,
}

impl std::fmt::Debug for MountedPlugins {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter
            .debug_struct("MountedPlugins")
            .field(
                "pending_leases",
                &self
                    .leases
                    .state
                    .lock()
                    .unwrap_or_else(std::sync::PoisonError::into_inner)
                    .leases
                    .len(),
            )
            .finish_non_exhaustive()
    }
}

impl MountedPlugins {
    fn from_parts(contributions: PluginContributions, leases: Vec<Box<dyn PluginLease>>) -> Self {
        Self {
            contributions,
            leases: Arc::new(LeaseSet {
                state: std::sync::Mutex::new(LeaseState {
                    leases,
                    holders: 1,
                    closing: false,
                }),
                notify: tokio::sync::Notify::new(),
            }),
        }
    }

    /// Prompt sections in order.
    pub fn prompts(&self) -> impl Iterator<Item = &PromptSection> {
        self.contributions.prompts.values()
    }
    /// Context contributors in mount order.
    pub fn contexts(&self) -> &[Arc<dyn ContextContributor>] {
        &self.contributions.contexts
    }
    /// Merged tool registry.
    pub fn tools(&self) -> &ToolRegistry {
        &self.contributions.tools
    }
    /// Hooks in mount order.
    pub fn hooks(&self) -> &[Arc<dyn Hook>] {
        &self.contributions.hooks
    }

    /// Release every lease in reverse order; failures are collected and retryable.
    pub async fn unmount(&mut self) -> Result<(), PluginError> {
        let leases = Arc::clone(&self.leases);
        {
            leases
                .state
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner)
                .closing = true;
        }
        loop {
            let notified = leases.notify.notified();
            let pending = {
                let mut state = leases
                    .state
                    .lock()
                    .unwrap_or_else(std::sync::PoisonError::into_inner);
                (state.holders == 1).then(|| std::mem::take(&mut state.leases))
            };
            if let Some(mut pending) = pending {
                let mut failed = Vec::new();
                let mut errors = Vec::new();
                while let Some(mut lease) = pending.pop() {
                    if let Err(error) = lease.unmount().await {
                        errors.push(error.to_string());
                        failed.push(lease);
                    }
                }
                leases
                    .state
                    .lock()
                    .unwrap_or_else(std::sync::PoisonError::into_inner)
                    .leases = failed;
                return if errors.is_empty() {
                    Ok(())
                } else {
                    Err(PluginError::Mount(format!(
                        "plugin cleanup failed: {}",
                        errors.join("; ")
                    )))
                };
            }
            notified.await;
        }
    }
}

impl Clone for MountedPlugins {
    fn clone(&self) -> Self {
        self.leases
            .state
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .holders += 1;
        Self {
            contributions: self.contributions.clone(),
            leases: Arc::clone(&self.leases),
        }
    }
}

impl Drop for MountedPlugins {
    fn drop(&mut self) {
        let mut state = self
            .leases
            .state
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        state.holders -= 1;
        if state.closing && state.holders <= 1 {
            self.leases.notify.notify_waiters();
        }
    }
}

struct LeaseSet {
    state: std::sync::Mutex<LeaseState>,
    notify: tokio::sync::Notify,
}

struct LeaseState {
    leases: Vec<Box<dyn PluginLease>>,
    holders: usize,
    closing: bool,
}

/// Registered plugins keyed by id.
#[derive(Default)]
pub struct PluginCatalog {
    plugins: BTreeMap<(String, String), Arc<dyn AgentPlugin>>,
}

impl std::fmt::Debug for PluginCatalog {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter
            .debug_struct("PluginCatalog")
            .field("count", &self.plugins.len())
            .finish()
    }
}

impl PluginCatalog {
    /// Register a plugin; duplicate ids never replace the first registration.
    pub fn register(&mut self, plugin: Arc<dyn AgentPlugin>) -> Result<(), PluginError> {
        plugin.manifest().validate()?;
        let id = plugin.manifest().id.clone();
        let version = plugin.manifest().version.clone();
        if self.plugins.contains_key(&(id.clone(), version.clone())) {
            return Err(PluginError::Conflict(format!("{id}@{version}")));
        }
        self.plugins.insert((id, version), plugin);
        Ok(())
    }

    /// Redacted descriptors of every plugin.
    pub fn descriptors(&self) -> impl Iterator<Item = PluginDescriptor> + '_ {
        self.plugins.values().map(|plugin| {
            let manifest = plugin.manifest();
            PluginDescriptor {
                id: manifest.id.clone(),
                version: manifest.version.clone(),
                dependencies: manifest.dependencies.clone(),
                permissions: manifest.permissions.clone(),
            }
        })
    }

    /// Mount the referenced plugins in dependency order, rolling back on failure.
    pub async fn mount(
        &self,
        context: &PluginMountContext,
        references: &BTreeMap<String, (String, Value)>,
    ) -> Result<MountedPlugins, PluginError> {
        let selected = self.select(references)?;
        let configs = references
            .iter()
            .map(|(id, (_, config))| (id.clone(), config.clone()))
            .collect();
        let borrowed = selected
            .iter()
            .map(|plugin| plugin.as_ref())
            .collect::<Vec<_>>();
        mount_plugins(context, &borrowed, &configs).await
    }

    /// Report references to unknown plugins, unsatisfied versions and invalid configs.
    pub fn validate_references(
        &self,
        references: &BTreeMap<String, (String, Value)>,
    ) -> Result<(), PluginError> {
        self.select(references).map(|_| ())
    }

    fn select(
        &self,
        references: &BTreeMap<String, (String, Value)>,
    ) -> Result<Vec<Arc<dyn AgentPlugin>>, PluginError> {
        if let Some(issue) =
            validate_plugin_references(&self.descriptors().collect::<Vec<_>>(), references)
                .into_iter()
                .next()
        {
            return Err(issue.error);
        }
        let mut selected = Vec::with_capacity(references.len());
        let mut configs = BTreeMap::new();
        for (id, (version, config)) in references {
            let plugin = if let Some(plugin) = self.plugins.get(&(id.clone(), version.clone())) {
                plugin
            } else if let Some(((_, actual), _)) = self
                .plugins
                .iter()
                .find(|((plugin_id, _), _)| plugin_id == id)
            {
                return Err(PluginError::DependencyVersion {
                    plugin: id.clone(),
                    dependency: id.clone(),
                    required: version.clone(),
                    actual: actual.clone(),
                });
            } else {
                return Err(PluginError::MissingDependency(format!("{id}@{version}")));
            };
            selected.push(Arc::clone(plugin));
            configs.insert(id.clone(), config.clone());
        }
        let borrowed = selected
            .iter()
            .map(|plugin| plugin.as_ref())
            .collect::<Vec<_>>();
        resolve_plugin_order(&borrowed)?;
        for plugin in &selected {
            let manifest = plugin.manifest();
            validate_json_schema(&manifest.config_schema, &configs[&manifest.id])
                .map_err(|error| PluginError::InvalidConfig(manifest.id.clone(), error))?;
        }
        Ok(selected)
    }
}

/// One problem found by [`PluginCatalog::validate_references`].
#[derive(Debug)]
pub struct PluginReferenceIssue {
    /// Plugin id the issue is about.
    pub plugin: String,
    /// Error message when the operation failed.
    pub error: PluginError,
}

/// Side-effect-free validation shared by profile inspection and runtime mount.
pub fn validate_plugin_references(
    descriptors: &[PluginDescriptor],
    references: &BTreeMap<String, (String, Value)>,
) -> Vec<PluginReferenceIssue> {
    let exact = descriptors
        .iter()
        .map(|descriptor| {
            (
                (descriptor.id.as_str(), descriptor.version.as_str()),
                descriptor,
            )
        })
        .collect::<BTreeMap<_, _>>();
    let mut issues = Vec::new();
    for (id, (version, _)) in references {
        if exact.contains_key(&(id.as_str(), version.as_str())) {
            continue;
        }
        let error = descriptors
            .iter()
            .find(|descriptor| descriptor.id == *id)
            .map_or_else(
                || PluginError::MissingDependency(format!("{id}@{version}")),
                |actual| PluginError::DependencyVersion {
                    plugin: id.clone(),
                    dependency: id.clone(),
                    required: version.clone(),
                    actual: actual.version.clone(),
                },
            );
        issues.push(PluginReferenceIssue {
            plugin: id.clone(),
            error,
        });
    }
    for (id, (version, _)) in references {
        let Some(descriptor) = exact.get(&(id.as_str(), version.as_str())) else {
            continue;
        };
        for (dependency, requirement) in &descriptor.dependencies {
            let error = match references.get(dependency).map(|(version, _)| version) {
                None => Some(PluginError::MissingDependency(format!(
                    "{id} -> {dependency}"
                ))),
                Some(actual) => match (VersionReq::parse(requirement), Version::parse(actual)) {
                    (Err(error), _) => Some(PluginError::InvalidManifest(format!(
                        "{id} dependency {dependency}: {error}"
                    ))),
                    (_, Err(error)) => Some(PluginError::InvalidManifest(error.to_string())),
                    (Ok(required), Ok(actual_version)) if !required.matches(&actual_version) => {
                        Some(PluginError::DependencyVersion {
                            plugin: id.clone(),
                            dependency: dependency.clone(),
                            required: requirement.clone(),
                            actual: actual.clone(),
                        })
                    }
                    _ => None,
                },
            };
            if let Some(error) = error {
                issues.push(PluginReferenceIssue {
                    plugin: id.clone(),
                    error,
                });
            }
        }
    }
    issues
}

/// Mount a complete Session plugin set atomically. Any validation or activation
/// failure tears down current and already-mounted plugins in reverse order.
pub async fn mount_plugins(
    context: &PluginMountContext,
    plugins: &[&dyn AgentPlugin],
    configs: &BTreeMap<String, Value>,
) -> Result<MountedPlugins, PluginError> {
    let ordered = resolve_plugin_order(plugins)?;
    let mut contributions = PluginContributions::default();
    let mut leases = Vec::with_capacity(ordered.len());
    for plugin in ordered {
        let manifest = plugin.manifest();
        let config = configs
            .get(&manifest.id)
            .cloned()
            .unwrap_or_else(|| Value::Object(Default::default()));
        if let Err(error) = validate_json_schema(&manifest.config_schema, &config) {
            return Err(rollback_error(
                PluginError::InvalidConfig(manifest.id.clone(), error),
                leases,
            )
            .await);
        }
        let mut additions = PluginContributions::default();
        let mut registrar = AgentRegistrar {
            owner: manifest,
            contributions: &mut additions,
        };
        leases.push(plugin.lease(context, &config));
        if let Err(error) = plugin.activate(context, &config, &mut registrar).await {
            return Err(rollback_error(error, leases).await);
        }
        if let Err(error) = contributions.merge(additions) {
            return Err(rollback_error(error, leases).await);
        }
    }
    Ok(MountedPlugins::from_parts(contributions, leases))
}

async fn rollback_error(cause: PluginError, leases: Vec<Box<dyn PluginLease>>) -> PluginError {
    let mut cleanup = MountedPlugins::from_parts(PluginContributions::default(), leases);
    match cleanup.unmount().await {
        Ok(()) => cause,
        Err(error) => PluginError::Rollback {
            cause: Box::new(cause),
            cleanup_error: error.to_string(),
            cleanup: Box::new(cleanup),
        },
    }
}

/// Topologically order the referenced plugins by their dependencies.
pub fn resolve_plugin_order<'a>(
    plugins: &'a [&'a dyn AgentPlugin],
) -> Result<Vec<&'a dyn AgentPlugin>, PluginError> {
    let by_id = plugins
        .iter()
        .map(|plugin| (plugin.manifest().id.as_str(), *plugin))
        .collect::<BTreeMap<_, _>>();
    if by_id.len() != plugins.len() {
        return Err(PluginError::Conflict("duplicate plugin id".into()));
    }
    let mut visiting = BTreeSet::new();
    let mut visited = BTreeSet::new();
    let mut ordered = Vec::with_capacity(plugins.len());
    fn visit<'a>(
        id: &'a str,
        by_id: &BTreeMap<&'a str, &'a dyn AgentPlugin>,
        visiting: &mut BTreeSet<&'a str>,
        visited: &mut BTreeSet<&'a str>,
        ordered: &mut Vec<&'a dyn AgentPlugin>,
    ) -> Result<(), PluginError> {
        if visited.contains(id) {
            return Ok(());
        }
        if !visiting.insert(id) {
            return Err(PluginError::DependencyCycle(id.into()));
        }
        let plugin = by_id
            .get(id)
            .ok_or_else(|| PluginError::MissingDependency(id.into()))?;
        plugin.manifest().validate()?;
        for (dependency, requirement) in &plugin.manifest().dependencies {
            if !by_id.contains_key(dependency.as_str()) {
                return Err(PluginError::MissingDependency(format!(
                    "{id} -> {dependency}"
                )));
            }
            let required = VersionReq::parse(requirement).map_err(|error| {
                PluginError::InvalidManifest(format!("{id} dependency {dependency}: {error}"))
            })?;
            let actual = Version::parse(&by_id[dependency.as_str()].manifest().version)
                .map_err(|error| PluginError::InvalidManifest(error.to_string()))?;
            if !required.matches(&actual) {
                return Err(PluginError::DependencyVersion {
                    plugin: id.into(),
                    dependency: dependency.clone(),
                    required: requirement.clone(),
                    actual: actual.to_string(),
                });
            }
            visit(dependency, by_id, visiting, visited, ordered)?;
        }
        visiting.remove(id);
        visited.insert(id);
        ordered.push(*plugin);
        Ok(())
    }
    for id in by_id.keys() {
        visit(id, &by_id, &mut visiting, &mut visited, &mut ordered)?;
    }
    Ok(ordered)
}

fn require(manifest: &PluginManifest, permission: PluginPermission) -> Result<(), PluginError> {
    manifest
        .permissions
        .contains(&permission)
        .then_some(())
        .ok_or_else(|| PluginError::Permission(format!("{} lacks {permission:?}", manifest.id)))
}

fn insert_once<T>(map: &mut BTreeMap<String, T>, key: String, value: T) -> Result<(), PluginError> {
    if map.contains_key(&key) {
        return Err(PluginError::Conflict(key));
    }
    map.insert(key, value);
    Ok(())
}

/// Plugin registration, validation or mount failure.
#[derive(Debug, thiserror::Error)]
pub enum PluginError {
    /// Invalid plugin manifest.
    #[error("invalid plugin manifest: {0}")]
    InvalidManifest(String),
    /// Missing plugin dependency.
    #[error("missing plugin dependency: {0}")]
    MissingDependency(String),
    /// Plugin dependency cycle at.
    #[error("plugin dependency cycle at {0}")]
    DependencyCycle(String),
    /// Plugin permission denied.
    #[error("plugin permission denied: {0}")]
    Permission(String),
    /// Plugin registration conflict.
    #[error("plugin registration conflict: {0}")]
    Conflict(String),
    /// Plugin mount failed.
    #[error("plugin mount failed: {0}")]
    Mount(String),
    /// `cause`; plugin rollback failed: `cleanup_error`.
    #[error("{cause}; plugin rollback failed: {cleanup_error}")]
    Rollback {
        /// The mount failure that triggered rollback.
        cause: Box<PluginError>,
        /// Why rollback failed.
        cleanup_error: String,
        /// Partially mounted state; call [`PluginError::retry_cleanup`].
        cleanup: Box<MountedPlugins>,
    },
    /// Plugin configuration is invalid.
    #[error("plugin {0} configuration is invalid: {1}")]
    InvalidConfig(String, String),
    /// Plugin `plugin` requires `dependency` `required`, found `actual`.
    #[error("plugin {plugin} requires {dependency} {required}, found {actual}")]
    DependencyVersion {
        /// Plugin declaring the dependency.
        plugin: String,
        /// Dependency id.
        dependency: String,
        /// Required semver range.
        required: String,
        /// Registered version.
        actual: String,
    },
}

impl PluginError {
    /// Retry a failed rollback; succeeds once every lease is released.
    pub async fn retry_cleanup(&mut self) -> Result<(), PluginError> {
        let Self::Rollback {
            cleanup,
            cleanup_error,
            ..
        } = self
        else {
            return Ok(());
        };
        match cleanup.unmount().await {
            Ok(()) => Ok(()),
            Err(error) => {
                *cleanup_error = error.to_string();
                Err(error)
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};

    use super::*;
    use crate::ToolExecutionContext;

    struct TestTool;
    #[async_trait]
    impl Tool for TestTool {
        fn name(&self) -> &str {
            "test.tool"
        }
        fn description(&self) -> &str {
            "test"
        }
        fn parameters(&self) -> Value {
            serde_json::json!({"type":"object"})
        }
        fn output_schema(&self) -> Value {
            serde_json::json!({"type":"object"})
        }
        async fn call(&self, arguments: Value) -> Result<Value, String> {
            Ok(arguments)
        }
    }

    struct TestHook;
    #[async_trait]
    impl Hook for TestHook {}

    struct Lease(Arc<AtomicUsize>);
    #[async_trait]
    impl PluginLease for Lease {
        async fn unmount(&mut self) -> Result<(), PluginError> {
            self.0.fetch_add(1, Ordering::SeqCst);
            Ok(())
        }
    }

    struct TestPlugin {
        manifest: PluginManifest,
        fail: bool,
        unmounted: Arc<AtomicUsize>,
    }

    struct RegisteringPlugin {
        manifest: PluginManifest,
        unmounted: Arc<AtomicUsize>,
    }

    #[async_trait]
    impl AgentPlugin for RegisteringPlugin {
        fn manifest(&self) -> &PluginManifest {
            &self.manifest
        }
        fn lease(&self, _: &PluginMountContext, _: &Value) -> Box<dyn PluginLease> {
            Box::new(Lease(Arc::clone(&self.unmounted)))
        }
        async fn activate(
            &self,
            _: &PluginMountContext,
            _: &Value,
            registrar: &mut AgentRegistrar<'_>,
        ) -> Result<(), PluginError> {
            registrar.tool(Arc::new(TestTool))?;
            Ok(())
        }
    }

    struct RetryLease {
        id: &'static str,
        fail_once: Arc<AtomicBool>,
        calls: Arc<std::sync::Mutex<Vec<&'static str>>>,
    }

    #[async_trait]
    impl PluginLease for RetryLease {
        async fn unmount(&mut self) -> Result<(), PluginError> {
            self.calls.lock().unwrap().push(self.id);
            if self.fail_once.swap(false, Ordering::SeqCst) {
                Err(PluginError::Mount(format!("{} failed", self.id)))
            } else {
                Ok(())
            }
        }
    }

    struct CleanupPlugin {
        manifest: PluginManifest,
        lease_id: &'static str,
        fail_activation: bool,
        fail_once: Arc<AtomicBool>,
        calls: Arc<std::sync::Mutex<Vec<&'static str>>>,
    }

    #[async_trait]
    impl AgentPlugin for CleanupPlugin {
        fn manifest(&self) -> &PluginManifest {
            &self.manifest
        }
        fn lease(&self, _: &PluginMountContext, _: &Value) -> Box<dyn PluginLease> {
            Box::new(RetryLease {
                id: self.lease_id,
                fail_once: Arc::clone(&self.fail_once),
                calls: Arc::clone(&self.calls),
            })
        }
        async fn activate(
            &self,
            _: &PluginMountContext,
            _: &Value,
            _: &mut AgentRegistrar<'_>,
        ) -> Result<(), PluginError> {
            if self.fail_activation {
                Err(PluginError::Mount(self.manifest.id.clone()))
            } else {
                Ok(())
            }
        }
    }
    #[async_trait]
    impl AgentPlugin for TestPlugin {
        fn manifest(&self) -> &PluginManifest {
            &self.manifest
        }
        fn lease(&self, _: &PluginMountContext, _: &Value) -> Box<dyn PluginLease> {
            Box::new(Lease(Arc::clone(&self.unmounted)))
        }
        async fn activate(
            &self,
            _context: &PluginMountContext,
            _config: &Value,
            _registrar: &mut AgentRegistrar<'_>,
        ) -> Result<(), PluginError> {
            if self.fail {
                Err(PluginError::Mount(self.manifest.id.clone()))
            } else {
                Ok(())
            }
        }
    }

    fn plugin(
        id: &str,
        dependencies: BTreeMap<String, String>,
        fail: bool,
        unmounted: Arc<AtomicUsize>,
    ) -> TestPlugin {
        TestPlugin {
            manifest: PluginManifest {
                id: id.into(),
                version: "1.0.0".into(),
                dependencies,
                config_schema: serde_json::json!({"type":"object"}),
                permissions: BTreeSet::new(),
            },
            fail,
            unmounted,
        }
    }

    #[tokio::test]
    async fn activation_failure_rolls_back_current_and_prior_leases() {
        let unmounted = Arc::new(AtomicUsize::new(0));
        let first = plugin("first", BTreeMap::new(), false, Arc::clone(&unmounted));
        let second = plugin(
            "second",
            BTreeMap::from([("first".into(), "^1".into())]),
            true,
            Arc::clone(&unmounted),
        );
        let context = PluginMountContext {
            tenant_id: "t".parse().unwrap(),
            subject_id: "s".parse().unwrap(),
            session_id: "session".parse().unwrap(),
            profile_revision_id: "profile".parse().unwrap(),
        };
        assert!(
            mount_plugins(&context, &[&second, &first], &BTreeMap::new())
                .await
                .is_err()
        );
        assert_eq!(unmounted.load(Ordering::SeqCst), 2);
    }

    #[tokio::test]
    async fn activation_failure_cleans_current_plugin_and_retries_failed_cleanup() {
        let calls = Arc::new(std::sync::Mutex::new(Vec::new()));
        let first = CleanupPlugin {
            manifest: plugin(
                "first",
                BTreeMap::new(),
                false,
                Arc::new(AtomicUsize::new(0)),
            )
            .manifest,
            lease_id: "first",
            fail_activation: false,
            fail_once: Arc::new(AtomicBool::new(false)),
            calls: Arc::clone(&calls),
        };
        let second = CleanupPlugin {
            manifest: plugin(
                "second",
                BTreeMap::from([("first".into(), "^1".into())]),
                false,
                Arc::new(AtomicUsize::new(0)),
            )
            .manifest,
            lease_id: "second",
            fail_activation: true,
            fail_once: Arc::new(AtomicBool::new(true)),
            calls: Arc::clone(&calls),
        };
        let context = PluginMountContext {
            tenant_id: "t".parse().unwrap(),
            subject_id: "s".parse().unwrap(),
            session_id: "rollback".parse().unwrap(),
            profile_revision_id: "profile".parse().unwrap(),
        };

        let mut error = mount_plugins(&context, &[&second, &first], &BTreeMap::new())
            .await
            .unwrap_err();
        assert!(matches!(error, PluginError::Rollback { .. }));
        assert_eq!(*calls.lock().unwrap(), vec!["second", "first"]);
        error.retry_cleanup().await.unwrap();
        assert_eq!(*calls.lock().unwrap(), vec!["second", "first", "second"]);
    }

    #[tokio::test]
    async fn registration_conflict_rolls_back_current_and_prior_leases() {
        let unmounted = Arc::new(AtomicUsize::new(0));
        let manifest = |id: &str, dependencies| PluginManifest {
            id: id.into(),
            version: "1.0.0".into(),
            dependencies,
            config_schema: serde_json::json!({"type":"object"}),
            permissions: BTreeSet::from([PluginPermission::Tool]),
        };
        let first = RegisteringPlugin {
            manifest: manifest("first", BTreeMap::new()),
            unmounted: Arc::clone(&unmounted),
        };
        let second = RegisteringPlugin {
            manifest: manifest("second", BTreeMap::from([("first".into(), "^1".into())])),
            unmounted: Arc::clone(&unmounted),
        };
        let context = PluginMountContext {
            tenant_id: "t".parse().unwrap(),
            subject_id: "s".parse().unwrap(),
            session_id: "session".parse().unwrap(),
            profile_revision_id: "profile".parse().unwrap(),
        };

        assert!(matches!(
            mount_plugins(&context, &[&second, &first], &BTreeMap::new()).await,
            Err(PluginError::Conflict(_))
        ));
        assert_eq!(unmounted.load(Ordering::SeqCst), 2);
    }

    #[tokio::test]
    async fn session_lease_unmounts_after_last_runtime_clone() {
        let unmounted = Arc::new(AtomicUsize::new(0));
        let plugin = plugin("session", BTreeMap::new(), false, Arc::clone(&unmounted));
        let context = PluginMountContext {
            tenant_id: "t".parse().unwrap(),
            subject_id: "s".parse().unwrap(),
            session_id: "session".parse().unwrap(),
            profile_revision_id: "profile".parse().unwrap(),
        };
        let mounted = mount_plugins(&context, &[&plugin], &BTreeMap::new())
            .await
            .unwrap();
        let runtime_clone = mounted.clone();
        let teardown = tokio::spawn(async move {
            let mut mounted = mounted;
            mounted.unmount().await
        });
        tokio::task::yield_now().await;
        assert_eq!(unmounted.load(Ordering::SeqCst), 0);
        drop(runtime_clone);
        teardown.await.unwrap().unwrap();
        assert_eq!(unmounted.load(Ordering::SeqCst), 1);
    }

    #[tokio::test]
    async fn cleanup_attempts_every_lease_and_retries_only_failures() {
        let calls = Arc::new(std::sync::Mutex::new(Vec::new()));
        let first = CleanupPlugin {
            manifest: plugin(
                "first-cleanup",
                BTreeMap::new(),
                false,
                Arc::new(AtomicUsize::new(0)),
            )
            .manifest,
            lease_id: "first",
            fail_activation: false,
            fail_once: Arc::new(AtomicBool::new(false)),
            calls: Arc::clone(&calls),
        };
        let second = CleanupPlugin {
            manifest: plugin(
                "second-cleanup",
                BTreeMap::from([("first-cleanup".into(), "^1".into())]),
                false,
                Arc::new(AtomicUsize::new(0)),
            )
            .manifest,
            lease_id: "second",
            fail_activation: false,
            fail_once: Arc::new(AtomicBool::new(true)),
            calls: Arc::clone(&calls),
        };
        let context = PluginMountContext {
            tenant_id: "t".parse().unwrap(),
            subject_id: "s".parse().unwrap(),
            session_id: "cleanup".parse().unwrap(),
            profile_revision_id: "profile".parse().unwrap(),
        };
        let mut mounted = mount_plugins(&context, &[&second, &first], &BTreeMap::new())
            .await
            .unwrap();
        assert!(mounted.unmount().await.is_err());
        assert_eq!(*calls.lock().unwrap(), vec!["second", "first"]);
        mounted.unmount().await.unwrap();
        assert_eq!(*calls.lock().unwrap(), vec!["second", "first", "second"]);
    }

    #[test]
    fn dependency_versions_and_cycles_fail_before_mount() {
        let counter = Arc::new(AtomicUsize::new(0));
        let base = plugin("base", BTreeMap::new(), false, Arc::clone(&counter));
        let incompatible = plugin(
            "consumer",
            BTreeMap::from([("base".into(), "^2".into())]),
            false,
            Arc::clone(&counter),
        );
        assert!(matches!(
            resolve_plugin_order(&[&base, &incompatible]),
            Err(PluginError::DependencyVersion { .. })
        ));
        let left = plugin(
            "left",
            BTreeMap::from([("right".into(), "^1".into())]),
            false,
            Arc::clone(&counter),
        );
        let right = plugin(
            "right",
            BTreeMap::from([("left".into(), "^1".into())]),
            false,
            counter,
        );
        assert!(matches!(
            resolve_plugin_order(&[&left, &right]),
            Err(PluginError::DependencyCycle(_))
        ));
    }

    #[test]
    fn descriptor_preflight_rejects_missing_and_incompatible_dependencies() {
        let descriptors = vec![
            PluginDescriptor {
                id: "base".into(),
                version: "1.0.0".into(),
                dependencies: BTreeMap::new(),
                permissions: BTreeSet::new(),
            },
            PluginDescriptor {
                id: "child".into(),
                version: "1.0.0".into(),
                dependencies: BTreeMap::from([("base".into(), "^2".into())]),
                permissions: BTreeSet::new(),
            },
        ];
        let references = BTreeMap::from([
            ("base".into(), ("1.0.0".into(), Value::Null)),
            ("child".into(), ("1.0.0".into(), Value::Null)),
        ]);

        let issues = validate_plugin_references(&descriptors, &references);
        assert_eq!(issues.len(), 1);
        assert!(matches!(
            issues[0].error,
            PluginError::DependencyVersion { .. }
        ));
    }

    #[test]
    fn descriptor_preflight_reports_every_missing_dependency() {
        let descriptors = vec![PluginDescriptor {
            id: "child".into(),
            version: "1.0.0".into(),
            dependencies: BTreeMap::from([
                ("first".into(), "^1".into()),
                ("second".into(), "^1".into()),
            ]),
            permissions: BTreeSet::new(),
        }];
        let references = BTreeMap::from([("child".into(), ("1.0.0".into(), Value::Null))]);

        let issues = validate_plugin_references(&descriptors, &references);

        assert_eq!(issues.len(), 2);
        assert!(issues.iter().all(|issue| issue.plugin == "child"
            && matches!(issue.error, PluginError::MissingDependency(_))));
    }

    #[tokio::test]
    async fn registrar_and_catalog_enforce_manifest_permissions_atomically() {
        let manifest = PluginManifest {
            id: "test.plugin".into(),
            version: "1.0.0".into(),
            dependencies: BTreeMap::new(),
            config_schema: serde_json::json!({"type":"object","required":["enabled"]}),
            permissions: BTreeSet::from([
                PluginPermission::Prompt,
                PluginPermission::Tool,
                PluginPermission::Hook,
            ]),
        };
        manifest.validate().unwrap();
        for invalid in [
            PluginManifest {
                id: "Bad Plugin".into(),
                ..manifest.clone()
            },
            PluginManifest {
                version: "latest".into(),
                ..manifest.clone()
            },
        ] {
            assert!(matches!(
                invalid.validate(),
                Err(PluginError::InvalidManifest(_))
            ));
        }

        let mut contributions = PluginContributions::default();
        {
            let mut registrar = AgentRegistrar {
                owner: &manifest,
                contributions: &mut contributions,
            };
            registrar.prompt("system", "prompt").unwrap();
            assert!(matches!(
                registrar.prompt("system", "again"),
                Err(PluginError::Conflict(_))
            ));
            registrar.tool(Arc::new(TestTool)).unwrap();
            assert!(matches!(
                registrar.tool(Arc::new(TestTool)),
                Err(PluginError::Conflict(_))
            ));
            registrar.hook(Arc::new(TestHook)).unwrap();
        }
        let prompt = contributions.prompts.values().next().unwrap();
        assert_eq!(prompt.id, "plugin.test.plugin.system");
        assert_eq!(prompt.source, "test.plugin");
        assert_eq!(prompt.version, "1.0.0");

        let denied = PluginManifest {
            permissions: BTreeSet::new(),
            ..manifest.clone()
        };
        let mut denied_contributions = PluginContributions::default();
        let mut denied_registrar = AgentRegistrar {
            owner: &denied,
            contributions: &mut denied_contributions,
        };
        assert!(matches!(
            denied_registrar.prompt("x", "x"),
            Err(PluginError::Permission(_))
        ));
        let context = ToolExecutionContext {
            request: crate::RequestContext {
                tenant_id: "tenant".parse().unwrap(),
                subject_id: "subject".parse().unwrap(),
                roles: Default::default(),
                locale: "en".into(),
                request_id: "request".parse().unwrap(),
                entitlements: Default::default(),
            },
            session_id: "session".parse().unwrap(),
            run_id: "run".parse().unwrap(),
            step: 1,
            call_id: "call".parse().unwrap(),
            source_event_seq: 1,
            interaction_resolution: None,
            cancellation: crate::CancellationToken::default(),
            deadline: std::time::Instant::now() + std::time::Duration::from_secs(15),
        };
        assert_eq!(
            TestHook
                .before_tool(&context, "test.tool", &serde_json::json!({}))
                .await,
            HookDecision::Continue
        );
        TestHook
            .after_tool(
                &context,
                "test.tool",
                &serde_json::json!({}),
                &serde_json::json!({}),
            )
            .await
            .unwrap();
        let unmounted = Arc::new(AtomicUsize::new(0));
        let plugin = Arc::new(TestPlugin {
            manifest: manifest.clone(),
            fail: false,
            unmounted: Arc::clone(&unmounted),
        });
        let mut catalog = PluginCatalog::default();
        catalog.register(plugin.clone()).unwrap();
        assert!(matches!(
            catalog.register(plugin),
            Err(PluginError::Conflict(_))
        ));
        let mount_context = PluginMountContext {
            tenant_id: "t".parse().unwrap(),
            subject_id: "s".parse().unwrap(),
            session_id: "session".parse().unwrap(),
            profile_revision_id: "profile".parse().unwrap(),
        };
        assert!(matches!(
            catalog
                .mount(
                    &mount_context,
                    &BTreeMap::from([("missing".into(), ("1.0.0".into(), serde_json::json!({})))])
                )
                .await,
            Err(PluginError::MissingDependency(_))
        ));
        assert!(matches!(
            catalog
                .mount(
                    &mount_context,
                    &BTreeMap::from([(
                        "test.plugin".into(),
                        ("2.0.0".into(), serde_json::json!({}))
                    )])
                )
                .await,
            Err(PluginError::DependencyVersion { .. })
        ));
        assert!(matches!(
            catalog
                .mount(
                    &mount_context,
                    &BTreeMap::from([(
                        "test.plugin".into(),
                        ("1.0.0".into(), serde_json::json!({}))
                    )])
                )
                .await,
            Err(PluginError::InvalidConfig(_, _))
        ));
        let mut mounted = catalog
            .mount(
                &mount_context,
                &BTreeMap::from([(
                    "test.plugin".into(),
                    ("1.0.0".into(), serde_json::json!({"enabled":true})),
                )]),
            )
            .await
            .unwrap();
        mounted.unmount().await.unwrap();
        assert_eq!(unmounted.load(Ordering::SeqCst), 1);
    }
}