yapi_core/agent_session/

extensions.rs

//! Extension dispatch: pi events, input and command handlers, and the
//! agent loop hooks that deliver loop events to extensions.

use std::sync::Arc;
use std::sync::atomic::Ordering;

use futures_util::future::BoxFuture;
use serde::Deserialize;
use serde_json::Value;
use tokio_util::sync::CancellationToken;
use yapi_agent::Tool;
use yapi_agent::hooks::AgentHooks;
use yapi_ai::registry::Auth;
use yapi_types::event::AgentEvent;
use yapi_types::extension_event::{ExtensionEvent, RegisteredMcpServer, UiPrompt};
use yapi_types::message::{
    Content, ContentBlock, ImageContent, Message, StopReason, SystemMessage, ToolResultMessage,
};
use yapi_types::model::Model;
use yapi_types::rpc::StreamingBehavior;
use yapi_types::rpc::{CommandSource, SlashCommand, SourceInfo};
use yapi_types::sync::{lock, write};
use yapi_types::system_prompt::SystemPromptOptions;

use crate::extensions::{
    Context, Extension, ExtensionUi, Mode, SessionActions, ToolRenderers, receives, tool_granted,
};
use crate::mcp::config::check_registration;
use crate::messages::convert_to_llm;
use crate::time::now_ms;

use super::{AgentSession, InputSource, Replacement, SessionChange, drain};

/// What a handler result that changes the session's conversation needs.
pub(super) const SESSION: &str = "the session grant";

/// The fields of a `tool_result` handler's result that change the result.
const TOOL_RESULT_KEYS: [&str; 5] = [
    "content",
    "details",
    "structuredContent",
    "isError",
    "usage",
];

pub(super) struct Hooks {
    pub(super) session: AgentSession,
    pub(super) cancel: CancellationToken,
    /// pi's turn index for extension events: turns since `agent_start`.
    pub(super) turn_index: std::sync::atomic::AtomicU64,
}

fn ui_prompt<'a>(kind: &'a str, title: Option<&'a str>) -> UiPrompt<'a> {
    UiPrompt {
        reason: "ui_prompt",
        kind,
        title,
    }
}

/// A loop event as pi's extension event (`_emitExtensionEvent` in pi's agent
/// session), when `wanted` takes its type, or `None` for events extensions
/// do not see.
fn extension_event(
    event: &AgentEvent,
    turn_index: u64,
    wanted: impl Fn(&str) -> bool,
) -> Option<Value> {
    let (kind, keys): (&str, &[&str]) = match event {
        AgentEvent::AgentStart => ("agent_start", &[]),
        AgentEvent::AgentEnd { .. } => ("agent_end", &["messages"]),
        AgentEvent::TurnStart => ("turn_start", &[]),
        // `turn_end` reaches extensions as a boundary, before the event.
        AgentEvent::MessageStart { .. } => ("message_start", &["message"]),
        AgentEvent::MessageUpdate { .. } => {
            ("message_update", &["message", "assistantMessageEvent"])
        }
        // `message_end` reaches extensions as the message is finished.
        AgentEvent::ToolExecutionStart { .. } => (
            "tool_execution_start",
            &["toolCallId", "toolName", "args", "parentToolCallId"],
        ),
        AgentEvent::ToolExecutionUpdate { .. } => (
            "tool_execution_update",
            &[
                "toolCallId",
                "toolName",
                "args",
                "partialResult",
                "parentToolCallId",
            ],
        ),
        AgentEvent::ToolExecutionEnd { .. } => (
            "tool_execution_end",
            &[
                "toolCallId",
                "toolName",
                "result",
                "isError",
                "parentToolCallId",
            ],
        ),
        _ => return None,
    };
    if !wanted(kind) {
        return None;
    }
    let value = serde_json::to_value(event).ok()?;
    let mut out = serde_json::Map::new();
    out.insert("type".into(), Value::String(kind.to_owned()));
    if kind == "turn_start" {
        out.insert("turnIndex".into(), Value::from(turn_index));
        out.insert("timestamp".into(), Value::from(now_ms()));
    }
    for key in keys {
        if let Some(field) = value.get(*key) {
            out.insert((*key).to_owned(), field.clone());
        }
    }
    Some(Value::Object(out))
}

impl AgentHooks for Hooks {
    fn on_event<'a>(&'a self, event: &'a AgentEvent) -> BoxFuture<'a, ()> {
        Box::pin(async move {
            let session = &self.session;
            if let AgentEvent::MessageStart { message } = event {
                session.dequeue_started(message);
            }
            // pi's interactive mode rewrites an aborted response's error before
            // the session records it.
            let rewritten = match event {
                AgentEvent::MessageEnd {
                    message: Message::Assistant(assistant),
                } if assistant.stop_reason == StopReason::Aborted
                    && lock(&session.inner.binding).1 == Mode::Tui =>
                {
                    let attempt = lock(&session.inner.recovery).retry_attempt;
                    let mut assistant = assistant.clone();
                    assistant.error_message = Some(if attempt > 0 {
                        let plural = if attempt > 1 { "s" } else { "" };
                        format!("Aborted after {attempt} retry attempt{plural}")
                    } else {
                        "Operation aborted".to_owned()
                    });
                    Some(AgentEvent::MessageEnd {
                        message: Message::Assistant(assistant),
                    })
                }
                _ => None,
            };
            let event = rewritten.as_ref().unwrap_or(event);
            if matches!(event, AgentEvent::AgentStart) {
                self.turn_index.store(0, Ordering::SeqCst);
            }
            let turn_index = self.turn_index.load(Ordering::SeqCst);
            if matches!(event, AgentEvent::TurnEnd { .. }) {
                self.turn_index.fetch_add(1, Ordering::SeqCst);
            }
            // As in pi, extensions see the event first, then listeners. The
            // session records it last, so entries extensions append come before
            // it.
            if let Some(extension_event) =
                extension_event(event, turn_index, |kind| session.has_handlers(kind))
            {
                session
                    .emit_extension_event(&extension_event, self.cancel.clone())
                    .await;
            }
            match event {
                AgentEvent::AgentEnd { messages, .. } => {
                    session.inner.nested.clear();
                    session.emit(&AgentEvent::AgentEnd {
                        messages: messages.clone(),
                        will_retry: session.will_retry_after(messages),
                    });
                }
                _ => session.emit(event),
            }
            match event {
                AgentEvent::MessageEnd { message } => {
                    let entry_id = session.with_session(|file| match message {
                        Message::Custom(custom) => file
                            .append_custom_message(
                                &custom.custom_type,
                                custom.content.clone(),
                                custom.display,
                                custom.details.clone(),
                            )
                            .ok(),
                        Message::System(_)
                        | Message::User(_)
                        | Message::Assistant(_)
                        | Message::ToolResult(_) => file.append_message(message.clone()).ok(),
                        _ => None,
                    });
                    match message {
                        Message::Assistant(assistant) => {
                            let finished_retry = {
                                let mut recovery = lock(&session.inner.recovery);
                                recovery.last_assistant = Some(((**assistant).clone(), entry_id));
                                recovery.turn_tool_results.clear();
                                if !matches!(
                                    assistant.stop_reason,
                                    StopReason::Error | StopReason::Length
                                ) {
                                    recovery.overflow_recovery_attempted = false;
                                }
                                if assistant.stop_reason != StopReason::Error {
                                    std::mem::take(&mut recovery.retry_attempt)
                                } else {
                                    0
                                }
                            };
                            if finished_retry > 0 {
                                session.emit(&AgentEvent::AutoRetryEnd {
                                    success: true,
                                    attempt: finished_retry,
                                    final_error: None,
                                });
                            }
                        }
                        Message::ToolResult(_) => {
                            if let Some(id) = entry_id {
                                lock(&session.inner.recovery).turn_tool_results.push(id);
                            }
                        }
                        _ => {}
                    }
                }
                AgentEvent::TurnEnd { .. } => {
                    let mut recovery = lock(&session.inner.recovery);
                    recovery.last_tool_results = std::mem::take(&mut recovery.turn_tool_results);
                }
                _ => {}
            }
            // Custom messages sent during the turn are appended once its handlers
            // ran, so those that turn_end handlers send join them, as in pi.
            if matches!(event, AgentEvent::TurnEnd { .. }) {
                session.flush_pending_custom();
            }
        })
    }

    fn transform_context(&self, messages: Vec<Message>) -> BoxFuture<'_, Vec<Message>> {
        Box::pin(async move {
            let session = &self.session;
            let mut messages = session
                .context_handlers(messages, self.cancel.clone())
                .await;
            // pi's hidden declarations projection: requests leave the tools
            // the loadout hides out of every declaration change.
            let hidden = lock(&session.inner.hidden_tools).clone();
            if !hidden.is_empty() {
                for message in &mut messages {
                    if let Message::System(system) = message {
                        system.tools_added = system.tools_added.take().and_then(|mut tools| {
                            tools.retain(|tool| !hidden.contains(&tool.name));
                            (!tools.is_empty()).then_some(tools)
                        });
                        system.tools_removed = system.tools_removed.take().and_then(|mut tools| {
                            tools.retain(|tool| !hidden.contains(&tool.name));
                            (!tools.is_empty()).then_some(tools)
                        });
                    }
                }
            }
            let forced = lock(&session.inner.forced_prompt).clone();
            if let Some(forced) = forced {
                // pi's forced prompt projection: one system message with the
                // forced text replaces the transcript's system messages.
                let current = yapi_ai::transcript::current_system_message(&messages);
                let head = Message::System(SystemMessage {
                    content: Content::Text(forced),
                    sections: None,
                    timestamp: current
                        .as_ref()
                        .map_or_else(now_ms, |system| system.timestamp),
                    tools_added: current.and_then(|system| system.tools_added),
                    tools_removed: None,
                });
                messages.retain(|message| !matches!(message, Message::System(_)));
                messages.insert(0, head);
            }
            messages
        })
    }

    fn convert_to_llm(&self, messages: Vec<Message>) -> Vec<Message> {
        let messages = convert_to_llm(messages);
        // Read on every request, so a change applies mid-session.
        let blocked = self
            .session
            .settings()
            .images
            .and_then(|images| images.block_images)
            == Some(true);
        if blocked {
            crate::messages::block_images(messages)
        } else {
            messages
        }
    }

    fn auth<'a>(&'a self, model: &'a Model) -> BoxFuture<'a, Auth> {
        // The registry is read under a short lock. Credential commands run outside it.
        let registry = self.session.registry();
        Box::pin(async move { registry.auth(model).await })
    }

    fn current_tools(&self) -> BoxFuture<'_, Option<Vec<Arc<dyn Tool>>>> {
        Box::pin(async move {
            let declared = self
                .session
                .inner
                .tools
                .with(|registry| registry.declared());
            self.session.refresh_loadouts(declared.clone()).await;
            Some(self.session.loadout(declared))
        })
    }

    fn before_tool_call<'a>(
        &'a self,
        call: yapi_agent::hooks::BeforeToolCall<'a>,
    ) -> BoxFuture<'a, Option<yapi_agent::hooks::Block>> {
        Box::pin(async move {
            let ctx = self.session.extension_context(self.cancel.clone());
            for extension in &self.session.inner.extensions {
                if let Some(reason) = extension
                    .tool_call(&ctx, &call.tool_call.name, &*call.args)
                    .await
                {
                    let granted = extension.grants().session;
                    return Some(yapi_agent::hooks::Block {
                        reason: self
                            .session
                            .applies(extension, "tool_call", granted, SESSION)
                            .then_some(reason),
                        terminate: false,
                    });
                }
            }
            let mut event = ExtensionEvent::ToolCall {
                tool_name: &call.tool_call.name,
                tool_call_id: &call.tool_call.id,
                parent_tool_call_id: call.parent_tool_call_id,
                input: call.args,
            }
            .to_value();
            // Handlers change the call's arguments by editing `event.input`,
            // which later handlers see, as in pi. The arguments are what the
            // model sees, and a handler changes them only when its grants
            // would also let it run the tool itself.
            let name = &call.tool_call.name;
            let tool = self
                .session
                .inner
                .tools
                .with(|registry| registry.get(name).cloned());
            for extension in self.session.handlers_of("tool_call") {
                let Some(result) = extension.handle(&ctx, &event).await else {
                    continue;
                };
                // The runtime reports the input whether or not a handler changed it.
                if let Some(input) = result
                    .get("input")
                    .filter(|input| input.is_object() && **input != event["input"])
                {
                    let grants = extension.grants();
                    let granted = grants.session
                        && tool
                            .as_ref()
                            .map_or(grants.is_all(), |tool| tool_granted(&grants, tool));
                    let missing =
                        format!("the session grant and the grants to run the {name} tool");
                    if self
                        .session
                        .applies(extension, "tool_call", granted, &missing)
                    {
                        event["input"] = input.clone();
                    }
                }
                // A block always applies, but its reason is text the model
                // reads, which needs the session grant.
                if result["block"] == true {
                    let granted = extension.grants().session;
                    let reason = result["reason"].as_str().filter(|_| {
                        self.session
                            .applies(extension, "tool_call", granted, SESSION)
                    });
                    return Some(yapi_agent::hooks::Block {
                        reason: reason.map(str::to_owned),
                        terminate: result["terminate"] == true,
                    });
                }
            }
            *call.args = event["input"].take();
            None
        })
    }

    fn after_tool_call<'a>(
        &'a self,
        call: yapi_agent::hooks::AfterToolCall<'a>,
    ) -> BoxFuture<'a, Option<yapi_agent::hooks::ResultPatch>> {
        Box::pin(async move {
            let mut patch = self.tool_result_handlers(&call).await;
            // After the handlers, so images they add are normalized too.
            let content = patch
                .as_ref()
                .and_then(|patch| patch.content.as_ref())
                .unwrap_or(&call.result.content);
            if content
                .iter()
                .any(|block| matches!(block, ContentBlock::Image(_)))
            {
                let content = content.clone();
                let (auto_resize, limits) = self.session.image_options();
                let normalized = tokio::task::spawn_blocking(move || {
                    crate::images::normalize_tool_result(&content, auto_resize, limits.as_ref())
                })
                .await
                .ok()
                .flatten();
                if let Some(content) = normalized {
                    // New images alone keep the structured content, as in pi.
                    patch
                        .get_or_insert_with(|| yapi_agent::hooks::ResultPatch {
                            structured_content: call.result.structured_content.clone(),
                            ..Default::default()
                        })
                        .content = Some(content);
                }
            }
            patch
        })
    }

    fn finish_turn<'a>(
        &'a self,
        message: &'a yapi_types::message::AssistantMessage,
        tool_results: &'a [ToolResultMessage],
    ) -> BoxFuture<'a, yapi_agent::hooks::TurnDecision> {
        Box::pin(async move {
            let turn_index = self.turn_index.load(Ordering::SeqCst);
            self.session
                .turn_end_boundary(message, tool_results, turn_index, self.cancel.clone())
                .await
        })
    }

    fn finish_message(&self, message: Message) -> BoxFuture<'_, Message> {
        Box::pin(async move {
            self.session
                .message_end_handlers(message, self.cancel.clone())
                .await
        })
    }

    fn complete_tool_result(&self, message: &mut ToolResultMessage) {
        let Some(summary) = self.session.inner.nested.take(&message.tool_call_id) else {
            return;
        };
        if summary.calls.is_some() {
            message.nested_calls = summary.calls;
        }
        if let Some(usage) = summary.usage {
            message.usage = Some(match &message.usage {
                Some(own) => own.combine(&usage),
                None => usage,
            });
        }
    }

    fn next_turn_messages(&self) -> BoxFuture<'_, Result<Option<Vec<Message>>, String>> {
        // As in pi, each turn continues from the session's transcript, with
        // what the session recorded since the last, such as custom messages
        // held while it ran and entries turn_end handlers committed, once a
        // context over the threshold is compacted.
        Box::pin(async move {
            self.session
                .compact_before_next_response(&self.cancel)
                .await?;
            Ok(Some(self.session.messages()))
        })
    }

    // As pi's agent, each drain reads the current mode, so a change applies
    // to the running agent.
    fn steering_messages(&self) -> BoxFuture<'_, Vec<Message>> {
        let messages = drain(&self.session.inner.steering, self.session.drains_all().0);
        Box::pin(async move { messages })
    }

    fn follow_up_messages(&self) -> BoxFuture<'_, Vec<Message>> {
        let messages = drain(&self.session.inner.follow_up, self.session.drains_all().1);
        Box::pin(async move { messages })
    }
}

impl Hooks {
    /// pi's `tool_result` event: the changes its handlers make, if any.
    async fn tool_result_handlers(
        &self,
        call: &yapi_agent::hooks::AfterToolCall<'_>,
    ) -> Option<yapi_agent::hooks::ResultPatch> {
        let handlers = self.session.handlers_of("tool_result");
        if handlers.is_empty() {
            return None;
        }
        let ctx = self.session.extension_context(self.cancel.clone());
        let mut event = ExtensionEvent::ToolResult {
            tool_name: &call.tool_call.name,
            tool_call_id: &call.tool_call.id,
            parent_tool_call_id: call.parent_tool_call_id,
            input: call.args,
            content: &call.result.content,
            details: call.result.details.as_ref(),
            structured_content: call.result.structured_content.as_ref(),
            is_error: call.is_error,
            usage: call.result.usage.as_ref(),
        }
        .to_value();
        let mut modified = false;
        for extension in handlers {
            let Some(result) = extension.handle(&ctx, &event).await else {
                continue;
            };
            let given = |key: &str| result.get(key).filter(|value| !value.is_null());
            let changes = TOOL_RESULT_KEYS.iter().any(|key| given(key).is_some());
            if changes
                && !self.session.applies(
                    extension,
                    "tool_result",
                    extension.grants().session,
                    SESSION,
                )
            {
                continue;
            }
            // As in pi, structured content not replaced along with the
            // content may no longer match it.
            if given("content").is_some()
                && given("structuredContent").is_none()
                && let Some(event) = event.as_object_mut()
            {
                event.shift_remove("structuredContent");
            }
            for key in TOOL_RESULT_KEYS {
                if let Some(value) = given(key)
                    && let Some(event) = event.as_object_mut()
                {
                    set_tool_result_key(event, key, value.clone());
                    modified = true;
                }
            }
        }
        modified.then(|| yapi_agent::hooks::ResultPatch {
            content: serde_json::from_value(event["content"].clone()).ok(),
            details: Some(event["details"].clone()).filter(|details| !details.is_null()),
            structured_content: event.get("structuredContent").cloned(),
            is_error: event["isError"].as_bool(),
            usage: serde_json::from_value(event["usage"].clone()).ok(),
            terminate: None,
        })
    }
}

/// Whether `extension` handles pi event `kind` and its grants let it hear it.
fn hears(extension: &Arc<dyn Extension>, kind: &str) -> bool {
    extension.handles(kind) && receives(kind, &extension.grants())
}

/// Sets `key` of a `tool_result` event where pi's event has it: pi's holds
/// `details` and `usage` even when undefined, after `content` and `isError`.
fn set_tool_result_key(event: &mut serde_json::Map<String, Value>, key: &str, value: Value) {
    if let Some(slot) = event.get_mut(key) {
        *slot = value;
        return;
    }
    let after = match key {
        "details" => "content",
        "usage" => "isError",
        _ => "",
    };
    let index = event
        .keys()
        .position(|name| name == after)
        .map_or(event.len(), |index| index + 1);
    event.shift_insert(index, key.to_owned(), value);
}

/// The servers in `registered` that runtimes other than `owner` registered.
fn others(
    registered: &indexmap::IndexMap<String, (u64, RegisteredMcpServer)>,
    owner: u64,
) -> impl Iterator<Item = &RegisteredMcpServer> + Clone {
    registered
        .values()
        .filter(move |(runtime, _)| *runtime != owner)
        .map(|(_, server)| server)
}

impl AgentSession {
    /// The context extensions get, with `cancel` for the operation at hand.
    pub(super) fn extension_context(&self, cancel: CancellationToken) -> Context {
        let (ui, mode, _) = lock(&self.inner.binding).clone();
        let (project_dir, project_trusted) = {
            let settings = lock(&self.inner.settings);
            (
                settings.project_dir().to_path_buf(),
                settings.project_trusted(),
            )
        };
        Context {
            cwd: self.inner.cwd.clone(),
            agent_dir: self.inner.agent_dir.clone(),
            project_dir,
            project_trusted,
            mode,
            ui,
            tools: self.inner.tools.clone(),
            cancel,
            session: self.downgrade(),
        }
    }

    /// The UI and mode extensions are bound to.
    pub fn extension_binding(&self) -> (Arc<dyn ExtensionUi>, Mode) {
        let (ui, mode, _) = lock(&self.inner.binding).clone();
        (ui, mode)
    }

    /// Whether a result of `extension`'s `event` handler applies: always
    /// when `granted`, else never, and the extension is warned once per
    /// event that the result needs `missing`, such as the session grant.
    pub(super) fn applies(
        &self,
        extension: &Arc<dyn Extension>,
        event: &str,
        granted: bool,
        missing: &str,
    ) -> bool {
        if !granted {
            self.dropped(extension, event, missing);
        }
        granted
    }

    /// Warns `extension`, once per event, that the session ignored a result
    /// of its `event` handler, which needs `missing`.
    pub(super) fn dropped(&self, extension: &Arc<dyn Extension>, event: &str, missing: &str) {
        let path = extension.source().path;
        if lock(&self.inner.dropped_results).insert((path.clone(), event.to_owned())) {
            extension.warn(&format!(
                "Ignored a {event} result of extension {path}: it needs {missing}"
            ));
        }
    }

    pub(super) fn handlers_of(&self, kind: &str) -> Vec<&Arc<dyn Extension>> {
        self.inner
            .extensions
            .iter()
            .filter(|extension| hears(extension, kind))
            .collect()
    }

    /// Whether any extension handles pi events of type `kind` and may hear them.
    pub fn has_handlers(&self, kind: &str) -> bool {
        self.inner
            .extensions
            .iter()
            .any(|extension| hears(extension, kind))
    }

    /// Delivers pi event `event` to every extension that handles it, in
    /// order, as pi's runner `emit` does. Only a `session_before_*` event has
    /// a result: the last one a handler gave, or the first that cancels,
    /// which stops delivery.
    pub async fn emit_extension_event(
        &self,
        event: &Value,
        cancel: CancellationToken,
    ) -> Option<Value> {
        let kind = event["type"].as_str().unwrap_or_default();
        let handlers = self.handlers_of(kind);
        if handlers.is_empty() {
            return None;
        }
        let ctx = self.extension_context(cancel);
        let before = kind.starts_with("session_before_");
        let mut result = None;
        for extension in handlers {
            let Some(mut next) = extension.handle(&ctx, event).await else {
                continue;
            };
            if !before {
                continue;
            }
            let cancelled = next["cancel"] == true;
            // Without the session grant, a handler may only cancel, such as
            // a compaction it would otherwise summarize itself. Anything else
            // counts as no result, so it never replaces a granted one.
            if !extension.grants().session {
                let changes = next
                    .as_object()
                    .is_some_and(|result| result.keys().any(|key| key != "cancel"));
                if changes {
                    self.dropped(extension, kind, SESSION);
                }
                if !cancelled {
                    continue;
                }
                next = serde_json::json!({"cancel": true});
            }
            result = Some(next);
            if cancelled {
                break;
            }
        }
        result
    }

    /// pi's `session_before_switch` or `session_before_fork`: whether an
    /// extension cancels `change`.
    pub async fn cancels(&self, change: &SessionChange) -> bool {
        let reason = change.reason().as_str();
        let event = match change {
            SessionChange::New => ExtensionEvent::SessionBeforeSwitch {
                reason,
                target_session_file: None,
            },
            SessionChange::Resume(target) => ExtensionEvent::SessionBeforeSwitch {
                reason,
                target_session_file: Some(target),
            },
            SessionChange::Fork { entry_id, at } => ExtensionEvent::SessionBeforeFork {
                entry_id,
                position: if *at { "at" } else { "before" },
            },
        };
        let result = self
            .emit_extension_event(&event.to_value(), CancellationToken::new())
            .await;
        result.is_some_and(|result| result["cancel"] == true)
    }

    /// pi's `message_end`: the message as handlers replaced it, each seeing
    /// the previous one's. A replacement of another role is reported and
    /// skipped, and a missing content of a user, assistant, tool result or
    /// custom message becomes empty.
    async fn message_end_handlers(&self, message: Message, cancel: CancellationToken) -> Message {
        let handlers = self.handlers_of("message_end");
        if handlers.is_empty() {
            return message;
        }
        let ctx = self.extension_context(cancel);
        let report = |extension: &Arc<dyn Extension>, error: &str| {
            ctx.ui
                .extension_error(&extension.source().path, "message_end", error, None);
        };
        let mut current = serde_json::json!(message);
        let mut replaced_by = None;
        for extension in handlers {
            let event = ExtensionEvent::MessageEnd { message: &current }.to_value();
            let Some(next) = extension
                .handle(&ctx, &event)
                .await
                .and_then(|mut result| result.get_mut("message").map(Value::take))
                .filter(|next| !next.is_null())
            else {
                continue;
            };
            if !self.applies(
                extension,
                "message_end",
                extension.grants().session,
                SESSION,
            ) {
                continue;
            }
            if next["role"] != current["role"] {
                report(
                    extension,
                    "message_end handlers must return a message with the same role",
                );
                continue;
            }
            current = next;
            replaced_by = Some(extension);
        }
        let Some(extension) = replaced_by else {
            return message;
        };
        if matches!(
            current["role"].as_str(),
            Some("user" | "assistant" | "toolResult" | "custom")
        ) && current["content"].is_null()
        {
            current["content"] = Value::Array(Vec::new());
        }
        serde_json::from_value(current).unwrap_or_else(|error| {
            report(extension, &error.to_string());
            message
        })
    }

    /// Passes `value` through the extensions handling `kind` as the event's
    /// `field`, each seeing the `field` the previous one returned.
    pub(super) async fn chain(
        &self,
        kind: &str,
        field: &str,
        value: Value,
        cancel: CancellationToken,
    ) -> Value {
        let handlers = self.handlers_of(kind);
        if handlers.is_empty() {
            return value;
        }
        let ctx = self.extension_context(cancel);
        let mut value = value;
        for extension in handlers {
            let event = serde_json::json!({"type": kind, field: value});
            if let Some(result) = extension.handle(&ctx, &event).await
                && let Some(next) = result.get(field)
                && *next != value
                && self.applies(extension, kind, extension.grants().session, SESSION)
            {
                value = next.clone();
            }
        }
        value
    }

    /// pi's `emitContext`: `context` handlers see the conversation without
    /// system messages, which come back as the current prompt when they
    /// change it. `context_with_system` handlers then see and return the
    /// whole transcript.
    async fn context_handlers(
        &self,
        messages: Vec<Message>,
        cancel: CancellationToken,
    ) -> Vec<Message> {
        let mut messages = messages;
        let ctx = self.extension_context(cancel);
        for extension in self.handlers_of("context") {
            let visible: Vec<Message> = messages
                .iter()
                .filter(|message| !matches!(message, Message::System(_)))
                .cloned()
                .collect();
            let event = ExtensionEvent::Context { messages: &visible }.to_value();
            if let Some(result) = extension.handle(&ctx, &event).await
                && let Ok(returned) =
                    serde_json::from_value::<Vec<Message>>(result["messages"].clone())
                && returned != visible
                && self.applies(extension, "context", extension.grants().session, SESSION)
            {
                let head = yapi_ai::transcript::current_system_message(&messages);
                messages = head
                    .map(Message::System)
                    .into_iter()
                    .chain(returned)
                    .collect();
            }
        }
        for extension in self.handlers_of("context_with_system") {
            let had_head = matches!(messages.first(), Some(Message::System(_)));
            let event = ExtensionEvent::ContextWithSystem {
                messages: &messages,
            }
            .to_value();
            if let Some(result) = extension.handle(&ctx, &event).await
                && let Ok(returned) =
                    serde_json::from_value::<Vec<Message>>(result["messages"].clone())
                && returned != messages
                && self.applies(
                    extension,
                    "context_with_system",
                    extension.grants().session,
                    SESSION,
                )
            {
                messages = returned;
            }
            if had_head && !matches!(messages.first(), Some(Message::System(_))) {
                ctx.ui.extension_error(
                    &extension.source().path,
                    "context_with_system",
                    "Handler removed the leading system message; the request has no prompt or initial tool declarations. Keep it at index 0 or replace a dropped prefix with getCurrentSystemMessage().",
                    None,
                );
            }
        }
        messages
    }

    /// pi's `before_provider_headers`: the headers after handlers changed
    /// them in place. A null value removes a header.
    pub(super) async fn before_provider_headers(
        &self,
        headers: yapi_ai::stream::RequestHeaders,
        cancel: CancellationToken,
    ) -> yapi_ai::stream::RequestHeaders {
        let value = serde_json::to_value(&headers).unwrap_or_default();
        let value = self
            .chain("before_provider_headers", "headers", value, cancel)
            .await;
        match value.as_object() {
            Some(headers) => headers
                .iter()
                .map(|(name, value)| (name.clone(), value.as_str().map(str::to_owned)))
                .collect(),
            None => headers,
        }
    }

    /// A blocking extension dialog of `kind` (`select`, `confirm`, `input`,
    /// `editor` or `custom`) with `title` opens. Extensions hear of the
    /// outermost as pi's `ui_prompt_start`.
    pub fn ui_prompt_opened(&self, kind: &str, title: Option<&str>) {
        let title = title.filter(|title| !title.is_empty()).map(str::to_owned);
        {
            let mut prompts = lock(&self.inner.ui_prompts);
            prompts.0 += 1;
            if prompts.0 > 1 {
                return;
            }
            prompts.1 = Some((kind.to_owned(), title.clone()));
        }
        self.announce(ExtensionEvent::UiPromptStart(ui_prompt(
            kind,
            title.as_deref(),
        )));
    }

    /// A blocking extension dialog closed. When it was the outermost,
    /// extensions hear of it as pi's `ui_prompt_end`.
    pub fn ui_prompt_closed(&self) {
        let (kind, title) = {
            let mut prompts = lock(&self.inner.ui_prompts);
            prompts.0 = prompts.0.saturating_sub(1);
            if prompts.0 > 0 {
                return;
            }
            match prompts.1.take() {
                Some(prompt) => prompt,
                None => return,
            }
        };
        self.announce(ExtensionEvent::UiPromptEnd(ui_prompt(
            &kind,
            title.as_deref(),
        )));
    }

    /// Delivers `event` to extensions in the background, after the events
    /// announced before it, as pi's runner delivers events it does not wait
    /// for. Outside a tokio runtime it waits for the next
    /// [`AgentSession::flush_announcements`].
    pub(super) fn announce(&self, event: ExtensionEvent<'_>) {
        let event = event.to_value();
        let kind = event["type"].as_str().unwrap_or_default();
        if !self.has_handlers(kind) {
            return;
        }
        lock(&self.inner.announcements).push_back(event);
        if let Ok(runtime) = tokio::runtime::Handle::try_current() {
            let session = self.clone();
            runtime.spawn(async move { session.flush_announcements().await });
        }
    }

    /// Sets the MCP servers that extension runtime `owner` registered, and
    /// announces `mcp_servers_change` when they changed after the extensions
    /// started. Before then, as runtimes bind, the reports hold the servers
    /// registered while extensions loaded, which the MCP extension reads
    /// when the session starts. Nothing changes when a server's name, or its
    /// tools' namespace, belongs to another runtime, and the error is pi's
    /// for that registration.
    pub fn set_mcp_servers(
        &self,
        owner: u64,
        servers: Vec<RegisteredMcpServer>,
    ) -> Result<(), String> {
        let changed = {
            let (registered, started) = &mut *lock(&self.inner.mcp_servers);
            for server in &servers {
                check_registration(others(registered, owner), &server.name)?;
            }
            let before = registered.clone();
            // As in pi's registry, a server registered again keeps its place.
            registered.retain(|name, (runtime, _)| {
                *runtime != owner || servers.iter().any(|server| server.name == *name)
            });
            for server in servers {
                registered.insert(server.name.clone(), (owner, server));
            }
            *started && *registered != before
        };
        if changed {
            let servers = self.mcp_servers();
            self.announce(ExtensionEvent::McpServersChange { servers: &servers });
            self.report_unhandled_mcp_servers();
        }
        Ok(())
    }

    /// pi's checks of server `name`, which an extension of runtime `owner`
    /// registers, against the servers of the other runtimes, without
    /// registering it.
    pub fn check_mcp_server(&self, owner: u64, name: &str) -> Result<(), String> {
        check_registration(others(&lock(&self.inner.mcp_servers).0, owner), name)
    }

    /// pi's `reportUnhandledMcpServers`: when no extension handles
    /// `mcp_servers_change`, nothing connects registered servers, and each
    /// is reported once as its extension's error.
    fn report_unhandled_mcp_servers(&self) {
        if self.has_handlers("mcp_servers_change") {
            return;
        }
        let (ui, _) = self.extension_binding();
        for server in self.mcp_servers() {
            if !lock(&self.inner.unhandled_mcp_servers).insert(server.name.clone()) {
                continue;
            }
            let error = format!(
                "MCP server \"{}\" is registered, but no loaded extension connects MCP servers; another extension may have replaced the built-in MCP support",
                server.name
            );
            ui.extension_error(&server.extension_path, "register_mcp_server", &error, None);
        }
    }

    /// pi's `getMcpServers`: the servers extensions registered, in
    /// registration order.
    pub fn mcp_servers(&self) -> Vec<RegisteredMcpServer> {
        lock(&self.inner.mcp_servers)
            .0
            .values()
            .map(|(_, server)| server.clone())
            .collect()
    }

    /// Records the extension at `path`, whose factory failed with `error` as
    /// the extensions loaded for this session. The factories of every
    /// session after the first run again then, as pi loads the extensions
    /// for each runtime.
    pub fn add_extension_load_error(&self, path: String, error: String) {
        lock(&self.inner.extension_load_errors).push((path, error));
    }

    /// The extensions that failed to load for this session, by path, with
    /// pi's message, which pi's interactive mode lists under
    /// `[Extension issues]`.
    pub fn extension_load_errors(&self) -> Vec<(String, String)> {
        lock(&self.inner.extension_load_errors).clone()
    }

    /// Records the extension at `path`, which was not loaded for this
    /// session, with pi's `warning`.
    pub fn add_extension_warning(&self, path: String, warning: String) {
        lock(&self.inner.extension_warnings).push((path, warning));
    }

    /// The extensions that were not loaded for this session, by path, with
    /// pi's warning, which pi's interactive mode lists under
    /// `[Extension issues]`.
    pub fn extension_warnings(&self) -> Vec<(String, String)> {
        lock(&self.inner.extension_warnings).clone()
    }

    /// Waits until the events announced so far have reached extensions.
    pub async fn flush_announcements(&self) {
        let _turn = self.inner.announcing.lock().await;
        loop {
            let next = lock(&self.inner.announcements).pop_front();
            let Some(event) = next else {
                return;
            };
            self.emit_extension_event(&event, CancellationToken::new())
                .await;
        }
    }

    /// pi's `input` event: `None` when a handler handled the input,
    /// otherwise the possibly transformed text and images.
    pub(super) async fn input_handlers(
        &self,
        text: &str,
        images: Vec<ImageContent>,
        source: InputSource,
        streaming: Option<StreamingBehavior>,
    ) -> Option<(String, Vec<ImageContent>)> {
        let handlers = self.handlers_of("input");
        if handlers.is_empty() {
            return Some((text.to_owned(), images));
        }
        let ctx = self.extension_context(CancellationToken::new());
        let mut text = text.to_owned();
        let mut images = images;
        for extension in handlers {
            let event = ExtensionEvent::Input {
                text: &text,
                images: &images,
                source: source.as_str(),
                streaming_behavior: streaming,
            }
            .to_value();
            let Some(result) = extension.handle(&ctx, &event).await else {
                continue;
            };
            match result["action"].as_str() {
                Some("handled") => return None,
                Some("transform")
                    if self.applies(extension, "input", extension.grants().session, SESSION) =>
                {
                    if let Some(next) = result["text"].as_str() {
                        next.clone_into(&mut text);
                    }
                    if let Ok(next) = serde_json::from_value(result["images"].clone()) {
                        images = next;
                    }
                }
                _ => {}
            }
        }
        Some((text, images))
    }

    /// pi's `before_agent_start` event for extensions handling it, with the
    /// run's prompt `options`, which each handler may edit for the next: the
    /// custom messages to send with the prompt, and the system prompt a
    /// handler forced, if any. The edited options become the run's, and
    /// edited `selectedTools` the active tools.
    pub(super) async fn before_agent_start_handlers(
        &self,
        prompt: &str,
        images: &[ImageContent],
        mut options: SystemPromptOptions,
    ) -> (Vec<Message>, Option<String>) {
        let handlers = self.handlers_of("before_agent_start");
        let mut messages = Vec::new();
        let selected = options.selected_tools.clone();
        let ctx = self.extension_context(CancellationToken::new());
        for extension in handlers {
            let system_prompt = self.render_prompt(&options).unwrap_or_default();
            let event = ExtensionEvent::BeforeAgentStart {
                prompt,
                images,
                system_prompt: &system_prompt,
                system_prompt_options: &options,
            }
            .to_value();
            let Some(result) = extension.handle(&ctx, &event).await else {
                continue;
            };
            // Without the session grant, no part of the result applies. The
            // runtime reports the options and the forced prompt whether or not
            // a handler changed them, so only a change is warned of.
            if !extension.grants().session {
                let edited = result["systemPromptOptions"].is_object()
                    && SystemPromptOptions::deserialize(&result["systemPromptOptions"])
                        .is_ok_and(|edited| edited != options);
                let forced = result["systemPrompt"]
                    .as_str()
                    .is_some_and(|prompt| Some(prompt) != options.force_system_prompt.as_deref());
                let added = result["messages"]
                    .as_array()
                    .is_some_and(|messages| !messages.is_empty());
                if edited || forced || added {
                    self.dropped(extension, "before_agent_start", SESSION);
                }
                continue;
            }
            for message in result["messages"].as_array().into_iter().flatten() {
                let content = match &message["content"] {
                    Value::Null => Value::Array(Vec::new()),
                    other => other.clone(),
                };
                messages.push(Message::Custom(yapi_types::message::CustomMessage {
                    custom_type: message["customType"]
                        .as_str()
                        .unwrap_or_default()
                        .to_owned(),
                    content: serde_json::from_value(content).unwrap_or(Content::Blocks(Vec::new())),
                    display: message["display"].as_bool().unwrap_or(false),
                    details: (!message["details"].is_null()).then(|| message["details"].clone()),
                    timestamp: now_ms(),
                }));
            }
            // An edit that is not pi's shape leaves the options as they were.
            if result["systemPromptOptions"].is_object()
                && let Ok(edited) = SystemPromptOptions::deserialize(&result["systemPromptOptions"])
            {
                options = edited;
            }
            if let Some(prompt) = result["systemPrompt"].as_str() {
                options.force_system_prompt = Some(prompt.to_owned());
            }
        }
        if options.selected_tools != selected {
            self.set_active_tools(options.selected_tools.clone());
        }
        let forced = options.force_system_prompt.take();
        *lock(&self.inner.run_options) = Some(options);
        (messages, forced)
    }

    /// Carries out a session change an extension command asks for: whether
    /// an extension cancelled it. Without a mode, nothing changes, as pi's
    /// default actions do.
    pub async fn session_action(
        &self,
        action: crate::extensions::SessionAction,
    ) -> Result<bool, String> {
        let actions = lock(&self.inner.binding).2.clone();
        match actions {
            Some(actions) => actions(action).await,
            None => Ok(false),
        }
    }

    /// Gives extensions their UI, mode and the `actions` that carry out the
    /// session changes their commands ask for, and starts them: pi's
    /// `bindExtensions`, which emits `session_start` with the reason this
    /// session replaced another and that session's file, or `startup` when
    /// `replaced` is `None`.
    pub async fn bind_extensions(
        &self,
        ui: Arc<dyn ExtensionUi>,
        mode: Mode,
        actions: Option<SessionActions>,
        replaced: Option<(Replacement, Option<String>)>,
    ) {
        *lock(&self.inner.binding) = (ui, mode, actions);
        let ctx = self.extension_context(CancellationToken::new());
        for extension in &self.inner.extensions {
            extension.session_start(&ctx).await;
        }
        lock(&self.inner.mcp_servers).1 = true;
        let (reason, previous) = match replaced {
            Some((reason, previous)) => (Some(reason), reason.reported(previous)),
            None => (None, None),
        };
        let event = ExtensionEvent::SessionStart {
            reason: reason.map_or("startup", Replacement::as_str),
            previous_session_file: previous.as_deref(),
        };
        self.emit_extension_event(&event.to_value(), CancellationToken::new())
            .await;
        self.report_unhandled_mcp_servers();
        let reload = reason == Some(Replacement::Reload);
        self.discover_resources(if reload { "reload" } else { "startup" })
            .await;
    }

    /// pi's `resources_discover`, after `session_start`: adds the skills,
    /// prompt templates and themes extensions name, each with its extension
    /// as source.
    async fn discover_resources(&self, reason: &str) {
        let handlers = self.handlers_of("resources_discover");
        if handlers.is_empty() {
            return;
        }
        let event = ExtensionEvent::ResourcesDiscover {
            cwd: &self.inner.cwd,
            reason,
        }
        .to_value();
        let ctx = self.extension_context(CancellationToken::new());
        let mut found: [Vec<SourceInfo>; 3] = Default::default();
        for extension in handlers {
            let Some(result) = extension.handle(&ctx, &event).await else {
                continue;
            };
            // The files it names load into the prompt and the commands.
            let grants = extension.grants();
            let names = ["skillPaths", "promptPaths", "themePaths"]
                .iter()
                .any(|kind| {
                    result[*kind]
                        .as_array()
                        .is_some_and(|paths| !paths.is_empty())
                });
            if names
                && !self.applies(
                    extension,
                    "resources_discover",
                    grants.session && grants.filesystem,
                    "the session and file grants",
                )
            {
                continue;
            }
            let origin = extension.source().path;
            let synthetic = origin.starts_with("builtin:") || origin.starts_with('<');
            let source = if synthetic {
                format!("extension:{}", origin.replace(['<', '>'], ""))
            } else {
                let name = std::path::Path::new(&origin)
                    .file_name()
                    .map(|name| name.to_string_lossy().into_owned())
                    .unwrap_or_default();
                let name = name
                    .strip_suffix(".ts")
                    .or_else(|| name.strip_suffix(".js"))
                    .unwrap_or(&name);
                format!("extension:{name}")
            };
            let base_dir = (!synthetic)
                .then(|| std::path::Path::new(&origin).parent())
                .flatten()
                .map(|dir| dir.display().to_string());
            for (kind, paths) in ["skillPaths", "promptPaths", "themePaths"]
                .iter()
                .zip(found.iter_mut())
            {
                for path in result[*kind].as_array().into_iter().flatten() {
                    let Some(path) = path.as_str() else { continue };
                    let path = crate::tools::path::resolve_to_cwd(path.trim(), &self.inner.cwd);
                    paths.push(SourceInfo {
                        path: path.display().to_string(),
                        source: source.clone(),
                        scope: "temporary".into(),
                        origin: "top-level".into(),
                        base_dir: base_dir.clone(),
                    });
                }
            }
        }
        // pi's `extendResources`: each kind extensions name paths for loads
        // again from its paths, merged with those it had.
        let [skills, prompts, themes] = found;
        let merged = |known: &mut Vec<SourceInfo>, found: Vec<SourceInfo>| {
            crate::resources::merge_sources(std::mem::take(known).into_iter().chain(found))
        };
        let mut resources = write(&self.inner.resources);
        let resources = &mut *resources;
        if !skills.is_empty() {
            resources.skill_sources = merged(&mut resources.skill_sources, skills);
            (resources.skills, resources.skill_diagnostics) =
                crate::resources::skills_from(&resources.skill_sources);
        }
        if !prompts.is_empty() {
            resources.template_sources = merged(&mut resources.template_sources, prompts);
            (resources.templates, resources.template_diagnostics) =
                crate::resources::templates_from(&resources.template_sources);
        }
        if !themes.is_empty() {
            resources.themes = merged(&mut resources.themes, themes);
        }
    }

    /// Stops extensions before the session ends: pi's `session_shutdown`
    /// with the reason `quit`.
    pub async fn shutdown(&self) {
        self.stop_extensions(ExtensionEvent::SessionShutdown {
            reason: "quit",
            target_session_file: None,
        })
        .await;
        self.release_resources().await;
    }

    /// Stops extensions before a session replaces this one for `reason`.
    /// `target` is the replacement's session file, which pi does not report
    /// on reload.
    pub async fn shutdown_for(&self, reason: Replacement, target: Option<String>) {
        let target = reason.reported(target);
        self.stop_extensions(ExtensionEvent::SessionShutdown {
            reason: reason.as_str(),
            target_session_file: target.as_deref(),
        })
        .await;
        if reason != Replacement::Reload {
            self.release_resources().await;
        }
    }

    async fn stop_extensions(&self, event: ExtensionEvent<'_>) {
        self.cancel_cache_warming();
        self.emit_extension_event(&event.to_value(), CancellationToken::new())
            .await;
        let ctx = self.extension_context(CancellationToken::new());
        for extension in &self.inner.extensions {
            extension.session_shutdown(&ctx).await;
        }
    }

    /// pi's `cleanupSessionResources`, run when a session ends or another
    /// replaces it, but not on reload: the session's cached sockets close
    /// and extensions release what they hold for it.
    async fn release_resources(&self) {
        let id = lock(&self.inner.state).session.id().to_owned();
        yapi_ai::api::close_codex_websockets(Some(&id));
        for extension in &self.inner.extensions {
            extension.end_session(&id).await;
        }
    }

    /// The session's extensions.
    pub fn extensions(&self) -> &[Arc<dyn Extension>] {
        &self.inner.extensions
    }

    /// The extension that draws tool `name`, and how.
    pub fn tool_renderer(&self, name: &str) -> Option<(Arc<dyn Extension>, ToolRenderers)> {
        self.inner.extensions.iter().find_map(|extension| {
            let renderers = extension.renderers().tools.get(name).copied()?;
            Some((extension.clone(), renderers))
        })
    }

    /// The extension that draws custom entries of type `custom_type`.
    pub fn entry_renderer(&self, custom_type: &str) -> Option<Arc<dyn Extension>> {
        self.inner
            .extensions
            .iter()
            .find(|extension| {
                extension
                    .renderers()
                    .entries
                    .iter()
                    .any(|kind| kind == custom_type)
            })
            .cloned()
    }

    /// The extensions that transform Markdown, in load order.
    pub fn markdown_transformers(&self) -> Vec<Arc<dyn Extension>> {
        self.inner
            .extensions
            .iter()
            .filter(|extension| extension.renderers().markdown)
            .cloned()
            .collect()
    }

    /// The extensions that render code blocks, in load order.
    pub fn code_renderers(&self) -> Vec<Arc<dyn Extension>> {
        self.inner
            .extensions
            .iter()
            .filter(|extension| extension.renderers().code)
            .cloned()
            .collect()
    }

    /// The extension that draws custom messages of type `custom_type`.
    pub fn message_renderer(&self, custom_type: &str) -> Option<Arc<dyn Extension>> {
        self.inner
            .extensions
            .iter()
            .find(|extension| {
                extension
                    .renderers()
                    .messages
                    .iter()
                    .any(|kind| kind == custom_type)
            })
            .cloned()
    }

    /// The shortcuts the session's extensions bind, given the keys of each
    /// built-in action, and pi's warnings about conflicts.
    pub fn extension_shortcuts(
        &self,
        builtin: &[(String, Vec<String>)],
    ) -> (
        Vec<crate::extensions::ShortcutBinding>,
        Vec<(String, String)>,
    ) {
        crate::extensions::resolve_shortcuts(&self.inner.extensions, builtin)
    }

    /// Runs an extension shortcut's handler and returns the handler's error.
    pub async fn run_shortcut(
        &self,
        binding: &crate::extensions::ShortcutBinding,
    ) -> Result<(), String> {
        let ctx = self.extension_context(CancellationToken::new());
        binding
            .extension
            .run_shortcut(&binding.registered, &ctx)
            .await
    }

    /// pi's `resolveRegisteredCommands`: the extensions' commands in load
    /// order, each with the name it is invoked by. A name registered more
    /// than once is invoked as `name:1`, `name:2` and so on.
    pub fn extension_commands(&self) -> Vec<crate::extensions::ResolvedCommand> {
        let all: Vec<(Arc<dyn Extension>, crate::extensions::Command)> = self
            .inner
            .extensions
            .iter()
            .flat_map(|extension| {
                extension
                    .commands()
                    .into_iter()
                    .map(move |command| (Arc::clone(extension), command))
            })
            .collect();
        let count = |name: &str| {
            all.iter()
                .filter(|(_, command)| command.name == name)
                .count()
        };
        let mut seen: std::collections::HashMap<String, usize> = std::collections::HashMap::new();
        let mut taken: std::collections::HashSet<String> = std::collections::HashSet::new();
        all.iter()
            .map(|(extension, command)| {
                let occurrence = {
                    let entry = seen.entry(command.name.clone()).or_default();
                    *entry += 1;
                    *entry
                };
                let mut invocation = if count(&command.name) > 1 {
                    format!("{}:{occurrence}", command.name)
                } else {
                    command.name.clone()
                };
                let mut suffix = occurrence;
                while taken.contains(&invocation) {
                    suffix += 1;
                    invocation = format!("{}:{suffix}", command.name);
                }
                taken.insert(invocation.clone());
                crate::extensions::ResolvedCommand {
                    invocation,
                    command: command.clone(),
                    extension: Arc::clone(extension),
                }
            })
            .collect()
    }

    /// pi's slash commands, as RPC's `get_commands` and `pi.getCommands()`
    /// list them: the extensions' commands, then the prompt templates and
    /// the skills, each with where it comes from.
    pub fn slash_commands(&self) -> Vec<SlashCommand> {
        let resources = self.resources();
        let mut commands: Vec<SlashCommand> = self
            .extension_commands()
            .into_iter()
            .map(|resolved| SlashCommand {
                name: resolved.invocation,
                // pi omits a description the command did not give.
                description: Some(resolved.command.description)
                    .filter(|description| !description.is_empty()),
                source: CommandSource::Extension,
                source_info: resolved.extension.source(),
            })
            .collect();
        commands.extend(
            resources
                .templates
                .into_iter()
                .map(|template| SlashCommand {
                    name: template.name,
                    description: Some(template.description),
                    source: CommandSource::Prompt,
                    source_info: template.source,
                }),
        );
        commands.extend(resources.skills.into_iter().map(|skill| SlashCommand {
            name: format!("skill:{}", skill.name),
            description: Some(skill.description),
            source: CommandSource::Skill,
            source_info: skill.source,
        }));
        commands
    }

    /// Whether `text` invokes an extension command.
    pub fn is_extension_command(&self, text: &str) -> bool {
        let Some(rest) = text.strip_prefix('/') else {
            return false;
        };
        let name = rest.split(' ').next().unwrap_or_default();
        self.extension_commands()
            .iter()
            .any(|command| command.invocation == name)
    }

    /// Runs `/name args` when an extension command is invoked by `name`, and
    /// returns whether it did. Errors are the extension's to report.
    pub(super) async fn run_extension_command(&self, text: &str) -> bool {
        let Some(rest) = text.strip_prefix('/') else {
            return false;
        };
        let (name, args) = rest.split_once(' ').unwrap_or((rest, ""));
        let Some(resolved) = self
            .extension_commands()
            .into_iter()
            .find(|command| command.invocation == name)
        else {
            return false;
        };
        let ctx = self.extension_context(CancellationToken::new());
        resolved
            .extension
            .run_command(&resolved.command.name, args, &ctx)
            .await;
        true
    }
}

#[cfg(test)]
mod tests {
    use serde_json::json;

    use super::*;

    /// A handler's `details` and `usage` take the places pi's event holds
    /// for them, as another extension runtime then sees the event.
    #[test]
    fn tool_result_key_placement() {
        let mut event = json!({"type": "tool_result", "content": [], "isError": false});
        let event = event.as_object_mut().unwrap();
        set_tool_result_key(event, "structuredContent", json!({}));
        set_tool_result_key(event, "usage", json!({}));
        set_tool_result_key(event, "details", json!({}));
        set_tool_result_key(event, "isError", json!(true));
        let keys: Vec<&String> = event.keys().collect();
        assert_eq!(
            keys,
            [
                "type",
                "content",
                "details",
                "isError",
                "usage",
                "structuredContent"
            ]
        );
        assert_eq!(event["isError"], true);
    }
}