agent_engine/extensions/runtime/

mod.rs

//! Extension runtime trait and registry.

pub mod process;
pub mod restart;

pub use restart::RestartPolicy;

use self::process::{
    NotificationFrame, ProviderCompleteParams, ProviderCompleteResult, ProviderStreamEvent,
};
use crate::extensions::commands::CommandOutputEvent;
use crate::extensions::hooks::events::{HookEvent, HookResult};
use crate::extensions::info::PluginInfo;
use crate::extensions::tasks::TaskEvent;
use async_trait::async_trait;
use serde_json::Value;

/// Streamed event delivered to a `command.invoke` caller.
#[derive(Debug, Clone, PartialEq)]
pub enum InvokeCommandEvent {
    /// Command output event matching the caller's `request_id`.
    Output(CommandOutputEvent),
    /// Spontaneous plugin task event (any `request_id`).
    Task(TaskEvent),
}

/// Health state for an extension handler.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ExtensionHealth {
    /// Manifest validated, process spawned, but `initialize` not yet completed.
    Loaded,
    /// Manifest failed validation — the extension never started.
    FailedValidation,
    /// `initialize` rpc failed — the extension started but couldn't capability-handshake.
    FailedInitialize,
    /// Healthy and serving requests.
    Running,
    /// Process restarting after transport failure but within restart budget.
    Restarting,
    /// Running, but at least one recent operation failed (e.g. hook timeout).
    Degraded,
    /// Permanent failure — restart budget exhausted or unrecoverable error.
    Failed,
}

impl ExtensionHealth {
    pub fn as_str(self) -> &'static str {
        match self {
            Self::Loaded => "loaded",
            Self::FailedValidation => "failed_validation",
            Self::FailedInitialize => "failed_initialize",
            Self::Running => "running",
            Self::Restarting => "restarting",
            Self::Degraded => "degraded",
            Self::Failed => "failed",
        }
    }
}

#[cfg(test)]
#[allow(clippy::items_after_test_module)]
mod health_tests {
    use super::ExtensionHealth;

    #[test]
    fn as_str_covers_all_variants() {
        assert_eq!(ExtensionHealth::Loaded.as_str(), "loaded");
        assert_eq!(
            ExtensionHealth::FailedValidation.as_str(),
            "failed_validation"
        );
        assert_eq!(
            ExtensionHealth::FailedInitialize.as_str(),
            "failed_initialize"
        );
        assert_eq!(ExtensionHealth::Running.as_str(), "running");
        assert_eq!(ExtensionHealth::Restarting.as_str(), "restarting");
        assert_eq!(ExtensionHealth::Degraded.as_str(), "degraded");
        assert_eq!(ExtensionHealth::Failed.as_str(), "failed");
    }
}

/// A point-in-time lifecycle observation, not a lease on future execution.
///
/// Generations are opaque and local to one handler instance. A caller must pin
/// both the handler identity and generation before dispatch, then require the
/// same generation with `Running` or `Degraded` health when accepting a reply
/// and before executing any deferred work. All other health states fail closed.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ExtensionLifecycle {
    pub generation: u64,
    pub health: ExtensionHealth,
}

/// Trait for extension runtimes that can handle hook events.
#[async_trait]
pub trait ExtensionHandler: Send + Sync {
    /// Unique identifier for this extension.
    fn id(&self) -> &str;

    /// Observe the actual serving runtime without waiting for locks or RPC I/O.
    ///
    /// `None` means unsupported or temporarily unavailable, never "healthy".
    /// Implementations must invalidate the generation on runtime/transport loss
    /// and before replacement, even if the handler itself remains the same Arc.
    /// Consumers retaining authority must fail closed on `None`, non-live health,
    /// or a changed generation, including while a previously valid reply waits
    /// to execute. Recheck at dispatch, reply acceptance, and each frontend tick;
    /// this snapshot cannot eliminate exit races after the final observation.
    fn lifecycle_snapshot(&self) -> Option<ExtensionLifecycle> {
        None
    }

    /// Handle a hook event. Returns the handler's decision.
    async fn handle(&self, event: &HookEvent) -> HookResult;

    /// Call an extension-provided tool.
    async fn call_tool(&self, _name: &str, _input: Value) -> Result<Value, String> {
        Err("extension runtime does not support tool.call".to_string())
    }

    /// Call an extension-provided tool on behalf of a conversation.
    ///
    /// `session_id` is the owning conversation id (`None` = worker/subagent
    /// or no conversation). Runtimes that put it on the wire override this;
    /// the default drops it and calls [`Self::call_tool`], so handlers that
    /// predate the field (and wrappers that fail closed) keep their exact
    /// behaviour.
    async fn call_tool_for(
        &self,
        name: &str,
        input: Value,
        session_id: Option<&str>,
    ) -> Result<Value, String> {
        let _ = session_id;
        self.call_tool(name, input).await
    }

    /// Complete a chat request through an extension-provided model provider.
    async fn provider_complete(
        &self,
        _params: ProviderCompleteParams,
    ) -> Result<ProviderCompleteResult, String> {
        Err("extension runtime does not support provider.complete".to_string())
    }

    /// Stream a chat request through an extension-provided model provider.
    ///
    /// The handler must forward `provider.stream.event` notifications to `sink`
    /// in order. The returned `ProviderCompleteResult` is the final aggregated
    /// response (so callers that don't need streaming can use it as the final
    /// state). Implementations that don't support streaming should return
    /// `Err("provider.stream is not supported by this extension")`.
    async fn provider_stream(
        &self,
        _params: ProviderCompleteParams,
        _sink: tokio::sync::mpsc::Sender<ProviderStreamEvent>,
    ) -> Result<ProviderCompleteResult, String> {
        Err("provider.stream is not supported by this extension".to_string())
    }

    /// Invoke a plugin-registered interactive slash command. The handler must
    /// forward `command.output` notifications matching `request_id` and any
    /// `task.*` notifications to `sink`. Returns the final response value.
    ///
    /// CP-11 fix-3: the sink is BOUNDED and metered
    /// ([`crate::extensions::invoke_output::InvokeEventSink`]). It must be
    /// paired with an EAGERLY CONCURRENT consumer — normally the collector
    /// half of [`crate::extensions::invoke_output::invoke_event_channel`],
    /// joined with this call by
    /// `ExtensionManager::invoke_command_collected` — which enforces
    /// invocation-local byte/event budgets at production time. A hostile
    /// `command.output` flood is therefore paced by awaited sends and
    /// bounded retention instead of parking aggregate bytes in an
    /// unbounded post-hoc queue (the pre-fix behavior). Handlers must
    /// treat a closed sink like the old closed-unbounded-sink case: stop
    /// forwarding, still complete the call.
    async fn invoke_command(
        &self,
        _command: &str,
        _args: Vec<String>,
        _request_id: &str,
        _sink: crate::extensions::invoke_output::InvokeEventSink,
    ) -> Result<Value, String> {
        Err("extension runtime does not support command.invoke".to_string())
    }

    /// Fetch optional plugin capability/build/model information.
    async fn get_info(&self) -> Result<PluginInfo, String> {
        Err("extension runtime does not support info.get".to_string())
    }

    /// Ask the plugin to supply sidecar spawn arguments. Used by the
    /// modality-neutral sidecar bootstrap path (see
    /// `crate::sidecar::spawn`); plugins that don't host a sidecar
    /// should leave the default `Err` in place. Core treats the
    /// returned [`SidecarSpawnArgs::args`] as opaque.
    async fn sidecar_spawn_args(&self) -> Result<crate::sidecar::spawn::SidecarSpawnArgs, String> {
        Err("extension runtime does not support sidecar.spawn_args".to_string())
    }

    /// Open a plugin-owned custom settings editor and return its initial render payload.
    async fn settings_editor_open(&self, _category: &str, _field: &str) -> Result<Value, String> {
        Err("extension runtime does not support settings.editor.open".to_string())
    }

    /// Forward a keypress to the active plugin-owned custom settings editor.
    async fn settings_editor_key(
        &self,
        _category: &str,
        _field: &str,
        _key: &str,
    ) -> Result<Value, String> {
        Err("extension runtime does not support settings.editor.key".to_string())
    }

    /// Ask the plugin to commit a custom editor value selected by the UI.
    async fn settings_editor_commit(
        &self,
        _category: &str,
        _field: &str,
        _value: Value,
    ) -> Result<Value, String> {
        Err("extension runtime does not support settings.editor.commit".to_string())
    }

    /// Gracefully shut down the extension.
    async fn shutdown(&self);

    /// Subscribe to raw JSON-RPC notifications from this extension.
    /// Returns a monotonic subscription id (for unsubscribing) and a
    /// receiver that yields every notification frame. The default
    /// implementation returns a channel whose sender is immediately
    /// dropped, so the receiver yields `None` right away — effectively
    /// a no-op for handler impls that don't support notifications.
    async fn subscribe_notifications(
        &self,
    ) -> (usize, tokio::sync::mpsc::Receiver<NotificationFrame>) {
        let (_tx, rx) = tokio::sync::mpsc::channel(1);
        (0, rx)
    }

    /// Current health state of this handler.
    async fn health(&self) -> ExtensionHealth {
        ExtensionHealth::Running
    }

    /// Number of transport restarts observed by this handler.
    async fn restart_count(&self) -> usize {
        0
    }
}

#[cfg(test)]
mod lifecycle_tests {
    use super::*;

    struct UnsupportedHandler;

    #[async_trait]
    impl ExtensionHandler for UnsupportedHandler {
        fn id(&self) -> &str {
            "unsupported-lifecycle"
        }

        async fn handle(&self, _event: &HookEvent) -> HookResult {
            unreachable!("snapshot must not dispatch a hook")
        }

        async fn shutdown(&self) {}
    }

    #[test]
    fn default_lifecycle_snapshot_is_unavailable_through_trait_object() {
        let handler: &dyn ExtensionHandler = &UnsupportedHandler;
        assert_eq!(handler.lifecycle_snapshot(), None);
    }
}

#[cfg(test)]
pub mod tests {
    use super::*;

    /// Minimal handler for unit tests. Reports generation 0, Running health.
    pub struct DummyHandler;

    #[async_trait]
    impl ExtensionHandler for DummyHandler {
        fn id(&self) -> &str {
            "dummy"
        }

        fn lifecycle_snapshot(&self) -> Option<ExtensionLifecycle> {
            Some(ExtensionLifecycle {
                health: ExtensionHealth::Running,
                generation: 0,
            })
        }

        async fn handle(&self, _event: &HookEvent) -> HookResult {
            HookResult::default()
        }

        async fn shutdown(&self) {}
    }
}