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};
pub(super) const SESSION: &str = "the session grant";
const BLOCK_REASON: &str = "the reason of a tool_call block";
const TOOL_RESULT_KEYS: [&str; 5] = [
"content",
"details",
"structuredContent",
"isError",
"usage",
];
pub(super) struct Hooks {
pub(super) session: AgentSession,
pub(super) cancel: CancellationToken,
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,
}
}
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", &[]),
AgentEvent::MessageStart { .. } => ("message_start", &["message"]),
AgentEvent::MessageUpdate { .. } => {
("message_update", &["message", "assistantMessageEvent"])
}
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);
}
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);
}
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);
}
_ => {}
}
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;
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 {
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);
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> {
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
.part_applies(extension, BLOCK_REASON, 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();
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;
};
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");
let what = "the input change of a tool_call result";
if self
.session
.part_applies(extension, what, granted, &missing)
{
event["input"] = input.clone();
}
}
if result["block"] == true {
let granted = extension.grants().session;
let reason = result["reason"].as_str().filter(|_| {
self.session
.part_applies(extension, BLOCK_REASON, 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;
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 {
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>> {
Box::pin(async move {
self.session
.compact_before_next_response(&self.cancel)
.await?;
Ok(Some(self.session.messages()))
})
}
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 {
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;
}
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,
})
}
}
fn hears(extension: &Arc<dyn Extension>, kind: &str) -> bool {
extension.handles(kind) && receives(kind, &extension.grants())
}
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);
}
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 {
pub(super) fn extension_context(&self, cancel: CancellationToken) -> Context {
let (ui, mode, _) = lock(&self.inner.binding).clone();
let (project_dir, project_trusted, project_only_mcp) = {
let settings = lock(&self.inner.settings);
let keep = crate::config::project_only().then(|| {
let global = settings.document(crate::settings::Scope::Global);
crate::packages::resolve::ProjectOnly::from_settings(global).mcp_servers
});
(
settings.project_dir().to_path_buf(),
settings.project_trusted(),
keep,
)
};
Context {
cwd: self.inner.cwd.clone(),
agent_dir: self.inner.agent_dir.clone(),
project_dir,
project_trusted,
project_only_mcp,
mode,
ui,
tools: self.inner.tools.clone(),
cancel,
session: self.downgrade(),
}
}
pub fn extension_binding(&self) -> (Arc<dyn ExtensionUi>, Mode) {
let (ui, mode, _) = lock(&self.inner.binding).clone();
(ui, mode)
}
pub(super) fn applies(
&self,
extension: &Arc<dyn Extension>,
event: &str,
granted: bool,
missing: &str,
) -> bool {
if !granted {
self.dropped(extension, event, missing);
}
granted
}
pub(super) fn part_applies(
&self,
extension: &Arc<dyn Extension>,
what: &str,
granted: bool,
missing: &str,
) -> bool {
if !granted {
self.ignored(extension, what, missing);
}
granted
}
pub(super) fn dropped(&self, extension: &Arc<dyn Extension>, event: &str, missing: &str) {
self.ignored(extension, &format!("a {event} result"), missing);
}
fn ignored(&self, extension: &Arc<dyn Extension>, what: &str, missing: &str) {
let source = extension.source();
if lock(&self.inner.dropped_results).insert((source.path.clone(), what.to_owned())) {
let project = self.project_dir();
let label = crate::chest::startup::entry_label(&source.source, &source.scope, &project);
let label = label.unwrap_or(source.path);
extension.warn(&format!(
"Ignored {what} of extension {label}: 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()
}
pub fn has_handlers(&self, kind: &str) -> bool {
self.inner
.extensions
.iter()
.any(|extension| hears(extension, kind))
}
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;
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
}
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)
}
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: ¤t }.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
})
}
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
}
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
}
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,
}
}
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(),
)));
}
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(),
)));
}
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 });
}
}
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();
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(())
}
pub fn check_mcp_server(&self, owner: u64, name: &str) -> Result<(), String> {
check_registration(others(&lock(&self.inner.mcp_servers).0, owner), name)
}
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);
}
}
pub fn mcp_servers(&self) -> Vec<RegisteredMcpServer> {
lock(&self.inner.mcp_servers)
.0
.values()
.map(|(_, server)| server.clone())
.collect()
}
pub fn add_extension_load_error(&self, path: String, error: String) {
lock(&self.inner.extension_load_errors).push((path, error));
}
pub fn extension_load_errors(&self) -> Vec<(String, String)> {
lock(&self.inner.extension_load_errors).clone()
}
pub fn add_extension_warning(&self, path: String, warning: String) {
lock(&self.inner.extension_warnings).push((path, warning));
}
pub fn extension_warnings(&self) -> Vec<(String, String)> {
lock(&self.inner.extension_warnings).clone()
}
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;
}
}
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))
}
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;
};
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(),
}));
}
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)
}
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),
}
}
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;
}
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;
};
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(),
});
}
}
}
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);
}
}
pub async fn shutdown(&self) {
self.stop_extensions(ExtensionEvent::SessionShutdown {
reason: "quit",
target_session_file: None,
})
.await;
self.release_resources().await;
}
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;
}
}
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;
}
}
pub fn extensions(&self) -> &[Arc<dyn Extension>] {
&self.inner.extensions
}
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))
})
}
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()
}
pub fn markdown_transformers(&self) -> Vec<Arc<dyn Extension>> {
self.inner
.extensions
.iter()
.filter(|extension| extension.renderers().markdown)
.cloned()
.collect()
}
pub fn code_renderers(&self) -> Vec<Arc<dyn Extension>> {
self.inner
.extensions
.iter()
.filter(|extension| extension.renderers().code)
.cloned()
.collect()
}
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()
}
pub fn extension_shortcuts(
&self,
builtin: &[(String, Vec<String>)],
) -> (
Vec<crate::extensions::ShortcutBinding>,
Vec<(String, String)>,
) {
crate::extensions::resolve_shortcuts(&self.inner.extensions, builtin)
}
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
}
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()
}
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,
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
}
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)
}
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::*;
#[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);
}
}