yapi_core/extensions/

mod.rs

//! Extensions: code that registers tools and commands with a session and
//! handles its events, as pi's extension runner (`core/extensions` in pi
//! `v1.0.0`) runs them.
//!
//! This is the part of pi's `ExtensionAPI` that the built-in extensions use.
//! The wasm host in `yapi-ext` implements [`Extension`] for loaded packages with
//! the same contract. Each session gets its own extension instances. A
//! replaced session receives `session_shutdown` before its successor starts.

pub mod codemode;
pub mod discovery;
mod events;
mod grants;
pub mod tool_search;
mod ui;

pub use events::{EventPolicy, delivers};
pub use grants::{CAPABILITIES, Grants, receives, tool_granted};
pub use ui::{
    ComponentHost, CustomOptions, DialogOptions, ExtensionUi, NoUi, NotifyKind, Placement,
    RemoteComponent, TerminalState, Widget, WorkingIndicator,
};

use std::path::PathBuf;
use std::sync::{Arc, Mutex};

use futures_util::future::BoxFuture;
use indexmap::IndexMap;
use serde_json::Value;
use tokio_util::sync::CancellationToken;
use yapi_ai::model_catalog::RefreshOptions;
use yapi_ai::registry::{Availability, ModelRegistry};
use yapi_types::autocomplete::ArgumentCompletions;
use yapi_types::model::Model;
use yapi_types::models::ModelDefinition;
use yapi_types::rpc::SourceInfo;

use crate::agent_session::WeakSession;
use crate::tools::RegisteredTool;
use crate::tools::registry::ToolRegistry;

/// pi's `ExtensionMode`: the mode a session runs in.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum Mode {
    /// The interactive terminal UI.
    Tui,
    /// RPC over stdin and stdout.
    Rpc,
    /// Print mode.
    #[default]
    Print,
    /// JSON event mode.
    Json,
}

/// pi's `ToolInfo`: a tool as extensions see it.
#[derive(Clone, Debug, PartialEq)]
pub struct ToolInfo {
    /// Name.
    pub name: String,
    /// Description.
    pub description: String,
    /// Argument schema.
    pub parameters: Value,
    /// How it reaches the model.
    pub exposure: crate::tools::Exposure,
    /// Its group.
    pub namespace: Option<crate::tools::Namespace>,
}

/// The session's tools, shared with its extensions.
#[derive(Clone, Default)]
pub struct Tools(Arc<Mutex<ToolRegistry>>);

impl Tools {
    /// Wraps a registry.
    pub fn new(registry: ToolRegistry) -> Tools {
        Tools(Arc::new(Mutex::new(registry)))
    }

    fn lock(&self) -> std::sync::MutexGuard<'_, ToolRegistry> {
        yapi_types::sync::lock(&self.0)
    }

    /// Runs `f` with the registry.
    pub fn with<T>(&self, f: impl FnOnce(&mut ToolRegistry) -> T) -> T {
        f(&mut self.lock())
    }

    /// pi's `registerTool`.
    pub fn register(&self, tool: RegisteredTool) {
        self.lock().register(tool);
    }

    /// pi's `getActiveTools`.
    pub fn active(&self) -> Vec<String> {
        self.lock().active()
    }

    /// pi's `setActiveTools`.
    pub fn set_active(&self, names: Vec<String>) {
        self.lock().set_active(names);
    }

    /// pi's `getAllTools`.
    pub fn all(&self) -> Vec<ToolInfo> {
        self.lock().all().iter().map(RegisteredTool::info).collect()
    }
}

/// What an extension can reach while handling an event or command.
#[derive(Clone)]
pub struct Context {
    /// The session's working directory.
    pub cwd: PathBuf,
    /// The agent directory.
    pub agent_dir: PathBuf,
    /// The directory the session's project files come from: the working
    /// directory unless the run names another.
    pub project_dir: PathBuf,
    /// Whether the project's resources are trusted.
    pub project_trusted: bool,
    /// In a project-only run, the patterns of the user's MCP servers it
    /// keeps, from the `projectOnly` setting. The user's others stay out.
    pub project_only_mcp: Option<Vec<String>>,
    /// The mode.
    pub mode: Mode,
    /// Dialogs and notifications.
    pub ui: Arc<dyn ExtensionUi>,
    /// The session's tools.
    pub tools: Tools,
    /// Cancelled when the current run is aborted.
    pub cancel: CancellationToken,
    /// The session, for extensions that act on it later.
    pub session: WeakSession,
}

/// A slash command an extension registered.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Command {
    /// Name without the slash.
    pub name: String,
    /// What it does.
    pub description: String,
}

/// pi's `ResolvedCommand`: an extension command with the name it is
/// invoked by.
#[derive(Clone)]
pub struct ResolvedCommand {
    /// `name`, or `name:N` when several extensions register `name`.
    pub invocation: String,
    /// The command as registered.
    pub command: Command,
    /// The extension that handles it.
    pub extension: Arc<dyn Extension>,
}

/// A key an extension binds with pi's `registerShortcut`.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Shortcut {
    /// The key id, such as `alt+u`.
    pub key: String,
    /// What it does, for `/hotkeys`.
    pub description: Option<String>,
}

/// An extension shortcut the interactive mode honors.
#[derive(Clone)]
pub struct ShortcutBinding {
    /// The key id, lowercased.
    pub key: String,
    /// The key id as the extension registered it.
    pub registered: String,
    /// What it does.
    pub description: Option<String>,
    /// The extension's path.
    pub path: String,
    /// The extension.
    pub extension: Arc<dyn Extension>,
}

/// pi's `RESERVED_KEYBINDINGS_FOR_EXTENSION_CONFLICTS`: actions whose keys
/// extensions may not take.
pub const RESERVED_ACTIONS: [&str; 18] = [
    "app.interrupt",
    "app.clear",
    "app.exit",
    "app.suspend",
    "app.thinking.cycle",
    "app.model.cycleForward",
    "app.model.cycleBackward",
    "app.model.select",
    "app.tools.expand",
    "app.thinking.toggle",
    "app.editor.external",
    "app.message.copy",
    "app.message.followUp",
    "tui.input.submit",
    "tui.select.confirm",
    "tui.select.cancel",
    "tui.input.copy",
    "tui.editor.deleteToLineEnd",
];

/// pi's `getShortcuts`: the shortcuts of `extensions` in load order, given
/// the keys bound to each built-in action. A key a reserved action uses is
/// refused. Another built-in key, or one an earlier extension took, goes to
/// the later extension. Returns the bindings and pi's warnings, each with the
/// path of the extension it concerns.
pub fn resolve_shortcuts(
    extensions: &[Arc<dyn Extension>],
    builtin: &[(String, Vec<String>)],
) -> (Vec<ShortcutBinding>, Vec<(String, String)>) {
    let mut taken: std::collections::HashMap<String, (&str, bool)> =
        std::collections::HashMap::new();
    for (action, keys) in builtin {
        let reserved = RESERVED_ACTIONS.contains(&action.as_str());
        for key in keys {
            let key = key.to_lowercase();
            if taken
                .get(&key)
                .is_some_and(|(_, existing)| *existing && !reserved)
            {
                continue;
            }
            taken.insert(key, (action.as_str(), reserved));
        }
    }
    let mut bindings: Vec<ShortcutBinding> = Vec::new();
    let mut warnings = Vec::new();
    for extension in extensions {
        let path = extension.source().path;
        for shortcut in extension.shortcuts() {
            let key = shortcut.key.to_lowercase();
            match taken.get(&key) {
                Some((_, true)) => {
                    warnings.push((
                        path.clone(),
                        format!(
                            "Extension shortcut '{}' from {path} conflicts with built-in shortcut. Skipping.",
                            shortcut.key
                        ),
                    ));
                    continue;
                }
                Some((action, false)) => warnings.push((
                    path.clone(),
                    format!(
                        "Extension shortcut conflict: '{}' is built-in shortcut for {action} and {path}. Using {path}.",
                        shortcut.key
                    ),
                )),
                None => {}
            }
            let binding = ShortcutBinding {
                key: key.clone(),
                registered: shortcut.key.clone(),
                description: shortcut.description,
                path: path.clone(),
                extension: extension.clone(),
            };
            // Like a `Map`, the later binding keeps the earlier one's place.
            match bindings.iter_mut().find(|existing| existing.key == key) {
                Some(existing) => {
                    warnings.push((
                        path.clone(),
                        format!(
                            "Extension shortcut conflict: '{}' registered by both {} and {path}. Using {path}.",
                            shortcut.key, existing.path
                        ),
                    ));
                    *existing = binding;
                }
                None => bindings.push(binding),
            }
        }
    }
    (bindings, warnings)
}

/// pi's `renderCall`, `renderResult` and `renderShell`: how an extension
/// draws one of its tools.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct ToolRenderers {
    /// It draws the call.
    pub call: bool,
    /// It draws the result.
    pub result: bool,
    /// It draws its own frame (`renderShell: "self"`).
    pub own_shell: bool,
}

/// pi's `ToolLoadout`: the tools a loadout hook sees.
#[derive(Clone, Default)]
pub struct Loadout {
    /// The active tools in activation order, with their own descriptions.
    pub declared: Vec<RegisteredTool>,
    /// The tools other tools may call, in registration order.
    pub callable: Vec<RegisteredTool>,
}

/// pi's `ToolLoadoutChanges`: what a loadout hook changes.
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct LoadoutChanges {
    /// Descriptions for declared tools, by name, replacing their own.
    pub descriptions: std::collections::HashMap<String, String>,
    /// Declared tools that stay callable but are left out of requests and
    /// of the prompt's tool list.
    pub hidden: std::collections::HashSet<String>,
}

/// What an extension draws in the transcript.
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct Renderers {
    /// Its tools, by name.
    pub tools: std::collections::HashMap<String, ToolRenderers>,
    /// Custom message types it draws (`registerMessageRenderer`).
    pub messages: Vec<String>,
    /// Custom entry types it draws (`registerEntryRenderer`).
    pub entries: Vec<String>,
    /// It transforms Markdown before it renders (`registerMarkdownTransformer`).
    pub markdown: bool,
    /// It renders code blocks (`pi.yapi.registerCodeRenderer`, yapi's own).
    pub code: bool,
}

/// An extension. Every hook has a no-op default.
pub trait Extension: Send + Sync {
    /// Where the extension comes from, as `get_commands` reports it.
    fn source(&self) -> SourceInfo;

    /// What the extension may reach. The session applies its event
    /// handlers' results only within these.
    fn grants(&self) -> Grants {
        Grants::default()
    }

    /// Reports `message`, a warning from the host about the extension, such
    /// as a handler result the session dropped for lack of a grant, with
    /// the extension's own output. Built-in extensions, which have every
    /// grant, report nothing.
    fn warn(&self, _message: &str) {}

    /// Registers the tools it offers from the start, before the session's
    /// active tools are chosen.
    fn load(&self, _tools: &Tools) {}

    /// The commands it handles.
    fn commands(&self) -> Vec<Command> {
        Vec::new()
    }

    /// Completions for the argument of `command`. While they are
    /// [`ArgumentCompletions::Pending`], the UI's `refresh_completions`
    /// follows once they are ready.
    fn complete(&self, _command: &str, _prefix: &str) -> ArgumentCompletions {
        ArgumentCompletions::Ready(None)
    }

    /// The keys it binds.
    fn shortcuts(&self) -> Vec<Shortcut> {
        Vec::new()
    }

    /// Runs the handler of shortcut `key` and returns the handler's error.
    fn run_shortcut<'a>(
        &'a self,
        _key: &'a str,
        _ctx: &'a Context,
    ) -> BoxFuture<'a, Result<(), String>> {
        Box::pin(async { Ok(()) })
    }

    /// Runs `command` with its argument text.
    fn run_command<'a>(
        &'a self,
        _command: &'a str,
        _args: &'a str,
        _ctx: &'a Context,
    ) -> BoxFuture<'a, ()> {
        Box::pin(async {})
    }

    /// Resolves once the extension code that is running has finished and
    /// delivered its results. pi runs extension code to completion before it
    /// reads the next input. Returns a count that grows with every call into
    /// the extension's code.
    fn settle(&self) -> BoxFuture<'_, u64> {
        Box::pin(async { 0 })
    }

    /// The session started or was bound to a mode.
    fn session_start<'a>(&'a self, _ctx: &'a Context) -> BoxFuture<'a, ()> {
        Box::pin(async {})
    }

    /// The session ends or is being replaced.
    fn session_shutdown<'a>(&'a self, _ctx: &'a Context) -> BoxFuture<'a, ()> {
        Box::pin(async {})
    }

    /// Session `session_id` ended or another replaced it, after its
    /// `session_shutdown`, but not on reload. pi disposes the session's
    /// resources then, with pi-ai's `cleanupSessionResources`.
    fn end_session<'a>(&'a self, _session_id: &'a str) -> BoxFuture<'a, ()> {
        Box::pin(async {})
    }

    /// A prompt is about to run. `sections` are extra system prompt sections
    /// by tag name, which the extension may change.
    fn before_agent_start<'a>(
        &'a self,
        _ctx: &'a Context,
        _sections: &'a mut IndexMap<String, String>,
    ) -> BoxFuture<'a, ()> {
        Box::pin(async {})
    }

    /// Whether it handles pi events of type `kind` through [`Extension::handle`].
    fn handles(&self, _kind: &str) -> bool {
        false
    }

    /// Runs its handlers for pi event `event`, a JSON object with a `type`,
    /// and returns their combined result as pi's runner combines one
    /// extension's handlers. The session combines results across extensions.
    fn handle<'a>(&'a self, _ctx: &'a Context, _event: &'a Value) -> BoxFuture<'a, Option<Value>> {
        Box::pin(async { None })
    }

    /// pi's `prepareLoadout` hook of a tool it registered: how the next
    /// request presents the declared tools. Runs before every request.
    fn prepare_loadout(&self, _loadout: &Loadout) -> LoadoutChanges {
        LoadoutChanges::default()
    }

    /// Runs before a request's [`Extension::prepare_loadout`] for
    /// `loadout`, for hooks that run elsewhere and answer later, so that
    /// `prepare_loadout` can give what they made of it.
    fn refresh_loadout<'a>(&'a self, _loadout: &'a Loadout) -> BoxFuture<'a, ()> {
        Box::pin(async {})
    }

    /// What it draws in the transcript.
    fn renderers(&self) -> Renderers {
        Renderers::default()
    }

    /// Builds the component for a transcript item it draws. `request` is
    /// `{"kind": "toolCall" | "toolResult", "name", "toolCallId", "args",
    /// "result", "options", "context"}`, `{"kind": "message", "key",
    /// "message", "options"}` or `{"kind": "entry", "key", "entry",
    /// "options"}`. `None` keeps the built-in rendering, which shows nothing
    /// for an entry.
    fn component<'a>(&'a self, _request: &'a Value) -> BoxFuture<'a, Option<RemoteComponent>> {
        Box::pin(async { None })
    }

    /// The Markdown its transformer makes of `request`, `{"markdown",
    /// "context"}` with pi's `MarkdownTransformContext`. `Ok(None)` keeps
    /// the Markdown, as when the transformer fails or returns no string.
    fn transform_markdown<'a>(
        &'a self,
        _request: &'a Value,
    ) -> BoxFuture<'a, Result<Option<String>, String>> {
        Box::pin(async { Ok(None) })
    }

    /// The lines its code renderer draws for `request`, `{"code", "lang",
    /// "width"}` with `lang` possibly `null`, as text with escape sequences.
    /// `Ok(None)` leaves the code to yapi, and an error is the renderer's.
    fn render_code<'a>(
        &'a self,
        _request: &'a Value,
    ) -> BoxFuture<'a, Result<Option<Vec<String>>, String>> {
        Box::pin(async { Ok(None) })
    }

    /// A tool is about to run. `Some(reason)` blocks it.
    fn tool_call<'a>(
        &'a self,
        _ctx: &'a Context,
        _tool: &'a str,
        _input: &'a Value,
    ) -> BoxFuture<'a, Option<String>> {
        Box::pin(async { None })
    }

    /// Runs `command` in `cwd` through the operations its `user_bash`
    /// handler returned as `handle`, sending output to `output`, until the
    /// command ends or `cancel` fires, and returns the exit code.
    fn run_bash<'a>(
        &'a self,
        _handle: &'a Value,
        _command: &'a str,
        _cwd: &'a std::path::Path,
        _output: tokio::sync::mpsc::UnboundedSender<Vec<u8>>,
        _cancel: CancellationToken,
    ) -> BoxFuture<'a, Result<Option<i32>, String>> {
        Box::pin(async { Err("This extension has no bash operations".to_owned()) })
    }

    /// Refreshes the model lists of the providers it registered with
    /// `refreshModels` that `options` selects, as pi does: first offline
    /// with the registry's runtime key or stored credential
    /// ([`ModelRegistry::read_credential`]), then, when `options` allow the network
    /// and a credential resolves, online with it. Then those with the legacy
    /// `oauth.modifyModels` project their models for the stored credential,
    /// and pi-ai `Provider`s check their API key and filter their models for
    /// it. Returns each step's update, or why it failed, by provider in order.
    fn refresh_models<'a>(
        &'a self,
        _registry: &'a ModelRegistry,
        _options: &'a RefreshOptions,
    ) -> BoxFuture<'a, Vec<(String, ModelList)>> {
        Box::pin(async { Vec::new() })
    }
}

/// pi's `newSession` `setup`: runs with the replacement session once it is
/// built, before its extensions start.
pub type Setup =
    Box<dyn FnOnce(crate::agent_session::AgentSession) -> BoxFuture<'static, ()> + Send>;

/// A session change an extension command asks for, one of pi's
/// `ExtensionCommandContextActions`.
pub enum SessionAction {
    /// `newSession`, linked to a parent session file when given.
    New {
        /// The parent session file.
        parent: Option<String>,
        /// What to write to the new session before it starts.
        setup: Option<Setup>,
    },
    /// `fork` at an entry. `at` keeps the entry, as `/clone` does.
    Fork {
        /// The entry.
        entry_id: String,
        /// Keep the entry.
        at: bool,
    },
    /// `navigateTree` to an entry.
    Tree {
        /// The entry.
        target_id: String,
        /// How.
        options: crate::agent_session::TreeNavigation,
    },
    /// `switchSession` to a session file.
    Switch {
        /// The file, as the extension gave it.
        path: String,
    },
    /// `reload`.
    Reload,
    /// The extension finished its `withSession` work after a replacement.
    /// The mode reports the change now, as pi does once it returns.
    Replaced,
}

/// How a mode carries out [`SessionAction`]s: whether an extension
/// cancelled the change, or why it failed.
pub type SessionActions =
    Arc<dyn Fn(SessionAction) -> BoxFuture<'static, Result<bool, String>> + Send + Sync>;

/// pi's `BashOperations`: operations a `user_bash` handler returned to run a
/// `!` command with.
#[derive(Clone)]
pub struct BashOperations {
    /// The extension whose handler returned them.
    pub extension: Arc<dyn Extension>,
    /// What the extension knows them by.
    pub handle: Value,
}

/// A provider's new models from a refresh, or why it failed.
pub type ModelList = Result<ModelUpdate, String>;

/// A provider's new models from a refresh.
#[derive(Debug)]
pub enum ModelUpdate {
    /// The list its `refreshModels` returned.
    Definitions(Vec<ModelDefinition>),
    /// Its chat models as its legacy `oauth.modifyModels` made them for the
    /// stored OAuth credential, or `None` without one.
    Modified(Option<Vec<Model>>),
    /// What its API key `check`, `filterModels` and `filterAllModels` make
    /// available for the stored credential.
    Availability(Availability),
}

/// The built-in extensions each session starts with.
pub fn builtins() -> Vec<Arc<dyn Extension>> {
    vec![
        Arc::new(tool_search::ToolSearchExtension),
        Arc::new(crate::mcp::extension::McpExtension::new()),
    ]
}

/// pi's source info for a built-in extension named `name`.
pub fn builtin_source(name: &str) -> SourceInfo {
    SourceInfo {
        path: format!("builtin:{name}"),
        source: "builtin".into(),
        scope: "temporary".into(),
        origin: "top-level".into(),
        base_dir: None,
    }
}