Skip to main content

yapi_core/agent_session/
extensions.rs

1//! Extension dispatch: pi events, input and command handlers, and the
2//! agent loop hooks that deliver loop events to extensions.
3
4use std::sync::Arc;
5use std::sync::atomic::Ordering;
6
7use futures_util::future::BoxFuture;
8use serde::Deserialize;
9use serde_json::Value;
10use tokio_util::sync::CancellationToken;
11use yapi_agent::Tool;
12use yapi_agent::hooks::AgentHooks;
13use yapi_ai::registry::Auth;
14use yapi_types::event::AgentEvent;
15use yapi_types::extension_event::{ExtensionEvent, RegisteredMcpServer, UiPrompt};
16use yapi_types::message::{
17    Content, ContentBlock, ImageContent, Message, StopReason, SystemMessage, ToolResultMessage,
18};
19use yapi_types::model::Model;
20use yapi_types::rpc::StreamingBehavior;
21use yapi_types::rpc::{CommandSource, SlashCommand, SourceInfo};
22use yapi_types::sync::{lock, write};
23use yapi_types::system_prompt::SystemPromptOptions;
24
25use crate::extensions::{
26    Context, Extension, ExtensionUi, Mode, SessionActions, ToolRenderers, receives, tool_granted,
27};
28use crate::mcp::config::check_registration;
29use crate::messages::convert_to_llm;
30use crate::time::now_ms;
31
32use super::{AgentSession, InputSource, Replacement, SessionChange, drain};
33
34/// What a handler result that changes the session's conversation needs.
35pub(super) const SESSION: &str = "the session grant";
36
37/// What the session drops from a tool_call block without [`SESSION`].
38const BLOCK_REASON: &str = "the reason of a tool_call block";
39
40/// The fields of a `tool_result` handler's result that change the result.
41const TOOL_RESULT_KEYS: [&str; 5] = [
42    "content",
43    "details",
44    "structuredContent",
45    "isError",
46    "usage",
47];
48
49pub(super) struct Hooks {
50    pub(super) session: AgentSession,
51    pub(super) cancel: CancellationToken,
52    /// pi's turn index for extension events: turns since `agent_start`.
53    pub(super) turn_index: std::sync::atomic::AtomicU64,
54}
55
56fn ui_prompt<'a>(kind: &'a str, title: Option<&'a str>) -> UiPrompt<'a> {
57    UiPrompt {
58        reason: "ui_prompt",
59        kind,
60        title,
61    }
62}
63
64/// A loop event as pi's extension event (`_emitExtensionEvent` in pi's agent
65/// session), when `wanted` takes its type, or `None` for events extensions
66/// do not see.
67fn extension_event(
68    event: &AgentEvent,
69    turn_index: u64,
70    wanted: impl Fn(&str) -> bool,
71) -> Option<Value> {
72    let (kind, keys): (&str, &[&str]) = match event {
73        AgentEvent::AgentStart => ("agent_start", &[]),
74        AgentEvent::AgentEnd { .. } => ("agent_end", &["messages"]),
75        AgentEvent::TurnStart => ("turn_start", &[]),
76        // `turn_end` reaches extensions as a boundary, before the event.
77        AgentEvent::MessageStart { .. } => ("message_start", &["message"]),
78        AgentEvent::MessageUpdate { .. } => {
79            ("message_update", &["message", "assistantMessageEvent"])
80        }
81        // `message_end` reaches extensions as the message is finished.
82        AgentEvent::ToolExecutionStart { .. } => (
83            "tool_execution_start",
84            &["toolCallId", "toolName", "args", "parentToolCallId"],
85        ),
86        AgentEvent::ToolExecutionUpdate { .. } => (
87            "tool_execution_update",
88            &[
89                "toolCallId",
90                "toolName",
91                "args",
92                "partialResult",
93                "parentToolCallId",
94            ],
95        ),
96        AgentEvent::ToolExecutionEnd { .. } => (
97            "tool_execution_end",
98            &[
99                "toolCallId",
100                "toolName",
101                "result",
102                "isError",
103                "parentToolCallId",
104            ],
105        ),
106        _ => return None,
107    };
108    if !wanted(kind) {
109        return None;
110    }
111    let value = serde_json::to_value(event).ok()?;
112    let mut out = serde_json::Map::new();
113    out.insert("type".into(), Value::String(kind.to_owned()));
114    if kind == "turn_start" {
115        out.insert("turnIndex".into(), Value::from(turn_index));
116        out.insert("timestamp".into(), Value::from(now_ms()));
117    }
118    for key in keys {
119        if let Some(field) = value.get(*key) {
120            out.insert((*key).to_owned(), field.clone());
121        }
122    }
123    Some(Value::Object(out))
124}
125
126impl AgentHooks for Hooks {
127    fn on_event<'a>(&'a self, event: &'a AgentEvent) -> BoxFuture<'a, ()> {
128        Box::pin(async move {
129            let session = &self.session;
130            if let AgentEvent::MessageStart { message } = event {
131                session.dequeue_started(message);
132            }
133            // pi's interactive mode rewrites an aborted response's error before
134            // the session records it.
135            let rewritten = match event {
136                AgentEvent::MessageEnd {
137                    message: Message::Assistant(assistant),
138                } if assistant.stop_reason == StopReason::Aborted
139                    && lock(&session.inner.binding).1 == Mode::Tui =>
140                {
141                    let attempt = lock(&session.inner.recovery).retry_attempt;
142                    let mut assistant = assistant.clone();
143                    assistant.error_message = Some(if attempt > 0 {
144                        let plural = if attempt > 1 { "s" } else { "" };
145                        format!("Aborted after {attempt} retry attempt{plural}")
146                    } else {
147                        "Operation aborted".to_owned()
148                    });
149                    Some(AgentEvent::MessageEnd {
150                        message: Message::Assistant(assistant),
151                    })
152                }
153                _ => None,
154            };
155            let event = rewritten.as_ref().unwrap_or(event);
156            if matches!(event, AgentEvent::AgentStart) {
157                self.turn_index.store(0, Ordering::SeqCst);
158            }
159            let turn_index = self.turn_index.load(Ordering::SeqCst);
160            if matches!(event, AgentEvent::TurnEnd { .. }) {
161                self.turn_index.fetch_add(1, Ordering::SeqCst);
162            }
163            // As in pi, extensions see the event first, then listeners. The
164            // session records it last, so entries extensions append come before
165            // it.
166            if let Some(extension_event) =
167                extension_event(event, turn_index, |kind| session.has_handlers(kind))
168            {
169                session
170                    .emit_extension_event(&extension_event, self.cancel.clone())
171                    .await;
172            }
173            match event {
174                AgentEvent::AgentEnd { messages, .. } => {
175                    session.inner.nested.clear();
176                    session.emit(&AgentEvent::AgentEnd {
177                        messages: messages.clone(),
178                        will_retry: session.will_retry_after(messages),
179                    });
180                }
181                _ => session.emit(event),
182            }
183            match event {
184                AgentEvent::MessageEnd { message } => {
185                    let entry_id = session.with_session(|file| match message {
186                        Message::Custom(custom) => file
187                            .append_custom_message(
188                                &custom.custom_type,
189                                custom.content.clone(),
190                                custom.display,
191                                custom.details.clone(),
192                            )
193                            .ok(),
194                        Message::System(_)
195                        | Message::User(_)
196                        | Message::Assistant(_)
197                        | Message::ToolResult(_) => file.append_message(message.clone()).ok(),
198                        _ => None,
199                    });
200                    match message {
201                        Message::Assistant(assistant) => {
202                            let finished_retry = {
203                                let mut recovery = lock(&session.inner.recovery);
204                                recovery.last_assistant = Some(((**assistant).clone(), entry_id));
205                                recovery.turn_tool_results.clear();
206                                if !matches!(
207                                    assistant.stop_reason,
208                                    StopReason::Error | StopReason::Length
209                                ) {
210                                    recovery.overflow_recovery_attempted = false;
211                                }
212                                if assistant.stop_reason != StopReason::Error {
213                                    std::mem::take(&mut recovery.retry_attempt)
214                                } else {
215                                    0
216                                }
217                            };
218                            if finished_retry > 0 {
219                                session.emit(&AgentEvent::AutoRetryEnd {
220                                    success: true,
221                                    attempt: finished_retry,
222                                    final_error: None,
223                                });
224                            }
225                        }
226                        Message::ToolResult(_) => {
227                            if let Some(id) = entry_id {
228                                lock(&session.inner.recovery).turn_tool_results.push(id);
229                            }
230                        }
231                        _ => {}
232                    }
233                }
234                AgentEvent::TurnEnd { .. } => {
235                    let mut recovery = lock(&session.inner.recovery);
236                    recovery.last_tool_results = std::mem::take(&mut recovery.turn_tool_results);
237                }
238                _ => {}
239            }
240            // Custom messages sent during the turn are appended once its handlers
241            // ran, so those that turn_end handlers send join them, as in pi.
242            if matches!(event, AgentEvent::TurnEnd { .. }) {
243                session.flush_pending_custom();
244            }
245        })
246    }
247
248    fn transform_context(&self, messages: Vec<Message>) -> BoxFuture<'_, Vec<Message>> {
249        Box::pin(async move {
250            let session = &self.session;
251            let mut messages = session
252                .context_handlers(messages, self.cancel.clone())
253                .await;
254            // pi's hidden declarations projection: requests leave the tools
255            // the loadout hides out of every declaration change.
256            let hidden = lock(&session.inner.hidden_tools).clone();
257            if !hidden.is_empty() {
258                for message in &mut messages {
259                    if let Message::System(system) = message {
260                        system.tools_added = system.tools_added.take().and_then(|mut tools| {
261                            tools.retain(|tool| !hidden.contains(&tool.name));
262                            (!tools.is_empty()).then_some(tools)
263                        });
264                        system.tools_removed = system.tools_removed.take().and_then(|mut tools| {
265                            tools.retain(|tool| !hidden.contains(&tool.name));
266                            (!tools.is_empty()).then_some(tools)
267                        });
268                    }
269                }
270            }
271            let forced = lock(&session.inner.forced_prompt).clone();
272            if let Some(forced) = forced {
273                // pi's forced prompt projection: one system message with the
274                // forced text replaces the transcript's system messages.
275                let current = yapi_ai::transcript::current_system_message(&messages);
276                let head = Message::System(SystemMessage {
277                    content: Content::Text(forced),
278                    sections: None,
279                    timestamp: current
280                        .as_ref()
281                        .map_or_else(now_ms, |system| system.timestamp),
282                    tools_added: current.and_then(|system| system.tools_added),
283                    tools_removed: None,
284                });
285                messages.retain(|message| !matches!(message, Message::System(_)));
286                messages.insert(0, head);
287            }
288            messages
289        })
290    }
291
292    fn convert_to_llm(&self, messages: Vec<Message>) -> Vec<Message> {
293        let messages = convert_to_llm(messages);
294        // Read on every request, so a change applies mid-session.
295        let blocked = self
296            .session
297            .settings()
298            .images
299            .and_then(|images| images.block_images)
300            == Some(true);
301        if blocked {
302            crate::messages::block_images(messages)
303        } else {
304            messages
305        }
306    }
307
308    fn auth<'a>(&'a self, model: &'a Model) -> BoxFuture<'a, Auth> {
309        // The registry is read under a short lock. Credential commands run outside it.
310        let registry = self.session.registry();
311        Box::pin(async move { registry.auth(model).await })
312    }
313
314    fn current_tools(&self) -> BoxFuture<'_, Option<Vec<Arc<dyn Tool>>>> {
315        Box::pin(async move {
316            let declared = self
317                .session
318                .inner
319                .tools
320                .with(|registry| registry.declared());
321            self.session.refresh_loadouts(declared.clone()).await;
322            Some(self.session.loadout(declared))
323        })
324    }
325
326    fn before_tool_call<'a>(
327        &'a self,
328        call: yapi_agent::hooks::BeforeToolCall<'a>,
329    ) -> BoxFuture<'a, Option<yapi_agent::hooks::Block>> {
330        Box::pin(async move {
331            let ctx = self.session.extension_context(self.cancel.clone());
332            for extension in &self.session.inner.extensions {
333                if let Some(reason) = extension
334                    .tool_call(&ctx, &call.tool_call.name, &*call.args)
335                    .await
336                {
337                    let granted = extension.grants().session;
338                    return Some(yapi_agent::hooks::Block {
339                        reason: self
340                            .session
341                            .part_applies(extension, BLOCK_REASON, granted, SESSION)
342                            .then_some(reason),
343                        terminate: false,
344                    });
345                }
346            }
347            let mut event = ExtensionEvent::ToolCall {
348                tool_name: &call.tool_call.name,
349                tool_call_id: &call.tool_call.id,
350                parent_tool_call_id: call.parent_tool_call_id,
351                input: call.args,
352            }
353            .to_value();
354            // Handlers change the call's arguments by editing `event.input`,
355            // which later handlers see, as in pi. The arguments are what the
356            // model sees, and a handler changes them only when its grants
357            // would also let it run the tool itself.
358            let name = &call.tool_call.name;
359            let tool = self
360                .session
361                .inner
362                .tools
363                .with(|registry| registry.get(name).cloned());
364            for extension in self.session.handlers_of("tool_call") {
365                let Some(result) = extension.handle(&ctx, &event).await else {
366                    continue;
367                };
368                // The runtime reports the input whether or not a handler changed it.
369                if let Some(input) = result
370                    .get("input")
371                    .filter(|input| input.is_object() && **input != event["input"])
372                {
373                    let grants = extension.grants();
374                    let granted = grants.session
375                        && tool
376                            .as_ref()
377                            .map_or(grants.is_all(), |tool| tool_granted(&grants, tool));
378                    let missing =
379                        format!("the session grant and the grants to run the {name} tool");
380                    let what = "the input change of a tool_call result";
381                    if self
382                        .session
383                        .part_applies(extension, what, granted, &missing)
384                    {
385                        event["input"] = input.clone();
386                    }
387                }
388                // A block always applies, but its reason is text the model
389                // reads, which needs the session grant.
390                if result["block"] == true {
391                    let granted = extension.grants().session;
392                    let reason = result["reason"].as_str().filter(|_| {
393                        self.session
394                            .part_applies(extension, BLOCK_REASON, granted, SESSION)
395                    });
396                    return Some(yapi_agent::hooks::Block {
397                        reason: reason.map(str::to_owned),
398                        terminate: result["terminate"] == true,
399                    });
400                }
401            }
402            *call.args = event["input"].take();
403            None
404        })
405    }
406
407    fn after_tool_call<'a>(
408        &'a self,
409        call: yapi_agent::hooks::AfterToolCall<'a>,
410    ) -> BoxFuture<'a, Option<yapi_agent::hooks::ResultPatch>> {
411        Box::pin(async move {
412            let mut patch = self.tool_result_handlers(&call).await;
413            // After the handlers, so images they add are normalized too.
414            let content = patch
415                .as_ref()
416                .and_then(|patch| patch.content.as_ref())
417                .unwrap_or(&call.result.content);
418            if content
419                .iter()
420                .any(|block| matches!(block, ContentBlock::Image(_)))
421            {
422                let content = content.clone();
423                let (auto_resize, limits) = self.session.image_options();
424                let normalized = tokio::task::spawn_blocking(move || {
425                    crate::images::normalize_tool_result(&content, auto_resize, limits.as_ref())
426                })
427                .await
428                .ok()
429                .flatten();
430                if let Some(content) = normalized {
431                    // New images alone keep the structured content, as in pi.
432                    patch
433                        .get_or_insert_with(|| yapi_agent::hooks::ResultPatch {
434                            structured_content: call.result.structured_content.clone(),
435                            ..Default::default()
436                        })
437                        .content = Some(content);
438                }
439            }
440            patch
441        })
442    }
443
444    fn finish_turn<'a>(
445        &'a self,
446        message: &'a yapi_types::message::AssistantMessage,
447        tool_results: &'a [ToolResultMessage],
448    ) -> BoxFuture<'a, yapi_agent::hooks::TurnDecision> {
449        Box::pin(async move {
450            let turn_index = self.turn_index.load(Ordering::SeqCst);
451            self.session
452                .turn_end_boundary(message, tool_results, turn_index, self.cancel.clone())
453                .await
454        })
455    }
456
457    fn finish_message(&self, message: Message) -> BoxFuture<'_, Message> {
458        Box::pin(async move {
459            self.session
460                .message_end_handlers(message, self.cancel.clone())
461                .await
462        })
463    }
464
465    fn complete_tool_result(&self, message: &mut ToolResultMessage) {
466        let Some(summary) = self.session.inner.nested.take(&message.tool_call_id) else {
467            return;
468        };
469        if summary.calls.is_some() {
470            message.nested_calls = summary.calls;
471        }
472        if let Some(usage) = summary.usage {
473            message.usage = Some(match &message.usage {
474                Some(own) => own.combine(&usage),
475                None => usage,
476            });
477        }
478    }
479
480    fn next_turn_messages(&self) -> BoxFuture<'_, Result<Option<Vec<Message>>, String>> {
481        // As in pi, each turn continues from the session's transcript, with
482        // what the session recorded since the last, such as custom messages
483        // held while it ran and entries turn_end handlers committed, once a
484        // context over the threshold is compacted.
485        Box::pin(async move {
486            self.session
487                .compact_before_next_response(&self.cancel)
488                .await?;
489            Ok(Some(self.session.messages()))
490        })
491    }
492
493    // As pi's agent, each drain reads the current mode, so a change applies
494    // to the running agent.
495    fn steering_messages(&self) -> BoxFuture<'_, Vec<Message>> {
496        let messages = drain(&self.session.inner.steering, self.session.drains_all().0);
497        Box::pin(async move { messages })
498    }
499
500    fn follow_up_messages(&self) -> BoxFuture<'_, Vec<Message>> {
501        let messages = drain(&self.session.inner.follow_up, self.session.drains_all().1);
502        Box::pin(async move { messages })
503    }
504}
505
506impl Hooks {
507    /// pi's `tool_result` event: the changes its handlers make, if any.
508    async fn tool_result_handlers(
509        &self,
510        call: &yapi_agent::hooks::AfterToolCall<'_>,
511    ) -> Option<yapi_agent::hooks::ResultPatch> {
512        let handlers = self.session.handlers_of("tool_result");
513        if handlers.is_empty() {
514            return None;
515        }
516        let ctx = self.session.extension_context(self.cancel.clone());
517        let mut event = ExtensionEvent::ToolResult {
518            tool_name: &call.tool_call.name,
519            tool_call_id: &call.tool_call.id,
520            parent_tool_call_id: call.parent_tool_call_id,
521            input: call.args,
522            content: &call.result.content,
523            details: call.result.details.as_ref(),
524            structured_content: call.result.structured_content.as_ref(),
525            is_error: call.is_error,
526            usage: call.result.usage.as_ref(),
527        }
528        .to_value();
529        let mut modified = false;
530        for extension in handlers {
531            let Some(result) = extension.handle(&ctx, &event).await else {
532                continue;
533            };
534            let given = |key: &str| result.get(key).filter(|value| !value.is_null());
535            let changes = TOOL_RESULT_KEYS.iter().any(|key| given(key).is_some());
536            if changes
537                && !self.session.applies(
538                    extension,
539                    "tool_result",
540                    extension.grants().session,
541                    SESSION,
542                )
543            {
544                continue;
545            }
546            // As in pi, structured content not replaced along with the
547            // content may no longer match it.
548            if given("content").is_some()
549                && given("structuredContent").is_none()
550                && let Some(event) = event.as_object_mut()
551            {
552                event.shift_remove("structuredContent");
553            }
554            for key in TOOL_RESULT_KEYS {
555                if let Some(value) = given(key)
556                    && let Some(event) = event.as_object_mut()
557                {
558                    set_tool_result_key(event, key, value.clone());
559                    modified = true;
560                }
561            }
562        }
563        modified.then(|| yapi_agent::hooks::ResultPatch {
564            content: serde_json::from_value(event["content"].clone()).ok(),
565            details: Some(event["details"].clone()).filter(|details| !details.is_null()),
566            structured_content: event.get("structuredContent").cloned(),
567            is_error: event["isError"].as_bool(),
568            usage: serde_json::from_value(event["usage"].clone()).ok(),
569            terminate: None,
570        })
571    }
572}
573
574/// Whether `extension` handles pi event `kind` and its grants let it hear it.
575fn hears(extension: &Arc<dyn Extension>, kind: &str) -> bool {
576    extension.handles(kind) && receives(kind, &extension.grants())
577}
578
579/// Sets `key` of a `tool_result` event where pi's event has it: pi's holds
580/// `details` and `usage` even when undefined, after `content` and `isError`.
581fn set_tool_result_key(event: &mut serde_json::Map<String, Value>, key: &str, value: Value) {
582    if let Some(slot) = event.get_mut(key) {
583        *slot = value;
584        return;
585    }
586    let after = match key {
587        "details" => "content",
588        "usage" => "isError",
589        _ => "",
590    };
591    let index = event
592        .keys()
593        .position(|name| name == after)
594        .map_or(event.len(), |index| index + 1);
595    event.shift_insert(index, key.to_owned(), value);
596}
597
598/// The servers in `registered` that runtimes other than `owner` registered.
599fn others(
600    registered: &indexmap::IndexMap<String, (u64, RegisteredMcpServer)>,
601    owner: u64,
602) -> impl Iterator<Item = &RegisteredMcpServer> + Clone {
603    registered
604        .values()
605        .filter(move |(runtime, _)| *runtime != owner)
606        .map(|(_, server)| server)
607}
608
609impl AgentSession {
610    /// The context extensions get, with `cancel` for the operation at hand.
611    pub(super) fn extension_context(&self, cancel: CancellationToken) -> Context {
612        let (ui, mode, _) = lock(&self.inner.binding).clone();
613        let (project_dir, project_trusted, project_only_mcp) = {
614            let settings = lock(&self.inner.settings);
615            let keep = crate::config::project_only().then(|| {
616                let global = settings.document(crate::settings::Scope::Global);
617                crate::packages::resolve::ProjectOnly::from_settings(global).mcp_servers
618            });
619            (
620                settings.project_dir().to_path_buf(),
621                settings.project_trusted(),
622                keep,
623            )
624        };
625        Context {
626            cwd: self.inner.cwd.clone(),
627            agent_dir: self.inner.agent_dir.clone(),
628            project_dir,
629            project_trusted,
630            project_only_mcp,
631            mode,
632            ui,
633            tools: self.inner.tools.clone(),
634            cancel,
635            session: self.downgrade(),
636        }
637    }
638
639    /// The UI and mode extensions are bound to.
640    pub fn extension_binding(&self) -> (Arc<dyn ExtensionUi>, Mode) {
641        let (ui, mode, _) = lock(&self.inner.binding).clone();
642        (ui, mode)
643    }
644
645    /// Whether a result of `extension`'s `event` handler applies: always
646    /// when `granted`, else never, and the extension is warned once per
647    /// event that the result needs `missing`, such as the session grant.
648    pub(super) fn applies(
649        &self,
650        extension: &Arc<dyn Extension>,
651        event: &str,
652        granted: bool,
653        missing: &str,
654    ) -> bool {
655        if !granted {
656            self.dropped(extension, event, missing);
657        }
658        granted
659    }
660
661    /// [`AgentSession::applies`] for part `what` of a handler's result, such
662    /// as the input change of a tool_call result, which the warning names.
663    pub(super) fn part_applies(
664        &self,
665        extension: &Arc<dyn Extension>,
666        what: &str,
667        granted: bool,
668        missing: &str,
669    ) -> bool {
670        if !granted {
671            self.ignored(extension, what, missing);
672        }
673        granted
674    }
675
676    /// Warns `extension`, once per event, that the session ignored a result
677    /// of its `event` handler, which needs `missing`.
678    pub(super) fn dropped(&self, extension: &Arc<dyn Extension>, event: &str, missing: &str) {
679        self.ignored(extension, &format!("a {event} result"), missing);
680    }
681
682    /// Warns `extension`, once per `what`, that the session ignored `what`,
683    /// which needs `missing`. A chest extension is named by its entry.
684    fn ignored(&self, extension: &Arc<dyn Extension>, what: &str, missing: &str) {
685        let source = extension.source();
686        if lock(&self.inner.dropped_results).insert((source.path.clone(), what.to_owned())) {
687            let project = self.project_dir();
688            let label = crate::chest::startup::entry_label(&source.source, &source.scope, &project);
689            let label = label.unwrap_or(source.path);
690            extension.warn(&format!(
691                "Ignored {what} of extension {label}: it needs {missing}"
692            ));
693        }
694    }
695
696    pub(super) fn handlers_of(&self, kind: &str) -> Vec<&Arc<dyn Extension>> {
697        self.inner
698            .extensions
699            .iter()
700            .filter(|extension| hears(extension, kind))
701            .collect()
702    }
703
704    /// Whether any extension handles pi events of type `kind` and may hear them.
705    pub fn has_handlers(&self, kind: &str) -> bool {
706        self.inner
707            .extensions
708            .iter()
709            .any(|extension| hears(extension, kind))
710    }
711
712    /// Delivers pi event `event` to every extension that handles it, in
713    /// order, as pi's runner `emit` does. Only a `session_before_*` event has
714    /// a result: the last one a handler gave, or the first that cancels,
715    /// which stops delivery.
716    pub async fn emit_extension_event(
717        &self,
718        event: &Value,
719        cancel: CancellationToken,
720    ) -> Option<Value> {
721        let kind = event["type"].as_str().unwrap_or_default();
722        let handlers = self.handlers_of(kind);
723        if handlers.is_empty() {
724            return None;
725        }
726        let ctx = self.extension_context(cancel);
727        let before = kind.starts_with("session_before_");
728        let mut result = None;
729        for extension in handlers {
730            let Some(mut next) = extension.handle(&ctx, event).await else {
731                continue;
732            };
733            if !before {
734                continue;
735            }
736            let cancelled = next["cancel"] == true;
737            // Without the session grant, a handler may only cancel, such as
738            // a compaction it would otherwise summarize itself. Anything else
739            // counts as no result, so it never replaces a granted one.
740            if !extension.grants().session {
741                let changes = next
742                    .as_object()
743                    .is_some_and(|result| result.keys().any(|key| key != "cancel"));
744                if changes {
745                    self.dropped(extension, kind, SESSION);
746                }
747                if !cancelled {
748                    continue;
749                }
750                next = serde_json::json!({"cancel": true});
751            }
752            result = Some(next);
753            if cancelled {
754                break;
755            }
756        }
757        result
758    }
759
760    /// pi's `session_before_switch` or `session_before_fork`: whether an
761    /// extension cancels `change`.
762    pub async fn cancels(&self, change: &SessionChange) -> bool {
763        let reason = change.reason().as_str();
764        let event = match change {
765            SessionChange::New => ExtensionEvent::SessionBeforeSwitch {
766                reason,
767                target_session_file: None,
768            },
769            SessionChange::Resume(target) => ExtensionEvent::SessionBeforeSwitch {
770                reason,
771                target_session_file: Some(target),
772            },
773            SessionChange::Fork { entry_id, at } => ExtensionEvent::SessionBeforeFork {
774                entry_id,
775                position: if *at { "at" } else { "before" },
776            },
777        };
778        let result = self
779            .emit_extension_event(&event.to_value(), CancellationToken::new())
780            .await;
781        result.is_some_and(|result| result["cancel"] == true)
782    }
783
784    /// pi's `message_end`: the message as handlers replaced it, each seeing
785    /// the previous one's. A replacement of another role is reported and
786    /// skipped, and a missing content of a user, assistant, tool result or
787    /// custom message becomes empty.
788    async fn message_end_handlers(&self, message: Message, cancel: CancellationToken) -> Message {
789        let handlers = self.handlers_of("message_end");
790        if handlers.is_empty() {
791            return message;
792        }
793        let ctx = self.extension_context(cancel);
794        let report = |extension: &Arc<dyn Extension>, error: &str| {
795            ctx.ui
796                .extension_error(&extension.source().path, "message_end", error, None);
797        };
798        let mut current = serde_json::json!(message);
799        let mut replaced_by = None;
800        for extension in handlers {
801            let event = ExtensionEvent::MessageEnd { message: &current }.to_value();
802            let Some(next) = extension
803                .handle(&ctx, &event)
804                .await
805                .and_then(|mut result| result.get_mut("message").map(Value::take))
806                .filter(|next| !next.is_null())
807            else {
808                continue;
809            };
810            if !self.applies(
811                extension,
812                "message_end",
813                extension.grants().session,
814                SESSION,
815            ) {
816                continue;
817            }
818            if next["role"] != current["role"] {
819                report(
820                    extension,
821                    "message_end handlers must return a message with the same role",
822                );
823                continue;
824            }
825            current = next;
826            replaced_by = Some(extension);
827        }
828        let Some(extension) = replaced_by else {
829            return message;
830        };
831        if matches!(
832            current["role"].as_str(),
833            Some("user" | "assistant" | "toolResult" | "custom")
834        ) && current["content"].is_null()
835        {
836            current["content"] = Value::Array(Vec::new());
837        }
838        serde_json::from_value(current).unwrap_or_else(|error| {
839            report(extension, &error.to_string());
840            message
841        })
842    }
843
844    /// Passes `value` through the extensions handling `kind` as the event's
845    /// `field`, each seeing the `field` the previous one returned.
846    pub(super) async fn chain(
847        &self,
848        kind: &str,
849        field: &str,
850        value: Value,
851        cancel: CancellationToken,
852    ) -> Value {
853        let handlers = self.handlers_of(kind);
854        if handlers.is_empty() {
855            return value;
856        }
857        let ctx = self.extension_context(cancel);
858        let mut value = value;
859        for extension in handlers {
860            let event = serde_json::json!({"type": kind, field: value});
861            if let Some(result) = extension.handle(&ctx, &event).await
862                && let Some(next) = result.get(field)
863                && *next != value
864                && self.applies(extension, kind, extension.grants().session, SESSION)
865            {
866                value = next.clone();
867            }
868        }
869        value
870    }
871
872    /// pi's `emitContext`: `context` handlers see the conversation without
873    /// system messages, which come back as the current prompt when they
874    /// change it. `context_with_system` handlers then see and return the
875    /// whole transcript.
876    async fn context_handlers(
877        &self,
878        messages: Vec<Message>,
879        cancel: CancellationToken,
880    ) -> Vec<Message> {
881        let mut messages = messages;
882        let ctx = self.extension_context(cancel);
883        for extension in self.handlers_of("context") {
884            let visible: Vec<Message> = messages
885                .iter()
886                .filter(|message| !matches!(message, Message::System(_)))
887                .cloned()
888                .collect();
889            let event = ExtensionEvent::Context { messages: &visible }.to_value();
890            if let Some(result) = extension.handle(&ctx, &event).await
891                && let Ok(returned) =
892                    serde_json::from_value::<Vec<Message>>(result["messages"].clone())
893                && returned != visible
894                && self.applies(extension, "context", extension.grants().session, SESSION)
895            {
896                let head = yapi_ai::transcript::current_system_message(&messages);
897                messages = head
898                    .map(Message::System)
899                    .into_iter()
900                    .chain(returned)
901                    .collect();
902            }
903        }
904        for extension in self.handlers_of("context_with_system") {
905            let had_head = matches!(messages.first(), Some(Message::System(_)));
906            let event = ExtensionEvent::ContextWithSystem {
907                messages: &messages,
908            }
909            .to_value();
910            if let Some(result) = extension.handle(&ctx, &event).await
911                && let Ok(returned) =
912                    serde_json::from_value::<Vec<Message>>(result["messages"].clone())
913                && returned != messages
914                && self.applies(
915                    extension,
916                    "context_with_system",
917                    extension.grants().session,
918                    SESSION,
919                )
920            {
921                messages = returned;
922            }
923            if had_head && !matches!(messages.first(), Some(Message::System(_))) {
924                ctx.ui.extension_error(
925                    &extension.source().path,
926                    "context_with_system",
927                    "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().",
928                    None,
929                );
930            }
931        }
932        messages
933    }
934
935    /// pi's `before_provider_headers`: the headers after handlers changed
936    /// them in place. A null value removes a header.
937    pub(super) async fn before_provider_headers(
938        &self,
939        headers: yapi_ai::stream::RequestHeaders,
940        cancel: CancellationToken,
941    ) -> yapi_ai::stream::RequestHeaders {
942        let value = serde_json::to_value(&headers).unwrap_or_default();
943        let value = self
944            .chain("before_provider_headers", "headers", value, cancel)
945            .await;
946        match value.as_object() {
947            Some(headers) => headers
948                .iter()
949                .map(|(name, value)| (name.clone(), value.as_str().map(str::to_owned)))
950                .collect(),
951            None => headers,
952        }
953    }
954
955    /// A blocking extension dialog of `kind` (`select`, `confirm`, `input`,
956    /// `editor` or `custom`) with `title` opens. Extensions hear of the
957    /// outermost as pi's `ui_prompt_start`.
958    pub fn ui_prompt_opened(&self, kind: &str, title: Option<&str>) {
959        let title = title.filter(|title| !title.is_empty()).map(str::to_owned);
960        {
961            let mut prompts = lock(&self.inner.ui_prompts);
962            prompts.0 += 1;
963            if prompts.0 > 1 {
964                return;
965            }
966            prompts.1 = Some((kind.to_owned(), title.clone()));
967        }
968        self.announce(ExtensionEvent::UiPromptStart(ui_prompt(
969            kind,
970            title.as_deref(),
971        )));
972    }
973
974    /// A blocking extension dialog closed. When it was the outermost,
975    /// extensions hear of it as pi's `ui_prompt_end`.
976    pub fn ui_prompt_closed(&self) {
977        let (kind, title) = {
978            let mut prompts = lock(&self.inner.ui_prompts);
979            prompts.0 = prompts.0.saturating_sub(1);
980            if prompts.0 > 0 {
981                return;
982            }
983            match prompts.1.take() {
984                Some(prompt) => prompt,
985                None => return,
986            }
987        };
988        self.announce(ExtensionEvent::UiPromptEnd(ui_prompt(
989            &kind,
990            title.as_deref(),
991        )));
992    }
993
994    /// Delivers `event` to extensions in the background, after the events
995    /// announced before it, as pi's runner delivers events it does not wait
996    /// for. Outside a tokio runtime it waits for the next
997    /// [`AgentSession::flush_announcements`].
998    pub(super) fn announce(&self, event: ExtensionEvent<'_>) {
999        let event = event.to_value();
1000        let kind = event["type"].as_str().unwrap_or_default();
1001        if !self.has_handlers(kind) {
1002            return;
1003        }
1004        lock(&self.inner.announcements).push_back(event);
1005        if let Ok(runtime) = tokio::runtime::Handle::try_current() {
1006            let session = self.clone();
1007            runtime.spawn(async move { session.flush_announcements().await });
1008        }
1009    }
1010
1011    /// Sets the MCP servers that extension runtime `owner` registered, and
1012    /// announces `mcp_servers_change` when they changed after the extensions
1013    /// started. Before then, as runtimes bind, the reports hold the servers
1014    /// registered while extensions loaded, which the MCP extension reads
1015    /// when the session starts. Nothing changes when a server's name, or its
1016    /// tools' namespace, belongs to another runtime, and the error is pi's
1017    /// for that registration.
1018    pub fn set_mcp_servers(
1019        &self,
1020        owner: u64,
1021        servers: Vec<RegisteredMcpServer>,
1022    ) -> Result<(), String> {
1023        let changed = {
1024            let (registered, started) = &mut *lock(&self.inner.mcp_servers);
1025            for server in &servers {
1026                check_registration(others(registered, owner), &server.name)?;
1027            }
1028            let before = registered.clone();
1029            // As in pi's registry, a server registered again keeps its place.
1030            registered.retain(|name, (runtime, _)| {
1031                *runtime != owner || servers.iter().any(|server| server.name == *name)
1032            });
1033            for server in servers {
1034                registered.insert(server.name.clone(), (owner, server));
1035            }
1036            *started && *registered != before
1037        };
1038        if changed {
1039            let servers = self.mcp_servers();
1040            self.announce(ExtensionEvent::McpServersChange { servers: &servers });
1041            self.report_unhandled_mcp_servers();
1042        }
1043        Ok(())
1044    }
1045
1046    /// pi's checks of server `name`, which an extension of runtime `owner`
1047    /// registers, against the servers of the other runtimes, without
1048    /// registering it.
1049    pub fn check_mcp_server(&self, owner: u64, name: &str) -> Result<(), String> {
1050        check_registration(others(&lock(&self.inner.mcp_servers).0, owner), name)
1051    }
1052
1053    /// pi's `reportUnhandledMcpServers`: when no extension handles
1054    /// `mcp_servers_change`, nothing connects registered servers, and each
1055    /// is reported once as its extension's error.
1056    fn report_unhandled_mcp_servers(&self) {
1057        if self.has_handlers("mcp_servers_change") {
1058            return;
1059        }
1060        let (ui, _) = self.extension_binding();
1061        for server in self.mcp_servers() {
1062            if !lock(&self.inner.unhandled_mcp_servers).insert(server.name.clone()) {
1063                continue;
1064            }
1065            let error = format!(
1066                "MCP server \"{}\" is registered, but no loaded extension connects MCP servers; another extension may have replaced the built-in MCP support",
1067                server.name
1068            );
1069            ui.extension_error(&server.extension_path, "register_mcp_server", &error, None);
1070        }
1071    }
1072
1073    /// pi's `getMcpServers`: the servers extensions registered, in
1074    /// registration order.
1075    pub fn mcp_servers(&self) -> Vec<RegisteredMcpServer> {
1076        lock(&self.inner.mcp_servers)
1077            .0
1078            .values()
1079            .map(|(_, server)| server.clone())
1080            .collect()
1081    }
1082
1083    /// Records the extension at `path`, whose factory failed with `error` as
1084    /// the extensions loaded for this session. The factories of every
1085    /// session after the first run again then, as pi loads the extensions
1086    /// for each runtime.
1087    pub fn add_extension_load_error(&self, path: String, error: String) {
1088        lock(&self.inner.extension_load_errors).push((path, error));
1089    }
1090
1091    /// The extensions that failed to load for this session, by path, with
1092    /// pi's message, which pi's interactive mode lists under
1093    /// `[Extension issues]`.
1094    pub fn extension_load_errors(&self) -> Vec<(String, String)> {
1095        lock(&self.inner.extension_load_errors).clone()
1096    }
1097
1098    /// Records the extension at `path`, which was not loaded for this
1099    /// session, with pi's `warning`.
1100    pub fn add_extension_warning(&self, path: String, warning: String) {
1101        lock(&self.inner.extension_warnings).push((path, warning));
1102    }
1103
1104    /// The extensions that were not loaded for this session, by path, with
1105    /// pi's warning, which pi's interactive mode lists under
1106    /// `[Extension issues]`.
1107    pub fn extension_warnings(&self) -> Vec<(String, String)> {
1108        lock(&self.inner.extension_warnings).clone()
1109    }
1110
1111    /// Waits until the events announced so far have reached extensions.
1112    pub async fn flush_announcements(&self) {
1113        let _turn = self.inner.announcing.lock().await;
1114        loop {
1115            let next = lock(&self.inner.announcements).pop_front();
1116            let Some(event) = next else {
1117                return;
1118            };
1119            self.emit_extension_event(&event, CancellationToken::new())
1120                .await;
1121        }
1122    }
1123
1124    /// pi's `input` event: `None` when a handler handled the input,
1125    /// otherwise the possibly transformed text and images.
1126    pub(super) async fn input_handlers(
1127        &self,
1128        text: &str,
1129        images: Vec<ImageContent>,
1130        source: InputSource,
1131        streaming: Option<StreamingBehavior>,
1132    ) -> Option<(String, Vec<ImageContent>)> {
1133        let handlers = self.handlers_of("input");
1134        if handlers.is_empty() {
1135            return Some((text.to_owned(), images));
1136        }
1137        let ctx = self.extension_context(CancellationToken::new());
1138        let mut text = text.to_owned();
1139        let mut images = images;
1140        for extension in handlers {
1141            let event = ExtensionEvent::Input {
1142                text: &text,
1143                images: &images,
1144                source: source.as_str(),
1145                streaming_behavior: streaming,
1146            }
1147            .to_value();
1148            let Some(result) = extension.handle(&ctx, &event).await else {
1149                continue;
1150            };
1151            match result["action"].as_str() {
1152                Some("handled") => return None,
1153                Some("transform")
1154                    if self.applies(extension, "input", extension.grants().session, SESSION) =>
1155                {
1156                    if let Some(next) = result["text"].as_str() {
1157                        next.clone_into(&mut text);
1158                    }
1159                    if let Ok(next) = serde_json::from_value(result["images"].clone()) {
1160                        images = next;
1161                    }
1162                }
1163                _ => {}
1164            }
1165        }
1166        Some((text, images))
1167    }
1168
1169    /// pi's `before_agent_start` event for extensions handling it, with the
1170    /// run's prompt `options`, which each handler may edit for the next: the
1171    /// custom messages to send with the prompt, and the system prompt a
1172    /// handler forced, if any. The edited options become the run's, and
1173    /// edited `selectedTools` the active tools.
1174    pub(super) async fn before_agent_start_handlers(
1175        &self,
1176        prompt: &str,
1177        images: &[ImageContent],
1178        mut options: SystemPromptOptions,
1179    ) -> (Vec<Message>, Option<String>) {
1180        let handlers = self.handlers_of("before_agent_start");
1181        let mut messages = Vec::new();
1182        let selected = options.selected_tools.clone();
1183        let ctx = self.extension_context(CancellationToken::new());
1184        for extension in handlers {
1185            let system_prompt = self.render_prompt(&options).unwrap_or_default();
1186            let event = ExtensionEvent::BeforeAgentStart {
1187                prompt,
1188                images,
1189                system_prompt: &system_prompt,
1190                system_prompt_options: &options,
1191            }
1192            .to_value();
1193            let Some(result) = extension.handle(&ctx, &event).await else {
1194                continue;
1195            };
1196            // Without the session grant, no part of the result applies. The
1197            // runtime reports the options and the forced prompt whether or not
1198            // a handler changed them, so only a change is warned of.
1199            if !extension.grants().session {
1200                let edited = result["systemPromptOptions"].is_object()
1201                    && SystemPromptOptions::deserialize(&result["systemPromptOptions"])
1202                        .is_ok_and(|edited| edited != options);
1203                let forced = result["systemPrompt"]
1204                    .as_str()
1205                    .is_some_and(|prompt| Some(prompt) != options.force_system_prompt.as_deref());
1206                let added = result["messages"]
1207                    .as_array()
1208                    .is_some_and(|messages| !messages.is_empty());
1209                if edited || forced || added {
1210                    self.dropped(extension, "before_agent_start", SESSION);
1211                }
1212                continue;
1213            }
1214            for message in result["messages"].as_array().into_iter().flatten() {
1215                let content = match &message["content"] {
1216                    Value::Null => Value::Array(Vec::new()),
1217                    other => other.clone(),
1218                };
1219                messages.push(Message::Custom(yapi_types::message::CustomMessage {
1220                    custom_type: message["customType"]
1221                        .as_str()
1222                        .unwrap_or_default()
1223                        .to_owned(),
1224                    content: serde_json::from_value(content).unwrap_or(Content::Blocks(Vec::new())),
1225                    display: message["display"].as_bool().unwrap_or(false),
1226                    details: (!message["details"].is_null()).then(|| message["details"].clone()),
1227                    timestamp: now_ms(),
1228                }));
1229            }
1230            // An edit that is not pi's shape leaves the options as they were.
1231            if result["systemPromptOptions"].is_object()
1232                && let Ok(edited) = SystemPromptOptions::deserialize(&result["systemPromptOptions"])
1233            {
1234                options = edited;
1235            }
1236            if let Some(prompt) = result["systemPrompt"].as_str() {
1237                options.force_system_prompt = Some(prompt.to_owned());
1238            }
1239        }
1240        if options.selected_tools != selected {
1241            self.set_active_tools(options.selected_tools.clone());
1242        }
1243        let forced = options.force_system_prompt.take();
1244        *lock(&self.inner.run_options) = Some(options);
1245        (messages, forced)
1246    }
1247
1248    /// Carries out a session change an extension command asks for: whether
1249    /// an extension cancelled it. Without a mode, nothing changes, as pi's
1250    /// default actions do.
1251    pub async fn session_action(
1252        &self,
1253        action: crate::extensions::SessionAction,
1254    ) -> Result<bool, String> {
1255        let actions = lock(&self.inner.binding).2.clone();
1256        match actions {
1257            Some(actions) => actions(action).await,
1258            None => Ok(false),
1259        }
1260    }
1261
1262    /// Gives extensions their UI, mode and the `actions` that carry out the
1263    /// session changes their commands ask for, and starts them: pi's
1264    /// `bindExtensions`, which emits `session_start` with the reason this
1265    /// session replaced another and that session's file, or `startup` when
1266    /// `replaced` is `None`.
1267    pub async fn bind_extensions(
1268        &self,
1269        ui: Arc<dyn ExtensionUi>,
1270        mode: Mode,
1271        actions: Option<SessionActions>,
1272        replaced: Option<(Replacement, Option<String>)>,
1273    ) {
1274        *lock(&self.inner.binding) = (ui, mode, actions);
1275        let ctx = self.extension_context(CancellationToken::new());
1276        for extension in &self.inner.extensions {
1277            extension.session_start(&ctx).await;
1278        }
1279        lock(&self.inner.mcp_servers).1 = true;
1280        let (reason, previous) = match replaced {
1281            Some((reason, previous)) => (Some(reason), reason.reported(previous)),
1282            None => (None, None),
1283        };
1284        let event = ExtensionEvent::SessionStart {
1285            reason: reason.map_or("startup", Replacement::as_str),
1286            previous_session_file: previous.as_deref(),
1287        };
1288        self.emit_extension_event(&event.to_value(), CancellationToken::new())
1289            .await;
1290        self.report_unhandled_mcp_servers();
1291        let reload = reason == Some(Replacement::Reload);
1292        self.discover_resources(if reload { "reload" } else { "startup" })
1293            .await;
1294    }
1295
1296    /// pi's `resources_discover`, after `session_start`: adds the skills,
1297    /// prompt templates and themes extensions name, each with its extension
1298    /// as source.
1299    async fn discover_resources(&self, reason: &str) {
1300        let handlers = self.handlers_of("resources_discover");
1301        if handlers.is_empty() {
1302            return;
1303        }
1304        let event = ExtensionEvent::ResourcesDiscover {
1305            cwd: &self.inner.cwd,
1306            reason,
1307        }
1308        .to_value();
1309        let ctx = self.extension_context(CancellationToken::new());
1310        let mut found: [Vec<SourceInfo>; 3] = Default::default();
1311        for extension in handlers {
1312            let Some(result) = extension.handle(&ctx, &event).await else {
1313                continue;
1314            };
1315            // The files it names load into the prompt and the commands.
1316            let grants = extension.grants();
1317            let names = ["skillPaths", "promptPaths", "themePaths"]
1318                .iter()
1319                .any(|kind| {
1320                    result[*kind]
1321                        .as_array()
1322                        .is_some_and(|paths| !paths.is_empty())
1323                });
1324            if names
1325                && !self.applies(
1326                    extension,
1327                    "resources_discover",
1328                    grants.session && grants.filesystem,
1329                    "the session and file grants",
1330                )
1331            {
1332                continue;
1333            }
1334            let origin = extension.source().path;
1335            let synthetic = origin.starts_with("builtin:") || origin.starts_with('<');
1336            let source = if synthetic {
1337                format!("extension:{}", origin.replace(['<', '>'], ""))
1338            } else {
1339                let name = std::path::Path::new(&origin)
1340                    .file_name()
1341                    .map(|name| name.to_string_lossy().into_owned())
1342                    .unwrap_or_default();
1343                let name = name
1344                    .strip_suffix(".ts")
1345                    .or_else(|| name.strip_suffix(".js"))
1346                    .unwrap_or(&name);
1347                format!("extension:{name}")
1348            };
1349            let base_dir = (!synthetic)
1350                .then(|| std::path::Path::new(&origin).parent())
1351                .flatten()
1352                .map(|dir| dir.display().to_string());
1353            for (kind, paths) in ["skillPaths", "promptPaths", "themePaths"]
1354                .iter()
1355                .zip(found.iter_mut())
1356            {
1357                for path in result[*kind].as_array().into_iter().flatten() {
1358                    let Some(path) = path.as_str() else { continue };
1359                    let path = crate::tools::path::resolve_to_cwd(path.trim(), &self.inner.cwd);
1360                    paths.push(SourceInfo {
1361                        path: path.display().to_string(),
1362                        source: source.clone(),
1363                        scope: "temporary".into(),
1364                        origin: "top-level".into(),
1365                        base_dir: base_dir.clone(),
1366                    });
1367                }
1368            }
1369        }
1370        // pi's `extendResources`: each kind extensions name paths for loads
1371        // again from its paths, merged with those it had.
1372        let [skills, prompts, themes] = found;
1373        let merged = |known: &mut Vec<SourceInfo>, found: Vec<SourceInfo>| {
1374            crate::resources::merge_sources(std::mem::take(known).into_iter().chain(found))
1375        };
1376        let mut resources = write(&self.inner.resources);
1377        let resources = &mut *resources;
1378        if !skills.is_empty() {
1379            resources.skill_sources = merged(&mut resources.skill_sources, skills);
1380            (resources.skills, resources.skill_diagnostics) =
1381                crate::resources::skills_from(&resources.skill_sources);
1382        }
1383        if !prompts.is_empty() {
1384            resources.template_sources = merged(&mut resources.template_sources, prompts);
1385            (resources.templates, resources.template_diagnostics) =
1386                crate::resources::templates_from(&resources.template_sources);
1387        }
1388        if !themes.is_empty() {
1389            resources.themes = merged(&mut resources.themes, themes);
1390        }
1391    }
1392
1393    /// Stops extensions before the session ends: pi's `session_shutdown`
1394    /// with the reason `quit`.
1395    pub async fn shutdown(&self) {
1396        self.stop_extensions(ExtensionEvent::SessionShutdown {
1397            reason: "quit",
1398            target_session_file: None,
1399        })
1400        .await;
1401        self.release_resources().await;
1402    }
1403
1404    /// Stops extensions before a session replaces this one for `reason`.
1405    /// `target` is the replacement's session file, which pi does not report
1406    /// on reload.
1407    pub async fn shutdown_for(&self, reason: Replacement, target: Option<String>) {
1408        let target = reason.reported(target);
1409        self.stop_extensions(ExtensionEvent::SessionShutdown {
1410            reason: reason.as_str(),
1411            target_session_file: target.as_deref(),
1412        })
1413        .await;
1414        if reason != Replacement::Reload {
1415            self.release_resources().await;
1416        }
1417    }
1418
1419    async fn stop_extensions(&self, event: ExtensionEvent<'_>) {
1420        self.cancel_cache_warming();
1421        self.emit_extension_event(&event.to_value(), CancellationToken::new())
1422            .await;
1423        let ctx = self.extension_context(CancellationToken::new());
1424        for extension in &self.inner.extensions {
1425            extension.session_shutdown(&ctx).await;
1426        }
1427    }
1428
1429    /// pi's `cleanupSessionResources`, run when a session ends or another
1430    /// replaces it, but not on reload: the session's cached sockets close
1431    /// and extensions release what they hold for it.
1432    async fn release_resources(&self) {
1433        let id = lock(&self.inner.state).session.id().to_owned();
1434        yapi_ai::api::close_codex_websockets(Some(&id));
1435        for extension in &self.inner.extensions {
1436            extension.end_session(&id).await;
1437        }
1438    }
1439
1440    /// The session's extensions.
1441    pub fn extensions(&self) -> &[Arc<dyn Extension>] {
1442        &self.inner.extensions
1443    }
1444
1445    /// The extension that draws tool `name`, and how.
1446    pub fn tool_renderer(&self, name: &str) -> Option<(Arc<dyn Extension>, ToolRenderers)> {
1447        self.inner.extensions.iter().find_map(|extension| {
1448            let renderers = extension.renderers().tools.get(name).copied()?;
1449            Some((extension.clone(), renderers))
1450        })
1451    }
1452
1453    /// The extension that draws custom entries of type `custom_type`.
1454    pub fn entry_renderer(&self, custom_type: &str) -> Option<Arc<dyn Extension>> {
1455        self.inner
1456            .extensions
1457            .iter()
1458            .find(|extension| {
1459                extension
1460                    .renderers()
1461                    .entries
1462                    .iter()
1463                    .any(|kind| kind == custom_type)
1464            })
1465            .cloned()
1466    }
1467
1468    /// The extensions that transform Markdown, in load order.
1469    pub fn markdown_transformers(&self) -> Vec<Arc<dyn Extension>> {
1470        self.inner
1471            .extensions
1472            .iter()
1473            .filter(|extension| extension.renderers().markdown)
1474            .cloned()
1475            .collect()
1476    }
1477
1478    /// The extensions that render code blocks, in load order.
1479    pub fn code_renderers(&self) -> Vec<Arc<dyn Extension>> {
1480        self.inner
1481            .extensions
1482            .iter()
1483            .filter(|extension| extension.renderers().code)
1484            .cloned()
1485            .collect()
1486    }
1487
1488    /// The extension that draws custom messages of type `custom_type`.
1489    pub fn message_renderer(&self, custom_type: &str) -> Option<Arc<dyn Extension>> {
1490        self.inner
1491            .extensions
1492            .iter()
1493            .find(|extension| {
1494                extension
1495                    .renderers()
1496                    .messages
1497                    .iter()
1498                    .any(|kind| kind == custom_type)
1499            })
1500            .cloned()
1501    }
1502
1503    /// The shortcuts the session's extensions bind, given the keys of each
1504    /// built-in action, and pi's warnings about conflicts.
1505    pub fn extension_shortcuts(
1506        &self,
1507        builtin: &[(String, Vec<String>)],
1508    ) -> (
1509        Vec<crate::extensions::ShortcutBinding>,
1510        Vec<(String, String)>,
1511    ) {
1512        crate::extensions::resolve_shortcuts(&self.inner.extensions, builtin)
1513    }
1514
1515    /// Runs an extension shortcut's handler and returns the handler's error.
1516    pub async fn run_shortcut(
1517        &self,
1518        binding: &crate::extensions::ShortcutBinding,
1519    ) -> Result<(), String> {
1520        let ctx = self.extension_context(CancellationToken::new());
1521        binding
1522            .extension
1523            .run_shortcut(&binding.registered, &ctx)
1524            .await
1525    }
1526
1527    /// pi's `resolveRegisteredCommands`: the extensions' commands in load
1528    /// order, each with the name it is invoked by. A name registered more
1529    /// than once is invoked as `name:1`, `name:2` and so on.
1530    pub fn extension_commands(&self) -> Vec<crate::extensions::ResolvedCommand> {
1531        let all: Vec<(Arc<dyn Extension>, crate::extensions::Command)> = self
1532            .inner
1533            .extensions
1534            .iter()
1535            .flat_map(|extension| {
1536                extension
1537                    .commands()
1538                    .into_iter()
1539                    .map(move |command| (Arc::clone(extension), command))
1540            })
1541            .collect();
1542        let count = |name: &str| {
1543            all.iter()
1544                .filter(|(_, command)| command.name == name)
1545                .count()
1546        };
1547        let mut seen: std::collections::HashMap<String, usize> = std::collections::HashMap::new();
1548        let mut taken: std::collections::HashSet<String> = std::collections::HashSet::new();
1549        all.iter()
1550            .map(|(extension, command)| {
1551                let occurrence = {
1552                    let entry = seen.entry(command.name.clone()).or_default();
1553                    *entry += 1;
1554                    *entry
1555                };
1556                let mut invocation = if count(&command.name) > 1 {
1557                    format!("{}:{occurrence}", command.name)
1558                } else {
1559                    command.name.clone()
1560                };
1561                let mut suffix = occurrence;
1562                while taken.contains(&invocation) {
1563                    suffix += 1;
1564                    invocation = format!("{}:{suffix}", command.name);
1565                }
1566                taken.insert(invocation.clone());
1567                crate::extensions::ResolvedCommand {
1568                    invocation,
1569                    command: command.clone(),
1570                    extension: Arc::clone(extension),
1571                }
1572            })
1573            .collect()
1574    }
1575
1576    /// pi's slash commands, as RPC's `get_commands` and `pi.getCommands()`
1577    /// list them: the extensions' commands, then the prompt templates and
1578    /// the skills, each with where it comes from.
1579    pub fn slash_commands(&self) -> Vec<SlashCommand> {
1580        let resources = self.resources();
1581        let mut commands: Vec<SlashCommand> = self
1582            .extension_commands()
1583            .into_iter()
1584            .map(|resolved| SlashCommand {
1585                name: resolved.invocation,
1586                // pi omits a description the command did not give.
1587                description: Some(resolved.command.description)
1588                    .filter(|description| !description.is_empty()),
1589                source: CommandSource::Extension,
1590                source_info: resolved.extension.source(),
1591            })
1592            .collect();
1593        commands.extend(
1594            resources
1595                .templates
1596                .into_iter()
1597                .map(|template| SlashCommand {
1598                    name: template.name,
1599                    description: Some(template.description),
1600                    source: CommandSource::Prompt,
1601                    source_info: template.source,
1602                }),
1603        );
1604        commands.extend(resources.skills.into_iter().map(|skill| SlashCommand {
1605            name: format!("skill:{}", skill.name),
1606            description: Some(skill.description),
1607            source: CommandSource::Skill,
1608            source_info: skill.source,
1609        }));
1610        commands
1611    }
1612
1613    /// Whether `text` invokes an extension command.
1614    pub fn is_extension_command(&self, text: &str) -> bool {
1615        let Some(rest) = text.strip_prefix('/') else {
1616            return false;
1617        };
1618        let name = rest.split(' ').next().unwrap_or_default();
1619        self.extension_commands()
1620            .iter()
1621            .any(|command| command.invocation == name)
1622    }
1623
1624    /// Runs `/name args` when an extension command is invoked by `name`, and
1625    /// returns whether it did. Errors are the extension's to report.
1626    pub(super) async fn run_extension_command(&self, text: &str) -> bool {
1627        let Some(rest) = text.strip_prefix('/') else {
1628            return false;
1629        };
1630        let (name, args) = rest.split_once(' ').unwrap_or((rest, ""));
1631        let Some(resolved) = self
1632            .extension_commands()
1633            .into_iter()
1634            .find(|command| command.invocation == name)
1635        else {
1636            return false;
1637        };
1638        let ctx = self.extension_context(CancellationToken::new());
1639        resolved
1640            .extension
1641            .run_command(&resolved.command.name, args, &ctx)
1642            .await;
1643        true
1644    }
1645}
1646
1647#[cfg(test)]
1648mod tests {
1649    use serde_json::json;
1650
1651    use super::*;
1652
1653    /// A handler's `details` and `usage` take the places pi's event holds
1654    /// for them, as another extension runtime then sees the event.
1655    #[test]
1656    fn tool_result_key_placement() {
1657        let mut event = json!({"type": "tool_result", "content": [], "isError": false});
1658        let event = event.as_object_mut().unwrap();
1659        set_tool_result_key(event, "structuredContent", json!({}));
1660        set_tool_result_key(event, "usage", json!({}));
1661        set_tool_result_key(event, "details", json!({}));
1662        set_tool_result_key(event, "isError", json!(true));
1663        let keys: Vec<&String> = event.keys().collect();
1664        assert_eq!(
1665            keys,
1666            [
1667                "type",
1668                "content",
1669                "details",
1670                "isError",
1671                "usage",
1672                "structuredContent"
1673            ]
1674        );
1675        assert_eq!(event["isError"], true);
1676    }
1677}